fix: Support transformer data parameterization (#3145)

Author: qidewenwhen

src/sagemaker/model.py

@@ -501,13 +501,14 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
             if is_pipeline_variable(self.model_data):
                 # model is not yet there, defer repacking to later during pipeline execution
                 if not isinstance(self.sagemaker_session, PipelineSession):
-                    # TODO: link the doc in the warning once ready
                     logging.warning(
                         "The model_data is a Pipeline variable of type %s, "
                         "which should be used under `PipelineSession` and "
                         "leverage `ModelStep` to create or register model. "
                         "Otherwise some functionalities e.g. "
-                        "runtime repack may be missing",
+                        "runtime repack may be missing. For more, see: "
+                        "https://sagemaker.readthedocs.io/en/stable/"
+                        "amazon_sagemaker_model_building_pipeline.html#model-step",
                         type(self.model_data),
                     )
                     return

src/sagemaker/pipeline.py

@@ -278,6 +278,7 @@ def register(
         description: Optional[str] = None,
         drift_check_baselines: Optional[DriftCheckBaselines] = None,
         customer_metadata_properties: Optional[Dict[str, str]] = None,
+        domain: Optional[str] = None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -305,6 +306,8 @@ def register(
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -335,6 +338,7 @@ def register(
             container_def_list=container_def,
             drift_check_baselines=drift_check_baselines,
             customer_metadata_properties=customer_metadata_properties,
+            domain=domain,
         )
 
         self.sagemaker_session.create_model_package_from_containers(**model_pkg_args)

src/sagemaker/tensorflow/model.py

@@ -362,13 +362,14 @@ def prepare_container_def(
             if isinstance(self.sagemaker_session, PipelineSession):
                 self.sagemaker_session.context.need_runtime_repack.add(id(self))
             else:
-                # TODO: link the doc in the warning once ready
                 logging.warning(
                     "The model_data is a Pipeline variable of type %s, "
                     "which should be used under `PipelineSession` and "
                     "leverage `ModelStep` to create or register model. "
                     "Otherwise some functionalities e.g. "
-                    "runtime repack may be missing",
+                    "runtime repack may be missing. For more, see: "
+                    "https://sagemaker.readthedocs.io/en/stable/"
+                    "amazon_sagemaker_model_building_pipeline.html#model-step",
                     type(self.model_data),
                 )
             model_data = self.model_data

src/sagemaker/transformer.py

@@ -191,7 +191,7 @@ def transform(
                 Only meaningful when wait is ``True`` (default: ``True``).
         """
         local_mode = self.sagemaker_session.local_mode
-        if not local_mode and not data.startswith("s3://"):
+        if not local_mode and not is_pipeline_variable(data) and not data.startswith("s3://"):
             raise ValueError("Invalid S3 URI: {}".format(data))
 
         if job_name is not None:

src/sagemaker/workflow/model_step.py

@@ -111,7 +111,6 @@ def __init__(
                 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
         from sagemaker.workflow.utilities import validate_step_args_input
 
         validate_step_args_input(
@@ -121,7 +120,8 @@ def __init__(
                 Session.create_model_package_from_containers.__name__,
             },
             error_message="The step_args of ModelStep must be obtained from model.create() "
-            "or model.register().",
+            "or model.register(). For more, see: https://sagemaker.readthedocs.io/en/stable/"
+            "amazon_sagemaker_model_building_pipeline.html#model-step",
         )
         if not (step_args.create_model_request is None) ^ (
             step_args.create_model_package_request is None

src/sagemaker/workflow/step_collections.py

@@ -253,12 +253,13 @@ def __init__(
         steps.append(register_model_step)
         self.steps = steps
 
-        # TODO: add public document link here once ready
         warnings.warn(
             (
                 "We are deprecating the use of RegisterModel. "
                 "Instead, please use the ModelStep, which simply takes in the step arguments "
-                "generated by model.register()."
+                "generated by model.register(). For more, see: "
+                "https://sagemaker.readthedocs.io/en/stable/"
+                "amazon_sagemaker_model_building_pipeline.html#model-step"
             ),
             DeprecationWarning,
         )

src/sagemaker/workflow/steps.py

@@ -444,12 +444,13 @@ def __init__(
 
         self._properties = Properties(path=f"Steps.{name}", shape_name="DescribeModelOutput")
 
-        # TODO: add public document link here once ready
         warnings.warn(
             (
                 "We are deprecating the use of CreateModelStep. "
                 "Instead, please use the ModelStep, which simply takes in the step arguments "
-                "generated by model.create()."
+                "generated by model.create(). For more, see: "
+                "https://sagemaker.readthedocs.io/en/stable/"
+                "amazon_sagemaker_model_building_pipeline.html#model-step"
             ),
             DeprecationWarning,
         )

tests/unit/sagemaker/workflow/test_transform_step.py

@@ -89,10 +89,10 @@ def test_transform_step_with_transformer(pipeline_session):
         output_path=f"s3://{pipeline_session.default_bucket()}/Transform",
         sagemaker_session=pipeline_session,
     )
-
-    transform_inputs = TransformInput(
-        data=f"s3://{pipeline_session.default_bucket()}/batch-data",
+    data = ParameterString(
+        name="Data", default_value=f"s3://{pipeline_session.default_bucket()}/batch-data"
     )
+    transform_inputs = TransformInput(data=data)
 
     with warnings.catch_warnings(record=True) as w:
         step_args = transformer.transform(
@@ -120,10 +120,11 @@ def test_transform_step_with_transformer(pipeline_session):
     pipeline = Pipeline(
         name="MyPipeline",
         steps=[step],
-        parameters=[model_name],
+        parameters=[model_name, data],
         sagemaker_session=pipeline_session,
     )
     step_args.args["ModelName"] = model_name.expr
+    step_args.args["TransformInput"]["DataSource"]["S3DataSource"]["S3Uri"] = data.expr
     assert json.loads(pipeline.definition())["Steps"][0] == {
         "Name": "MyTransformStep",
         "Type": "Transform",