Support optional input channels in local mode.

Author: nadiaya

CHANGELOG.rst

@@ -2,6 +2,11 @@
 CHANGELOG
 =========
 
+1.14.1.dev
+==========
+
+* enhancement: Local Mode: support optional input channels
+
 1.14.0
 ======
 

src/sagemaker/local/entities.py

@@ -47,22 +47,23 @@ def __init__(self, container):
         self.end_time = None
 
     def start(self, input_data_config, output_data_config, hyperparameters, job_name):
-        for channel in input_data_config:
-            if channel['DataSource'] and 'S3DataSource' in channel['DataSource']:
-                data_distribution = channel['DataSource']['S3DataSource']['S3DataDistributionType']
-                data_uri = channel['DataSource']['S3DataSource']['S3Uri']
-            elif channel['DataSource'] and 'FileDataSource' in channel['DataSource']:
-                data_distribution = channel['DataSource']['FileDataSource']['FileDataDistributionType']
-                data_uri = channel['DataSource']['FileDataSource']['FileUri']
-            else:
-                raise ValueError('Need channel[\'DataSource\'] to have [\'S3DataSource\'] or [\'FileDataSource\']')
-
-            # use a single Data URI - this makes handling S3 and File Data easier down the stack
-            channel['DataUri'] = data_uri
-
-            if data_distribution != 'FullyReplicated':
-                raise RuntimeError('DataDistribution: %s is not currently supported in Local Mode' %
-                                   data_distribution)
+        if input_data_config:
+            for channel in input_data_config:
+                if channel['DataSource'] and 'S3DataSource' in channel['DataSource']:
+                    data_distribution = channel['DataSource']['S3DataSource']['S3DataDistributionType']
+                    data_uri = channel['DataSource']['S3DataSource']['S3Uri']
+                elif channel['DataSource'] and 'FileDataSource' in channel['DataSource']:
+                    data_distribution = channel['DataSource']['FileDataSource']['FileDataDistributionType']
+                    data_uri = channel['DataSource']['FileDataSource']['FileUri']
+                else:
+                    raise ValueError('Need channel[\'DataSource\'] to have [\'S3DataSource\'] or [\'FileDataSource\']')
+
+                # use a single Data URI - this makes handling S3 and File Data easier down the stack
+                channel['DataUri'] = data_uri
+
+                if data_distribution != 'FullyReplicated':
+                    raise RuntimeError('DataDistribution: %s is not currently supported in Local Mode' %
+                                       data_distribution)
 
         self.start = datetime.datetime.now()
         self.state = self._TRAINING

src/sagemaker/local/image.py

@@ -262,15 +262,17 @@ def write_config_files(self, host, hyperparameters, input_data_config):
             'hosts': self.hosts
         }
 
-        print(input_data_config)
         json_input_data_config = {}
-        for c in input_data_config:
-            channel_name = c['ChannelName']
-            json_input_data_config[channel_name] = {
-                'TrainingInputMode': 'File'
-            }
-            if 'ContentType' in c:
-                json_input_data_config[channel_name]['ContentType'] = c['ContentType']
+        if input_data_config:
+            print(input_data_config)
+
+            for c in input_data_config:
+                channel_name = c['ChannelName']
+                json_input_data_config[channel_name] = {
+                    'TrainingInputMode': 'File'
+                }
+                if 'ContentType' in c:
+                    json_input_data_config[channel_name]['ContentType'] = c['ContentType']
 
         _write_json_file(os.path.join(config_path, 'hyperparameters.json'), hyperparameters)
         _write_json_file(os.path.join(config_path, 'resourceconfig.json'), resource_config)
@@ -285,14 +287,16 @@ def _prepare_training_volumes(self, data_dir, input_data_config, hyperparameters
         # Set up the channels for the containers. For local data we will
         # mount the local directory to the container. For S3 Data we will download the S3 data
         # first.
-        for channel in input_data_config:
-            uri = channel['DataUri']
-            channel_name = channel['ChannelName']
-            channel_dir = os.path.join(data_dir, channel_name)
-            os.mkdir(channel_dir)
-
-            data_source = sagemaker.local.data.get_data_source_instance(uri, self.sagemaker_session)
-            volumes.append(_Volume(data_source.get_root_dir(), channel=channel_name))
+
+        if input_data_config:
+            for channel in input_data_config:
+                uri = channel['DataUri']
+                channel_name = channel['ChannelName']
+                channel_dir = os.path.join(data_dir, channel_name)
+                os.mkdir(channel_dir)
+
+                data_source = sagemaker.local.data.get_data_source_instance(uri, self.sagemaker_session)
+                volumes.append(_Volume(data_source.get_root_dir(), channel=channel_name))
 
         # If there is a training script directory and it is a local directory,
         #  mount it to the container.

src/sagemaker/local/local_session.py

@@ -53,8 +53,8 @@ def __init__(self, sagemaker_session=None):
         """
         self.sagemaker_session = sagemaker_session or LocalSession()
 
-    def create_training_job(self, TrainingJobName, AlgorithmSpecification, InputDataConfig, OutputDataConfig,
-                            ResourceConfig, **kwargs):
+    def create_training_job(self, TrainingJobName, AlgorithmSpecification, OutputDataConfig,
+                            ResourceConfig, InputDataConfig=None, **kwargs):
         """
         Create a training job in Local Mode
         Args:

tests/unit/test_local_session.py

@@ -61,8 +61,8 @@ def test_create_training_job(train, LocalSession):
     resource_config = {'InstanceType': 'local', 'InstanceCount': instance_count}
     hyperparameters = {'a': 1, 'b': 'bee'}
 
-    local_sagemaker_client.create_training_job('my-training-job', algo_spec, input_data_config,
-                                               output_data_config, resource_config, HyperParameters=hyperparameters)
+    local_sagemaker_client.create_training_job('my-training-job', algo_spec, output_data_config, resource_config,
+                                               InputDataConfig=input_data_config, HyperParameters=hyperparameters)
 
     expected = {
         'ResourceConfig': {'InstanceCount': instance_count},
@@ -111,8 +111,8 @@ def test_create_training_job_invalid_data_source(train, LocalSession):
     hyperparameters = {'a': 1, 'b': 'bee'}
 
     with pytest.raises(ValueError):
-        local_sagemaker_client.create_training_job('my-training-job', algo_spec, input_data_config,
-                                                   output_data_config, resource_config, HyperParameters=hyperparameters)
+        local_sagemaker_client.create_training_job('my-training-job', algo_spec, output_data_config, resource_config,
+                                                   InputDataConfig=input_data_config, HyperParameters=hyperparameters)
 
 
 @patch('sagemaker.local.image._SageMakerContainer.train', return_value="/some/path/to/model")
@@ -141,8 +141,8 @@ def test_create_training_job_not_fully_replicated(train, LocalSession):
     hyperparameters = {'a': 1, 'b': 'bee'}
 
     with pytest.raises(RuntimeError):
-        local_sagemaker_client.create_training_job('my-training-job', algo_spec, input_data_config,
-                                                   output_data_config, resource_config, HyperParameters=hyperparameters)
+        local_sagemaker_client.create_training_job('my-training-job', algo_spec, output_data_config, resource_config,
+                                                   InputDataConfig=input_data_config, HyperParameters=hyperparameters)
 
 
 @patch('sagemaker.local.local_session.LocalSession')