Skip to content

Commit 8fde2ce

Browse files
authored
fix: Populate service_account and network in PipelineJob instead of pipeline_spec (#658)
1 parent d97da41 commit 8fde2ce

File tree

2 files changed

+8
-2
lines changed

2 files changed

+8
-2
lines changed

google/cloud/aiplatform/pipeline_jobs.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -250,10 +250,10 @@ def run(
250250
Optional. Whether to execute this method synchronously. If False, this method will unblock and it will be executed in a concurrent Future.
251251
"""
252252
if service_account:
253-
self._gca_resource.pipeline_spec.service_account = service_account
253+
self._gca_resource.service_account = service_account
254254

255255
if network:
256-
self._gca_resource.pipeline_spec.network = network
256+
self._gca_resource.network = network
257257

258258
_LOGGER.log_create_with_lro(self.__class__)
259259

tests/unit/aiplatform/test_pipeline_jobs.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,8 @@ def mock_pipeline_service_create():
8484
name=_TEST_PIPELINE_JOB_NAME,
8585
state=gca_pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_SUCCEEDED,
8686
create_time=_TEST_PIPELINE_CREATE_TIME,
87+
service_account=_TEST_SERVICE_ACCOUNT,
88+
network=_TEST_NETWORK,
8789
)
8890
yield mock_create_pipeline_job
8991

@@ -93,6 +95,8 @@ def make_pipeline_job(state):
9395
name=_TEST_PIPELINE_JOB_NAME,
9496
state=state,
9597
create_time=_TEST_PIPELINE_CREATE_TIME,
98+
service_account=_TEST_SERVICE_ACCOUNT,
99+
network=_TEST_NETWORK,
96100
)
97101

98102

@@ -239,6 +243,8 @@ def test_run_call_pipeline_service_create(
239243
"root": _TEST_PIPELINE_JOB_SPEC["pipelineSpec"]["root"],
240244
},
241245
runtime_config=runtime_config,
246+
service_account=_TEST_SERVICE_ACCOUNT,
247+
network=_TEST_NETWORK,
242248
)
243249

244250
mock_pipeline_service_create.assert_called_once_with(

0 commit comments

Comments
 (0)