change: Add override_pipeline_parameter_var decorator to give grace period to update invalid pipeline var args (aws#3180)

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

Author: 2 people

src/sagemaker/image_uris.py

@@ -24,13 +24,16 @@
 from sagemaker.spark import defaults
 from sagemaker.jumpstart import artifacts
 from sagemaker.workflow import is_pipeline_variable
+from sagemaker.workflow.utilities import override_pipeline_parameter_var
 
 logger = logging.getLogger(__name__)
 
 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
 HUGGING_FACE_FRAMEWORK = "huggingface"
 
 
+# TODO: we should remove this decorator later
+@override_pipeline_parameter_var
 def retrieve(
     framework,
     region,

src/sagemaker/processing.py

@@ -1195,7 +1195,9 @@ def __init__(
             source (str): The source for the output.
             destination (str): The destination of the output. If a destination
                 is not provided, one will be generated:
-                "s3://<default-bucket-name>/<job-name>/output/<output-name>".
+                "s3://<default-bucket-name>/<job-name>/output/<output-name>"
+                (Note: this does not apply when used with
+                :class:`~sagemaker.workflow.steps.ProcessingStep`).
             output_name (str): The name of the output. If a name
                 is not provided, one will be generated (eg. "output-1").
             s3_upload_mode (str): Valid options are "EndOfJob" or "Continuous".

src/sagemaker/workflow/utilities.py

@@ -13,30 +13,38 @@
 """Utilities to support workflow."""
 from __future__ import absolute_import
 
+import inspect
+import logging
+from functools import wraps
 from pathlib import Path
-from typing import List, Sequence, Union, Set
+from typing import List, Sequence, Union, Set, TYPE_CHECKING
 import hashlib
 from urllib.parse import unquote, urlparse
 from _hashlib import HASH as Hash
 
+from sagemaker.workflow.parameters import Parameter
 from sagemaker.workflow.pipeline_context import _StepArguments
-from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.entities import (
     Entity,
     RequestType,
 )
 
+if TYPE_CHECKING:
+    from sagemaker.workflow.step_collections import StepCollection
+
 BUF_SIZE = 65536  # 64KiB
 
 
-def list_to_request(entities: Sequence[Union[Entity, StepCollection]]) -> List[RequestType]:
+def list_to_request(entities: Sequence[Union[Entity, "StepCollection"]]) -> List[RequestType]:
     """Get the request structure for list of entities.
 
     Args:
         entities (Sequence[Entity]): A list of entities.
     Returns:
         list: A request structure for a workflow service call.
     """
+    from sagemaker.workflow.step_collections import StepCollection
+
     request_dicts = []
     for entity in entities:
         if isinstance(entity, Entity):
@@ -151,3 +159,41 @@ def validate_step_args_input(
         raise TypeError(error_message)
     if step_args.caller_name not in expected_caller:
         raise ValueError(error_message)
+
+
+def override_pipeline_parameter_var(func):
+    """A decorator to override pipeline Parameters passed into a function
+
+    This is a temporary decorator to override pipeline Parameter objects with their default value
+    and display warning information to instruct users to update their code.
+
+    This decorator can help to give a grace period for users to update their code when
+    we make changes to explicitly prevent passing any pipeline variables to a function.
+
+    We should remove this decorator after the grace period.
+    """
+    warning_msg_template = (
+        "%s should not be a pipeline variable (%s). "
+        "The default_value of this Parameter object will be used to override it. "
+        "Please remove this pipeline variable and use python primitives instead."
+    )
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        params = inspect.signature(func).parameters
+        args = list(args)
+        for i, (arg_name, _) in enumerate(params.items()):
+            if i >= len(args):
+                break
+            if isinstance(args[i], Parameter):
+                logging.warning(warning_msg_template, arg_name, type(args[i]))
+                args[i] = args[i].default_value
+        args = tuple(args)
+
+        for arg_name, value in kwargs.items():
+            if isinstance(value, Parameter):
+                logging.warning(warning_msg_template, arg_name, type(value))
+                kwargs[arg_name] = value.default_value
+        return func(*args, **kwargs)
+
+    return wrapper

tests/unit/sagemaker/image_uris/test_retrieve.py

@@ -19,6 +19,7 @@
 from mock import patch
 
 from sagemaker import image_uris
+from sagemaker.workflow.functions import Join
 from sagemaker.workflow.parameters import ParameterString
 
 BASE_CONFIG = {
@@ -721,16 +722,50 @@ def test_retrieve_huggingface(config_for_framework):
 
 
 def test_retrieve_with_pipeline_variable():
+    kwargs = dict(
+        framework="tensorflow",
+        version="1.15",
+        py_version="py3",
+        instance_type="ml.m5.xlarge",
+        region="us-east-1",
+        image_scope="training",
+    )
+    # instance_type is plain string which should not break anything
+    image_uris.retrieve(**kwargs)
+
+    # instance_type is parameter string with not None default value
+    # which should not break anything
+    kwargs["instance_type"] = ParameterString(
+        name="TrainingInstanceType",
+        default_value="ml.m5.xlarge",
+    )
+    image_uris.retrieve(**kwargs)
+
+    # instance_type is parameter string without default value
+    # (equivalent to pass in None to instance_type field)
+    # which should fail due to empty instance type check
+    kwargs["instance_type"] = ParameterString(name="TrainingInstanceType")
     with pytest.raises(Exception) as error:
-        image_uris.retrieve(
-            framework="tensorflow",
-            version="1.15",
-            py_version="py3",
-            instance_type=ParameterString(
-                name="TrainingInstanceType",
-                default_value="ml.m5.xlarge",
-            ),
-            region="us-east-1",
-            image_scope="training",
-        )
+        image_uris.retrieve(**kwargs)
+    assert "Empty SageMaker instance type" in str(error.value)
+
+    # instance_type is other types of pipeline variable
+    # which should break loudly
+    kwargs["instance_type"] = Join(on="", values=["a", "b"])
+    with pytest.raises(Exception) as error:
+        image_uris.retrieve(**kwargs)
     assert "instance_type should not be a pipeline variable" in str(error.value)
+
+    # instance_type (ParameterString) is given as args rather than kwargs
+    # which should not break anything
+    image_uris.retrieve(
+        "tensorflow",
+        "us-east-1",
+        "1.15",
+        "py3",
+        ParameterString(
+            name="TrainingInstanceType",
+            default_value="ml.m5.xlarge",
+        ),
+        image_scope="training",
+    )