fix: log pipeline completion and raise pipeline failures (#523)

Author: sasha-gitg

google/cloud/aiplatform/pipeline_jobs.py

@@ -44,6 +44,10 @@
  ]
 )
 
+_PIPELINE_ERROR_STATES = set(
+ [gca_pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_FAILED]
+)
+
 # Vertex AI Pipelines service API job name relative name prefix pattern.
 _JOB_NAME_PATTERN = "{parent}/pipelineJobs/{job_id}"
 
@@ -311,6 +315,13 @@ def _block_until_complete(self):
  previous_time = current_time
  time.sleep(wait)
 
+ # Error is only populated when the job state is
+ # JOB_STATE_FAILED or JOB_STATE_CANCELLED.
+ if self._gca_resource.state in _PIPELINE_ERROR_STATES:
+ raise RuntimeError("Job failed with:\n%s" % self._gca_resource.error)
+ else:
+ _LOGGER.log_action_completed_against_resource("run", "completed", self)
+
  def cancel(self) -> None:
  """Starts asynchronous cancellation on the PipelineJob. The server
  makes a best effort to cancel the job, but success is not guaranteed.

tests/unit/aiplatform/test_pipeline_jobs.py

@@ -133,6 +133,26 @@ def mock_pipeline_service_get():
  yield mock_get_pipeline_job
 
 
+@pytest.fixture
+def mock_pipeline_service_get_with_fail():
+ with mock.patch.object(
+ pipeline_service_client_v1beta1.PipelineServiceClient, "get_pipeline_job"
+ ) as mock_get_pipeline_job:
+ mock_get_pipeline_job.side_effect = [
+ make_pipeline_job(
+ gca_pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_RUNNING
+ ),
+ make_pipeline_job(
+ gca_pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_RUNNING
+ ),
+ make_pipeline_job(
+ gca_pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_FAILED
+ ),
+ ]
+
+ yield mock_get_pipeline_job
+
+
 @pytest.fixture
 def mock_pipeline_service_cancel():
  with mock.patch.object(
@@ -269,3 +289,33 @@ def test_cancel_pipeline_job_without_running(
  job.cancel()
 
  assert e.match(regexp=r"PipelineJob has not been launched")
+
+ @pytest.mark.usefixtures(
+ "mock_pipeline_service_create",
+ "mock_pipeline_service_get_with_fail",
+ "mock_load_json",
+ )
+ @pytest.mark.parametrize("sync", [True, False])
+ def test_pipeline_failure_raises(self, sync):
+ aiplatform.init(
+ project=_TEST_PROJECT,
+ staging_bucket=_TEST_GCS_BUCKET_NAME,
+ location=_TEST_LOCATION,
+ credentials=_TEST_CREDENTIALS,
+ )
+
+ job = pipeline_jobs.PipelineJob(
+ display_name=_TEST_PIPELINE_JOB_ID,
+ template_path=_TEST_TEMPLATE_PATH,
+ job_id=_TEST_PIPELINE_JOB_ID,
+ parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
+ enable_caching=True,
+ )
+
+ with pytest.raises(RuntimeError):
+ job.run(
+ service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK, sync=sync,
+ )
+
+ if not sync:
+ job.wait()