Address Owen's feedback.

In general boto_session is not None anymore. no_internet renamed to
local_code. Added more unit tests to bring coverage up.

Author: Ignacio Quintero

src/sagemaker/estimator.py

@@ -20,7 +20,7 @@
 from six import with_metaclass, string_types
 
 from sagemaker.fw_utils import tar_and_upload_dir, parse_s3_url, UploadedCode, validate_source_dir
-from sagemaker.local.local_session import LocalSession, file_input
+from sagemaker.local import LocalSession, file_input
 
 from sagemaker.model import Model
 from sagemaker.model import (SCRIPT_PARAM_NAME, DIR_PARAM_NAME, CLOUDWATCH_METRICS_PARAM_NAME,
@@ -155,10 +155,10 @@ def fit(self, inputs, wait=True, logs=True, job_name=None):
             self._current_job_name = name_from_base(base_name)
 
         # 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
+        # For Local Mode with local_code=True we don't need an explicit output_path
         if self.output_path is None:
-            no_internet = get_config_value('local.no_internet', self.sagemaker_session.config)
-            if self.sagemaker_session.local_mode and no_internet:
+            local_code = get_config_value('local.local_code', self.sagemaker_session.config)
+            if self.sagemaker_session.local_mode and local_code:
                 self.output_path = ''
             else:
                 self.output_path = 's3://{}/'.format(self.sagemaker_session.default_bucket())
@@ -610,10 +610,10 @@ def fit(self, inputs, wait=True, logs=True, job_name=None):
         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
+        # if we are in local mode with local_code=True. We want the container to just
         # mount the source dir instead of uploading to S3.
-        no_internet = get_config_value('local.no_internet', self.sagemaker_session.config)
-        if self.sagemaker_session.local_mode and no_internet:
+        local_code = get_config_value('local.local_code', self.sagemaker_session.config)
+        if self.sagemaker_session.local_mode and local_code:
             # 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)
@@ -632,7 +632,7 @@ def fit(self, inputs, wait=True, logs=True, job_name=None):
         self._hyperparameters[CLOUDWATCH_METRICS_PARAM_NAME] = self.enable_cloudwatch_metrics
         self._hyperparameters[CONTAINER_LOG_LEVEL_PARAM_NAME] = self.container_log_level
         self._hyperparameters[JOB_NAME_PARAM_NAME] = self._current_job_name
-        self._hyperparameters[SAGEMAKER_REGION_PARAM_NAME] = self.sagemaker_session.region_name
+        self._hyperparameters[SAGEMAKER_REGION_PARAM_NAME] = self.sagemaker_session.boto_region_name
         super(Framework, self).fit(inputs, wait, logs, self._current_job_name)
 
     def _stage_user_code_in_s3(self):

src/sagemaker/local/__init__.py

@@ -10,3 +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.
+from .local_session import (file_input, LocalSession, LocalSagemakerRuntimeClient,
+                            LocalSagemakerClient)
+
+__all__ = [file_input, LocalSession, LocalSagemakerClient, LocalSagemakerRuntimeClient]

src/sagemaker/local/image.py

@@ -71,9 +71,6 @@ def __init__(self, instance_type, instance_count, image, sagemaker_session=None)
         self.hosts = ['{}-{}-{}'.format(CONTAINER_PREFIX, i, suffix) for i in range(1, self.instance_count + 1)]
         self.container_root = None
         self.container = None
-        # set the local config. This is optional and will use reasonable defaults
-        # if not present.
-        self.local_config = get_config_value('local', self.sagemaker_session.config)
 
     def train(self, input_data_config, hyperparameters):
         """Run a training job locally using docker-compose.
@@ -86,7 +83,10 @@ 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'))
+        # A shared directory for all the containers. It is only mounted if the training script is
+        # Local.
+        shared_dir = os.path.join(self.container_root, 'shared')
+        os.mkdir(shared_dir)
 
         data_dir = self._create_tmp_folder()
         volumes = []
@@ -123,7 +123,8 @@ def train(self, input_data_config, hyperparameters):
         parsed_uri = urlparse(training_dir)
         if parsed_uri.scheme == 'file':
             volumes.append(_Volume(parsed_uri.path, '/opt/ml/code'))
-            volumes.append(_Volume(os.path.join(self.container_root, 'shared'), '/opt/ml/shared'))
+            # Also mount a directory that all the containers can access.
+            volumes.append(_Volume(shared_dir, '/opt/ml/shared'))
 
         # Create the configuration files for each container that we will create
         # Each container will map the additional local volumes (if any).
@@ -144,6 +145,7 @@ def train(self, input_data_config, hyperparameters):
         # lots of data downloaded from S3. This doesn't delete any local
         # data that was just mounted to the container.
         _delete_tree(data_dir)
+        _delete_tree(shared_dir)
         # Also free the container config files.
         for host in self.hosts:
             container_config_path = os.path.join(self.container_root, host)

src/sagemaker/local/local_session.py

@@ -173,20 +173,13 @@ def __init__(self, boto_session=None):
             logger.warning("Windows Support for Local Mode is Experimental")
 
     def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
-        """Initialize a boto session for this Local SageMaker Session."""
-        if get_config_value('local.no_internet', self.config):
-            # if no_internet is set to True in the config file then we won't create a boto_session
-            # this will make any component that defaults to using S3 utilize a local file instead.
-            self.boto_session = None
-            self.region_name = get_config_value('local.region_name', self.config)
-            if self.region_name is None:
-                raise ValueError('Must setup region_name in the sagemaker config file. See <Link to Readme Here>')
-        else:
-            self.boto_session = boto_session or boto3.Session()
-            self.region_name = self.boto_session.region_name
+        """Initialize this Local SageMaker Session."""
+
+        self.boto_session = boto_session or boto3.Session()
+        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.')
+        if self._region_name is None:
+            raise ValueError('Must setup local AWS configuration with a region supported by SageMaker.')
 
         self.sagemaker_client = LocalSagemakerClient(self)
         self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self.config)

src/sagemaker/model.py

@@ -160,21 +160,21 @@ def prepare_container_def(self, instance_type):
         Returns:
             dict[str, str]: A container definition object usable with the CreateModel API.
         """
-        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))
+        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)
 
     def _upload_code(self, key_prefix):
-        self.uploaded_code = tar_and_upload_dir(session=self.sagemaker_session.boto_session,
-                                                bucket=self.bucket or self.sagemaker_session.default_bucket(),
-                                                s3_key_prefix=key_prefix,
-                                                script=self.entry_point,
-                                                directory=self.source_dir)
+        local_code = get_config_value('local.local_code', self.sagemaker_session.config)
+        if self.sagemaker_session.local_mode and local_code:
+            self.uploaded_code = None
+        else:
+            self.uploaded_code = tar_and_upload_dir(session=self.sagemaker_session.boto_session,
+                                                    bucket=self.bucket or self.sagemaker_session.default_bucket(),
+                                                    s3_key_prefix=key_prefix,
+                                                    script=self.entry_point,
+                                                    directory=self.source_dir)
 
     def _framework_env_vars(self):
         if self.uploaded_code:
@@ -189,5 +189,5 @@ def _framework_env_vars(self):
             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
+            SAGEMAKER_REGION_PARAM_NAME.upper(): self.sagemaker_session.boto_region_name
         }

src/sagemaker/mxnet/estimator.py

@@ -65,7 +65,7 @@ def train_image(self):
         Returns:
             str: The URI of the Docker image.
         """
-        return create_image_uri(self.sagemaker_session.region_name, self.__framework_name__,
+        return create_image_uri(self.sagemaker_session.boto_region_name, self.__framework_name__,
                                 self.train_instance_type, framework_version=self.framework_version,
                                 py_version=self.py_version)
 

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_name = self.boto_session.region_name
-        if self.region_name 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')
@@ -101,7 +101,7 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
 
     @property
     def boto_region_name(self):
-        return self.boto_session.region_name
+        return self._region_name
 
     def upload_data(self, path, bucket=None, key_prefix='data'):
         """Upload local file or directory to S3.

src/sagemaker/tensorflow/estimator.py

@@ -279,7 +279,7 @@ def train_image(self):
         Returns:
             str: The URI of the Docker image.
         """
-        return create_image_uri(self.sagemaker_session.region_name, self.__framework_name__,
+        return create_image_uri(self.sagemaker_session.boto_region_name, self.__framework_name__,
                                 self.train_instance_type, self.framework_version, py_version=self.py_version)
 
     def create_model(self, model_server_workers=None):
@@ -306,8 +306,8 @@ def hyperparameters(self):
         hyperparameters = super(TensorFlow, self).hyperparameters()
 
         if not self.checkpoint_path:
-            no_internet = get_config_value('local.no_internet', self.sagemaker_session.config)
-            if self.sagemaker_session.local_mode and no_internet:
+            local_code = get_config_value('local.local_code', self.sagemaker_session.config)
+            if self.sagemaker_session.local_mode and local_code:
                 self.checkpoint_path = '/opt/ml/shared/checkpoints'
             else:
                 self.checkpoint_path = os.path.join(self.output_path,

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, get_config_value
+from sagemaker.utils import name_from_image
 
 
 class TensorFlowPredictor(RealTimePredictor):
@@ -83,18 +83,13 @@ def prepare_container_def(self, instance_type):
         """
         deploy_image = self.image
         if not deploy_image:
-            region_name = self.sagemaker_session.region_name
+            region_name = self.sagemaker_session.boto_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)
 
-        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)
+        self._upload_code(deploy_key_prefix)
         deploy_env = dict(self.env)
         deploy_env.update(self._framework_env_vars())
 

tests/unit/test_estimator.py

@@ -17,6 +17,7 @@
 from mock import Mock, patch
 
 from sagemaker.estimator import Estimator, Framework, _TrainingJob
+from sagemaker.local import LocalSession
 from sagemaker.session import s3_input
 from sagemaker.model import FrameworkModel
 from sagemaker.predictor import RealTimePredictor
@@ -89,12 +90,12 @@ 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,
-               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',
+    sms = Mock(name='sagemaker_session', boto_session=boto_mock,
+               boto_region_name=REGION, config=None, local_mode=False)
+    sms.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
+    sms.sagemaker_client.describe_training_job = Mock(name='describe_training_job',
                                                       return_value=DESCRIBE_TRAINING_JOB_RESULT)
-    return ims
+    return sms
 
 
 def test_sagemaker_s3_uri_invalid(sagemaker_session):
@@ -142,6 +143,24 @@ def test_invalid_custom_code_bucket(sagemaker_session):
 }
 
 
+def test_local_code_location():
+    config = {
+        'local': {
+            'local_code': True,
+            'region': 'us-west-2'
+        }
+    }
+    sms = Mock(name='sagemaker_session', boto_session=None,
+               boto_region_name=REGION, config=config, local_mode=True)
+    t = DummyFramework(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sms,
+                       train_instance_count=1, train_instance_type='local',
+                       base_job_name=IMAGE_NAME, hyperparameters={123: [456], 'learning_rate': 0.1})
+
+    t.fit('file:///data/file')
+    assert t.source_dir == DATA_DIR
+    assert t.entry_point == 'dummy_script.py'
+
+
 @patch('time.strftime', return_value=TIMESTAMP)
 def test_start_new_convert_hyperparameters_to_str(strftime, sagemaker_session):
     uri = 'bucket/mydata'

tests/unit/test_image.py

@@ -49,6 +49,9 @@
                    'b': 'bee',
                    'sagemaker_submit_directory': json.dumps('s3://my_bucket/code')}
 
+LOCAL_CODE_HYPERPARAMETERS = {'a': 1,
+                              'b': 2,
+                              'sagemaker_submit_directory': json.dumps('file:///tmp/code')}
 
 @pytest.fixture()
 def sagemaker_session():
@@ -216,6 +219,37 @@ def test_train(_download_folder, _cleanup, _execute_and_stream_output, LocalSess
                 assert config['services'][h]['command'] == 'train'
 
 
+@patch('sagemaker.local.local_session.LocalSession')
+@patch('sagemaker.local.image._execute_and_stream_output')
+@patch('sagemaker.local.image._SageMakerContainer._cleanup')
+@patch('sagemaker.local.image._SageMakerContainer._download_folder')
+def test_train_local_code(_download_folder, _cleanup, _execute_and_stream_output,
+                          _local_session, tmpdir, sagemaker_session):
+    directories = [str(tmpdir.mkdir('container-root')), str(tmpdir.mkdir('data'))]
+    with patch('sagemaker.local.image._SageMakerContainer._create_tmp_folder',
+               side_effect=directories):
+        instance_count = 2
+        image = 'my-image'
+        sagemaker_container = _SageMakerContainer('local', instance_count, image,
+                                                  sagemaker_session=sagemaker_session)
+
+        sagemaker_container.train(INPUT_DATA_CONFIG, LOCAL_CODE_HYPERPARAMETERS)
+
+        docker_compose_file = os.path.join(sagemaker_container.container_root,
+                                           'docker-compose.yaml')
+        shared_folder_path = os.path.join(sagemaker_container.container_root, 'shared')
+
+        with open(docker_compose_file, 'r') as f:
+            config = yaml.load(f)
+            assert len(config['services']) == instance_count
+            for h in sagemaker_container.hosts:
+                assert config['services'][h]['image'] == image
+                assert config['services'][h]['command'] == 'train'
+                volumes = config['services'][h]['volumes']
+                assert '%s:/opt/ml/code' % '/tmp/code' in volumes
+                assert '%s:/opt/ml/shared' % shared_folder_path in volumes
+
+
 @patch('sagemaker.local.image._HostingContainer.up')
 @patch('shutil.copy')
 @patch('shutil.copytree')
@@ -244,6 +278,38 @@ def test_serve(up, copy, copytree, tmpdir, sagemaker_session):
                 assert config['services'][h]['command'] == 'serve'
 
 
+@patch('sagemaker.local.image._HostingContainer.up')
+@patch('shutil.copy')
+@patch('shutil.copytree')
+def test_serve_local_code(up, copy, copytree, tmpdir, sagemaker_session):
+
+    with patch('sagemaker.local.image._SageMakerContainer._create_tmp_folder',
+               return_value=str(tmpdir.mkdir('container-root'))):
+
+        image = 'my-image'
+        sagemaker_container = _SageMakerContainer('local', 1, image, sagemaker_session=sagemaker_session)
+        primary_container = {'ModelDataUrl': '/some/model/path',
+                             'Environment': {'env1': 1,
+                                             'env2': 'b',
+                                             'SAGEMAKER_SUBMIT_DIRECTORY': 'file:///tmp/code'
+                                             }
+                             }
+
+        sagemaker_container.serve(primary_container)
+        docker_compose_file = os.path.join(sagemaker_container.container_root,
+                                           'docker-compose.yaml')
+
+        with open(docker_compose_file, 'r') as f:
+            config = yaml.load(f)
+
+            for h in sagemaker_container.hosts:
+                assert config['services'][h]['image'] == image
+                assert config['services'][h]['command'] == 'serve'
+
+                volumes = config['services'][h]['volumes']
+                assert '%s:/opt/ml/code' % '/tmp/code' in volumes
+
+
 @patch('os.makedirs')
 def test_download_folder(makedirs):
     boto_mock = Mock(name='boto_session')

tests/unit/test_lda.py

@@ -37,8 +37,7 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name='boto_session', region_name=REGION)
-    sms = Mock(name='sagemaker_session', boto_session=boto_mock,
-               region_name=REGION, config=None, local_mode=False)
+    sms = Mock(name='sagemaker_session', boto_session=boto_mock, 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',