fix: Prevent passing PipelineVariable object into image_uris.retrieve (#3054)

Author: qidewenwhen

src/sagemaker/image_uris.py

@@ -23,6 +23,7 @@
 from sagemaker.jumpstart.utils import is_jumpstart_model_input
 from sagemaker.spark import defaults
 from sagemaker.jumpstart import artifacts
+from sagemaker.workflow import is_pipeline_variable
 
 logger = logging.getLogger(__name__)
 
@@ -104,11 +105,17 @@ def retrieve(
 
     Raises:
         NotImplementedError: If the scope is not supported.
-        ValueError: If the combination of arguments specified is not supported.
+        ValueError: If the combination of arguments specified is not supported or
+            any PipelineVariable object is passed in.
         VulnerableJumpStartModelError: If any of the dependencies required by the script have
             known security vulnerabilities.
         DeprecatedJumpStartModelError: If the version of the model is deprecated.
     """
+    args = dict(locals())
+    for name, val in args.items():
+        if is_pipeline_variable(val):
+            raise ValueError("%s should not be a pipeline variable (%s)" % (name, type(val)))
+
     if is_jumpstart_model_input(model_id, model_version):
         return artifacts._retrieve_image_uri(
             model_id,

src/sagemaker/workflow/entities.py

@@ -78,7 +78,11 @@ def __add__(self, other: Union[Expression, PrimitiveType]):
 
     def __str__(self):
         """Override built-in String function for PipelineVariable"""
-        raise TypeError("Pipeline variables do not support __str__ operation.")
+        raise TypeError(
+            "Pipeline variables do not support __str__ operation. "
+            "Please use `.to_string()` to convert it to string type in execution time"
+            "or use `.expr` to translate it to Json for display purpose in Python SDK."
+        )
 
     def __int__(self):
         """Override built-in Integer function for PipelineVariable"""

tests/integ/sagemaker/workflow/test_model_create_and_registration.py

@@ -90,10 +90,12 @@ def test_conditional_pytorch_training_model_registration(
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
     good_enough_input = ParameterInteger(name="GoodEnoughInput", default_value=1)
     in_condition_input = ParameterString(name="Foo", default_value="Foo")
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=entry_point,
         role=role,
@@ -153,7 +155,6 @@ def test_conditional_pytorch_training_model_registration(
             in_condition_input,
             good_enough_input,
             instance_count,
-            instance_type,
         ],
         steps=[step_train, step_cond],
         sagemaker_session=sagemaker_session,
@@ -259,8 +260,10 @@ def test_sklearn_xgboost_sip_model_registration(
     prefix = "sip"
     bucket_name = sagemaker_session.default_bucket()
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
 
+    # The instance_type should not be a pipeline variable
+    # since it is used to retrieve image_uri in compile time (PySDK)
     sklearn_processor = SKLearnProcessor(
         role=role,
         instance_type=instance_type,
@@ -331,6 +334,8 @@ def test_sklearn_xgboost_sip_model_registration(
     source_dir = base_dir
     code_location = "s3://{0}/{1}/code".format(bucket_name, prefix)
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     estimator = XGBoost(
         entry_point=entry_point,
         source_dir=source_dir,
@@ -416,7 +421,6 @@ def test_sklearn_xgboost_sip_model_registration(
             train_data_path_param,
             val_data_path_param,
             model_path_param,
-            instance_type,
             instance_count,
             output_path_param,
         ],
@@ -462,7 +466,7 @@ def test_model_registration_with_drift_check_baselines(
     pipeline_name,
 ):
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
 
     # upload model data to s3
     model_local_path = os.path.join(DATA_DIR, "mxnet_mnist/model.tar.gz")
@@ -550,6 +554,9 @@ def test_model_registration_with_drift_check_baselines(
         ),
     )
     customer_metadata_properties = {"key1": "value1"}
+
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     estimator = XGBoost(
         entry_point="training.py",
         source_dir=os.path.join(DATA_DIR, "sip"),
@@ -579,7 +586,6 @@ def test_model_registration_with_drift_check_baselines(
         parameters=[
             model_uri_param,
             metrics_uri_param,
-            instance_type,
             instance_count,
         ],
         steps=[step_register],
@@ -667,9 +673,11 @@ def test_model_registration_with_model_repack(
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
     good_enough_input = ParameterInteger(name="GoodEnoughInput", default_value=1)
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=entry_point,
         role=role,
@@ -724,7 +732,7 @@ def test_model_registration_with_model_repack(
 
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[good_enough_input, instance_count, instance_type],
+        parameters=[good_enough_input, instance_count],
         steps=[step_cond],
         sagemaker_session=sagemaker_session,
     )
@@ -767,8 +775,10 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     tensorflow_estimator = TensorFlow(
         entry_point=entry_point,
         role=role,
@@ -809,10 +819,7 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
 
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[
-            instance_count,
-            instance_type,
-        ],
+        parameters=[instance_count],
         steps=[step_train, step_register_model],
         sagemaker_session=sagemaker_session,
     )

tests/integ/sagemaker/workflow/test_model_steps.py

@@ -83,15 +83,16 @@ def test_pytorch_training_model_registration_and_creation_without_custom_inferen
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=entry_point,
         role=role,
         framework_version="1.5.0",
         py_version="py3",
         instance_count=instance_count,
-        instance_type=instance_type,
+        instance_type="ml.m5.xlarge",
         sagemaker_session=pipeline_session,
     )
     train_step_args = pytorch_estimator.fit(inputs=inputs)
@@ -140,7 +141,7 @@ def test_pytorch_training_model_registration_and_creation_without_custom_inferen
     )
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[instance_count, instance_type],
+        parameters=[instance_count],
         steps=[step_train, step_model_regis, step_model_create, step_fail],
         sagemaker_session=pipeline_session,
     )
@@ -203,15 +204,16 @@ def test_pytorch_training_model_registration_and_creation_with_custom_inference(
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=entry_point,
         role=role,
         framework_version="1.5.0",
         py_version="py3",
         instance_count=instance_count,
-        instance_type=instance_type,
+        instance_type="ml.m5.xlarge",
         sagemaker_session=pipeline_session,
         output_kms_key=kms_key,
     )
@@ -267,7 +269,7 @@ def test_pytorch_training_model_registration_and_creation_with_custom_inference(
     )
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[instance_count, instance_type],
+        parameters=[instance_count],
         steps=[step_train, step_model_regis, step_model_create, step_fail],
         sagemaker_session=pipeline_session,
     )
@@ -400,7 +402,6 @@ def test_model_registration_with_drift_check_baselines_and_model_metrics(
     pipeline_name,
 ):
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
 
     # upload model data to s3
     model_local_path = os.path.join(DATA_DIR, "mxnet_mnist/model.tar.gz")
@@ -488,10 +489,12 @@ def test_model_registration_with_drift_check_baselines_and_model_metrics(
         ),
     )
     customer_metadata_properties = {"key1": "value1"}
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     estimator = XGBoost(
         entry_point="training.py",
         source_dir=os.path.join(DATA_DIR, "sip"),
-        instance_type=instance_type,
+        instance_type="ml.m5.xlarge",
         instance_count=instance_count,
         framework_version="0.90-2",
         sagemaker_session=pipeline_session,
@@ -524,7 +527,6 @@ def test_model_registration_with_drift_check_baselines_and_model_metrics(
         parameters=[
             model_uri_param,
             metrics_uri_param,
-            instance_type,
             instance_count,
         ],
         steps=[step_model_register],
@@ -606,13 +608,14 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
     )
     inputs = TrainingInput(s3_data=input_path)
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     tensorflow_estimator = TensorFlow(
         entry_point=entry_point,
         role=role,
         instance_count=instance_count,
-        instance_type=instance_type,
+        instance_type="ml.m5.xlarge",
         framework_version=tf_full_version,
         py_version=tf_full_py_version,
         sagemaker_session=pipeline_session,
@@ -645,10 +648,7 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
     )
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[
-            instance_count,
-            instance_type,
-        ],
+        parameters=[instance_count],
         steps=[step_train, step_register_model],
         sagemaker_session=pipeline_session,
     )

tests/integ/sagemaker/workflow/test_retry.py

@@ -31,10 +31,7 @@
     _REGISTER_MODEL_NAME_BASE,
     _CREATE_MODEL_NAME_BASE,
 )
-from sagemaker.workflow.parameters import (
-    ParameterInteger,
-    ParameterString,
-)
+from sagemaker.workflow.parameters import ParameterInteger
 from sagemaker.pytorch.estimator import PyTorch
 from sagemaker.workflow.pipeline import Pipeline
 from sagemaker.workflow.retry import (
@@ -185,9 +182,11 @@ def test_model_registration_with_model_repack(
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
     good_enough_input = ParameterInteger(name="GoodEnoughInput", default_value=1)
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=entry_point,
         role=role,
@@ -252,7 +251,7 @@ def test_model_registration_with_model_repack(
     )
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[good_enough_input, instance_count, instance_type],
+        parameters=[good_enough_input, instance_count],
         steps=[step_train, step_cond],
         sagemaker_session=pipeline_session,
     )

tests/integ/sagemaker/workflow/test_training_steps.py

@@ -59,7 +59,7 @@ def test_training_job_with_debugger_and_profiler(
     pytorch_training_latest_py_version,
 ):
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
 
     rules = [
         Rule.sagemaker(rule_configs.vanishing_gradient()),
@@ -78,6 +78,8 @@ def test_training_job_with_debugger_and_profiler(
     )
     inputs = TrainingInput(s3_data=input_path)
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=script_path,
         role="SageMakerRole",
@@ -98,7 +100,7 @@ def test_training_job_with_debugger_and_profiler(
 
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[instance_count, instance_type],
+        parameters=[instance_count],
         steps=[step_train],
         sagemaker_session=sagemaker_session,
     )

tests/integ/sagemaker/workflow/test_tuning_steps.py

@@ -91,8 +91,10 @@ def test_tuning_single_algo(
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=entry_point,
         role=role,
@@ -167,7 +169,7 @@ def test_tuning_single_algo(
 
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[instance_count, instance_type, min_batch_size, max_batch_size],
+        parameters=[instance_count, min_batch_size, max_batch_size],
         steps=[step_tune, step_best_model, step_second_best_model],
         sagemaker_session=pipeline_session,
     )
@@ -221,10 +223,12 @@ def test_tuning_multi_algos(
     )
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
 
     input_data = f"s3://sagemaker-sample-data-{region_name}/processing/census/census-income.csv"
 
+    # The instance_type should not be a pipeline variable
+    # since it is used to retrieve image_uri in compile time (PySDK)
     sklearn_processor = SKLearnProcessor(
         framework_version="0.20.0",
         instance_type=instance_type,
@@ -259,6 +263,8 @@ def test_tuning_multi_algos(
     json_get_hp = JsonGet(
         step_name=step_process.name, property_file=property_file, json_path="train_size"
     )
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=entry_point,
         role=role,
@@ -307,7 +313,7 @@ def test_tuning_multi_algos(
 
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[instance_count, instance_type, min_batch_size, max_batch_size],
+        parameters=[instance_count, min_batch_size, max_batch_size, static_hp_1],
         steps=[step_process, step_tune],
         sagemaker_session=sagemaker_session,
     )