change: Add more validations for pipeline step new interfaces (#3119)

Co-authored-by: Dewen Qi <[email protected]>

Author: qidewenwhen

src/sagemaker/estimator.py

@@ -1340,6 +1340,11 @@ def register(
         Returns:
             str: A string of SageMaker Model Package ARN.
         """
+        if isinstance(self.sagemaker_session, PipelineSession):
+            raise TypeError(
+                "estimator.register does not support PipelineSession at this moment. "
+                "Please use model.register with PipelineSession if you're using the ModelStep."
+            )
         default_name = name_from_base(self.base_job_name)
         model_name = model_name or default_name
         if compile_model_family is not None:

src/sagemaker/session.py

@@ -591,7 +591,7 @@ def submit(request):
             LOGGER.debug("train request: %s", json.dumps(request, indent=4))
             self.sagemaker_client.create_training_job(**request)
 
-        self._intercept_create_request(train_request, submit)
+        self._intercept_create_request(train_request, submit, self.train.__name__)
 
     def _get_train_request(  # noqa: C901
         self,
@@ -922,7 +922,7 @@ def submit(request):
             LOGGER.debug("process request: %s", json.dumps(request, indent=4))
             self.sagemaker_client.create_processing_job(**request)
 
-        self._intercept_create_request(process_request, submit)
+        self._intercept_create_request(process_request, submit, self.process.__name__)
 
     def _get_process_request(
         self,
@@ -2099,7 +2099,7 @@ def submit(request):
             LOGGER.debug("tune request: %s", json.dumps(request, indent=4))
             self.sagemaker_client.create_hyper_parameter_tuning_job(**request)
 
-        self._intercept_create_request(tune_request, submit)
+        self._intercept_create_request(tune_request, submit, self.create_tuning_job.__name__)
 
     def _get_tuning_request(
         self,
@@ -2569,7 +2569,7 @@ def submit(request):
             LOGGER.debug("Transform request: %s", json.dumps(request, indent=4))
             self.sagemaker_client.create_transform_job(**request)
 
-        self._intercept_create_request(transform_request, submit)
+        self._intercept_create_request(transform_request, submit, self.transform.__name__)
 
     def _create_model_request(
         self,

src/sagemaker/workflow/model_step.py

@@ -16,10 +16,10 @@
 import logging
 from typing import Union, List, Dict, Optional
 
-from sagemaker import Model, PipelineModel
+from sagemaker import Model, PipelineModel, Session
 from sagemaker.workflow._utils import _RegisterModelStep, _RepackModelStep
 from sagemaker.workflow.pipeline_context import PipelineSession, _ModelStepArguments
-from sagemaker.workflow.retry import RetryPolicy
+from sagemaker.workflow.retry import RetryPolicy, SageMakerJobStepRetryPolicy
 from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import Step, CreateModelStep
 
@@ -57,17 +57,72 @@ def __init__(
                 If a listed `Step` name does not exist, an error is returned (default: None).
             retry_policies (List[RetryPolicy] or Dict[str, List[RetryPolicy]]): The list of retry
                 policies for the `ModelStep` (default: None).
+
+                If a list of retry policies is provided, it would be applied to all steps in the
+                `ModelStep` collection. Note: in this case, `SageMakerJobStepRetryPolicy`
+                is not allowed, since create/register model step does not support it.
+                Please find the example below:
+
+                .. code:: python
+
+                    ModelStep(
+                        ...
+                        retry_policies=[
+                            StepRetryPolicy(...),
+                        ],
+                    )
+
+                If a dict is provided, it can specify different retry policies for different
+                types of steps in the `ModelStep` collection. Similarly,
+                `SageMakerJobStepRetryPolicy` is not allowed for create/register model step.
+                See examples below:
+
+                .. code:: python
+
+                    ModelStep(
+                        ...
+                        retry_policies=dict(
+                           register_model_retry_policies=[
+                               StepRetryPolicy(...),
+                           ],
+                           repack_model_retry_policies=[
+                               SageMakerJobStepRetryPolicy(...),
+                           ],
+                        )
+                    )
+
+                or
+
+                .. code:: python
+
+                    ModelStep(
+                        ...
+                        retry_policies=dict(
+                           create_model_retry_policies=[
+                               StepRetryPolicy(...),
+                           ],
+                           repack_model_retry_policies=[
+                               SageMakerJobStepRetryPolicy(...),
+                           ],
+                        )
+                    )
+
             display_name (str): The display name of the `ModelStep`.
                 The display name provides better UI readability. (default: None).
             description (str): The description of the `ModelStep` (default: None).
         """
         # TODO: add a doc link in error message once ready
-        if not isinstance(step_args, _ModelStepArguments):
-            raise TypeError(
-                "To correctly configure a ModelStep, "
-                "the step_args must be a `_ModelStepArguments` object generated by "
-                ".create() or .register()."
-            )
+        from sagemaker.workflow.utilities import validate_step_args_input
+
+        validate_step_args_input(
+            step_args=step_args,
+            expected_caller={
+                Session.create_model.__name__,
+                Session.create_model_package_from_containers.__name__,
+            },
+            error_message="The step_args of ModelStep must be obtained from model.create() "
+            "or model.register().",
+        )
         if not (step_args.create_model_request is None) ^ (
             step_args.create_model_package_request is None
         ):
@@ -93,7 +148,22 @@ def __init__(
         self._create_model_args = self.step_args.create_model_request
         self._register_model_args = self.step_args.create_model_package_request
         self._need_runtime_repack = self.step_args.need_runtime_repack
+        self._assign_and_validate_retry_policies(retry_policies)
+
+        if self._need_runtime_repack:
+            self._append_repack_model_step()
+        if self._register_model_args:
+            self._append_register_model_step()
+        else:
+            self._append_create_model_step()
 
+    def _assign_and_validate_retry_policies(self, retry_policies):
+        """Assign and validate retry policies according to each kind of sub steps
+
+        Args:
+            retry_policies (List[RetryPolicy] or Dict[str, List[RetryPolicy]]): The list of retry
+                policies for the `ModelStep`.
+        """
         if isinstance(retry_policies, dict):
             self._create_model_retry_policies = retry_policies.get(
                 _CREATE_MODEL_RETRY_POLICIES, None
@@ -109,12 +179,23 @@ def __init__(
             self._register_model_retry_policies = retry_policies
             self._repack_model_retry_policies = retry_policies
 
-        if self._need_runtime_repack:
-            self._append_repack_model_step()
-        if self._register_model_args:
-            self._append_register_model_step()
-        else:
-            self._append_create_model_step()
+        self._validate_sagemaker_job_step_retry_policy()
+
+    def _validate_sagemaker_job_step_retry_policy(self):
+        """Validate SageMakerJobStepRetryPolicy
+
+        Validate that SageMakerJobStepRetryPolicy is not assigning to create/register model step.
+        """
+        retry_policies = set(
+            (self._create_model_retry_policies or []) + (self._register_model_retry_policies or [])
+        )
+        for policy in retry_policies:
+            if not isinstance(policy, SageMakerJobStepRetryPolicy):
+                continue
+            raise ValueError(
+                "SageMakerJobStepRetryPolicy is not allowed for a create/register"
+                " model step. Please use StepRetryPolicy instead"
+            )
 
     def _append_register_model_step(self):
         """Create and append a `_RegisterModelStep`"""

src/sagemaker/workflow/pipeline_context.py

@@ -15,14 +15,45 @@
 
 import warnings
 import inspect
-from typing import Dict
 from functools import wraps
+from typing import Dict, Optional
 
 from sagemaker.session import Session, SessionSettings
 
 
-class _ModelStepArguments:
-    """Step arguments entity for ModelStep"""
+class _StepArguments:
+    """Step arguments entity for `Step`"""
+
+    def __init__(self, caller_name: str = None):
+        """Create a `_StepArguments`
+
+        Args:
+            caller_name (str): The name of the caller function which is intercepted by the
+                PipelineSession to get the step arguments.
+        """
+        self.caller_name = caller_name
+
+
+class _JobStepArguments(_StepArguments):
+    """Step arguments entity for job step types
+
+    Job step types include: TrainingStep, ProcessingStep, TuningStep, TransformStep
+    """
+
+    def __init__(self, caller_name: str, args: dict):
+        """Create a `_JobStepArguments`
+
+        Args:
+            caller_name (str): The name of the caller function which is intercepted by the
+                PipelineSession to get the step arguments.
+            args (dict): The arguments to be used for composing the SageMaker API request.
+        """
+        super(_JobStepArguments, self).__init__(caller_name)
+        self.args = args
+
+
+class _ModelStepArguments(_StepArguments):
+    """Step arguments entity for `ModelStep`"""
 
     def __init__(self, model):
         """Create a `_ModelStepArguments`
@@ -31,6 +62,7 @@ def __init__(self, model):
             model (Model or PipelineModel): A `sagemaker.model.Model`
                 or `sagemaker.pipeline.PipelineModel` instance
         """
+        super(_ModelStepArguments, self).__init__()
         self.model = model
         self.create_model_package_request = None
         self.create_model_request = None
@@ -87,9 +119,8 @@ def context(self):
         return self._context
 
     @context.setter
-    def context(self, args: Dict):
-        # TODO: we should use _StepArguments type to formalize non-ModelStep args
-        self._context = args
+    def context(self, value: Optional[_StepArguments] = None):
+        self._context = value
 
     def _intercept_create_request(self, request: Dict, create, func_name: str = None):
         """This function intercepts the create job request
@@ -99,12 +130,14 @@ def _intercept_create_request(self, request: Dict, create, func_name: str = None
             create (functor): a functor calls the sagemaker client create method
             func_name (str): the name of the function needed intercepting
         """
-        if func_name == "create_model":
+        if func_name == self.create_model.__name__:
             self.context.create_model_request = request
-        elif func_name == "create_model_package_from_containers":
+            self.context.caller_name = func_name
+        elif func_name == self.create_model_package_from_containers.__name__:
             self.context.create_model_package_request = request
+            self.context.caller_name = func_name
         else:
-            self.context = request
+            self.context = _JobStepArguments(func_name, request)
 
     def init_step_arguments(self, model):
         """Create a `_ModelStepArguments` (if not exist) as pipeline context

src/sagemaker/workflow/retry.py

@@ -148,6 +148,10 @@ def to_request(self) -> RequestType:
         request["ExceptionType"] = [e.value for e in self.exception_types]
         return request
 
+    def __hash__(self):
+        """Hash function for StepRetryPolicy types"""
+        return hash(tuple(self.to_request()))
+
 
 class SageMakerJobStepRetryPolicy(RetryPolicy):
     """RetryPolicy for exception thrown by SageMaker Job.
@@ -202,3 +206,7 @@ def to_request(self) -> RequestType:
         request = super().to_request()
         request["ExceptionType"] = [e.value for e in self.exception_type_list]
         return request
+
+    def __hash__(self):
+        """Hash function for SageMakerJobStepRetryPolicy types"""
+        return hash(tuple(self.to_request()))