Add Local Mode support

CHANGELOG.rst

@@ -5,6 +5,7 @@ CHANGELOG
 1.1.dev4
 ========
 * feature: Frameworks: Use more idiomatic ECR repository naming scheme
+* feature: Add Support for Local Mode
 
 1.1.3
 ========

setup.py

@@ -32,7 +32,8 @@ def read(fname):
       ],
 
       # Declare minimal set for installation
-      install_requires=['boto3>=1.4.8', 'numpy>=1.9.0', 'protobuf>=3.1', 'scipy>=1.0.0'],
+      install_requires=['boto3>=1.4.8', 'numpy>=1.9.0', 'protobuf>=3.1', 'scipy>=1.0.0', 'urllib3>=1.2',
+                        'PyYAML>=3.2'],
 
       extras_require={
           'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'pytest-xdist',

src/sagemaker/__init__.py

@@ -23,6 +23,7 @@
 
 from sagemaker.model import Model
 from sagemaker.predictor import RealTimePredictor
+from sagemaker.local_session import LocalSession
 from sagemaker.session import Session
 from sagemaker.session import container_def
 from sagemaker.session import production_variant
@@ -34,5 +35,5 @@
            LinearLearnerModel, LinearLearnerPredictor,
            LDA, LDAModel, LDAPredictor,
            FactorizationMachines, FactorizationMachinesModel, FactorizationMachinesPredictor,
-           Model, NTM, NTMModel, NTMPredictor, RealTimePredictor, Session,
+           Model, NTM, NTMModel, NTMPredictor, RealTimePredictor, LocalSession, Session,
            container_def, s3_input, production_variant, get_execution_role]

src/sagemaker/estimator.py

@@ -21,6 +21,7 @@
 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.local_session import LocalSession
 from sagemaker.model import Model
 from sagemaker.model import (SCRIPT_PARAM_NAME, DIR_PARAM_NAME, CLOUDWATCH_METRICS_PARAM_NAME,
                              CONTAINER_LOG_LEVEL_PARAM_NAME, JOB_NAME_PARAM_NAME, SAGEMAKER_REGION_PARAM_NAME)
@@ -78,7 +79,17 @@ def __init__(self, role, train_instance_count, train_instance_type,
         self.train_volume_size = train_volume_size
         self.train_max_run = train_max_run
         self.input_mode = input_mode
-        self.sagemaker_session = sagemaker_session or Session()
+
+        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 = LocalSession()
+        else:
+            self.local_mode = False
+            self.sagemaker_session = sagemaker_session or Session()
+
         self.base_job_name = base_job_name
         self._current_job_name = None
         self.output_path = output_path
@@ -303,7 +314,7 @@ def start_new(cls, estimator, inputs):
         """Create a new Amazon SageMaker training job from the estimator.
 
         Args:
-            estimator (sagemaker.estimator.Framework): Estimator object created by the user.
+            estimator (sagemaker.estimator.EstimatorBase): Estimator object created by the user.
             inputs (str): Parameters used when called  :meth:`~sagemaker.estimator.EstimatorBase.fit`.
 
         Returns:

src/sagemaker/fw_utils.py

@@ -45,19 +45,27 @@ def create_image_uri(region, framework, instance_type, framework_version, py_ver
         str: The appropriate image URI based on the given parameters.
     """
 
-    if not instance_type.startswith('ml.'):
+    # Handle Local Mode
+    if instance_type.startswith('local'):
+        if instance_type == 'local':
+            device_type = 'cpu'
+        elif instance_type == 'local_gpu':
+            device_type = 'gpu'
+
+    elif not instance_type.startswith('ml.'):
         raise ValueError('{} is not a valid SageMaker instance type. See: '
                          'https://aws.amazon.com/sagemaker/pricing/instance-types/'.format(instance_type))
-    family = instance_type.split('.')[1]
-
-    # For some frameworks, we have optimized images for specific families, e.g c5 or p3. In those cases,
-    # we use the family name in the image tag. In other cases, we use 'cpu' or 'gpu'.
-    if family in optimized_families:
-        device_type = family
-    elif family[0] in ['g', 'p']:
-        device_type = 'gpu'
     else:
-        device_type = 'cpu'
+        family = instance_type.split('.')[1]
+
+        # For some frameworks, we have optimized images for specific families, e.g c5 or p3. In those cases,
+        # we use the family name in the image tag. In other cases, we use 'cpu' or 'gpu'.
+        if family in optimized_families:
+            device_type = family
+        elif family[0] in ['g', 'p']:
+            device_type = 'gpu'
+        else:
+            device_type = 'cpu'
 
     tag = "{}-{}-{}".format(framework_version, device_type, py_version)
     return "{}.dkr.ecr.{}.amazonaws.com/sagemaker-{}:{}" \