Add VPC config to estimator for training job creation

CreateTraningJob api supports vpc. This change adds vpc config as an optional
argument to Estimator.

Author: icywang86rui

CHANGELOG.rst

@@ -6,6 +6,7 @@ CHANGELOG
 ========
 
 * bug-fix: Estimators: Fix serialization of single records
+* enhancement: Enable VPC config in training job creation
 
 1.9.0
 =====

src/sagemaker/estimator.py

@@ -46,7 +46,8 @@ class EstimatorBase(with_metaclass(ABCMeta, object)):
 
     def __init__(self, role, train_instance_count, train_instance_type,
                  train_volume_size=30, train_max_run=24 * 60 * 60, input_mode='File',
-                 output_path=None, output_kms_key=None, base_job_name=None, sagemaker_session=None, tags=None):
+                 output_path=None, output_kms_key=None, base_job_name=None, sagemaker_session=None, tags=None,
+                 subnets=None, security_group_ids=None):
         """Initialize an ``EstimatorBase`` instance.
 
         Args:
@@ -99,6 +100,10 @@ def __init__(self, role, train_instance_count, train_instance_type,
         self.output_kms_key = output_kms_key
         self.latest_training_job = None
 
+        # VPC configurations
+        self.subnets = subnets
+        self.security_group_ids = security_group_ids
+
     @abstractmethod
     def train_image(self):
         """Return the Docker image to use for training.
@@ -398,8 +403,9 @@ def start_new(cls, estimator, inputs):
         estimator.sagemaker_session.train(image=estimator.train_image(), input_mode=estimator.input_mode,
                                           input_config=config['input_config'], role=config['role'],
                                           job_name=estimator._current_job_name, output_config=config['output_config'],
-                                          resource_config=config['resource_config'], hyperparameters=hyperparameters,
-                                          stop_condition=config['stop_condition'], tags=estimator.tags)
+                                          resource_config=config['resource_config'], vpc_config=config['vpc_config'],
+                                          hyperparameters=hyperparameters, stop_condition=config['stop_condition'],
+                                          tags=estimator.tags)
 
         return cls(estimator.sagemaker_session, estimator._current_job_name)
 

src/sagemaker/job.py

@@ -59,12 +59,14 @@ def _load_config(inputs, estimator):
                                                         estimator.train_instance_type,
                                                         estimator.train_volume_size)
         stop_condition = _Job._prepare_stop_condition(estimator.train_max_run)
+        vpc_config = _Job._prepare_vpc_config(estimator.subnets, estimator.security_group_ids)
 
         return {'input_config': input_config,
                 'role': role,
                 'output_config': output_config,
                 'resource_config': resource_config,
-                'stop_condition': stop_condition}
+                'stop_condition': stop_condition,
+                'vpc_config': vpc_config}
 
     @staticmethod
     def _format_inputs_to_input_config(inputs):
@@ -143,6 +145,13 @@ def _prepare_resource_config(instance_count, instance_type, volume_size):
                 'InstanceType': instance_type,
                 'VolumeSizeInGB': volume_size}
 
+    @staticmethod
+    def _prepare_vpc_config(subnets, security_group_ids):
+        if subnets is None or security_group_ids is None:
+            return None
+        return {'Subnets': subnets,
+                'SecurityGroupIds': security_group_ids}
+
     @staticmethod
     def _prepare_stop_condition(max_run):
         return {'MaxRuntimeInSeconds': max_run}

src/sagemaker/session.py

@@ -202,7 +202,7 @@ def default_bucket(self):
         return self._default_bucket
 
     def train(self, image, input_mode, input_config, role, job_name, output_config,
-              resource_config, hyperparameters, stop_condition, tags):
+              resource_config, vpc_config, hyperparameters, stop_condition, tags):
         """Create an Amazon SageMaker training job.
 
         Args:
@@ -259,6 +259,9 @@ def train(self, image, input_mode, input_config, role, job_name, output_config,
         if tags is not None:
             train_request['Tags'] = tags
 
+        if vpc_config is not None:
+            train_request['VpcConfig'] = vpc_config
+
         LOGGER.info('Creating training-job with name: {}'.format(job_name))
         LOGGER.debug('train request: {}'.format(json.dumps(train_request, indent=4)))
         self.sagemaker_client.create_training_job(**train_request)

tests/integ/test_tf.py

@@ -22,6 +22,8 @@
 from tests.integ.timeout import timeout_and_delete_endpoint_by_name, timeout
 
 DATA_PATH = os.path.join(DATA_DIR, 'iris', 'data')
+VPC_SUBNETS = ['subnet-06b8537735fac3757']
+VPC_SECURITY_GROUP_IDS = ['sg-0a1008de6e1f384c3']
 
 
 @pytest.mark.continuous_testing
@@ -98,10 +100,17 @@ def test_failed_tf_training(sagemaker_session, tf_full_version):
                                hyperparameters={'input_tensor_name': 'inputs'},
                                train_instance_count=1,
                                train_instance_type='ml.c4.xlarge',
-                               sagemaker_session=sagemaker_session)
+                               sagemaker_session=sagemaker_session,
+                               subnets=VPC_SUBNETS,
+                               security_group_ids=VPC_SECURITY_GROUP_IDS)
 
         inputs = estimator.sagemaker_session.upload_data(path=DATA_PATH, key_prefix='integ-test-data/tf-failure')
 
         with pytest.raises(ValueError) as e:
             estimator.fit(inputs)
         assert 'This failure is expected' in str(e.value)
+
+        job_desc = estimator.sagemaker_session.sagemaker_client.describe_training_job(
+            TrainingJobName=estimator.latest_training_job.name)
+        assert VPC_SUBNETS == job_desc['VpcConfig']['Subnets']
+        assert VPC_SECURITY_GROUP_IDS == job_desc['VpcConfig']['SecurityGroupIds']

tests/unit/test_chainer.py

@@ -122,6 +122,7 @@ def _create_train_job(version):
             'MaxRuntimeInSeconds': 24 * 60 * 60
         },
         'tags': None,
+        'vpc_config': {'SecurityGroupIds': None, 'Subnets': None}
     }
 
 

tests/unit/test_estimator.py

@@ -658,6 +658,7 @@ def test_unsupported_type_in_dict():
     },
     'stop_condition': {'MaxRuntimeInSeconds': 86400},
     'tags': None,
+    'vpc_config': {'SecurityGroupIds': None, 'Subnets': None}
 }
 
 HYPERPARAMS = {'x': 1, 'y': 'hello'}

tests/unit/test_mxnet.py

@@ -95,6 +95,7 @@ def _create_train_job(version):
             'MaxRuntimeInSeconds': 24 * 60 * 60
         },
         'tags': None,
+        'vpc_config': {'SecurityGroupIds': None, 'Subnets': None}
     }
 
 

tests/unit/test_pytorch.py

@@ -112,7 +112,8 @@ def _create_train_job(version):
         'stop_condition': {
             'MaxRuntimeInSeconds': 24 * 60 * 60
         },
-        'tags': None
+        'tags': None,
+        'vpc_config': {'SecurityGroupIds': None, 'Subnets': None}
     }
 
 

tests/unit/test_session.py

@@ -176,6 +176,7 @@ def test_s3_input_all_arguments():
 MAX_TIME = 3 * 60 * 60
 JOB_NAME = 'jobname'
 TAGS = [{'Name': 'some-tag', 'Value': 'value-for-tag'}]
+VPC_CONFIG = {'Subnets': 'subnet', 'SecurityGroupIds': 'sgi-blahblah'}
 
 DEFAULT_EXPECTED_TRAIN_JOB_ARGS = {
     'OutputDataConfig': {
@@ -259,7 +260,7 @@ def test_train_pack_to_request(sagemaker_session):
 
     sagemaker_session.train(image=IMAGE, input_mode='File', input_config=in_config, role=EXPANDED_ROLE,
                             job_name=JOB_NAME, output_config=out_config, resource_config=resource_config,
-                            hyperparameters=None, stop_condition=stop_cond, tags=None)
+                            hyperparameters=None, stop_condition=stop_cond, tags=None, vpc_config=None)
 
     assert sagemaker_session.sagemaker_client.method_calls[0] == (
         'create_training_job', (), DEFAULT_EXPECTED_TRAIN_JOB_ARGS)
@@ -322,12 +323,13 @@ def test_train_pack_to_request_with_optional_params(sagemaker_session):
 
     sagemaker_session.train(image=IMAGE, input_mode='File', input_config=in_config, role=EXPANDED_ROLE,
                             job_name=JOB_NAME, output_config=out_config, resource_config=resource_config,
-                            hyperparameters=hyperparameters, stop_condition=stop_cond, tags=TAGS)
+                            hyperparameters=hyperparameters, stop_condition=stop_cond, tags=TAGS, vpc_config=VPC_CONFIG)
 
     _, _, actual_train_args = sagemaker_session.sagemaker_client.method_calls[0]
 
     assert actual_train_args['HyperParameters'] == hyperparameters
     assert actual_train_args['Tags'] == TAGS
+    assert actual_train_args['VpcConfig'] == VPC_CONFIG
 
 
 def test_transform_pack_to_request(sagemaker_session):

tests/unit/test_tf_estimator.py

@@ -103,6 +103,7 @@ def _create_train_job(tf_version):
             'MaxRuntimeInSeconds': 24 * 60 * 60
         },
         'tags': None,
+        'vpc_config': {'SecurityGroupIds': None, 'Subnets': None}
     }