change: allow serving script to be defined for deploy() and transformer() with frameworks (#944)

Author: laurenyu

src/sagemaker/chainer/estimator.py

@@ -155,7 +155,13 @@ def hyperparameters(self):
         return hyperparameters
 
     def create_model(
-        self, model_server_workers=None, role=None, vpc_config_override=VPC_CONFIG_DEFAULT
+        self,
+        model_server_workers=None,
+        role=None,
+        vpc_config_override=VPC_CONFIG_DEFAULT,
+        entry_point=None,
+        source_dir=None,
+        dependencies=None,
     ):
         """Create a SageMaker ``ChainerModel`` object that can be deployed to an
         ``Endpoint``.
@@ -171,17 +177,24 @@ def create_model(
                 the model. Default: use subnets and security groups from this Estimator.
                 * 'Subnets' (list[str]): List of subnet ids.
                 * 'SecurityGroupIds' (list[str]): List of security group ids.
+            entry_point (str): Path (absolute or relative) to the local Python source file which should be executed
+                as the entry point to training. If not specified, the training entry point is used.
+            source_dir (str): Path (absolute or relative) to a directory with any other serving
+                source code dependencies aside from the entry point file.
+                If not specified, the model source directory from training is used.
+            dependencies (list[str]): A list of paths to directories (absolute or relative) with
+                any additional libraries that will be exported to the container.
+                If not specified, the dependencies from training are used.
 
         Returns:
             sagemaker.chainer.model.ChainerModel: A SageMaker ``ChainerModel``
             object. See :func:`~sagemaker.chainer.model.ChainerModel` for full details.
         """
-        role = role or self.role
         return ChainerModel(
             self.model_data,
-            role,
-            self.entry_point,
-            source_dir=self._model_source_dir(),
+            role or self.role,
+            entry_point or self.entry_point,
+            source_dir=(source_dir or self._model_source_dir()),
             enable_cloudwatch_metrics=self.enable_cloudwatch_metrics,
             name=self._current_job_name,
             container_log_level=self.container_log_level,
@@ -192,7 +205,7 @@ def create_model(
             image=self.image_name,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
-            dependencies=self.dependencies,
+            dependencies=(dependencies or self.dependencies),
         )
 
     @classmethod

src/sagemaker/estimator.py

@@ -1,4 +1,4 @@
-# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2017-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"). You
 # may not use this file except in compliance with the License. A copy of
@@ -1524,6 +1524,7 @@ def transformer(
         role=None,
         model_server_workers=None,
         volume_kms_key=None,
+        entry_point=None,
     ):
         """Return a ``Transformer`` that uses a SageMaker Model based on the
         training job. It reuses the SageMaker Session and base job name used by
@@ -1561,11 +1562,19 @@ def transformer(
                 worker per vCPU.
             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 should be executed
+                as the entry point to training. If not specified, the training entry point is used.
+
+        Returns:
+            sagemaker.transformer.Transformer: a ``Transformer`` object that can be used to start a
+                SageMaker Batch Transform job.
         """
         role = role or self.role
 
         if self.latest_training_job is not None:
-            model = self.create_model(role=role, model_server_workers=model_server_workers)
+            model = self.create_model(
+                role=role, model_server_workers=model_server_workers, entry_point=entry_point
+            )
 
             container_def = model.prepare_container_def(instance_type)
             model_name = model.name or name_from_image(container_def["Image"])

src/sagemaker/mxnet/estimator.py

@@ -135,7 +135,13 @@ def _configure_distribution(self, distributions):
             self._hyperparameters[self.LAUNCH_PS_ENV_NAME] = enabled
 
     def create_model(
-        self, model_server_workers=None, role=None, vpc_config_override=VPC_CONFIG_DEFAULT
+        self,
+        model_server_workers=None,
+        role=None,
+        vpc_config_override=VPC_CONFIG_DEFAULT,
+        entry_point=None,
+        source_dir=None,
+        dependencies=None,
     ):
         """Create a SageMaker ``MXNetModel`` object that can be deployed to an
         ``Endpoint``.
@@ -151,17 +157,24 @@ def create_model(
                 the model. Default: use subnets and security groups from this Estimator.
                 * 'Subnets' (list[str]): List of subnet ids.
                 * 'SecurityGroupIds' (list[str]): List of security group ids.
+            entry_point (str): Path (absolute or relative) to the local Python source file which should be executed
+                as the entry point to training. If not specified, the training entry point is used.
+            source_dir (str): Path (absolute or relative) to a directory with any other serving
+                source code dependencies aside from the entry point file.
+                If not specified, the model source directory from training is used.
+            dependencies (list[str]): A list of paths to directories (absolute or relative) with
+                any additional libraries that will be exported to the container.
+                If not specified, the dependencies from training are used.
 
         Returns:
             sagemaker.mxnet.model.MXNetModel: A SageMaker ``MXNetModel`` object.
             See :func:`~sagemaker.mxnet.model.MXNetModel` for full details.
         """
-        role = role or self.role
         return MXNetModel(
             self.model_data,
-            role,
-            self.entry_point,
-            source_dir=self._model_source_dir(),
+            role or self.role,
+            entry_point or self.entry_point,
+            source_dir=(source_dir or self._model_source_dir()),
             enable_cloudwatch_metrics=self.enable_cloudwatch_metrics,
             name=self._current_job_name,
             container_log_level=self.container_log_level,
@@ -172,7 +185,7 @@ def create_model(
             model_server_workers=model_server_workers,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
-            dependencies=self.dependencies,
+            dependencies=(dependencies or self.dependencies),
         )
 
     @classmethod

src/sagemaker/pytorch/estimator.py

@@ -108,7 +108,13 @@ def __init__(
         self.py_version = py_version
 
     def create_model(
-        self, model_server_workers=None, role=None, vpc_config_override=VPC_CONFIG_DEFAULT
+        self,
+        model_server_workers=None,
+        role=None,
+        vpc_config_override=VPC_CONFIG_DEFAULT,
+        entry_point=None,
+        source_dir=None,
+        dependencies=None,
     ):
         """Create a SageMaker ``PyTorchModel`` object that can be deployed to an
         ``Endpoint``.
@@ -124,17 +130,24 @@ def create_model(
                 the model. Default: use subnets and security groups from this Estimator.
                 * 'Subnets' (list[str]): List of subnet ids.
                 * 'SecurityGroupIds' (list[str]): List of security group ids.
+            entry_point (str): Path (absolute or relative) to the local Python source file which should be executed
+                as the entry point to training. If not specified, the training entry point is used.
+            source_dir (str): Path (absolute or relative) to a directory with any other serving
+                source code dependencies aside from the entry point file.
+                If not specified, the model source directory from training is used.
+            dependencies (list[str]): A list of paths to directories (absolute or relative) with
+                any additional libraries that will be exported to the container.
+                If not specified, the dependencies from training are used.
 
         Returns:
             sagemaker.pytorch.model.PyTorchModel: A SageMaker ``PyTorchModel``
             object. See :func:`~sagemaker.pytorch.model.PyTorchModel` for full details.
         """
-        role = role or self.role
         return PyTorchModel(
             self.model_data,
-            role,
-            self.entry_point,
-            source_dir=self._model_source_dir(),
+            role or self.role,
+            entry_point or self.entry_point,
+            source_dir=(source_dir or self._model_source_dir()),
             enable_cloudwatch_metrics=self.enable_cloudwatch_metrics,
             name=self._current_job_name,
             container_log_level=self.container_log_level,
@@ -145,7 +158,7 @@ def create_model(
             model_server_workers=model_server_workers,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
-            dependencies=self.dependencies,
+            dependencies=(dependencies or self.dependencies),
         )
 
     @classmethod

src/sagemaker/tensorflow/estimator.py

@@ -470,8 +470,12 @@ def create_model(
         role=None,
         vpc_config_override=VPC_CONFIG_DEFAULT,
         endpoint_type=None,
+        entry_point=None,
+        source_dir=None,
+        dependencies=None,
     ):
-        """Create a SageMaker ``TensorFlowModel`` object that can be deployed to an ``Endpoint``.
+        """Create a ``Model`` 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
@@ -482,27 +486,55 @@ def create_model(
                 Default: use subnets and security groups from this Estimator.
                 * 'Subnets' (list[str]): List of subnet ids.
                 * 'SecurityGroupIds' (list[str]): List of security group ids.
-            endpoint_type: Optional. Selects the software stack used by the inference server.
+            endpoint_type (str): Optional. Selects the software stack used by the inference server.
                 If  not specified, the model will be configured to use the default
                 SageMaker model server. If 'tensorflow-serving', the model will be configured to
                 use the SageMaker Tensorflow Serving container.
+            entry_point (str): Path (absolute or relative) to the local Python source file which should be executed
+                as the entry point to training. If not specified and ``endpoint_type`` is 'tensorflow-serving',
+                no entry point is used. If ``endpoint_type`` is also ``None``, then the training entry point is used.
+            source_dir (str): Path (absolute or relative) to a directory with any other serving
+                source code dependencies aside from the entry point file. If not specified and
+                ``endpoint_type`` is 'tensorflow-serving', no source_dir is used. If ``endpoint_type`` is also ``None``,
+                then the model source directory from training is used.
+            dependencies (list[str]): A list of paths to directories (absolute or relative) with
+                any additional libraries that will be exported to the container.
+                If not specified and ``endpoint_type`` is 'tensorflow-serving', ``dependencies`` is set to ``None``.
+                If ``endpoint_type`` is also ``None``, then the dependencies from training are used.
 
         Returns:
-            sagemaker.tensorflow.model.TensorFlowModel: A SageMaker ``TensorFlowModel`` object.
-            See :func:`~sagemaker.tensorflow.model.TensorFlowModel` for full details.
+            sagemaker.tensorflow.model.TensorFlowModel or sagemaker.tensorflow.serving.Model: A ``Model`` object.
+                See :class:`~sagemaker.tensorflow.serving.Model` or :class:`~sagemaker.tensorflow.model.TensorFlowModel`
+                for full details.
         """
-
         role = role or self.role
+
         if endpoint_type == "tensorflow-serving" or self._script_mode_enabled():
-            return self._create_tfs_model(role=role, vpc_config_override=vpc_config_override)
+            return self._create_tfs_model(
+                role=role,
+                vpc_config_override=vpc_config_override,
+                entry_point=entry_point,
+                source_dir=source_dir,
+                dependencies=dependencies,
+            )
 
         return self._create_default_model(
             model_server_workers=model_server_workers,
             role=role,
             vpc_config_override=vpc_config_override,
+            entry_point=entry_point,
+            source_dir=source_dir,
+            dependencies=dependencies,
         )
 
-    def _create_tfs_model(self, role=None, vpc_config_override=VPC_CONFIG_DEFAULT):
+    def _create_tfs_model(
+        self,
+        role=None,
+        vpc_config_override=VPC_CONFIG_DEFAULT,
+        entry_point=None,
+        source_dir=None,
+        dependencies=None,
+    ):
         """Placeholder docstring"""
         return Model(
             model_data=self.model_data,
@@ -513,15 +545,26 @@ def _create_tfs_model(self, role=None, vpc_config_override=VPC_CONFIG_DEFAULT):
             framework_version=utils.get_short_version(self.framework_version),
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
+            entry_point=entry_point,
+            source_dir=source_dir,
+            dependencies=dependencies,
         )
 
-    def _create_default_model(self, model_server_workers, role, vpc_config_override):
+    def _create_default_model(
+        self,
+        model_server_workers,
+        role,
+        vpc_config_override,
+        entry_point=None,
+        source_dir=None,
+        dependencies=None,
+    ):
         """Placeholder docstring"""
         return TensorFlowModel(
             self.model_data,
             role,
-            self.entry_point,
-            source_dir=self._model_source_dir(),
+            entry_point or self.entry_point,
+            source_dir=source_dir or self._model_source_dir(),
             enable_cloudwatch_metrics=self.enable_cloudwatch_metrics,
             env={"SAGEMAKER_REQUIREMENTS": self.requirements_file},
             image=self.image_name,
@@ -533,7 +576,7 @@ def _create_default_model(self, model_server_workers, role, vpc_config_override)
             model_server_workers=model_server_workers,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
-            dependencies=self.dependencies,
+            dependencies=dependencies or self.dependencies,
         )
 
     def hyperparameters(self):
@@ -625,6 +668,7 @@ def transformer(
         model_server_workers=None,
         volume_kms_key=None,
         endpoint_type=None,
+        entry_point=None,
     ):
         """Return a ``Transformer`` that uses a SageMaker Model based on the training job. It reuses the
         SageMaker Session and base job name used by the Estimator.
@@ -656,6 +700,9 @@ def transformer(
                 SageMaker model server.
                 If 'tensorflow-serving', the model will be configured to
                 use the SageMaker Tensorflow Serving container.
+            entry_point (str): Path (absolute or relative) to the local Python source file which should be executed
+                as the entry point to training. If not specified and ``endpoint_type`` is 'tensorflow-serving',
+                no entry point is used. If ``endpoint_type`` is also ``None``, then the training entry point is used.
         """
 
         role = role or self.role
@@ -664,6 +711,7 @@ def transformer(
             role=role,
             vpc_config_override=VPC_CONFIG_DEFAULT,
             endpoint_type=endpoint_type,
+            entry_point=entry_point,
         )
         return model.transformer(
             instance_count,

tests/data/tfs/tfs-test-entrypoint-with-handler/inference.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"). You
 # may not use this file except in compliance with the License. A copy of
@@ -10,19 +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.
-
-"""Exports a toy TensorFlow model.
-Exports a TensorFlow model to /opt/ml/model/
-This graph calculates,
-  y = a*x + b
-where a and b are variables with a=0.5 and b=2.
-"""
 import json
-import shutil
-
-
-def save_model():
-    shutil.copytree("/opt/ml/code/123", "/opt/ml/model/123")
 
 
 def input_handler(data, context):
@@ -36,7 +24,3 @@ def output_handler(data, context):
     response_content_type = context.accept_header
     prediction = data.content
     return prediction, response_content_type
-
-
-if __name__ == "__main__":
-    save_model()

tests/data/tfs/tfs-test-entrypoint-with-handler/training.py

@@ -0,0 +1,28 @@
+# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# 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.
+
+"""Exports a toy TensorFlow model.
+Exports a TensorFlow model to /opt/ml/model/
+This graph calculates,
+  y = a*x + b
+where a and b are variables with a=0.5 and b=2.
+"""
+import shutil
+
+
+def save_model():
+    shutil.copytree("/opt/ml/code/123", "/opt/ml/model/123")
+
+
+if __name__ == "__main__":
+    save_model()