feat: Add preflight validations to PipelineJob submit and run methods.

PiperOrigin-RevId: 651504412

Author: vertex-sdk-bot

google/cloud/aiplatform/pipeline_jobs.py

@@ -305,6 +305,7 @@ def run(
         reserved_ip_ranges: Optional[List[str]] = None,
         sync: Optional[bool] = True,
         create_request_timeout: Optional[float] = None,
+        enable_preflight_validations: Optional[bool] = False,
     ) -> None:
         """Run this configured PipelineJob and monitor the job until completion.
 
@@ -325,6 +326,8 @@ def run(
                 Optional. Whether to execute this method synchronously. If False, this method will unblock and it will be executed in a concurrent Future.
             create_request_timeout (float):
                 Optional. The timeout for the create request in seconds.
+            enable_preflight_validations (bool):
+                Optional. Whether to enable preflight validations for the PipelineJob.
         """
         network = network or initializer.global_config.network
 
@@ -334,6 +337,7 @@ def run(
             reserved_ip_ranges=reserved_ip_ranges,
             sync=sync,
             create_request_timeout=create_request_timeout,
+            enable_preflight_validations=enable_preflight_validations,
         )
 
     @base.optional_sync()
@@ -344,6 +348,7 @@ def _run(
         reserved_ip_ranges: Optional[List[str]] = None,
         sync: Optional[bool] = True,
         create_request_timeout: Optional[float] = None,
+        enable_preflight_validations: Optional[bool] = False,
     ) -> None:
         """Helper method to ensure network synchronization and to run
         the configured PipelineJob and monitor the job until completion.
@@ -363,12 +368,15 @@ def _run(
                 Optional. Whether to execute this method synchronously. If False, this method will unblock and it will be executed in a concurrent Future.
             create_request_timeout (float):
                 Optional. The timeout for the create request in seconds.
+            enable_preflight_validations (bool):
+                Optional. Whether to enable preflight validations for the PipelineJob.
         """
         self.submit(
             service_account=service_account,
             network=network,
             reserved_ip_ranges=reserved_ip_ranges,
             create_request_timeout=create_request_timeout,
+            enable_preflight_validations=enable_preflight_validations,
         )
 
         self._block_until_complete()
@@ -402,6 +410,7 @@ def submit(
         create_request_timeout: Optional[float] = None,
         *,
         experiment: Optional[Union[str, experiment_resources.Experiment]] = None,
+        enable_preflight_validations: Optional[bool] = False,
     ) -> None:
         """Run this configured PipelineJob.
 
@@ -432,6 +441,8 @@ def submit(
 
                 Pipeline parameters will be associated as parameters to the
                 current Experiment Run.
+            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
@@ -471,12 +482,47 @@ def submit(
 
         _LOGGER.log_create_with_lro(self.__class__)
 
-        self._gca_resource = self.api_client.create_pipeline_job(
-            parent=self._parent,
-            pipeline_job=self._gca_resource,
-            pipeline_job_id=self.job_id,
-            timeout=create_request_timeout,
-        )
+        if enable_preflight_validations:
+            self._gca_resource.preflight_validations = True
+
+        def extract_error_messages(error_string):
+            """
+            Extracts error messages from a string containing structured errors.
+
+            Args:
+                error_string: The string containing the error data.
+
+            Returns:
+                A list of formatted error messages.
+            """
+
+            message_pattern = (
+                r"CreatePipelineJobApiErrorDetail\"\n.*message=(.*),\ cause=null"
+            )
+
+            matches = re.findall(message_pattern, error_string)
+
+            formatted_errors = [
+                f"{i+1}. {message}" for i, message in enumerate(matches)
+            ]
+
+            return formatted_errors
+
+        try:
+            self._gca_resource = self.api_client.create_pipeline_job(
+                parent=self._parent,
+                pipeline_job=self._gca_resource,
+                pipeline_job_id=self.job_id,
+                timeout=create_request_timeout,
+            )
+        except Exception as e:
+            preflight_validations_error_messages = extract_error_messages(str(e))
+            if preflight_validations_error_messages:
+                raise Exception(
+                    "PipelineJob Preflight validations failed with the following errors:\n"
+                    + "\n".join(preflight_validations_error_messages)
+                ) from e
+            raise
 
         _LOGGER.log_create_complete_with_getter(
             self.__class__, self._gca_resource, "pipeline_job"

tests/unit/aiplatform/test_pipeline_jobs.py

@@ -242,6 +242,22 @@
 )
 
 
+@pytest.fixture
+def mock_pipeline_service_create_with_preflight_validations():
+    with mock.patch.object(
+        pipeline_service_client.PipelineServiceClient, "create_pipeline_job"
+    ) as mock_create_pipeline_job:
+        mock_create_pipeline_job.return_value = gca_pipeline_job.PipelineJob(
+            name=_TEST_PIPELINE_JOB_NAME,
+            state=gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED,
+            create_time=_TEST_PIPELINE_CREATE_TIME,
+            service_account=_TEST_SERVICE_ACCOUNT,
+            network=_TEST_NETWORK,
+            reserved_ip_ranges=_TEST_RESERVED_IP_RANGES,
+        )
+        yield mock_create_pipeline_job
+
+
 @pytest.fixture
 def mock_pipeline_service_create():
     with mock.patch.object(
@@ -2267,3 +2283,89 @@ def test_submit_v1beta1_pipeline_job_returns_response(
         job.submit()
 
         assert mock_pipeline_v1beta1_service_create.call_count == 1
+
+    @pytest.mark.parametrize(
+        "job_spec",
+        [_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_SPEC_YAML, _TEST_PIPELINE_JOB],
+    )
+    @pytest.mark.parametrize("sync", [True, False])
+    def test_run_call_pipeline_service_run_with_preflight_validations(
+        self,
+        mock_pipeline_service_create_with_preflight_validations,
+        mock_pipeline_service_get,
+        mock_pipeline_bucket_exists,
+        job_spec,
+        mock_load_yaml_and_json,
+        sync,
+    ):
+        import yaml
+
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            staging_bucket=_TEST_GCS_BUCKET_NAME,
+            location=_TEST_LOCATION,
+            credentials=_TEST_CREDENTIALS,
+            service_account=_TEST_SERVICE_ACCOUNT,
+            network=_TEST_NETWORK,
+        )
+
+        job = pipeline_jobs.PipelineJob(
+            display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
+            template_path=_TEST_TEMPLATE_PATH,
+            job_id=_TEST_PIPELINE_JOB_ID,
+            parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
+            input_artifacts=_TEST_PIPELINE_INPUT_ARTIFACTS,
+            enable_caching=True,
+        )
+
+        job.run(
+            reserved_ip_ranges=_TEST_RESERVED_IP_RANGES,
+            sync=sync,
+            create_request_timeout=None,
+            enable_preflight_validations=True,
+        )
+
+        if not sync:
+            job.wait()
+
+        expected_runtime_config_dict = {
+            "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME,
+            "parameterValues": _TEST_PIPELINE_PARAMETER_VALUES,
+            "inputArtifacts": {"vertex_model": {"artifactId": "456"}},
+        }
+        runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb
+        json_format.ParseDict(expected_runtime_config_dict, runtime_config)
+
+        job_spec = yaml.safe_load(job_spec)
+        pipeline_spec = job_spec.get("pipelineSpec") or job_spec
+
+        # Construct expected request
+        expected_gapic_pipeline_job = gca_pipeline_job.PipelineJob(
+            display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
+            pipeline_spec={
+                "components": {},
+                "pipelineInfo": pipeline_spec["pipelineInfo"],
+                "root": pipeline_spec["root"],
+                "schemaVersion": "2.1.0",
+            },
+            runtime_config=runtime_config,
+            service_account=_TEST_SERVICE_ACCOUNT,
+            network=_TEST_NETWORK,
+            reserved_ip_ranges=_TEST_RESERVED_IP_RANGES,
+            preflight_validations=True,
+        )
+
+        mock_pipeline_service_create_with_preflight_validations.assert_called_once_with(
+            parent=_TEST_PARENT,
+            pipeline_job=expected_gapic_pipeline_job,
+            pipeline_job_id=_TEST_PIPELINE_JOB_ID,
+            timeout=None,
+        )
+
+        mock_pipeline_service_get.assert_called_with(
+            name=_TEST_PIPELINE_JOB_NAME, retry=base._DEFAULT_RETRY
+        )
+
+        assert job._gca_resource == make_pipeline_job(
+            gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED
+        )