From 0f422f549a7d6ff86c9936a957ac8d9b2444a6d5 Mon Sep 17 00:00:00 2001 From: Sasha Sobran Date: Thu, 26 Aug 2021 09:37:49 -0400 Subject: [PATCH 1/2] fix: Populate service_account and network in PipelineJob instead of pipeline_spec --- google/cloud/aiplatform/pipeline_jobs.py | 4 ++-- tests/unit/aiplatform/test_pipeline_jobs.py | 6 ++++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index bb98b1f0d5..5c948d17b9 100644 --- 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 index 0e3eddbf22..2570c74a55 100644 --- 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( From af85adfe8be5c48a2f5e588e2ef553cb8dde7bf5 Mon Sep 17 00:00:00 2001 From: Sasha Sobran Date: Thu, 26 Aug 2021 10:05:52 -0400 Subject: [PATCH 2/2] chore:lint --- tests/unit/aiplatform/test_pipeline_jobs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index 2570c74a55..bc39f4d9d9 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -85,7 +85,7 @@ def mock_pipeline_service_create(): state=gca_pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_SUCCEEDED, create_time=_TEST_PIPELINE_CREATE_TIME, service_account=_TEST_SERVICE_ACCOUNT, - network=_TEST_NETWORK + network=_TEST_NETWORK, ) yield mock_create_pipeline_job @@ -96,7 +96,7 @@ def make_pipeline_job(state): state=state, create_time=_TEST_PIPELINE_CREATE_TIME, service_account=_TEST_SERVICE_ACCOUNT, - network=_TEST_NETWORK + network=_TEST_NETWORK, ) @@ -244,7 +244,7 @@ def test_run_call_pipeline_service_create( }, runtime_config=runtime_config, service_account=_TEST_SERVICE_ACCOUNT, - network=_TEST_NETWORK + network=_TEST_NETWORK, ) mock_pipeline_service_create.assert_called_once_with(