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",