diff --git a/google/cloud/aiplatform/preview/pipelinejob/pipeline_jobs.py b/google/cloud/aiplatform/preview/pipelinejob/pipeline_jobs.py index b0eb7293e4..969dbce254 100644 --- a/google/cloud/aiplatform/preview/pipelinejob/pipeline_jobs.py +++ b/google/cloud/aiplatform/preview/pipelinejob/pipeline_jobs.py @@ -438,6 +438,8 @@ def submit( reserved_ip_ranges: Optional[List[str]] = None, create_request_timeout: Optional[float] = None, job_id: Optional[str] = None, + *, + enable_preflight_validations: Optional[bool] = False, ) -> None: """Run this configured PipelineJob. @@ -464,6 +466,8 @@ def submit( Optional. The ID to use for the PipelineJob, which will become the final component of the PipelineJob name. If not provided, an ID will be automatically generated. + enable_preflight_validations (bool): + Optional. Whether to enable preflight validations for the PipelineJob. """ network = network or initializer.global_config.network service_account = service_account or initializer.global_config.service_account @@ -475,6 +479,9 @@ def submit( if network: gca_resouce.network = network + if enable_preflight_validations: + gca_resouce.preflight_validations = True + if reserved_ip_ranges: gca_resouce.reserved_ip_ranges = reserved_ip_ranges user_project = initializer.global_config.project diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index d7addc86a9..d1252d94ce 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -2415,6 +2415,46 @@ def test_submit_v1beta1_pipeline_job_returns_response( assert mock_pipeline_v1beta1_service_create.call_count == 1 + @pytest.mark.parametrize("job_spec", [_TEST_PIPELINE_SPEC_JSON]) + @pytest.mark.parametrize("enable_preflight_validations", [True, False]) + @pytest.mark.parametrize("sync", [True, False]) + @mock.patch.object(pipeline_jobs, "_JOB_WAIT_TIME", 0.05) + @mock.patch.object(pipeline_jobs, "_LOG_WAIT_TIME", 0.05) + def test_run_v1beta1_pipeline_job_passes_preflight_validations( + self, + mock_load_yaml_and_json, + job_spec, + mock_pipeline_v1beta1_service_create, + mock_pipeline_service_get, + enable_preflight_validations, + sync, + ): + aiplatform.init( + project=_TEST_PROJECT, + staging_bucket=_TEST_GCS_BUCKET_NAME, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + + job = preview_pipeline_jobs._PipelineJob( + display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, + template_path=_TEST_TEMPLATE_PATH, + job_id=_TEST_PIPELINE_JOB_ID, + ) + + # run() comes from the GA PipelineJob and passes + # enable_preflight_validations on to submit() + job.run(sync=sync, enable_preflight_validations=enable_preflight_validations) + + if not sync: + job.wait() + + assert mock_pipeline_v1beta1_service_create.call_count == 1 + request = mock_pipeline_v1beta1_service_create.call_args.kwargs["request"] + assert ( + request.pipeline_job.preflight_validations is enable_preflight_validations + ) + @pytest.mark.usefixtures( "mock_pipeline_v1beta1_service_create", )