From 429a2946e941076fe6deed0626497ea4888dc50a Mon Sep 17 00:00:00 2001 From: abhishekacharya200 Date: Thu, 22 Oct 2026 21:11:43 +1100 Subject: [PATCH] fix: use the job's project, location and credentials in preview PipelineJob submit and rerun The preview _PipelineJob.submit() and rerun() built the request parent and the v1beta1 client from the aiplatform.init() project and location, and created the client without the job's credentials. submit() also dropped create_request_timeout. A job created with its own project, location or credentials was therefore sent to the global project and region with the global credentials. Use the job's parent, location and credentials the way the GA PipelineJob.submit() does, and pass create_request_timeout to create_pipeline_job. --- .../preview/pipelinejob/pipeline_jobs.py | 24 ++---- tests/unit/aiplatform/test_pipeline_jobs.py | 84 +++++++++++++++++++ 2 files changed, 93 insertions(+), 15 deletions(-) diff --git a/google/cloud/aiplatform/preview/pipelinejob/pipeline_jobs.py b/google/cloud/aiplatform/preview/pipelinejob/pipeline_jobs.py index b0eb7293e4..55185893c2 100644 --- a/google/cloud/aiplatform/preview/pipelinejob/pipeline_jobs.py +++ b/google/cloud/aiplatform/preview/pipelinejob/pipeline_jobs.py @@ -477,14 +477,10 @@ def submit( if reserved_ip_ranges: gca_resouce.reserved_ip_ranges = reserved_ip_ranges - user_project = initializer.global_config.project - user_location = initializer.global_config.location - parent = initializer.global_config.common_location_path( - project=user_project, location=user_location - ) client = self._instantiate_client( - location=user_location, + location=self.location, + credentials=self.credentials, appended_user_agent=["preview-pipeline-job-submit"], ) v1beta1_client = client.select_version(compat.V1BETA1) @@ -492,12 +488,14 @@ def submit( _LOGGER.log_create_with_lro(self.__class__) request = aiplatform_v1beta1.CreatePipelineJobRequest( - parent=parent, + parent=self._parent, pipeline_job=self._v1_beta1_pipeline_job, pipeline_job_id=job_id or self.job_id, ) - response = v1beta1_client.create_pipeline_job(request=request) + response = v1beta1_client.create_pipeline_job( + request=request, timeout=create_request_timeout + ) self._gca_resource = response @@ -561,14 +559,10 @@ def rerun( if reserved_ip_ranges: gca_resouce.reserved_ip_ranges = reserved_ip_ranges - user_project = initializer.global_config.project - user_location = initializer.global_config.location - parent = initializer.global_config.common_location_path( - project=user_project, location=user_location - ) client = self._instantiate_client( - location=user_location, + location=self.location, + credentials=self.credentials, appended_user_agent=["preview-pipeline-job-submit"], ) v1beta1_client = client.select_version(compat.V1BETA1) @@ -627,7 +621,7 @@ def rerun( ) request = aiplatform_v1beta1.CreatePipelineJobRequest( - parent=parent, + parent=self._parent, pipeline_job=self._v1_beta1_pipeline_job, pipeline_job_id=job_id, ) diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index d7addc86a9..e45704610d 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -70,6 +70,9 @@ _TEST_PROJECT = "test-project" _TEST_LOCATION = "us-central1" +# a project and location for the job itself, different from aiplatform.init +_TEST_JOB_PROJECT = "test-job-project" +_TEST_JOB_LOCATION = "europe-west4" _TEST_PIPELINE_JOB_DISPLAY_NAME = "sample-pipeline-job-display-name" _TEST_PIPELINE_JOB_DISPLAY_NAME_2 = "sample-pipeline-job-display-name-2" _TEST_PIPELINE_JOB_ID = "sample-test-pipeline-202111111" @@ -2498,6 +2501,87 @@ def test_rerun_v1beta1_pipeline_job_returns_response( assert mock_pipeline_v1beta1_service_get.call_count == 1 assert mock_pipeline_v1beta1_service_create.call_count == 2 + def _make_preview_job_in_another_project(self, job_credentials): + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + staging_bucket=_TEST_GCS_BUCKET_NAME, + credentials=_TEST_CREDENTIALS, + ) + return preview_pipeline_jobs._PipelineJob( + display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, + template_path=_TEST_TEMPLATE_PATH, + job_id=_TEST_PIPELINE_JOB_ID, + project=_TEST_JOB_PROJECT, + location=_TEST_JOB_LOCATION, + credentials=job_credentials, + ) + + def _assert_created_with_the_jobs_project_location_and_credentials( + self, mock_create, job_credentials + ): + # autospec records the client instance that made the call + client = mock_create.call_args.args[0] + assert client._transport._host == ( + f"{_TEST_JOB_LOCATION}-aiplatform.googleapis.com:443" + ) + assert client._transport._credentials == job_credentials + assert mock_create.call_args.kwargs["request"].parent == ( + f"projects/{_TEST_JOB_PROJECT}/locations/{_TEST_JOB_LOCATION}" + ) + + @pytest.mark.parametrize("job_spec", [_TEST_PIPELINE_SPEC_JSON]) + def test_submit_v1beta1_pipeline_job_uses_the_jobs_project_location_and_credentials( + self, + mock_load_yaml_and_json, + job_spec, + ): + job_credentials = auth_credentials.AnonymousCredentials() + job = self._make_preview_job_in_another_project(job_credentials) + + with mock.patch.object( + v1beta1_pipeline_service.PipelineServiceClient, + "create_pipeline_job", + autospec=True, + return_value=v1beta1_pipeline_job.PipelineJob(name=_TEST_PIPELINE_JOB_NAME), + ) as mock_create: + job.submit(create_request_timeout=180.0) + + self._assert_created_with_the_jobs_project_location_and_credentials( + mock_create, job_credentials + ) + assert mock_create.call_args.kwargs["timeout"] == 180.0 + + @pytest.mark.usefixtures("mock_pipeline_v1beta1_service_get") + @pytest.mark.parametrize("job_spec", [_TEST_PIPELINE_SPEC_JSON]) + def test_rerun_v1beta1_pipeline_job_uses_the_jobs_project_location_and_credentials( + self, + mock_load_yaml_and_json, + job_spec, + ): + job_credentials = auth_credentials.AnonymousCredentials() + job = self._make_preview_job_in_another_project(job_credentials) + + with mock.patch.object( + v1beta1_pipeline_service.PipelineServiceClient, + "create_pipeline_job", + autospec=True, + return_value=v1beta1_pipeline_job.PipelineJob(name=_TEST_PIPELINE_JOB_NAME), + ) as mock_create: + job.rerun( + original_pipelinejob_name=_TEST_PIPELINE_JOB_NAME, + pipeline_task_rerun_configs=[ + aiplatform_v1beta1.PipelineTaskRerunConfig( + task_name="task-name", + task_id=100, + ) + ], + ) + + self._assert_created_with_the_jobs_project_location_and_credentials( + mock_create, job_credentials + ) + @pytest.mark.usefixtures( "mock_pipeline_v1beta1_service_create", "mock_pipeline_v1beta1_service_get_with_task_details",