fix: Include deprecated JsonGet into PipelineVariable

Author: Dewen Qi

src/sagemaker/estimator.py

@@ -74,7 +74,7 @@
     get_config_value,
     name_from_base,
 )
-from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.workflow import is_pipeline_variable
 
 logger = logging.getLogger(__name__)
 
@@ -600,7 +600,7 @@ def _json_encode_hyperparameters(hyperparameters: Dict[str, Any]) -> Dict[str, A
         current_hyperparameters = hyperparameters
         if current_hyperparameters is not None:
             hyperparameters = {
-                str(k): (v.to_string() if isinstance(v, PipelineVariable) else json.dumps(v))
+                str(k): (v.to_string() if is_pipeline_variable(v) else json.dumps(v))
                 for (k, v) in current_hyperparameters.items()
             }
         return hyperparameters
@@ -1811,7 +1811,7 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
         current_hyperparameters = estimator.hyperparameters()
         if current_hyperparameters is not None:
             hyperparameters = {
-                str(k): (v.to_string() if isinstance(v, PipelineVariable) else str(v))
+                str(k): (v.to_string() if is_pipeline_variable(v) else str(v))
                 for (k, v) in current_hyperparameters.items()
             }
 

src/sagemaker/model.py

@@ -37,7 +37,7 @@
 from sagemaker.utils import unique_name_from_base
 from sagemaker.async_inference import AsyncInferenceConfig
 from sagemaker.predictor_async import AsyncPredictor
-from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.workflow import is_pipeline_variable
 
 LOGGER = logging.getLogger("sagemaker")
 
@@ -444,7 +444,7 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
             )
 
         if repack and self.model_data is not None and self.entry_point is not None:
-            if isinstance(self.model_data, PipelineVariable):
+            if is_pipeline_variable(self.model_data):
                 # model is not yet there, defer repacking to later during pipeline execution
                 return
 

src/sagemaker/parameter.py

@@ -15,7 +15,7 @@
 
 import json
 
-from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.workflow import is_pipeline_variable
 
 
 class ParameterRange(object):
@@ -72,10 +72,10 @@ def as_tuning_range(self, name):
         return {
             "Name": name,
             "MinValue": str(self.min_value)
-            if not isinstance(self.min_value, PipelineVariable)
+            if not is_pipeline_variable(self.min_value)
             else self.min_value.to_string(),
             "MaxValue": str(self.max_value)
-            if not isinstance(self.max_value, PipelineVariable)
+            if not is_pipeline_variable(self.max_value)
             else self.max_value.to_string(),
             "ScalingType": self.scaling_type,
         }
@@ -110,9 +110,7 @@ def __init__(self, values):  # pylint: disable=super-init-not-called
                 This input will be converted into a list of strings.
         """
         values = values if isinstance(values, list) else [values]
-        self.values = [
-            str(v) if not isinstance(v, PipelineVariable) else v.to_string() for v in values
-        ]
+        self.values = [str(v) if not is_pipeline_variable(v) else v.to_string() for v in values]
 
     def as_tuning_range(self, name):
         """Represent the parameter range as a dictionary.

src/sagemaker/processing.py

@@ -34,14 +34,14 @@
 from sagemaker.local import LocalSession
 from sagemaker.utils import base_name_from_image, get_config_value, name_from_base
 from sagemaker.session import Session
+from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.properties import Properties
 from sagemaker.workflow.parameters import Parameter
-from sagemaker.workflow.entities import Expression, PipelineVariable
+from sagemaker.workflow.entities import Expression
 from sagemaker.dataset_definition.inputs import S3Input, DatasetDefinition
 from sagemaker.apiutils._base_types import ApiObject
 from sagemaker.s3 import S3Uploader
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -233,7 +233,7 @@ def _normalize_args(
             kms_key (str): The ARN of the KMS key that is used to encrypt the
                 user code file (default: None).
         """
-        if code and isinstance(code, PipelineVariable):
+        if code and is_pipeline_variable(code):
             raise ValueError(
                 "code argument has to be a valid S3 URI or local file path "
                 + "rather than a pipeline variable"

src/sagemaker/tensorflow/model.py

@@ -21,7 +21,7 @@
 from sagemaker.deprecations import removed_kwargs
 from sagemaker.predictor import Predictor
 from sagemaker.serializers import JSONSerializer
-from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.workflow import is_pipeline_variable
 
 
 class TensorFlowPredictor(Predictor):
@@ -333,7 +333,7 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
 
         # If self.model_data is pipeline variable, model is not yet there.
         # So defer repacking to later during pipeline execution
-        if self.entry_point and not isinstance(self.model_data, PipelineVariable):
+        if self.entry_point and not is_pipeline_variable(self.model_data):
             key_prefix = sagemaker.fw_utils.model_code_key_prefix(
                 self.key_prefix, self.name, image_uri
             )

src/sagemaker/tuner.py

@@ -38,10 +38,10 @@
     IntegerParameter,
     ParameterRange,
 )
-from sagemaker.workflow.entities import PipelineVariable
 
 from sagemaker.session import Session
 from sagemaker.utils import base_from_name, base_name_from_image, name_from_base
+from sagemaker.workflow import is_pipeline_variable
 
 AMAZON_ESTIMATOR_MODULE = "sagemaker"
 AMAZON_ESTIMATOR_CLS_NAMES = {
@@ -362,7 +362,7 @@ def _prepare_static_hyperparameters(
         """Prepare static hyperparameters for one estimator before tuning."""
         # Remove any hyperparameter that will be tuned
         static_hyperparameters = {
-            str(k): str(v) if not isinstance(v, PipelineVariable) else v.to_string()
+            str(k): str(v) if not is_pipeline_variable(v) else v.to_string()
             for (k, v) in estimator.hyperparameters().items()
         }
         for hyperparameter_name in hyperparameter_ranges.keys():

src/sagemaker/workflow/__init__.py

@@ -20,3 +20,18 @@
 from sagemaker.workflow.properties import Properties
 
 PipelineNonPrimitiveInputTypes = Union[ExecutionVariable, Expression, Parameter, Properties]
+
+
+def is_pipeline_variable(var: object) -> bool:
+    """Check if the variable is a pipeline variable
+
+    Args:
+        var (object): The variable to be verified.
+    Returns:
+         bool: True if it is, False otherwise.
+    """
+
+    # Currently Expression is on top of all kinds of pipeline variables
+    # as well as PipelineExperimentConfigProperty and PropertyFile
+    # TODO: We should deprecate the Expression and replace it with PipelineVariable
+    return isinstance(var, Expression)

src/sagemaker/workflow/clarify_check_step.py

@@ -37,8 +37,8 @@
 from sagemaker.model_monitor.model_monitoring import _MODEL_MONITOR_S3_PATH
 from sagemaker.processing import ProcessingInput, ProcessingOutput, ProcessingJob
 from sagemaker.utils import name_from_base
-from sagemaker.workflow import PipelineNonPrimitiveInputTypes
-from sagemaker.workflow.entities import RequestType, PipelineVariable
+from sagemaker.workflow import PipelineNonPrimitiveInputTypes, is_pipeline_variable
+from sagemaker.workflow.entities import RequestType
 from sagemaker.workflow.properties import Properties
 from sagemaker.workflow.steps import Step, StepTypeEnum, CacheConfig
 from sagemaker.workflow.check_job_config import CheckJobConfig
@@ -193,16 +193,15 @@ def __init__(
                 + "DataBiasCheckConfig, ModelBiasCheckConfig or ModelExplainabilityCheckConfig"
             )
 
-        if isinstance(
-            clarify_check_config.data_config.s3_analysis_config_output_path, PipelineVariable
-        ):
+        if is_pipeline_variable(clarify_check_config.data_config.s3_analysis_config_output_path):
             raise RuntimeError(
                 "s3_analysis_config_output_path cannot be of type "
                 + "ExecutionVariable/Expression/Parameter/Properties"
             )
 
-        if not clarify_check_config.data_config.s3_analysis_config_output_path and isinstance(
-            clarify_check_config.data_config.s3_output_path, PipelineVariable
+        if (
+            not clarify_check_config.data_config.s3_analysis_config_output_path
+            and is_pipeline_variable(clarify_check_config.data_config.s3_output_path)
         ):
             raise RuntimeError(
                 "`s3_output_path` cannot be of type ExecutionVariable/Expression/Parameter"

src/sagemaker/workflow/condition_step.py

@@ -26,8 +26,8 @@
 from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.utilities import list_to_request
 from sagemaker.workflow.entities import (
-    Expression,
     RequestType,
+    PipelineVariable,
 )
 from sagemaker.workflow.properties import (
     Properties,
@@ -95,7 +95,7 @@ def properties(self):
 
 
 @attr.s
-class JsonGet(Expression):  # pragma: no cover
+class JsonGet(PipelineVariable):  # pragma: no cover
     """Get JSON properties from PropertyFiles.
 
     Attributes:

src/sagemaker/workflow/conditions.py

@@ -22,13 +22,13 @@
 
 import attr
 
+from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.entities import (
     DefaultEnumMeta,
     Entity,
     Expression,
     PrimitiveType,
     RequestType,
-    PipelineVariable,
 )
 from sagemaker.workflow.execution_variables import ExecutionVariable
 from sagemaker.workflow.parameters import Parameter
@@ -262,6 +262,6 @@ def primitive_or_expr(
     Returns:
         Either the expression of the value or the primitive value.
     """
-    if isinstance(value, PipelineVariable):
+    if is_pipeline_variable(value):
         return value.expr
     return value

src/sagemaker/workflow/pipeline.py

@@ -339,16 +339,20 @@ def interpolate(
     Args:
         request_obj (RequestType): The request dict.
         callback_output_to_step_map (Dict[str, str]): A dict of output name -> step name.
+        lambda_output_to_step_map (Dict[str, str]): A dict of output name -> step name.
 
     Returns:
         RequestType: The request dict with Parameter values replaced by their expression.
     """
-    request_obj_copy = deepcopy(request_obj)
-    return _interpolate(
-        request_obj_copy,
-        callback_output_to_step_map=callback_output_to_step_map,
-        lambda_output_to_step_map=lambda_output_to_step_map,
-    )
+    try:
+        request_obj_copy = deepcopy(request_obj)
+        return _interpolate(
+            request_obj_copy,
+            callback_output_to_step_map=callback_output_to_step_map,
+            lambda_output_to_step_map=lambda_output_to_step_map,
+        )
+    except TypeError as type_err:
+        raise TypeError("Not able to interpolate Pipeline definition: %s" % type_err)
 
 
 def _interpolate(

src/sagemaker/workflow/quality_check_step.py

@@ -22,9 +22,9 @@
 from sagemaker import s3
 from sagemaker.model_monitor import ModelMonitor
 from sagemaker.processing import ProcessingOutput, ProcessingJob, Processor, ProcessingInput
-from sagemaker.workflow import PipelineNonPrimitiveInputTypes
+from sagemaker.workflow import PipelineNonPrimitiveInputTypes, is_pipeline_variable
 
-from sagemaker.workflow.entities import RequestType, PipelineVariable
+from sagemaker.workflow.entities import RequestType
 from sagemaker.workflow.properties import (
     Properties,
 )
@@ -279,7 +279,7 @@ def _generate_baseline_job_inputs(self):
                 _CONTAINER_BASE_PATH, _CONTAINER_INPUT_PATH, _BASELINE_DATASET_INPUT_NAME
             )
         )
-        if isinstance(baseline_dataset, PipelineVariable):
+        if is_pipeline_variable(baseline_dataset):
             baseline_dataset_input = ProcessingInput(
                 source=self.quality_check_config.baseline_dataset,
                 destination=baseline_dataset_des,

tests/integ/sagemaker/workflow/test_workflow_with_clarify.py

@@ -33,8 +33,8 @@
 from sagemaker.processing import ProcessingInput, ProcessingOutput
 from sagemaker.session import get_execution_role
 from sagemaker.workflow.conditions import ConditionLessThanOrEqualTo
-from sagemaker.workflow.condition_step import ConditionStep
-from sagemaker.workflow.functions import JsonGet
+from sagemaker.workflow.condition_step import ConditionStep, JsonGet
+
 from sagemaker.workflow.parameters import (
     ParameterInteger,
     ParameterString,
@@ -237,8 +237,9 @@ def test_workflow_with_clarify(
             property_files=[property_file],
         )
 
+        # Keep the deprecated JsonGet in test to verify it's compatible with new changes
         cond_left = JsonGet(
-            step_name=step_process.name,
+            step=step_process,
             property_file="BiasOutput",
             json_path="post_training_bias_metrics.facets.F1[0].metrics[0].value",
         )