change: Add PipelineVariable annotation in framework models

Author: Dewen Qi

src/sagemaker/chainer/model.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 import logging
+from typing import Optional, Union
 
 import sagemaker
 from sagemaker import image_uris
@@ -27,6 +28,8 @@
 from sagemaker.deserializers import NumpyDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.serializers import NumpySerializer
+from sagemaker.utils import to_string
+from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger("sagemaker")
 
@@ -75,14 +78,14 @@ class ChainerModel(FrameworkModel):
 
     def __init__(
         self,
-        model_data,
-        role,
-        entry_point,
-        image_uri=None,
-        framework_version=None,
-        py_version=None,
-        predictor_cls=ChainerPredictor,
-        model_server_workers=None,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        entry_point: str,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        framework_version: Optional[str] = None,
+        py_version: Optional[str] = None,
+        predictor_cls: callable = ChainerPredictor,
+        model_server_workers: Optional[Union[int, PipelineVariable]] = None,
         **kwargs
     ):
         """Initialize an ChainerModel.
@@ -180,7 +183,9 @@ def prepare_container_def(
         deploy_env.update(self._script_mode_env_vars())
 
         if self.model_server_workers:
-            deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
+            deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = to_string(
+                self.model_server_workers
+            )
         return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
 
     def serving_image_uri(

src/sagemaker/estimator.py

@@ -73,6 +73,7 @@
     build_dict,
     get_config_value,
     name_from_base,
+    to_string,
 )
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.pipeline_context import (
@@ -1848,10 +1849,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 is_pipeline_variable(v) else str(v))
-                for (k, v) in current_hyperparameters.items()
-            }
+            hyperparameters = {str(k): to_string(v) for (k, v) in current_hyperparameters.items()}
 
         train_args = config.copy()
         train_args["input_mode"] = estimator.input_mode

src/sagemaker/huggingface/model.py

@@ -14,18 +14,23 @@
 from __future__ import absolute_import
 
 import logging
+from typing import Optional, Union, List, Dict
 
 import sagemaker
-from sagemaker import image_uris
+from sagemaker import image_uris, ModelMetrics
 from sagemaker.deserializers import JSONDeserializer
+from sagemaker.drift_check_baselines import DriftCheckBaselines
 from sagemaker.fw_utils import (
     model_code_key_prefix,
     validate_version_or_image_args,
 )
+from sagemaker.metadata_properties import MetadataProperties
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.predictor import Predictor
 from sagemaker.serializers import JSONSerializer
 from sagemaker.session import Session
+from sagemaker.utils import to_string
+from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger("sagemaker")
 
@@ -92,16 +97,16 @@ class HuggingFaceModel(FrameworkModel):
 
     def __init__(
         self,
-        role,
-        model_data=None,
-        entry_point=None,
-        transformers_version=None,
-        tensorflow_version=None,
-        pytorch_version=None,
-        py_version=None,
-        image_uri=None,
-        predictor_cls=HuggingFacePredictor,
-        model_server_workers=None,
+        role: str,
+        model_data: Optional[Union[str, PipelineVariable]] = None,
+        entry_point: Optional[str] = None,
+        transformers_version: Optional[str] = None,
+        tensorflow_version: Optional[str] = None,
+        pytorch_version: Optional[str] = None,
+        py_version: Optional[str] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        predictor_cls: callable = HuggingFacePredictor,
+        model_server_workers: Optional[Union[int, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize a HuggingFaceModel.
@@ -291,21 +296,21 @@ def deploy(
 
     def register(
         self,
-        content_types,
-        response_types,
-        inference_instances=None,
-        transform_instances=None,
-        model_package_name=None,
-        model_package_group_name=None,
-        image_uri=None,
-        model_metrics=None,
-        metadata_properties=None,
-        marketplace_cert=False,
-        approval_status=None,
-        description=None,
-        drift_check_baselines=None,
-        customer_metadata_properties=None,
-        domain=None,
+        content_types: List[Union[str, PipelineVariable]],
+        response_types: List[Union[str, PipelineVariable]],
+        inference_instances: Optional[List[Union[str, PipelineVariable]]] = None,
+        transform_instances: Optional[List[Union[str, PipelineVariable]]] = None,
+        model_package_name: Optional[Union[str, PipelineVariable]] = None,
+        model_package_group_name: Optional[Union[str, PipelineVariable]] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        model_metrics: Optional[ModelMetrics] = None,
+        metadata_properties: Optional[MetadataProperties] = None,
+        marketplace_cert: bool = False,
+        approval_status: Optional[Union[str, PipelineVariable]] = None,
+        description: Optional[str] = None,
+        drift_check_baselines: Optional[DriftCheckBaselines] = None,
+        customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        domain: Optional[Union[str, PipelineVariable]] = None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -409,7 +414,9 @@ def prepare_container_def(
         deploy_env.update(self._script_mode_env_vars())
 
         if self.model_server_workers:
-            deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
+            deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = to_string(
+                self.model_server_workers
+            )
         return sagemaker.container_def(
             deploy_image, self.repacked_model_data or self.model_data, deploy_env
         )

src/sagemaker/multidatamodel.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 import os
+from typing import Union, Optional
 
 from six.moves.urllib.parse import urlparse
 
@@ -22,6 +23,8 @@
 from sagemaker.deprecations import removed_kwargs
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.utils import pop_out_unused_kwarg
+from sagemaker.workflow.entities import PipelineVariable
 
 MULTI_MODEL_CONTAINER_MODE = "MultiModel"
 
@@ -34,12 +37,12 @@ class MultiDataModel(Model):
 
     def __init__(
         self,
-        name,
-        model_data_prefix,
-        model=None,
-        image_uri=None,
-        role=None,
-        sagemaker_session=None,
+        name: str,
+        model_data_prefix: str,
+        model: Model = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        role: Optional[str] = None,
+        sagemaker_session: Optional[Session] = None,
         **kwargs,
     ):
         """Initialize a ``MultiDataModel``.
@@ -106,6 +109,7 @@ def __init__(
 
         # Set the ``Model`` parameters if the model parameter is not specified
         if not self.model:
+            pop_out_unused_kwarg("model_data", kwargs, self.model_data_prefix)
             super(MultiDataModel, self).__init__(
                 image_uri,
                 self.model_data_prefix,
@@ -115,7 +119,9 @@ def __init__(
                 **kwargs,
             )
 
-    def prepare_container_def(self, instance_type=None, accelerator_type=None):
+    def prepare_container_def(
+        self, instance_type=None, accelerator_type=None, serverless_inference_config=None
+    ):
         """Return a container definition set.
 
         Definition set includes MultiModel mode, model data and other parameters

src/sagemaker/mxnet/model.py

@@ -14,21 +14,26 @@
 from __future__ import absolute_import
 
 import logging
+from typing import Union, Optional, List, Dict
 
 import packaging.version
 
 import sagemaker
-from sagemaker import image_uris
+from sagemaker import image_uris, ModelMetrics
 from sagemaker.deserializers import JSONDeserializer
+from sagemaker.drift_check_baselines import DriftCheckBaselines
 from sagemaker.fw_utils import (
     model_code_key_prefix,
     python_deprecation_warning,
     validate_version_or_image_args,
 )
+from sagemaker.metadata_properties import MetadataProperties
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.mxnet import defaults
 from sagemaker.predictor import Predictor
 from sagemaker.serializers import JSONSerializer
+from sagemaker.utils import to_string
+from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger("sagemaker")
 
@@ -77,14 +82,14 @@ class MXNetModel(FrameworkModel):
 
     def __init__(
         self,
-        model_data,
-        role,
-        entry_point,
-        framework_version=None,
-        py_version=None,
-        image_uri=None,
-        predictor_cls=MXNetPredictor,
-        model_server_workers=None,
+        model_data: Union[str, PipelineVariable],
+        role: str,
+        entry_point: str,
+        framework_version: str = _LOWEST_MMS_VERSION,
+        py_version: Optional[str] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        predictor_cls: callable = MXNetPredictor,
+        model_server_workers: Optional[Union[int, PipelineVariable]] = None,
         **kwargs
     ):
         """Initialize an MXNetModel.
@@ -144,21 +149,21 @@ def __init__(
 
     def register(
         self,
-        content_types,
-        response_types,
-        inference_instances=None,
-        transform_instances=None,
-        model_package_name=None,
-        model_package_group_name=None,
-        image_uri=None,
-        model_metrics=None,
-        metadata_properties=None,
-        marketplace_cert=False,
-        approval_status=None,
-        description=None,
-        drift_check_baselines=None,
-        customer_metadata_properties=None,
-        domain=None,
+        content_types: List[Union[str, PipelineVariable]],
+        response_types: List[Union[str, PipelineVariable]],
+        inference_instances: Optional[List[Union[str, PipelineVariable]]] = None,
+        transform_instances: Optional[List[Union[str, PipelineVariable]]] = None,
+        model_package_name: Optional[Union[str, PipelineVariable]] = None,
+        model_package_group_name: Optional[Union[str, PipelineVariable]] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        model_metrics: Optional[ModelMetrics] = None,
+        metadata_properties: Optional[MetadataProperties] = None,
+        marketplace_cert: bool = False,
+        approval_status: Optional[Union[str, PipelineVariable]] = None,
+        description: Optional[str] = None,
+        drift_check_baselines: Optional[DriftCheckBaselines] = None,
+        customer_metadata_properties: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        domain: Optional[Union[str, PipelineVariable]] = None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -262,7 +267,9 @@ def prepare_container_def(
         deploy_env.update(self._script_mode_env_vars())
 
         if self.model_server_workers:
-            deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
+            deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = to_string(
+                self.model_server_workers
+            )
         return sagemaker.container_def(
             deploy_image, self.repacked_model_data or self.model_data, deploy_env
         )

src/sagemaker/parameter.py

@@ -15,7 +15,7 @@
 
 import json
 
-from sagemaker.workflow import is_pipeline_variable
+from sagemaker.utils import to_string
 
 
 class ParameterRange(object):
@@ -71,12 +71,8 @@ def as_tuning_range(self, name):
         """
         return {
             "Name": name,
-            "MinValue": str(self.min_value)
-            if not is_pipeline_variable(self.min_value)
-            else self.min_value.to_string(),
-            "MaxValue": str(self.max_value)
-            if not is_pipeline_variable(self.max_value)
-            else self.max_value.to_string(),
+            "MinValue": to_string(self.min_value),
+            "MaxValue": to_string(self.max_value),
             "ScalingType": self.scaling_type,
         }
 
@@ -110,7 +106,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 is_pipeline_variable(v) else v.to_string() for v in values]
+        self.values = [to_string(v) for v in values]
 
     def as_tuning_range(self, name):
         """Represent the parameter range as a dictionary.