Allow Local Mode to work without a boto session.

This change works when ~/.sagemaker/config.yaml has

local:
   no_internet: True

It depends on the container image supporting a local training script
instead of an s3 location.

Author: Ignacio Quintero

src/sagemaker/estimator.py

@@ -14,13 +14,12 @@
 
 import json
 import logging
+import os
 from abc import ABCMeta
 from abc import abstractmethod
 from six import with_metaclass, string_types
 
-from sagemaker.fw_utils import tar_and_upload_dir
-from sagemaker.fw_utils import parse_s3_url
-from sagemaker.fw_utils import UploadedCode
+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.model import Model
@@ -30,7 +29,7 @@
 from sagemaker.predictor import RealTimePredictor
 from sagemaker.session import Session
 from sagemaker.session import s3_input
-from sagemaker.utils import base_name_from_image, name_from_base
+from sagemaker.utils import base_name_from_image, name_from_base, get_config_value
 
 
 class EstimatorBase(with_metaclass(ABCMeta, object)):
@@ -158,9 +157,14 @@ 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 output_path was specified we use it otherwise initialize here
+        # 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:
-            self.output_path = 's3://{}/'.format(self.sagemaker_session.default_bucket())
+            if self.local_mode and get_config_value('local.no_internet',
+                                                    self.sagemaker_session.config):
+                self.output_path = ''
+            else:
+                self.output_path = 's3://{}/'.format(self.sagemaker_session.default_bucket())
 
         self.latest_training_job = _TrainingJob.start_new(self, inputs)
         if wait:
@@ -604,27 +608,54 @@ 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 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):
+            # 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 self.local_mode and get_config_value('local.no_internet', self.sagemaker_session.config):
+            code_dir = 'file://' + self.source_dir
+            script = self.entry_point
+        else:
+            self.uploaded_code = self._stage_user_code_in_s3()
+            code_dir = self.uploaded_code.s3_prefix
+            script = self.uploaded_code.script_name
+
+        # Modify hyperparameters in-place to point to the right code directory and script URIs
+        self._hyperparameters[DIR_PARAM_NAME] = code_dir
+        self._hyperparameters[SCRIPT_PARAM_NAME] = script
+        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
+        super(Framework, self).fit(inputs, wait, logs, self._current_job_name)
+
+    def _stage_user_code_in_s3(self):
+        """ Upload the user training script to s3 and return the location.
+
+        Returns: s3 uri
+
+        """
         if self.code_location is None:
             code_bucket = self.sagemaker_session.default_bucket()
             code_s3_prefix = '{}/source'.format(self._current_job_name)
         else:
             code_bucket, key_prefix = parse_s3_url(self.code_location)
             code_s3_prefix = '{}/{}/source'.format(key_prefix, self._current_job_name)
 
-        self.uploaded_code = tar_and_upload_dir(session=self.sagemaker_session.boto_session,
-                                                bucket=code_bucket,
-                                                s3_key_prefix=code_s3_prefix,
-                                                script=self.entry_point,
-                                                directory=self.source_dir)
-
-        # Modify hyperparameters in-place to add the URLs to the uploaded code.
-        self._hyperparameters[DIR_PARAM_NAME] = self.uploaded_code.s3_prefix
-        self._hyperparameters[SCRIPT_PARAM_NAME] = self.uploaded_code.script_name
-        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.boto_session.region_name
-        super(Framework, self).fit(inputs, wait, logs, self._current_job_name)
+        return tar_and_upload_dir(session=self.sagemaker_session.boto_session,
+                                  bucket=code_bucket,
+                                  s3_key_prefix=code_s3_prefix,
+                                  script=self.entry_point,
+                                  directory=self.source_dir)
 
     def hyperparameters(self):
         """Return the hyperparameters as a dictionary to use for training.

src/sagemaker/fw_utils.py

@@ -68,6 +68,27 @@ def create_image_uri(region, framework, instance_type, framework_version, py_ver
         .format(account, region, framework, tag)
 
 
+def validate_source_dir(script, directory):
+    """Validate that the source directory exists and it contains the user script
+
+    Args:
+        script (str):  Script filename.
+        directory (str): Directory containing the source file.
+
+    Raises:
+        ValueError: If ``directory`` does not exist, is not a directory, or does not contain ``script``.
+    """
+    if directory:
+        if not os.path.exists(directory):
+            raise ValueError('"{}" does not exist.'.format(directory))
+        if not os.path.isdir(directory):
+            raise ValueError('"{}" is not a directory.'.format(directory))
+        if script not in os.listdir(directory):
+            raise ValueError('No file named "{}" was found in directory "{}".'.format(script, directory))
+
+    return True
+
+
 def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory):
     """Pack and upload source files to S3 only if directory is empty or local.
 
@@ -83,21 +104,13 @@ def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory):
 
     Returns:
         sagemaker.fw_utils.UserCode: An object with the S3 bucket and key (S3 prefix) and script name.
-
-    Raises:
-        ValueError: If ``directory`` does not exist, is not a directory, or does not contain ``script``.
     """
     if directory:
         if directory.lower().startswith("s3://"):
             return UploadedCode(s3_prefix=directory, script_name=os.path.basename(script))
-        if not os.path.exists(directory):
-            raise ValueError('"{}" does not exist.'.format(directory))
-        if not os.path.isdir(directory):
-            raise ValueError('"{}" is not a directory.'.format(directory))
-        if script not in os.listdir(directory):
-            raise ValueError('No file named "{}" was found in directory "{}".'.format(script, directory))
-        script_name = script
-        source_files = [os.path.join(directory, name) for name in os.listdir(directory)]
+        else:
+            script_name = script
+            source_files = [os.path.join(directory, name) for name in os.listdir(directory)]
     else:
         # If no directory is specified, the script parameter needs to be a valid relative path.
         os.path.exists(script)

src/sagemaker/local/image.py

@@ -29,6 +29,9 @@
 
 import yaml
 
+import sagemaker
+from sagemaker.utils import get_config_value
+
 CONTAINER_PREFIX = "algo"
 DOCKER_COMPOSE_FILENAME = 'docker-compose.yaml'
 
@@ -70,9 +73,7 @@ def __init__(self, instance_type, instance_count, image, sagemaker_session=None)
         self.container = None
         # set the local config. This is optional and will use reasonable defaults
         # if not present.
-        self.local_config = None
-        if self.sagemaker_session.config and 'local' in self.sagemaker_session.config:
-            self.local_config = self.sagemaker_session.config['local']
+        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.
@@ -116,6 +117,13 @@ def train(self, input_data_config, hyperparameters):
             else:
                 raise ValueError('Unknown URI scheme {}'.format(parsed_uri.scheme))
 
+        # If the training script directory is a local directory, mount it to the container.
+        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'))
+
         # Create the configuration files for each container that we will create
         # Each container will map the additional local volumes (if any).
         for host in self.hosts:
@@ -366,8 +374,9 @@ def _create_docker_host(self, host, environment, optml_subdirs, command, volumes
             }
         }
 
-        serving_port = 8080 if self.local_config is None else self.local_config.get('serving_port', 8080)
         if command == 'serve':
+            serving_port = get_config_value('local.serving_port',
+                                            self.sagemaker_session.config) or 8080
             host_config.update({
                 'ports': [
                     '%s:8080' % serving_port
@@ -377,9 +386,9 @@ def _create_docker_host(self, host, environment, optml_subdirs, command, volumes
         return host_config
 
     def _create_tmp_folder(self):
-        root_dir = None
-        if self.local_config and 'container_root' in self.local_config:
-            root_dir = os.path.abspath(self.local_config['container_root'])
+        root_dir = get_config_value('local.container_root', self.sagemaker_session.config)
+        if root_dir:
+            root_dir = os.path.abspath(root_dir)
 
         dir = tempfile.mkdtemp(dir=root_dir)
 

src/sagemaker/local/local_session.py

@@ -15,11 +15,13 @@
 import platform
 import time
 
+import boto3
 import urllib3
 from botocore.exceptions import ClientError
 
 from sagemaker.local.image import _SageMakerContainer
 from sagemaker.session import Session
+from sagemaker.utils import get_config_value
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.WARNING)
@@ -115,9 +117,7 @@ def create_endpoint(self, EndpointName, EndpointConfigName):
 
         i = 0
         http = urllib3.PoolManager()
-        serving_port = 8080
-        if self.sagemaker_session.config and 'local' in self.sagemaker_session.config:
-            serving_port = self.sagemaker_session.config['local'].get('serving_port', 8080)
+        serving_port = get_config_value('local.serving_port', self.sagemaker_session.config) or 8080
         endpoint_url = "http://localhost:%s/ping" % serving_port
         while True:
             i += 1
@@ -153,8 +153,8 @@ def __init__(self, config=None):
         """
         self.http = urllib3.PoolManager()
         self.serving_port = 8080
-        if config and 'local' in config:
-            self.serving_port = config['local'].get('serving_port', 8080)
+        self.config = config
+        self.serving_port = get_config_value('local.serving_port', config) or 8080
 
     def invoke_endpoint(self, Body, EndpointName, ContentType, Accept):
         url = "http://localhost:%s/invocations" % self.serving_port
@@ -171,6 +171,23 @@ def __init__(self, boto_session=None):
 
         if platform.system() == 'Windows':
             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
+
+            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/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.boto_session.region_name, self.__framework_name__,
+        return create_image_uri(self.sagemaker_session.region_name, self.__framework_name__,
                                 self.train_instance_type, framework_version=self.framework_version,
                                 py_version=self.py_version)
 

src/sagemaker/session.py

@@ -70,10 +70,25 @@ def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_c
                 If not provided, one will be created using this instance's ``boto_session``.
         """
         self._default_bucket = None
+
+        sagemaker_config_file = os.path.join(os.path.expanduser('~'), '.sagemaker', 'config.yaml')
+        if os.path.exists(sagemaker_config_file):
+            self.config = yaml.load(open(sagemaker_config_file, 'r'))
+        else:
+            self.config = None
+
+        self._initialize(boto_session, sagemaker_client, sagemaker_runtime_client)
+
+    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
+        """Initialize this SageMaker Session.
+
+        Creates or uses a boto_session, sagemaker_client and sagemaker_runtime_client.
+        Sets the region_name.
+        """
         self.boto_session = boto_session or boto3.Session()
 
-        region = self.boto_session.region_name
-        if region is None:
+        self.region = self.boto_session.region_name
+        if self.region 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')
@@ -82,12 +97,6 @@ def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_c
         self.sagemaker_runtime_client = sagemaker_runtime_client or self.boto_session.client('runtime.sagemaker')
         prepend_user_agent(self.sagemaker_runtime_client)
 
-        sagemaker_config_file = os.path.join(os.path.expanduser('~'), '.sagemaker', 'config.yaml')
-        if os.path.exists(sagemaker_config_file):
-            self.config = yaml.load(open(sagemaker_config_file, 'r'))
-        else:
-            self.config = None
-
     @property
     def boto_region_name(self):
         return self.boto_session.region_name

src/sagemaker/tensorflow/estimator.py

@@ -278,7 +278,7 @@ def train_image(self):
         Returns:
             str: The URI of the Docker image.
         """
-        return create_image_uri(self.sagemaker_session.boto_session.region_name, self.__framework_name__,
+        return create_image_uri(self.sagemaker_session.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):

src/sagemaker/utils.py

@@ -76,3 +76,17 @@ def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
 
     return wrapper
+
+
+def get_config_value(key_path, config):
+    if config is None:
+        return None
+
+    current_section = config
+    for key in key_path.split('.'):
+        if key in current_section:
+            current_section = current_section[key]
+        else:
+            return None
+
+    return current_section

tests/unit/test_estimator.py

@@ -89,7 +89,7 @@ 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)
+    ims = Mock(name='sagemaker_session', boto_session=boto_mock, region_name=REGION)
     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)

tests/unit/test_fw_utils.py

@@ -14,7 +14,7 @@
 from mock import Mock
 import os
 from sagemaker.fw_utils import create_image_uri, framework_name_from_image, framework_version_from_tag
-from sagemaker.fw_utils import tar_and_upload_dir, parse_s3_url, UploadedCode
+from sagemaker.fw_utils import tar_and_upload_dir, parse_s3_url, UploadedCode, validate_source_dir
 import pytest
 
 
@@ -110,13 +110,11 @@ def test_tar_and_upload_dir_is_not_directory(sagemaker_session):
     assert 'is not a directory' in str(error)
 
 
-def test_tar_and_upload_dir_file_not_in_dir(sagemaker_session):
-    bucket = 'mybucker'
-    s3_key_prefix = 'something/source'
+def test_validate_source_dir_file_not_in_dir():
     script = ' !@#$%^&*() .myscript. !@#$%^&*() '
     directory = '.'
     with pytest.raises(ValueError) as error:
-        tar_and_upload_dir(sagemaker_session, bucket, s3_key_prefix, script, directory)
+        validate_source_dir(script, directory)
     assert 'No file named' in str(error)