fix: Populate service_account and network in PipelineJob instead of p…

…ipeline_spec (#658)
googleapis · Aug 26, 2021 · 8fde2ce · 8fde2ce
1 parent d97da41
commit 8fde2ce
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py
@@ -250,10 +250,10 @@ def run(
                 Optional. Whether to execute this method synchronously. If False, this method will unblock and it will be executed in a concurrent Future.
         """
         if service_account:
-            self._gca_resource.pipeline_spec.service_account = service_account
+            self._gca_resource.service_account = service_account
 
         if network:
-            self._gca_resource.pipeline_spec.network = network
+            self._gca_resource.network = network
 
         _LOGGER.log_create_with_lro(self.__class__)
 

diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py
@@ -84,6 +84,8 @@ def mock_pipeline_service_create():
             name=_TEST_PIPELINE_JOB_NAME,
             state=gca_pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_SUCCEEDED,
             create_time=_TEST_PIPELINE_CREATE_TIME,
+            service_account=_TEST_SERVICE_ACCOUNT,
+            network=_TEST_NETWORK,
         )
         yield mock_create_pipeline_job
 
@@ -93,6 +95,8 @@ def make_pipeline_job(state):
         name=_TEST_PIPELINE_JOB_NAME,
         state=state,
         create_time=_TEST_PIPELINE_CREATE_TIME,
+        service_account=_TEST_SERVICE_ACCOUNT,
+        network=_TEST_NETWORK,
     )
 
 
@@ -239,6 +243,8 @@ def test_run_call_pipeline_service_create(
                 "root": _TEST_PIPELINE_JOB_SPEC["pipelineSpec"]["root"],
             },
             runtime_config=runtime_config,
+            service_account=_TEST_SERVICE_ACCOUNT,
+            network=_TEST_NETWORK,
         )
 
         mock_pipeline_service_create.assert_called_once_with(