Fix unit tests, changes for serving.

This fixes the unit tests that were broken.
Also adds the remaining work that allows serving for TF to work.
Pending: MXNet.

Author: Ignacio Quintero

src/sagemaker/estimator.py

@@ -82,13 +82,10 @@ def __init__(self, role, train_instance_count, train_instance_type,
         self.input_mode = input_mode
 
         if self.train_instance_type in ('local', 'local_gpu'):
-            self.local_mode = True
             if self.train_instance_type == 'local_gpu' and self.train_instance_count > 1:
                 raise RuntimeError("Distributed Training in Local GPU is not supported")
-
             self.sagemaker_session = sagemaker_session or LocalSession()
         else:
-            self.local_mode = False
             self.sagemaker_session = sagemaker_session or Session()
 
         self.base_job_name = base_job_name
@@ -160,8 +157,8 @@ def fit(self, inputs, wait=True, logs=True, job_name=None):
         # if output_path was specified we use it otherwise initialize here.
         # For Local Mode with no_internet=True we don't need an explicit output_path
         if self.output_path is None:
-            if self.local_mode and get_config_value('local.no_internet',
-                                                    self.sagemaker_session.config):
+            no_internet = get_config_value('local.no_internet', self.sagemaker_session.config)
+            if self.sagemaker_session.local_mode and no_internet:
                 self.output_path = ''
             else:
                 self.output_path = 's3://{}/'.format(self.sagemaker_session.default_bucket())
@@ -327,7 +324,7 @@ def start_new(cls, estimator, inputs):
             sagemaker.estimator.Framework: Constructed object that captures all information about the started job.
         """
 
-        local_mode = estimator.local_mode
+        local_mode = estimator.sagemaker_session.local_mode
 
         # Allow file:// input only in local mode
         if isinstance(inputs, str) and inputs.startswith('file://'):
@@ -608,19 +605,20 @@ def fit(self, inputs, wait=True, logs=True, job_name=None):
             base_name = self.base_job_name or base_name_from_image(self.train_image())
             self._current_job_name = name_from_base(base_name)
 
-        # if there is no source dir, use the directory containing the entry point.
-        if self.source_dir is None:
-            self.source_dir = os.path.dirname(self.entry_point)
-        self.entry_point = os.path.basename(self.entry_point)
-
         # validate source dir will raise a ValueError if there is something wrong with the
         # source directory. We are intentionally not handling it because this is a critical error.
         if self.source_dir and not self.source_dir.lower().startswith('s3://'):
             validate_source_dir(self.entry_point, self.source_dir)
 
         # if we are in local mode with no_internet=True. We want the container to just
         # mount the source dir instead of uploading to S3.
-        if self.local_mode and get_config_value('local.no_internet', self.sagemaker_session.config):
+        no_internet = get_config_value('local.no_internet', self.sagemaker_session.config)
+        if self.sagemaker_session.local_mode and no_internet:
+            # if there is no source dir, use the directory containing the entry point.
+            if self.source_dir is None:
+                self.source_dir = os.path.dirname(self.entry_point)
+            self.entry_point = os.path.basename(self.entry_point)
+
             code_dir = 'file://' + self.source_dir
             script = self.entry_point
         else:

src/sagemaker/local/image.py

@@ -86,6 +86,7 @@ def train(self, input_data_config, hyperparameters):
         """
         self.container_root = self._create_tmp_folder()
         os.mkdir(os.path.join(self.container_root, 'output'))
+        os.mkdir(os.path.join(self.container_root, 'shared'))
 
         data_dir = self._create_tmp_folder()
         volumes = []
@@ -121,8 +122,8 @@ def train(self, input_data_config, hyperparameters):
         training_dir = json.loads(hyperparameters[sagemaker.estimator.DIR_PARAM_NAME])
         parsed_uri = urlparse(training_dir)
         if parsed_uri.scheme == 'file':
-            print('appended Volume')
             volumes.append(_Volume(parsed_uri.path, '/opt/ml/code'))
+            volumes.append(_Volume(os.path.join(self.container_root, 'shared'), '/opt/ml/shared'))
 
         # Create the configuration files for each container that we will create
         # Each container will map the additional local volumes (if any).
@@ -179,7 +180,16 @@ def serve(self, primary_container):
 
         _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image)
 
-        self._generate_compose_file('serve', additional_env_vars=env_vars)
+        # If the user script was passed as a file:// mount it to the container.
+        script_dir = primary_container['Environment'][sagemaker.estimator.DIR_PARAM_NAME.upper()]
+        parsed_uri = urlparse(script_dir)
+        volumes = []
+        if parsed_uri.scheme == 'file':
+            volumes.append(_Volume(parsed_uri.path, '/opt/ml/code'))
+
+        self._generate_compose_file('serve',
+                                    additional_env_vars=env_vars,
+                                    additional_volumes=volumes)
         compose_command = self._compose()
         self.container = _HostingContainer(compose_command)
         self.container.up()
@@ -574,6 +584,10 @@ def _ecr_login_if_needed(boto_session, image):
     if _check_output('docker images -q %s' % image).strip():
         return
 
+    if not boto_session:
+        raise RuntimeError('A boto session is required to login to ECR.'
+                           'Please pull the image: %s manually.' % image)
+
     ecr = boto_session.client('ecr')
     auth = ecr.get_authorization_token(registryIds=[image.split('.')[0]])
     authorization_data = auth['authorizationData'][0]

src/sagemaker/local/local_session.py

@@ -190,6 +190,7 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
 
         self.sagemaker_client = LocalSagemakerClient(self)
         self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self.config)
+        self.local_mode = True
 
     def logs_for_job(self, job_name, wait=False, poll=5):
         # override logs_for_job() as it doesn't need to perform any action

src/sagemaker/model.py

@@ -16,7 +16,7 @@
 
 from sagemaker.fw_utils import tar_and_upload_dir, parse_s3_url
 from sagemaker.session import Session
-from sagemaker.utils import name_from_image
+from sagemaker.utils import name_from_image, get_config_value
 
 
 class Model(object):
@@ -160,7 +160,11 @@ def prepare_container_def(self, instance_type):
         Returns:
             dict[str, str]: A container definition object usable with the CreateModel API.
         """
-        self._upload_code(self.key_prefix or self.name or name_from_image(self.image))
+        no_internet = get_config_value('local.no_internet', self.sagemaker_session.config)
+        if self.sagemaker_session.local_mode and no_internet:
+            self.uploaded_code = None
+        else:
+            self._upload_code(self.key_prefix or self.name or name_from_image(self.image))
         deploy_env = dict(self.env)
         deploy_env.update(self._framework_env_vars())
         return sagemaker.container_def(self.image, self.model_data, deploy_env)
@@ -173,8 +177,17 @@ def _upload_code(self, key_prefix):
                                                 directory=self.source_dir)
 
     def _framework_env_vars(self):
-        return {SCRIPT_PARAM_NAME.upper(): self.uploaded_code.script_name,
-                DIR_PARAM_NAME.upper(): self.uploaded_code.s3_prefix,
-                CLOUDWATCH_METRICS_PARAM_NAME.upper(): str(self.enable_cloudwatch_metrics).lower(),
-                CONTAINER_LOG_LEVEL_PARAM_NAME.upper(): str(self.container_log_level),
-                SAGEMAKER_REGION_PARAM_NAME.upper(): self.sagemaker_session.boto_session.region_name}
+        if self.uploaded_code:
+            script_name = self.uploaded_code.script_name
+            dir_name = self.uploaded_code.s3_prefix
+        else:
+            script_name = self.entry_point
+            dir_name = 'file://' + self.source_dir
+
+        return {
+            SCRIPT_PARAM_NAME.upper(): script_name,
+            DIR_PARAM_NAME.upper(): dir_name,
+            CLOUDWATCH_METRICS_PARAM_NAME.upper(): str(self.enable_cloudwatch_metrics).lower(),
+            CONTAINER_LOG_LEVEL_PARAM_NAME.upper(): str(self.container_log_level),
+            SAGEMAKER_REGION_PARAM_NAME.upper(): self.sagemaker_session.region_name
+        }

src/sagemaker/session.py

@@ -87,8 +87,8 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
         """
         self.boto_session = boto_session or boto3.Session()
 
-        self.region = self.boto_session.region_name
-        if self.region is None:
+        self.region_name = self.boto_session.region_name
+        if self.region_name is None:
             raise ValueError('Must setup local AWS configuration with a region supported by SageMaker.')
 
         self.sagemaker_client = sagemaker_client or self.boto_session.client('sagemaker')
@@ -97,6 +97,8 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
         self.sagemaker_runtime_client = sagemaker_runtime_client or self.boto_session.client('runtime.sagemaker')
         prepend_user_agent(self.sagemaker_runtime_client)
 
+        self.local_mode = False
+
     @property
     def boto_region_name(self):
         return self.boto_session.region_name

src/sagemaker/tensorflow/estimator.py

@@ -20,6 +20,7 @@
 
 from sagemaker.estimator import Framework
 from sagemaker.fw_utils import create_image_uri, framework_name_from_image, framework_version_from_tag
+from sagemaker.utils import get_config_value
 
 from sagemaker.tensorflow.defaults import TF_VERSION
 from sagemaker.tensorflow.model import TensorFlowModel
@@ -305,7 +306,12 @@ def hyperparameters(self):
         hyperparameters = super(TensorFlow, self).hyperparameters()
 
         if not self.checkpoint_path:
-            self.checkpoint_path = os.path.join(self.output_path, self._current_job_name, 'checkpoints')
+            no_internet = get_config_value('local.no_internet', self.sagemaker_session.config)
+            if self.sagemaker_session.local_mode and no_internet:
+                self.checkpoint_path = '/opt/ml/shared/checkpoints'
+            else:
+                self.checkpoint_path = os.path.join(self.output_path,
+                                                    self._current_job_name, 'checkpoints')
 
         additional_hyperparameters = {'checkpoint_path': self.checkpoint_path,
                                       'training_steps': self.training_steps,

src/sagemaker/tensorflow/model.py

@@ -16,7 +16,7 @@
 from sagemaker.predictor import RealTimePredictor
 from sagemaker.tensorflow.defaults import TF_VERSION
 from sagemaker.tensorflow.predictor import tf_json_serializer, tf_json_deserializer
-from sagemaker.utils import name_from_image
+from sagemaker.utils import name_from_image, get_config_value
 
 
 class TensorFlowPredictor(RealTimePredictor):
@@ -83,11 +83,18 @@ def prepare_container_def(self, instance_type):
         """
         deploy_image = self.image
         if not deploy_image:
-            region_name = self.sagemaker_session.boto_session.region_name
+            region_name = self.sagemaker_session.region_name
             deploy_image = create_image_uri(region_name, self.__framework_name__, instance_type,
                                             self.framework_version, self.py_version)
+
         deploy_key_prefix = self.key_prefix or self.name or name_from_image(deploy_image)
-        self._upload_code(deploy_key_prefix)
+
+        no_internet = get_config_value('local.no_internet', self.sagemaker_session.config)
+        print('no_internet: %s  local_mode: %s' % (no_internet, self.sagemaker_session.local_mode))
+        if self.sagemaker_session.local_mode and no_internet:
+            self.uploaded_code = None
+        else:
+            self._upload_code(deploy_key_prefix)
         deploy_env = dict(self.env)
         deploy_env.update(self._framework_env_vars())
 

tests/component/test_mxnet_estimator.py

@@ -34,7 +34,9 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name='boto_session', region_name=REGION)
-    ims = Mock(name='sagemaker_session', boto_session=boto_mock)
+    ims = Mock(name='sagemaker_session', boto_session=boto_mock,
+               config=None, local_mode=False, region_name=REGION)
+
     ims.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     ims.expand_role = Mock(name="expand_role", return_value=ROLE)
     ims.sagemaker_client.describe_training_job = Mock(return_value={'ModelArtifacts':

tests/component/test_tf_estimator.py

@@ -34,7 +34,8 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name='boto_session', region_name=REGION)
-    ims = Mock(name='sagemaker_session', boto_session=boto_mock)
+    ims = Mock(name='sagemaker_session', boto_session=boto_mock, config=None,
+               local_mode=False, region_name=REGION)
     ims.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     ims.expand_role = Mock(name="expand_role", return_value=ROLE)
     ims.sagemaker_client.describe_training_job = Mock(return_value={'ModelArtifacts':
@@ -62,6 +63,7 @@ def test_deploy(sagemaker_session, tf_version):
          {'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
           'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
           'SAGEMAKER_SUBMIT_DIRECTORY': SOURCE_DIR,
+          'SAGEMAKER_REQUIREMENTS': '',
           'SAGEMAKER_REGION': REGION,
           'SAGEMAKER_PROGRAM': SCRIPT},
          'Image': image,

tests/unit/test_amazon_estimator.py

@@ -28,7 +28,8 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name='boto_session', region_name=REGION)
-    sms = Mock(name='sagemaker_session', boto_session=boto_mock)
+    sms = Mock(name='sagemaker_session', boto_session=boto_mock,
+               region_name=REGION, config=None, local_mode=False)
     sms.boto_region_name = REGION
     sms.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     returned_job_description = {'AlgorithmSpecification': {'TrainingInputMode': 'File',

tests/unit/test_estimator.py

@@ -89,7 +89,8 @@ def create_predictor(self, endpoint_name):
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name='boto_session', region_name=REGION)
-    ims = Mock(name='sagemaker_session', boto_session=boto_mock, region_name=REGION)
+    ims = Mock(name='sagemaker_session', boto_session=boto_mock, region_name=REGION,
+               config=None, local_mode=False)
     ims.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     ims.sagemaker_client.describe_training_job = Mock(name='describe_training_job',
                                                       return_value=DESCRIBE_TRAINING_JOB_RESULT)
@@ -533,18 +534,26 @@ def test_generic_to_deploy(sagemaker_session):
 
 
 @patch('sagemaker.estimator.LocalSession')
-def test_local_mode(sagemaker_session):
-    e = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, 'local', output_path='s3://bucket/prefix',
-                  sagemaker_session=sagemaker_session)
-    assert e.local_mode is True
+@patch('sagemaker.estimator.Session')
+def test_local_mode(session_class, local_session_class):
+    local_session = Mock()
+    local_session.local_mode = True
+
+    session = Mock()
+    session.local_mode = False
+
+    local_session_class.return_value = local_session
+    session_class.return_value = session
+
+    e = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, 'local')
+    print(e.sagemaker_session.local_mode)
+    assert e.sagemaker_session.local_mode is True
 
-    e2 = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, 'local_gpu', output_path='s3://bucket/prefix',
-                   sagemaker_session=sagemaker_session)
-    assert e2.local_mode is True
+    e2 = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, 'local_gpu')
+    assert e2.sagemaker_session.local_mode is True
 
-    e3 = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path='s3://bucket/prefix',
-                   sagemaker_session=sagemaker_session)
-    assert e3.local_mode is False
+    e3 = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE)
+    assert e3.sagemaker_session.local_mode is False
 
 
 @patch('sagemaker.estimator.LocalSession')

tests/unit/test_fm.py

@@ -39,7 +39,8 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name='boto_session', region_name=REGION)
-    sms = Mock(name='sagemaker_session', boto_session=boto_mock)
+    sms = Mock(name='sagemaker_session', boto_session=boto_mock,
+               region_name=REGION, config=None, local_mode=False)
     sms.boto_region_name = REGION
     sms.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     sms.sagemaker_client.describe_training_job = Mock(name='describe_training_job',

tests/unit/test_fw_utils.py

@@ -90,23 +90,19 @@ def test_tar_and_upload_dir_s3(sagemaker_session):
     assert result == UploadedCode('s3://m', 'mnist.py')
 
 
-def test_tar_and_upload_dir_does_not_exits(sagemaker_session):
-    bucket = 'mybucker'
-    s3_key_prefix = 'something/source'
+def test_validate_source_dir_does_not_exits(sagemaker_session):
     script = 'mnist.py'
     directory = ' !@#$%^&*()path probably in not there.!@#$%^&*()'
     with pytest.raises(ValueError) as error:
-        tar_and_upload_dir(sagemaker_session, bucket, s3_key_prefix, script, directory)
+        validate_source_dir(script, directory)
     assert 'does not exist' in str(error)
 
 
-def test_tar_and_upload_dir_is_not_directory(sagemaker_session):
-    bucket = 'mybucker'
-    s3_key_prefix = 'something/source'
+def test_validate_source_dir_is_not_directory(sagemaker_session):
     script = 'mnist.py'
     directory = inspect.getfile(inspect.currentframe())
     with pytest.raises(ValueError) as error:
-        tar_and_upload_dir(sagemaker_session, bucket, s3_key_prefix, script, directory)
+        validate_source_dir(script, directory)
     assert 'is not a directory' in str(error)
 
 

tests/unit/test_image.py

@@ -226,7 +226,12 @@ def test_serve(up, copy, copytree, tmpdir, sagemaker_session):
 
         image = 'my-image'
         sagemaker_container = _SageMakerContainer('local', 1, image, sagemaker_session=sagemaker_session)
-        primary_container = {'ModelDataUrl': '/some/model/path', 'Environment': {'env1': 1, 'env2': 'b'}}
+        primary_container = {'ModelDataUrl': '/some/model/path',
+                             'Environment': {'env1': 1,
+                                             'env2': 'b',
+                                             'SAGEMAKER_SUBMIT_DIRECTORY': 's3://some/path'
+                                             }
+                             }
 
         sagemaker_container.serve(primary_container)
         docker_compose_file = os.path.join(sagemaker_container.container_root, 'docker-compose.yaml')

tests/unit/test_kmeans.py

@@ -38,7 +38,8 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name='boto_session', region_name=REGION)
-    sms = Mock(name='sagemaker_session', boto_session=boto_mock)
+    sms = Mock(name='sagemaker_session', boto_session=boto_mock,
+               region_name=REGION, config=None, local_mode=False)
     sms.boto_region_name = REGION
     sms.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     sms.sagemaker_client.describe_training_job = Mock(name='describe_training_job',

tests/unit/test_lda.py

@@ -37,7 +37,8 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name='boto_session', region_name=REGION)
-    sms = Mock(name='sagemaker_session', boto_session=boto_mock)
+    sms = Mock(name='sagemaker_session', boto_session=boto_mock,
+               region_name=REGION, config=None, local_mode=False)
     sms.boto_region_name = REGION
     sms.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     sms.sagemaker_client.describe_training_job = Mock(name='describe_training_job',