Fix script mode image async attach and deploy (#692)

Enable framework_name_from_image to parse script mode image URIs. Currently
attach to a script job will generate an estimator fails to deploy because the
sdk thinks this is an custom image and use the training image for hosting.

Author: icywang86rui

src/sagemaker/chainer/estimator.py

@@ -157,7 +157,7 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
                 init_params[argument[len('sagemaker_'):]] = value
 
         image_name = init_params.pop('image')
-        framework, py_version, tag = framework_name_from_image(image_name)
+        framework, py_version, tag, _ = framework_name_from_image(image_name)
 
         if not framework:
             # If we were unable to parse the framework name from the image it is not one of our

src/sagemaker/fw_utils.py

@@ -211,29 +211,31 @@ def framework_name_from_image(image_name):
             str: The framework name
             str: The Python version
             str: The image tag
+            str: If the image is script mode
     """
     sagemaker_pattern = re.compile(r'^(\d+)(\.)dkr(\.)ecr(\.)(.+)(\.)amazonaws.com(/)(.*:.*)$')
     sagemaker_match = sagemaker_pattern.match(image_name)
     if sagemaker_match is None:
-        return None, None, None
+        return None, None, None, None
     else:
         # extract framework, python version and image tag
         # We must support both the legacy and current image name format.
         name_pattern = re.compile(
-            r'^sagemaker(?:-rl)?-(tensorflow|mxnet|chainer|pytorch|scikit-learn):(.*)-(.*?)-(py2|py3)$')
+            r'^sagemaker(?:-rl)?-(tensorflow|mxnet|chainer|pytorch|scikit-learn)(?:-)?(scriptmode)?:(.*)-(.*?)-(py2|py3)$') # noqa
         legacy_name_pattern = re.compile(
             r'^sagemaker-(tensorflow|mxnet)-(py2|py3)-(cpu|gpu):(.*)$')
 
         name_match = name_pattern.match(sagemaker_match.group(8))
         legacy_match = legacy_name_pattern.match(sagemaker_match.group(8))
 
         if name_match is not None:
-            fw, ver, device, py = name_match.group(1), name_match.group(2), name_match.group(3), name_match.group(4)
-            return fw, py, '{}-{}-{}'.format(ver, device, py)
+            fw, scriptmode, ver, device, py = name_match.group(1), name_match.group(2), name_match.group(3),\
+                name_match.group(4), name_match.group(5)
+            return fw, py, '{}-{}-{}'.format(ver, device, py), scriptmode
         elif legacy_match is not None:
-            return legacy_match.group(1), legacy_match.group(2), legacy_match.group(4)
+            return legacy_match.group(1), legacy_match.group(2), legacy_match.group(4), None
         else:
-            return None, None, None
+            return None, None, None, None
 
 
 def framework_version_from_tag(image_tag):

src/sagemaker/mxnet/estimator.py

@@ -132,7 +132,7 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
         """
         init_params = super(MXNet, cls)._prepare_init_params_from_job_description(job_details, model_channel_name)
         image_name = init_params.pop('image')
-        framework, py_version, tag = framework_name_from_image(image_name)
+        framework, py_version, tag, _ = framework_name_from_image(image_name)
 
         if not framework:
             # If we were unable to parse the framework name from the image it is not one of our

src/sagemaker/pytorch/estimator.py

@@ -114,7 +114,7 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
         """
         init_params = super(PyTorch, cls)._prepare_init_params_from_job_description(job_details, model_channel_name)
         image_name = init_params.pop('image')
-        framework, py_version, tag = framework_name_from_image(image_name)
+        framework, py_version, tag, _ = framework_name_from_image(image_name)
 
         if not framework:
             # If we were unable to parse the framework name from the image it is not one of our

src/sagemaker/rl/estimator.py

@@ -253,7 +253,7 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
             ._prepare_init_params_from_job_description(job_details, model_channel_name)
 
         image_name = init_params.pop('image')
-        framework, _, tag = fw_utils.framework_name_from_image(image_name)
+        framework, _, tag, _ = fw_utils.framework_name_from_image(image_name)
 
         if not framework:
             # If we were unable to parse the framework name from the image it is not one of our

src/sagemaker/sklearn/estimator.py

@@ -136,7 +136,7 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
         init_params = super(SKLearn, cls)._prepare_init_params_from_job_description(job_details)
 
         image_name = init_params.pop('image')
-        framework, py_version, _ = framework_name_from_image(image_name)
+        framework, py_version, _, _ = framework_name_from_image(image_name)
         init_params['py_version'] = py_version
 
         if framework and framework != cls.__framework_name__:

src/sagemaker/tensorflow/estimator.py

@@ -356,13 +356,16 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
                 init_params[argument] = value
 
         image_name = init_params.pop('image')
-        framework, py_version, tag = fw.framework_name_from_image(image_name)
+        framework, py_version, tag, script_mode = fw.framework_name_from_image(image_name)
         if not framework:
             # If we were unable to parse the framework name from the image it is not one of our
             # officially supported images, in this case just add the image to the init params.
             init_params['image_name'] = image_name
             return init_params
 
+        if script_mode:
+            init_params['script_mode'] = True
+
         init_params['py_version'] = py_version
 
         # We switched image tagging scheme from regular image version (e.g. '1.0') to more expressive

tests/data/tensorflow_mnist/mnist.py

@@ -187,6 +187,6 @@ def serving_input_fn():
     tf.estimator.train_and_evaluate(mnist_classifier, train_spec, eval_spec)
 
     if args.current_host == args.hosts[0]:
-        mnist_classifier.export_savedmodel(args.model_dir, serving_input_fn)
+        mnist_classifier.export_savedmodel('/opt/ml/model', serving_input_fn)
 
     tf_logger.info('====== Training finished =========')

tests/integ/test_tf_script_mode.py

@@ -12,8 +12,10 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import numpy as np
 import os
 import pytest
+import time
 
 import boto3
 from sagemaker.tensorflow import TensorFlow
@@ -40,7 +42,7 @@ def test_mnist(sagemaker_session, instance_type):
                            train_instance_type=instance_type,
                            sagemaker_session=sagemaker_session,
                            py_version='py3',
-                           framework_version='1.11',
+                           framework_version=TensorFlow.LATEST_VERSION,
                            base_job_name='test-tf-sm-mnist')
     inputs = estimator.sagemaker_session.upload_data(
         path=os.path.join(RESOURCE_PATH, 'data'),
@@ -49,7 +51,7 @@ def test_mnist(sagemaker_session, instance_type):
     with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
         estimator.fit(inputs)
     _assert_s3_files_exist(estimator.model_dir,
-                           ['graph.pbtxt', 'model.ckpt-0.index', 'model.ckpt-0.meta', 'saved_model.pb'])
+                           ['graph.pbtxt', 'model.ckpt-0.index', 'model.ckpt-0.meta'])
 
 
 @pytest.mark.canary_quick
@@ -63,7 +65,7 @@ def test_mnist_distributed(sagemaker_session, instance_type):
                            sagemaker_session=sagemaker_session,
                            py_version=integ.PYTHON_VERSION,
                            script_mode=True,
-                           framework_version='1.11',
+                           framework_version=TensorFlow.LATEST_VERSION,
                            distributions=PARAMETER_SERVER_DISTRIBUTION,
                            base_job_name='test-tf-sm-mnist')
     inputs = estimator.sagemaker_session.upload_data(
@@ -73,7 +75,32 @@ def test_mnist_distributed(sagemaker_session, instance_type):
     with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
         estimator.fit(inputs)
     _assert_s3_files_exist(estimator.model_dir,
-                           ['graph.pbtxt', 'model.ckpt-0.index', 'model.ckpt-0.meta', 'saved_model.pb'])
+                           ['graph.pbtxt', 'model.ckpt-0.index', 'model.ckpt-0.meta'])
+
+
+def test_mnist_async(sagemaker_session):
+    estimator = TensorFlow(entry_point=SCRIPT,
+                           role='SageMakerRole',
+                           train_instance_count=1,
+                           train_instance_type='ml.c5.4xlarge',
+                           sagemaker_session=sagemaker_session,
+                           py_version='py3',
+                           framework_version=TensorFlow.LATEST_VERSION,
+                           base_job_name='test-tf-sm-mnist')
+    inputs = estimator.sagemaker_session.upload_data(
+        path=os.path.join(RESOURCE_PATH, 'data'),
+        key_prefix='scriptmode/mnist')
+    estimator.fit(inputs, wait=False)
+    training_job_name = estimator.latest_training_job.name
+    time.sleep(20)
+    endpoint_name = training_job_name
+    with timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+        estimator = TensorFlow.attach(training_job_name=training_job_name, sagemaker_session=sagemaker_session)
+        predictor = estimator.deploy(initial_instance_count=1, instance_type='ml.c4.xlarge',
+                                     endpoint_name=endpoint_name)
+
+        result = predictor.predict(np.zeros(784))
+        print('predict result: {}'.format(result))
 
 
 def _assert_s3_files_exist(s3_url, files):

tests/unit/test_fw_utils.py

@@ -354,45 +354,50 @@ def walk():
 
 def test_framework_name_from_image_mxnet():
     image_name = '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.1-gpu-py3'
-    assert ('mxnet', 'py3', '1.1-gpu-py3') == fw_utils.framework_name_from_image(image_name)
+    assert ('mxnet', 'py3', '1.1-gpu-py3', None) == fw_utils.framework_name_from_image(image_name)
 
 
 def test_framework_name_from_image_tf():
     image_name = '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tensorflow:1.6-cpu-py2'
-    assert ('tensorflow', 'py2', '1.6-cpu-py2') == fw_utils.framework_name_from_image(image_name)
+    assert ('tensorflow', 'py2', '1.6-cpu-py2', None) == fw_utils.framework_name_from_image(image_name)
+
+
+def test_framework_name_from_image_tf_scriptmode():
+    image_name = '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tensorflow-scriptmode:1.12-cpu-py3'
+    assert ('tensorflow', 'py3', '1.12-cpu-py3', 'scriptmode') == fw_utils.framework_name_from_image(image_name)
 
 
 def test_framework_name_from_image_rl():
     image_name = '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-rl-mxnet:toolkit1.1-gpu-py3'
-    assert ('mxnet', 'py3', 'toolkit1.1-gpu-py3') == fw_utils.framework_name_from_image(image_name)
+    assert ('mxnet', 'py3', 'toolkit1.1-gpu-py3', None) == fw_utils.framework_name_from_image(image_name)
 
 
 def test_legacy_name_from_framework_image():
     image_name = '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-py3-gpu:2.5.6-gpu-py2'
-    framework, py_ver, tag = fw_utils.framework_name_from_image(image_name)
+    framework, py_ver, tag, _ = fw_utils.framework_name_from_image(image_name)
     assert framework == 'mxnet'
     assert py_ver == 'py3'
     assert tag == '2.5.6-gpu-py2'
 
 
 def test_legacy_name_from_wrong_framework():
-    framework, py_ver, tag = fw_utils.framework_name_from_image(
+    framework, py_ver, tag, _ = fw_utils.framework_name_from_image(
         '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-myown-py2-gpu:1')
     assert framework is None
     assert py_ver is None
     assert tag is None
 
 
 def test_legacy_name_from_wrong_python():
-    framework, py_ver, tag = fw_utils.framework_name_from_image(
+    framework, py_ver, tag, _ = fw_utils.framework_name_from_image(
         '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-myown-py4-gpu:1')
     assert framework is None
     assert py_ver is None
     assert tag is None
 
 
 def test_legacy_name_from_wrong_device():
-    framework, py_ver, tag = fw_utils.framework_name_from_image(
+    framework, py_ver, tag, _ = fw_utils.framework_name_from_image(
         '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-myown-py4-gpu:1')
     assert framework is None
     assert py_ver is None
@@ -401,7 +406,7 @@ def test_legacy_name_from_wrong_device():
 
 def test_legacy_name_from_image_any_tag():
     image_name = '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tensorflow-py2-cpu:any-tag'
-    framework, py_ver, tag = fw_utils.framework_name_from_image(image_name)
+    framework, py_ver, tag, _ = fw_utils.framework_name_from_image(image_name)
     assert framework == 'tensorflow'
     assert py_ver == 'py2'
     assert tag == 'any-tag'