feat: add support for failure_policy in PipelineJob (#1452)

* implement failure_policy * add tests * raise valueerror if failure_policy is invalid

Author: chongyouquan

google/cloud/aiplatform/compat/__init__.py

@@ -96,6 +96,7 @@
  types.model_evaluation_slice = types.model_evaluation_slice_v1beta1
  types.model_service = types.model_service_v1beta1
  types.operation = types.operation_v1beta1
+ types.pipeline_failure_policy = types.pipeline_failure_policy_v1beta1
  types.pipeline_job = types.pipeline_job_v1beta1
  types.pipeline_service = types.pipeline_service_v1beta1
  types.pipeline_state = types.pipeline_state_v1beta1
@@ -180,6 +181,7 @@
  types.model_evaluation_slice = types.model_evaluation_slice_v1
  types.model_service = types.model_service_v1
  types.operation = types.operation_v1
+ types.pipeline_failure_policy = types.pipeline_failure_policy_v1
  types.pipeline_job = types.pipeline_job_v1
  types.pipeline_service = types.pipeline_service_v1
  types.pipeline_state = types.pipeline_state_v1

google/cloud/aiplatform/compat/types/__init__.py

@@ -63,6 +63,7 @@
  model_evaluation_slice as model_evaluation_slice_v1beta1,
  model_service as model_service_v1beta1,
  operation as operation_v1beta1,
+ pipeline_failure_policy as pipeline_failure_policy_v1beta1,
  pipeline_job as pipeline_job_v1beta1,
  pipeline_service as pipeline_service_v1beta1,
  pipeline_state as pipeline_state_v1beta1,
@@ -126,6 +127,7 @@
  model_evaluation_slice as model_evaluation_slice_v1,
  model_service as model_service_v1,
  operation as operation_v1,
+ pipeline_failure_policy as pipeline_failure_policy_v1,
  pipeline_job as pipeline_job_v1,
  pipeline_service as pipeline_service_v1,
  pipeline_state as pipeline_state_v1,
@@ -191,6 +193,7 @@
  model_evaluation_slice_v1,
  model_service_v1,
  operation_v1,
+ pipeline_failure_policy_v1beta1,
  pipeline_job_v1,
  pipeline_service_v1,
  pipeline_state_v1,
@@ -253,6 +256,7 @@
  model_evaluation_slice_v1beta1,
  model_service_v1beta1,
  operation_v1beta1,
+ pipeline_failure_policy_v1beta1,
  pipeline_job_v1beta1,
  pipeline_service_v1beta1,
  pipeline_state_v1beta1,

google/cloud/aiplatform/pipeline_jobs.py

@@ -119,6 +119,7 @@ def __init__(
  credentials: Optional[auth_credentials.Credentials] = None,
  project: Optional[str] = None,
  location: Optional[str] = None,
+ failure_policy: Optional[str] = None,
  ):
  """Retrieves a PipelineJob resource and instantiates its
  representation.
@@ -173,6 +174,15 @@ def __init__(
  location (str):
  Optional. Location to create PipelineJob. If not set,
  location set in aiplatform.init will be used.
+ failure_policy (str):
+ Optional. The failure policy - "slow" or "fast".
+ Currently, the default of a pipeline is that the pipeline will continue to
+ run until no more tasks can be executed, also known as
+ PIPELINE_FAILURE_POLICY_FAIL_SLOW (corresponds to "slow").
+ However, if a pipeline is set to
+ PIPELINE_FAILURE_POLICY_FAIL_FAST (corresponds to "fast"),
+ it will stop scheduling any new tasks when a task has failed. Any
+ scheduled tasks will continue to completion.
 
  Raises:
  ValueError: If job_id or labels have incorrect format.
@@ -219,6 +229,7 @@ def __init__(
  )
  builder.update_pipeline_root(pipeline_root)
  builder.update_runtime_parameters(parameter_values)
+ builder.update_failure_policy(failure_policy)
  runtime_config_dict = builder.build()
 
  runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb

google/cloud/aiplatform/utils/pipeline_utils.py

@@ -17,6 +17,7 @@
 import copy
 import json
 from typing import Any, Dict, Mapping, Optional, Union
+from google.cloud.aiplatform.compat.types import pipeline_failure_policy
 import packaging.version
 
 
@@ -32,6 +33,7 @@ def __init__(
  schema_version: str,
  parameter_types: Mapping[str, str],
  parameter_values: Optional[Dict[str, Any]] = None,
+ failure_policy: Optional[pipeline_failure_policy.PipelineFailurePolicy] = None,
  ):
  """Creates a PipelineRuntimeConfigBuilder object.
 
@@ -44,11 +46,20 @@ def __init__(
  Required. The mapping from pipeline parameter name to its type.
  parameter_values (Dict[str, Any]):
  Optional. The mapping from runtime parameter name to its value.
+ failure_policy (pipeline_failure_policy.PipelineFailurePolicy):
+ Optional. Represents the failure policy of a pipeline. Currently, the
+ default of a pipeline is that the pipeline will continue to
+ run until no more tasks can be executed, also known as
+ PIPELINE_FAILURE_POLICY_FAIL_SLOW. However, if a pipeline is
+ set to PIPELINE_FAILURE_POLICY_FAIL_FAST, it will stop
+ scheduling any new tasks when a task has failed. Any
+ scheduled tasks will continue to completion.
  """
  self._pipeline_root = pipeline_root
  self._schema_version = schema_version
  self._parameter_types = parameter_types
  self._parameter_values = copy.deepcopy(parameter_values or {})
+ self._failure_policy = failure_policy
 
  @classmethod
  def from_job_spec_json(
@@ -80,7 +91,14 @@ def from_job_spec_json(
 
  pipeline_root = runtime_config_spec.get("gcsOutputDirectory")
  parameter_values = _parse_runtime_parameters(runtime_config_spec)
- return cls(pipeline_root, schema_version, parameter_types, parameter_values)
+ failure_policy = runtime_config_spec.get("failurePolicy")
+ return cls(
+ pipeline_root,
+ schema_version,
+ parameter_types,
+ parameter_values,
+ failure_policy,
+ )
 
  def update_pipeline_root(self, pipeline_root: Optional[str]) -> None:
  """Updates pipeline_root value.
@@ -111,6 +129,24 @@ def update_runtime_parameters(
  parameters[k] = json.dumps(v)
  self._parameter_values.update(parameters)
 
+ def update_failure_policy(self, failure_policy: Optional[str] = None) -> None:
+ """Merges runtime failure policy.
+
+ Args:
+ failure_policy (str):
+ Optional. The failure policy - "slow" or "fast".
+
+ Raises:
+ ValueError: if failure_policy is not valid.
+ """
+ if failure_policy:
+ if failure_policy in _FAILURE_POLICY_TO_ENUM_VALUE:
+ self._failure_policy = _FAILURE_POLICY_TO_ENUM_VALUE[failure_policy]
+ else:
+ raise ValueError(
+ f'failure_policy should be either "slow" or "fast", but got: "{failure_policy}".'
+ )
+
  def build(self) -> Dict[str, Any]:
  """Build a RuntimeConfig proto.
 
@@ -128,7 +164,8 @@ def build(self) -> Dict[str, Any]:
  parameter_values_key = "parameterValues"
  else:
  parameter_values_key = "parameters"
- return {
+
+ runtime_config = {
  "gcsOutputDirectory": self._pipeline_root,
  parameter_values_key: {
  k: self._get_vertex_value(k, v)
@@ -137,6 +174,11 @@ def build(self) -> Dict[str, Any]:
  },
  }
 
+ if self._failure_policy:
+ runtime_config["failurePolicy"] = self._failure_policy
+
+ return runtime_config
+
  def _get_vertex_value(
  self, name: str, value: Union[int, float, str, bool, list, dict]
  ) -> Union[int, float, str, bool, list, dict]:
@@ -205,3 +247,10 @@ def _parse_runtime_parameters(
  else:
  raise TypeError("Got unknown type of value: {}".format(value))
  return result
+
+
+_FAILURE_POLICY_TO_ENUM_VALUE = {
+ "slow": pipeline_failure_policy.PipelineFailurePolicy.PIPELINE_FAILURE_POLICY_FAIL_SLOW,
+ "fast": pipeline_failure_policy.PipelineFailurePolicy.PIPELINE_FAILURE_POLICY_FAIL_FAST,
+ None: pipeline_failure_policy.PipelineFailurePolicy.PIPELINE_FAILURE_POLICY_UNSPECIFIED,
+}

tests/unit/aiplatform/test_pipeline_jobs.py

@@ -30,6 +30,7 @@
 from google.cloud.aiplatform import base
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform import pipeline_jobs
+from google.cloud.aiplatform.compat.types import pipeline_failure_policy
 from google.cloud import storage
 from google.protobuf import json_format
 
@@ -621,6 +622,99 @@ def test_run_call_pipeline_service_create_with_timeout_not_explicitly_set(
  timeout=None,
  )
 
+ @pytest.mark.parametrize(
+ "job_spec",
+ [_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_SPEC_YAML, _TEST_PIPELINE_JOB],
+ )
+ @pytest.mark.parametrize(
+ "failure_policy",
+ [
+ (
+ "slow",
+ pipeline_failure_policy.PipelineFailurePolicy.PIPELINE_FAILURE_POLICY_FAIL_SLOW,
+ ),
+ (
+ "fast",
+ pipeline_failure_policy.PipelineFailurePolicy.PIPELINE_FAILURE_POLICY_FAIL_FAST,
+ ),
+ ],
+ )
+ @pytest.mark.parametrize("sync", [True, False])
+ def test_run_call_pipeline_service_create_with_failure_policy(
+ self,
+ mock_pipeline_service_create,
+ mock_pipeline_service_get,
+ job_spec,
+ mock_load_yaml_and_json,
+ failure_policy,
+ 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_DISPLAY_NAME,
+ template_path=_TEST_TEMPLATE_PATH,
+ job_id=_TEST_PIPELINE_JOB_ID,
+ parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
+ enable_caching=True,
+ failure_policy=failure_policy[0],
+ )
+
+ job.run(
+ service_account=_TEST_SERVICE_ACCOUNT,
+ network=_TEST_NETWORK,
+ sync=sync,
+ create_request_timeout=None,
+ )
+
+ if not sync:
+ job.wait()
+
+ expected_runtime_config_dict = {
+ "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME,
+ "parameterValues": _TEST_PIPELINE_PARAMETER_VALUES,
+ "failurePolicy": failure_policy[1],
+ }
+ 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,
+ )
+
+ mock_pipeline_service_create.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
+ )
+
  @pytest.mark.parametrize(
  "job_spec",
  [

tests/unit/aiplatform/test_utils.py

@@ -28,6 +28,7 @@
 from google.api_core import client_options, gapic_v1
 from google.cloud import aiplatform
 from google.cloud.aiplatform import compat, utils
+from google.cloud.aiplatform.compat.types import pipeline_failure_policy
 from google.cloud.aiplatform.utils import pipeline_utils, tensorboard_utils, yaml_utils
 from google.cloud.aiplatform_v1.services.model_service import (
  client as model_service_client_v1,
@@ -454,7 +455,22 @@ def test_pipeline_utils_runtime_config_builder_with_no_op_updates(self):
  expected_runtime_config = self.SAMPLE_JOB_SPEC["runtimeConfig"]
  assert expected_runtime_config == actual_runtime_config
 
- def test_pipeline_utils_runtime_config_builder_with_merge_updates(self):
+ @pytest.mark.parametrize(
+ "failure_policy",
+ [
+ (
+ "slow",
+ pipeline_failure_policy.PipelineFailurePolicy.PIPELINE_FAILURE_POLICY_FAIL_SLOW,
+ ),
+ (
+ "fast",
+ pipeline_failure_policy.PipelineFailurePolicy.PIPELINE_FAILURE_POLICY_FAIL_FAST,
+ ),
+ ],
+ )
+ def test_pipeline_utils_runtime_config_builder_with_merge_updates(
+ self, failure_policy
+ ):
  my_builder = pipeline_utils.PipelineRuntimeConfigBuilder.from_job_spec_json(
  self.SAMPLE_JOB_SPEC
  )
@@ -468,6 +484,7 @@ def test_pipeline_utils_runtime_config_builder_with_merge_updates(self):
  "bool_param": True,
  }
  )
+ my_builder.update_failure_policy(failure_policy[0])
  actual_runtime_config = my_builder.build()
 
  expected_runtime_config = {
@@ -481,9 +498,21 @@ def test_pipeline_utils_runtime_config_builder_with_merge_updates(self):
  "list_param": {"stringValue": "[1, 2, 3]"},
  "bool_param": {"stringValue": "true"},
  },
+ "failurePolicy": failure_policy[1],
  }
  assert expected_runtime_config == actual_runtime_config
 
+ def test_pipeline_utils_runtime_config_builder_invalid_failure_policy(self):
+ my_builder = pipeline_utils.PipelineRuntimeConfigBuilder.from_job_spec_json(
+ self.SAMPLE_JOB_SPEC
+ )
+ with pytest.raises(ValueError) as e:
+ my_builder.update_failure_policy("slo")
+
+ assert e.match(
+ regexp=r'failure_policy should be either "slow" or "fast", but got: "slo".'
+ )
+
  def test_pipeline_utils_runtime_config_builder_parameter_not_found(self):
  my_builder = pipeline_utils.PipelineRuntimeConfigBuilder.from_job_spec_json(
  self.SAMPLE_JOB_SPEC