feature: network isolation mode in training (aws#791)

* feature: network isolation mode in training

* feature: network isolation mode in tar support training

* change: documentation and check describe training job network isolation

* doc update

* doc update, remove inference section

* sourcedir

* type error fix

Author: wiltonwu

doc/overview.rst

@@ -559,6 +559,33 @@ Likewise, when you create ``Transformer`` from the ``Estimator`` using ``transfo
     # Transform Job container instances will run in your VPC
     mxnet_vpc_transformer.transform('s3://my-bucket/batch-transform-input')
 
+Secure Training with Network Isolation (Internet-Free) Mode
+-------------------------------------------------------------------------
+You can enable network isolation mode when running training and inference on Amazon SageMaker.
+
+For more information about Amazon SageMaker network isolation mode, see the `SageMaker documentation on network isolation or internet-free mode <https://docs.aws.amazon.com/sagemaker/latest/dg/mkt-algo-model-internet-free.html>`__.
+
+To train a model in network isolation mode, set the optional parameter ``enable_network_isolation`` to ``True`` in any network isolation supported Framework Estimator.
+
+.. code:: python
+
+    # set the enable_network_isolation parameter to True
+    sklearn_estimator = SKLearn('sklearn-train.py',
+                                train_instance_type='ml.m4.xlarge',
+                                framework_version='0.20.0',
+                                hyperparameters = {'epochs': 20, 'batch-size': 64, 'learning-rate': 0.1},
+                                enable_network_isolation=True)
+
+    # SageMaker Training Job will in the container without   any inbound or outbound network calls during runtime
+    sklearn_estimator.fit({'train': 's3://my-data-bucket/path/to/my/training/data',
+                            'test': 's3://my-data-bucket/path/to/my/test/data'})
+
+When this training job is created, the SageMaker Python SDK will upload the files in ``entry_point``, ``source_dir``, and ``dependencies`` to S3 as a compressed ``sourcedir.tar.gz`` file (``'s3://mybucket/sourcedir.tar.gz'``).
+
+A new training job channel, named ``code``, will be added with that S3 URI.  Before the training docker container is initialized, the ``sourcedir.tar.gz`` will be downloaded from S3 to the ML storage volume like any other offline input channel.
+
+Once the training job begins, the training container will look at the offline input ``code`` channel to install dependencies and run the entry script. This isolates the training container, so no inbound or outbound network calls can be made.
+
 
 FAQ
 ---

src/sagemaker/estimator.py

@@ -117,6 +117,8 @@ def __init__(self, role, train_instance_count, train_instance_type,
         self.metric_definitions = metric_definitions
         self.model_uri = model_uri
         self.model_channel_name = model_channel_name
+        self.code_uri = None
+        self.code_channel_name = 'code'
 
         if self.train_instance_type in ('local', 'local_gpu'):
             if self.train_instance_type == 'local_gpu' and self.train_instance_count > 1:
@@ -773,9 +775,11 @@ class Framework(EstimatorBase):
     LAUNCH_MPI_ENV_NAME = 'sagemaker_mpi_enabled'
     MPI_NUM_PROCESSES_PER_HOST = 'sagemaker_mpi_num_of_processes_per_host'
     MPI_CUSTOM_MPI_OPTIONS = 'sagemaker_mpi_custom_mpi_options'
+    CONTAINER_CODE_CHANNEL_SOURCEDIR_PATH = '/opt/ml/input/data/code/sourcedir.tar.gz'
 
     def __init__(self, entry_point, source_dir=None, hyperparameters=None, enable_cloudwatch_metrics=False,
-                 container_log_level=logging.INFO, code_location=None, image_name=None, dependencies=None, **kwargs):
+                 container_log_level=logging.INFO, code_location=None, image_name=None, dependencies=None,
+                 enable_network_isolation=False, **kwargs):
         """Base class initializer. Subclasses which override ``__init__`` should invoke ``super()``
 
         Args:
@@ -784,6 +788,21 @@ def __init__(self, entry_point, source_dir=None, hyperparameters=None, enable_cl
             source_dir (str): Path (absolute or relative) to a directory with any other training
                 source code dependencies aside from the entry point file (default: None). Structure within this
                 directory are preserved when training on Amazon SageMaker.
+            hyperparameters (dict): Hyperparameters that will be used for training (default: None).
+                The hyperparameters are made accessible as a dict[str, str] to the training code on SageMaker.
+                For convenience, this accepts other types for keys and values, but ``str()`` will be called
+                to convert them before training.
+            enable_cloudwatch_metrics (bool): [DEPRECATED] Now there are cloudwatch metrics emitted by all SageMaker
+                training jobs. This will be ignored for now and removed in a further release.
+            container_log_level (int): Log level to use within the container (default: logging.INFO).
+                Valid values are defined in the Python logging module.
+            code_location (str): The S3 prefix URI where custom code will be uploaded (default: None).
+                The code file uploaded in S3 is 'code_location/source/sourcedir.tar.gz'.
+                If not specified, the default code location is s3://default_bucket/job-name/. And code file
+                uploaded to S3 is s3://default_bucket/job-name/source/sourcedir.tar.gz
+            image_name (str): An alternate image name to use instead of the official Sagemaker image
+                for the framework. This is useful to run one of the Sagemaker supported frameworks
+                with an image containing custom dependencies.
             dependencies (list[str]): A list of paths to directories (absolute or relative) with
                 any additional libraries that will be exported to the container (default: []).
                 The library folders will be copied to SageMaker in the same folder where the entrypoint is copied.
@@ -800,21 +819,11 @@ def __init__(self, entry_point, source_dir=None, hyperparameters=None, enable_cl
                     >>>     |------ common
                     >>>     |------ virtual-env
 
-            hyperparameters (dict): Hyperparameters that will be used for training (default: None).
-                The hyperparameters are made accessible as a dict[str, str] to the training code on SageMaker.
-                For convenience, this accepts other types for keys and values, but ``str()`` will be called
-                to convert them before training.
-            enable_cloudwatch_metrics (bool): [DEPRECATED] Now there are cloudwatch metrics emitted by all SageMaker
-                training jobs. This will be ignored for now and removed in a further release.
-            container_log_level (int): Log level to use within the container (default: logging.INFO).
-                Valid values are defined in the Python logging module.
-            code_location (str): The S3 prefix URI where custom code will be uploaded (default: None).
-                The code file uploaded in S3 is 'code_location/source/sourcedir.tar.gz'.
-                If not specified, the default code location is s3://default_bucket/job-name/. And code file
-                uploaded to S3 is s3://default_bucket/job-name/source/sourcedir.tar.gz
-            image_name (str): An alternate image name to use instead of the official Sagemaker image
-                for the framework. This is useful to run one of the Sagemaker supported frameworks
-                with an image containing custom dependencies.
+            enable_network_isolation (bool): Specifies whether container will run in network isolation mode. Network
+                isolation mode restricts the container access to outside networks (such as the internet). The container
+                does not make any inbound or outbound network calls. If True, a channel named "code" will be created
+                for any user entry script for training. The user entry script, files in source_dir (if specified), and
+                dependencies will be uploaded in a tar to S3. Also known as internet-free mode (default: `False`).
             **kwargs: Additional kwargs passed to the ``EstimatorBase`` constructor.
         """
         super(Framework, self).__init__(**kwargs)
@@ -830,9 +839,18 @@ def __init__(self, entry_point, source_dir=None, hyperparameters=None, enable_cl
         self.container_log_level = container_log_level
         self.code_location = code_location
         self.image_name = image_name
+        self._enable_network_isolation = enable_network_isolation
 
         self._hyperparameters = hyperparameters or {}
 
+    def enable_network_isolation(self):
+        """Return True if this Estimator can use network isolation to run.
+
+        Returns:
+            bool: Whether this Estimator can use network isolation or not.
+        """
+        return self._enable_network_isolation
+
     def _prepare_for_training(self, job_name=None):
         """Set hyperparameters needed for training. This method will also validate ``source_dir``.
 
@@ -858,6 +876,11 @@ def _prepare_for_training(self, job_name=None):
 
             code_dir = 'file://' + self.source_dir
             script = self.entry_point
+        elif self.enable_network_isolation() and self.entry_point:
+            self.uploaded_code = self._stage_user_code_in_s3()
+            code_dir = self.CONTAINER_CODE_CHANNEL_SOURCEDIR_PATH
+            script = self.uploaded_code.script_name
+            self.code_uri = self.uploaded_code.s3_prefix
         else:
             self.uploaded_code = self._stage_user_code_in_s3()
             code_dir = self.uploaded_code.s3_prefix
@@ -881,12 +904,12 @@ def _stage_user_code_in_s3(self):
 
         if self.code_location is None and local_mode:
             code_bucket = self.sagemaker_session.default_bucket()
-            code_s3_prefix = '{}/source'.format(self._current_job_name)
+            code_s3_prefix = '{}/{}'.format(self._current_job_name, 'source')
             kms_key = None
 
         elif self.code_location is None:
             code_bucket, _ = parse_s3_url(self.output_path)
-            code_s3_prefix = '{}/source'.format(self._current_job_name)
+            code_s3_prefix = '{}/{}'.format(self._current_job_name, 'source')
             kms_key = self.output_kms_key
         else:
             code_bucket, key_prefix = parse_s3_url(self.code_location)

src/sagemaker/job.py

@@ -60,12 +60,21 @@ def _load_config(inputs, estimator, expand_role=True, validate_uri=True):
         stop_condition = _Job._prepare_stop_condition(estimator.train_max_run)
         vpc_config = estimator.get_vpc_config()
 
-        model_channel = _Job._prepare_model_channel(input_config, estimator.model_uri, estimator.model_channel_name,
-                                                    validate_uri)
+        model_channel = _Job._prepare_channel(input_config, estimator.model_uri, estimator.model_channel_name,
+                                              validate_uri, content_type='application/x-sagemaker-model',
+                                              input_mode='File')
         if model_channel:
             input_config = [] if input_config is None else input_config
             input_config.append(model_channel)
 
+        if estimator.enable_network_isolation():
+            code_channel = _Job._prepare_channel(input_config, estimator.code_uri, estimator.code_channel_name,
+                                                 validate_uri)
+
+            if code_channel:
+                input_config = [] if input_config is None else input_config
+                input_config.append(code_channel)
+
         return {'input_config': input_config,
                 'role': role,
                 'output_config': output_config,
@@ -110,16 +119,16 @@ def _convert_input_to_channel(channel_name, channel_s3_input):
         return channel_config
 
     @staticmethod
-    def _format_string_uri_input(uri_input, validate_uri=True):
+    def _format_string_uri_input(uri_input, validate_uri=True, content_type=None, input_mode=None):
         if isinstance(uri_input, str) and validate_uri and uri_input.startswith('s3://'):
-            return s3_input(uri_input)
+            return s3_input(uri_input, content_type=content_type, input_mode=input_mode)
         elif isinstance(uri_input, str) and validate_uri and uri_input.startswith('file://'):
             return file_input(uri_input)
         elif isinstance(uri_input, str) and validate_uri:
-            raise ValueError('Training input data must be a valid S3 or FILE URI: must start with "s3://" or '
-                             '"file://"')
+            raise ValueError('URI input {} must be a valid S3 or FILE URI: must start with "s3://" or '
+                             '"file://"'.format(uri_input))
         elif isinstance(uri_input, str):
-            return s3_input(uri_input)
+            return s3_input(uri_input, content_type=content_type, input_mode=input_mode)
         elif isinstance(uri_input, s3_input):
             return uri_input
         elif isinstance(uri_input, file_input):
@@ -128,21 +137,22 @@ def _format_string_uri_input(uri_input, validate_uri=True):
             raise ValueError('Cannot format input {}. Expecting one of str, s3_input, or file_input'.format(uri_input))
 
     @staticmethod
-    def _prepare_model_channel(input_config, model_uri=None, model_channel_name=None, validate_uri=True):
-        if not model_uri:
+    def _prepare_channel(input_config, channel_uri=None, channel_name=None, validate_uri=True, content_type=None,
+                         input_mode=None):
+        if not channel_uri:
             return
-        elif not model_channel_name:
-            raise ValueError('Expected a pre-trained model channel name if a model URL is specified.')
+        elif not channel_name:
+            raise ValueError('Expected a channel name if a channel URI {} is specified'.format(channel_uri))
 
         if input_config:
-            for channel in input_config:
-                if channel['ChannelName'] == model_channel_name:
-                    raise ValueError('Duplicate channels not allowed.')
+            for existing_channel in input_config:
+                if existing_channel['ChannelName'] == channel_name:
+                    raise ValueError('Duplicate channel {} not allowed.'.format(channel_name))
 
-        model_input = _Job._format_model_uri_input(model_uri, validate_uri)
-        model_channel = _Job._convert_input_to_channel(model_channel_name, model_input)
+        channel_input = _Job._format_string_uri_input(channel_uri, validate_uri, content_type, input_mode)
+        channel = _Job._convert_input_to_channel(channel_name, channel_input)
 
-        return model_channel
+        return channel
 
     @staticmethod
     def _format_model_uri_input(model_uri, validate_uri=True):

tests/integ/test_sklearn_train.py

@@ -55,6 +55,33 @@ def test_training_with_additional_hyperparameters(sagemaker_session, sklearn_ful
         return sklearn.latest_training_job.name
 
 
+@pytest.mark.skipif(PYTHON_VERSION != 'py3', reason="Scikit-learn image supports only python 3.")
+def test_training_with_network_isolation(sagemaker_session, sklearn_full_version):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        script_path = os.path.join(DATA_DIR, 'sklearn_mnist', 'mnist.py')
+        data_path = os.path.join(DATA_DIR, 'sklearn_mnist')
+
+        sklearn = SKLearn(entry_point=script_path,
+                          role='SageMakerRole',
+                          train_instance_type="ml.c4.xlarge",
+                          framework_version=sklearn_full_version,
+                          py_version=PYTHON_VERSION,
+                          sagemaker_session=sagemaker_session,
+                          hyperparameters={'epochs': 1},
+                          enable_network_isolation=True)
+
+        train_input = sklearn.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
+                                                            key_prefix='integ-test-data/sklearn_mnist/train')
+        test_input = sklearn.sagemaker_session.upload_data(path=os.path.join(data_path, 'test'),
+                                                           key_prefix='integ-test-data/sklearn_mnist/test')
+        job_name = unique_name_from_base('test-sklearn-hp')
+
+        sklearn.fit({'train': train_input, 'test': test_input}, job_name=job_name)
+        assert sagemaker_session.sagemaker_client \
+            .describe_training_job(TrainingJobName=job_name)['EnableNetworkIsolation']
+        return sklearn.latest_training_job.name
+
+
 @pytest.mark.canary_quick
 @pytest.mark.regional_testing
 @pytest.mark.skipif(PYTHON_VERSION != 'py3', reason="Scikit-learn image supports only python 3.")

tests/unit/test_fw_utils.py

@@ -175,27 +175,26 @@ def test_unoptimized_gpu_family():
 
 
 def test_tar_and_upload_dir_s3(sagemaker_session):
-    bucket = 'mybucker'
+    bucket = 'mybucket'
     s3_key_prefix = 'something/source'
     script = 'mnist.py'
     directory = 's3://m'
     result = fw_utils.tar_and_upload_dir(sagemaker_session, bucket, s3_key_prefix, script, directory)
+
     assert result == fw_utils.UploadedCode('s3://m', 'mnist.py')
 
 
 @patch('sagemaker.utils')
 def test_tar_and_upload_dir_s3_with_kms(utils, sagemaker_session):
+    bucket = 'mybucket'
+    s3_key_prefix = 'something/source'
+    script = 'mnist.py'
+    kms_key = 'kms-key'
+    result = fw_utils.tar_and_upload_dir(sagemaker_session, bucket, s3_key_prefix, script, kms_key=kms_key)
 
-    result = fw_utils.tar_and_upload_dir(sagemaker_session,
-                                         'mybucker',
-                                         'something/source',
-                                         'mnist.py',
-                                         kms_key='kms-key')
-
-    assert result == fw_utils.UploadedCode('s3://mybucker/something/source/sourcedir.tar.gz',
-                                           'mnist.py')
+    assert result == fw_utils.UploadedCode('s3://{}/{}/sourcedir.tar.gz'.format(bucket, s3_key_prefix), script)
 
-    extra_args = {'ServerSideEncryption': 'aws:kms', 'SSEKMSKeyId': 'kms-key'}
+    extra_args = {'ServerSideEncryption': 'aws:kms', 'SSEKMSKeyId': kms_key}
     obj = sagemaker_session.resource('s3').Object('', '')
     obj.upload_file.assert_called_with(utils.create_tar_file(), ExtraArgs=extra_args)