Skip to content

Commit 3f56ae7

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
chore: Add tensorboard to custom job samples.
PiperOrigin-RevId: 585763909
1 parent 4a71c8c commit 3f56ae7

6 files changed

+24
-0
lines changed

samples/model-builder/create_training_pipeline_custom_container_job_sample.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ def create_training_pipeline_custom_container_job_sample(
3636
validation_fraction_split: float = 0.1,
3737
test_fraction_split: float = 0.1,
3838
sync: bool = True,
39+
tensorboard_resource_name: Optional[str] = None,
40+
service_account: Optional[str] = None,
3941
):
4042
aiplatform.init(project=project, location=location, staging_bucket=staging_bucket)
4143

@@ -60,6 +62,8 @@ def create_training_pipeline_custom_container_job_sample(
6062
validation_fraction_split=validation_fraction_split,
6163
test_fraction_split=test_fraction_split,
6264
sync=sync,
65+
tensorboard=tensorboard_resource_name,
66+
service_account=service_account,
6367
)
6468

6569
model.wait()

samples/model-builder/create_training_pipeline_custom_container_job_test.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,8 @@ def test_create_training_pipeline_custom_container_job_sample(
4242
training_fraction_split=constants.TRAINING_FRACTION_SPLIT,
4343
validation_fraction_split=constants.VALIDATION_FRACTION_SPLIT,
4444
test_fraction_split=constants.TEST_FRACTION_SPLIT,
45+
tensorboard_resource_name=constants.TENSORBOARD_NAME,
46+
service_account=constants.SERVICE_ACCOUNT,
4547
)
4648

4749
mock_sdk_init.assert_called_once_with(
@@ -67,4 +69,6 @@ def test_create_training_pipeline_custom_container_job_sample(
6769
validation_fraction_split=constants.VALIDATION_FRACTION_SPLIT,
6870
test_fraction_split=constants.TEST_FRACTION_SPLIT,
6971
sync=True,
72+
tensorboard=constants.TENSORBOARD_NAME,
73+
service_account=constants.SERVICE_ACCOUNT,
7074
)

samples/model-builder/create_training_pipeline_custom_job_sample.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ def create_training_pipeline_custom_job_sample(
3737
validation_fraction_split: float = 0.1,
3838
test_fraction_split: float = 0.1,
3939
sync: bool = True,
40+
tensorboard_resource_name: Optional[str] = None,
41+
service_account: Optional[str] = None,
4042
):
4143
aiplatform.init(project=project, location=location, staging_bucket=staging_bucket)
4244

@@ -62,6 +64,8 @@ def create_training_pipeline_custom_job_sample(
6264
validation_fraction_split=validation_fraction_split,
6365
test_fraction_split=test_fraction_split,
6466
sync=sync,
67+
tensorboard=tensorboard_resource_name,
68+
service_account=service_account,
6569
)
6670

6771
model.wait()

samples/model-builder/create_training_pipeline_custom_job_test.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ def test_create_training_pipeline_custom_job_sample(
4343
training_fraction_split=constants.TRAINING_FRACTION_SPLIT,
4444
validation_fraction_split=constants.VALIDATION_FRACTION_SPLIT,
4545
test_fraction_split=constants.TEST_FRACTION_SPLIT,
46+
tensorboard_resource_name=constants.TENSORBOARD_NAME,
47+
service_account=constants.SERVICE_ACCOUNT,
4648
)
4749

4850
mock_sdk_init.assert_called_once_with(
@@ -68,4 +70,6 @@ def test_create_training_pipeline_custom_job_sample(
6870
validation_fraction_split=constants.VALIDATION_FRACTION_SPLIT,
6971
test_fraction_split=constants.TEST_FRACTION_SPLIT,
7072
sync=True,
73+
tensorboard=constants.TENSORBOARD_NAME,
74+
service_account=constants.SERVICE_ACCOUNT,
7175
)

samples/model-builder/create_training_pipeline_custom_package_job_sample.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ def create_training_pipeline_custom_package_job_sample(
3838
validation_fraction_split: float = 0.1,
3939
test_fraction_split: float = 0.1,
4040
sync: bool = True,
41+
tensorboard_resource_name: Optional[str] = None,
42+
service_account: Optional[str] = None,
4143
):
4244
aiplatform.init(project=project, location=location, staging_bucket=staging_bucket)
4345

@@ -64,6 +66,8 @@ def create_training_pipeline_custom_package_job_sample(
6466
validation_fraction_split=validation_fraction_split,
6567
test_fraction_split=test_fraction_split,
6668
sync=sync,
69+
tensorboard=tensorboard_resource_name,
70+
service_account=service_account,
6771
)
6872

6973
model.wait()

samples/model-builder/create_training_pipeline_custom_package_job_test.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ def test_create_training_pipeline_custom_package_job_sample(
4444
training_fraction_split=constants.TRAINING_FRACTION_SPLIT,
4545
validation_fraction_split=constants.VALIDATION_FRACTION_SPLIT,
4646
test_fraction_split=constants.TEST_FRACTION_SPLIT,
47+
tensorboard_resource_name=constants.TENSORBOARD_NAME,
48+
service_account=constants.SERVICE_ACCOUNT,
4749
)
4850

4951
mock_sdk_init.assert_called_once_with(
@@ -72,4 +74,6 @@ def test_create_training_pipeline_custom_package_job_sample(
7274
validation_fraction_split=constants.VALIDATION_FRACTION_SPLIT,
7375
test_fraction_split=constants.TEST_FRACTION_SPLIT,
7476
sync=True,
77+
tensorboard=constants.TENSORBOARD_NAME,
78+
service_account=constants.SERVICE_ACCOUNT,
7579
)

0 commit comments

Comments
 (0)