breaking: rename sagemaker.tensorflow.serving.Model/Predictor to sagemaker.tensorflow.model.TensorFlowModel/Predictor

This commit also changes two attributes to
match the other framework model classes:

* FRAMEWORK_NAME --> __framework_name__
* _framework_version --> framework_version

Author: laurenyu

src/sagemaker/cli/tensorflow.py

@@ -68,9 +68,9 @@ def create_model(self, model_url):
         Args:
             model_url:
         """
-        from sagemaker.tensorflow.serving import Model
+        from sagemaker.tensorflow.model import TensorFlowModel
 
-        return Model(
+        return TensorFlowModel(
             model_data=model_url,
             role=self.role_name,
             entry_point=self.script,

src/sagemaker/rl/estimator.py

@@ -21,6 +21,7 @@
 import sagemaker.fw_utils as fw_utils
 from sagemaker.model import FrameworkModel, SAGEMAKER_OUTPUT_LOCATION
 from sagemaker.mxnet.model import MXNetModel
+from sagemaker.tensorflow.model import TensorFlowModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
 logger = logging.getLogger("sagemaker")
@@ -90,7 +91,7 @@ def __init__(
         :meth:`~sagemaker.amazon.estimator.Framework.deploy` creates a hosted
         SageMaker endpoint and based on the specified framework returns an
         :class:`~sagemaker.amazon.mxnet.model.MXNetPredictor` or
-        :class:`~sagemaker.amazon.tensorflow.serving.Predictor` instance that
+        :class:`~sagemaker.amazon.tensorflow.model.TensorFlowPredictor` instance that
         can be used to perform inference against the hosted model.
 
         Technical documentation on preparing RLEstimator scripts for
@@ -205,15 +206,15 @@ def create_model(
             sagemaker.model.FrameworkModel: Depending on input parameters returns
                 one of the following:
 
-                * :class:`~sagemaker.model.FrameworkModel` - if ``image_name`` was specified
+                * :class:`~sagemaker.model.FrameworkModel` - if ``image_name`` is specified
                     on the estimator;
-                * :class:`~sagemaker.mxnet.MXNetModel` - if ``image_name`` wasn't specified and
-                    MXNet was used as the RL backend;
-                * :class:`~sagemaker.tensorflow.serving.Model` - if ``image_name`` wasn't specified
-                    and TensorFlow was used as the RL backend.
+                * :class:`~sagemaker.mxnet.MXNetModel` - if ``image_name`` isn't specified and
+                    MXNet is used as the RL backend;
+                * :class:`~sagemaker.tensorflow.model.TensorFlowModel` - if ``image_name`` isn't
+                    specified and TensorFlow is used as the RL backend.
 
         Raises:
-            ValueError: If image_name was not specified and framework enum is not valid.
+            ValueError: If image_name is not specified and framework enum is not valid.
         """
         base_args = dict(
             model_data=self.model_data,
@@ -252,9 +253,7 @@ def create_model(
             )
 
         if self.framework == RLFramework.TENSORFLOW.value:
-            from sagemaker.tensorflow.serving import Model as tfsModel
-
-            return tfsModel(framework_version=self.framework_version, **base_args)
+            return TensorFlowModel(framework_version=self.framework_version, **base_args)
         if self.framework == RLFramework.MXNET.value:
             return MXNetModel(
                 framework_version=self.framework_version, py_version=PYTHON_VERSION, **extended_args

src/sagemaker/tensorflow/__init__.py

@@ -14,3 +14,4 @@
 from __future__ import absolute_import
 
 from sagemaker.tensorflow.estimator import TensorFlow  # noqa: F401 (imported but unused)
+from sagemaker.tensorflow.model import TensorFlowModel, TensorFlowPredictor  # noqa: F401

src/sagemaker/tensorflow/estimator.py

@@ -23,7 +23,7 @@
 from sagemaker.estimator import Framework
 import sagemaker.fw_utils as fw
 from sagemaker.tensorflow import defaults
-from sagemaker.tensorflow.serving import Model
+from sagemaker.tensorflow.model import TensorFlowModel
 from sagemaker.transformer import Transformer
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
@@ -249,12 +249,13 @@ def create_model(
         dependencies=None,
         **kwargs
     ):
-        """Create a ``Model`` object that can be used for creating SageMaker model entities,
-        deploying to a SageMaker endpoint, or starting SageMaker Batch Transform jobs.
+        """Create a ``TensorFlowModel`` object that can be used for creating
+        SageMaker model entities, deploying to a SageMaker endpoint, or
+        starting SageMaker Batch Transform jobs.
 
         Args:
-            role (str): The ``ExecutionRoleArn`` IAM Role ARN for the ``Model``, which is also
-                used during transform jobs. If not specified, the role from the Estimator is used.
+            role (str): The ``TensorFlowModel``, which is also used during transform jobs.
+                If not specified, the role from the Estimator is used.
             vpc_config_override (dict[str, list[str]]): Optional override for VpcConfig set on the
                 model. Default: use subnets and security groups from this Estimator.
 
@@ -267,11 +268,12 @@ def create_model(
                 source code dependencies aside from the entry point file (default: None).
             dependencies (list[str]): A list of paths to directories (absolute or relative) with
                 any additional libraries that will be exported to the container (default: None).
-            **kwargs: Additional kwargs passed to :class:`~sagemaker.tensorflow.serving.Model`.
+            **kwargs: Additional kwargs passed to
+                :class:`~sagemaker.tensorflow.model.TensorFlowModel`.
 
         Returns:
-            sagemaker.tensorflow.serving.Model: A ``Model`` object.
-                See :class:`~sagemaker.tensorflow.serving.Model` for full details.
+            sagemaker.tensorflow.model.TensorFlowModel: A ``TensorFlowModel`` object.
+                See :class:`~sagemaker.tensorflow.model.TensorFlowModel` for full details.
         """
         if "image" not in kwargs:
             kwargs["image"] = self.image_name
@@ -282,7 +284,7 @@ def create_model(
         if "enable_network_isolation" not in kwargs:
             kwargs["enable_network_isolation"] = self.enable_network_isolation()
 
-        return Model(
+        return TensorFlowModel(
             model_data=self.model_data,
             role=role or self.role,
             container_log_level=self.container_log_level,
@@ -418,9 +420,8 @@ def transformer(
                 container in MB.
             tags (list[dict]): List of tags for labeling a transform job. If none specified, then
                 the tags used for the training job are used for the transform job.
-            role (str): The ``ExecutionRoleArn`` IAM Role ARN for the ``Model``, which is also
-                used during transform jobs. If not specified, the role from the Estimator will be
-                used.
+            role (str): The IAM Role ARN for the ``TensorFlowModel``, which is also used
+                during transform jobs. If not specified, the role from the Estimator is used.
             volume_kms_key (str): Optional. KMS key ID for encrypting the volume attached to the ML
                 compute instance (default: None).
             entry_point (str): Path (absolute or relative) to the local Python source file which

src/sagemaker/tensorflow/serving.py renamed to src/sagemaker/tensorflow/model.py

@@ -10,7 +10,7 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-"""Placeholder docstring"""
+"""Classes for using TensorFlow on Amazon SageMaker for inference."""
 from __future__ import absolute_import
 
 import logging
@@ -22,7 +22,7 @@
 from sagemaker.tensorflow.defaults import TF_VERSION
 
 
-class Predictor(sagemaker.RealTimePredictor):
+class TensorFlowPredictor(sagemaker.RealTimePredictor):
     """A ``RealTimePredictor`` implementation for inference against TensorFlow
     Serving endpoints.
     """
@@ -37,7 +37,7 @@ def __init__(
         model_name=None,
         model_version=None,
     ):
-        """Initialize a ``TFSPredictor``. See ``sagemaker.RealTimePredictor``
+        """Initialize a ``TensorFlowPredictor``. See :class:`~sagemaker.predictor.RealTimePredictor`
         for more info about parameters.
 
         Args:
@@ -61,7 +61,7 @@ def __init__(
                 that should handle the request. If not specified, the latest
                 version of the model will be used.
         """
-        super(Predictor, self).__init__(
+        super(TensorFlowPredictor, self).__init__(
             endpoint_name, sagemaker_session, serializer, deserializer, content_type
         )
 
@@ -115,13 +115,13 @@ def predict(self, data, initial_args=None):
             else:
                 args["CustomAttributes"] = self._model_attributes
 
-        return super(Predictor, self).predict(data, args)
+        return super(TensorFlowPredictor, self).predict(data, args)
 
 
-class Model(sagemaker.model.FrameworkModel):
-    """Placeholder docstring"""
+class TensorFlowModel(sagemaker.model.FrameworkModel):
+    """A ``FrameworkModel`` implementation for inference with TensorFlow Serving."""
 
-    FRAMEWORK_NAME = "tensorflow-serving"
+    __framework_name__ = "tensorflow-serving"
     LOG_LEVEL_PARAM_NAME = "SAGEMAKER_TFS_NGINX_LOGLEVEL"
     LOG_LEVEL_MAP = {
         logging.DEBUG: "debug",
@@ -140,7 +140,7 @@ def __init__(
         image=None,
         framework_version=TF_VERSION,
         container_log_level=None,
-        predictor_cls=Predictor,
+        predictor_cls=TensorFlowPredictor,
         **kwargs
     ):
         """Initialize a Model.
@@ -171,15 +171,15 @@ def __init__(
             :class:`~sagemaker.model.FrameworkModel` and
             :class:`~sagemaker.model.Model`.
         """
-        super(Model, self).__init__(
+        super(TensorFlowModel, self).__init__(
             model_data=model_data,
             role=role,
             image=image,
             predictor_cls=predictor_cls,
             entry_point=entry_point,
             **kwargs
         )
-        self._framework_version = framework_version
+        self.framework_version = framework_version
         self._container_log_level = container_log_level
 
     def deploy(
@@ -196,10 +196,10 @@ def deploy(
     ):
 
         if accelerator_type and not self._eia_supported():
-            msg = "The TensorFlow version %s doesn't support EIA." % self._framework_version
-
+            msg = "The TensorFlow version %s doesn't support EIA." % self.framework_version
             raise AttributeError(msg)
-        return super(Model, self).deploy(
+
+        return super(TensorFlowModel, self).deploy(
             initial_instance_count=initial_instance_count,
             instance_type=instance_type,
             accelerator_type=accelerator_type,
@@ -213,7 +213,7 @@ def deploy(
 
     def _eia_supported(self):
         """Return true if TF version is EIA enabled"""
-        return [int(s) for s in self._framework_version.split(".")][:2] <= self.LATEST_EIA_VERSION
+        return [int(s) for s in self.framework_version.split(".")][:2] <= self.LATEST_EIA_VERSION
 
     def prepare_container_def(self, instance_type, accelerator_type=None):
         """
@@ -249,12 +249,12 @@ def _get_container_env(self):
         if not self._container_log_level:
             return self.env
 
-        if self._container_log_level not in Model.LOG_LEVEL_MAP:
+        if self._container_log_level not in self.LOG_LEVEL_MAP:
             logging.warning("ignoring invalid container log level: %s", self._container_log_level)
             return self.env
 
         env = dict(self.env)
-        env[Model.LOG_LEVEL_PARAM_NAME] = Model.LOG_LEVEL_MAP[self._container_log_level]
+        env[self.LOG_LEVEL_PARAM_NAME] = self.LOG_LEVEL_MAP[self._container_log_level]
         return env
 
     def _get_image_uri(self, instance_type, accelerator_type=None):
@@ -269,9 +269,9 @@ def _get_image_uri(self, instance_type, accelerator_type=None):
         region_name = self.sagemaker_session.boto_region_name
         return create_image_uri(
             region_name,
-            Model.FRAMEWORK_NAME,
+            self.__framework_name__,
             instance_type,
-            self._framework_version,
+            self.framework_version,
             accelerator_type=accelerator_type,
         )
 

tests/unit/sagemaker/tensorflow/test_estimator.py

@@ -21,7 +21,7 @@
 import pytest
 
 from sagemaker.estimator import _TrainingJob
-from sagemaker.tensorflow import defaults, serving, TensorFlow
+from sagemaker.tensorflow import defaults, model, TensorFlow
 from tests.unit import DATA_DIR
 
 SCRIPT_FILE = "dummy_script.py"
@@ -188,7 +188,7 @@ def test_create_model(sagemaker_session, tf_version):
     model = tf.create_model()
 
     assert model.sagemaker_session == sagemaker_session
-    assert model._framework_version == tf_version
+    assert model.framework_version == tf_version
     assert model.entry_point is None
     assert model.role == ROLE
     assert model.name == job_name
@@ -372,17 +372,17 @@ def test_script_mode_create_model(sagemaker_session):
     )
     tf._prepare_for_training()  # set output_path and job name as if training happened
 
-    model = tf.create_model()
+    tf_model = tf.create_model()
 
-    assert isinstance(model, serving.Model)
+    assert isinstance(tf_model, model.TensorFlowModel)
 
-    assert model.model_data == tf.model_data
-    assert model.role == tf.role
-    assert model.name == tf._current_job_name
-    assert model.container_log_level == tf.container_log_level
-    assert model._framework_version == "1.11"
-    assert model.sagemaker_session == sagemaker_session
-    assert model.enable_network_isolation()
+    assert tf_model.model_data == tf.model_data
+    assert tf_model.role == tf.role
+    assert tf_model.name == tf._current_job_name
+    assert tf_model.container_log_level == tf.container_log_level
+    assert tf_model.framework_version == "1.11"
+    assert tf_model.sagemaker_session == sagemaker_session
+    assert tf_model.enable_network_isolation()
 
 
 @patch("time.strftime", return_value=TIMESTAMP)