Enable setting VPC config when creating/deploying models (#412)

* Create SageMaker Models including optional VpcConfig

Models created from Estimator or Training Job use the same VpcConfig

* Adjust naming, preparing, and passing vpc_config in base classes Estimator, Model, Session

* Add optional vpc_config parameter for Amazon/Frameworks models

* Improve integ test coverage of VPC features

Configure subnets and security groups to support multi-node VPC jobs and endpoints
Implement more VPC integ tests covering several Estimators and Predictors

* Update changelog

* Add vpc_utils with constants and utilities for handling VpcConfig

* Rename/refactor parameter vpc_config_override for models created from estimators

* Consolidate VPC integ tests to a single TensorFlow multi-instance case

* Tweaks to vpc_utils: make VPC_CONFIG_DEFAULT immutable and rename validate to sanitize

* Add integ test for transform with vpc config

* Update README with SageMaker VPC documentation and examples

* Remove tests/integ/test_vpc.py

Author: mmeidl

CHANGELOG.rst

@@ -2,6 +2,12 @@
 CHANGELOG
 =========
 
+=========
+1.11.2dev
+=========
+* enhancement: Enable setting VPC config when creating/deploying models
+
+=======
 1.11.1
 ======
 

README.rst

@@ -36,7 +36,8 @@ Table of Contents
 8. `BYO Docker Containers with SageMaker Estimators <#byo-docker-containers-with-sagemaker-estimators>`__
 9. `SageMaker Automatic Model Tuning <#sagemaker-automatic-model-tuning>`__
 10. `SageMaker Batch Transform <#sagemaker-batch-transform>`__
-11. `BYO Model <#byo-model>`__
+11. `Secure Training and Inference with VPC <#secure-training-and-inference-with-vpc>`__
+12. `BYO Model <#byo-model>`__
 
 
 Installing the SageMaker Python SDK
@@ -458,6 +459,71 @@ You can also specify other attributes of your data, such as the content type.
 For more details about what can be specified here, see `API docs <https://sagemaker.readthedocs.io/en/latest/transformer.html#sagemaker.transformer.Transformer.transform>`__.
 
 
+Secure Training and Inference with VPC
+--------------------------------------
+
+Amazon SageMaker allows you to control network traffic to and from model container instances using Amazon Virtual Private Cloud (VPC).
+You can configure SageMaker to use your own private VPC in order to further protect and monitor traffic.
+
+For more information about Amazon SageMaker VPC features, and guidelines for configuring your VPC,
+see the following documentation:
+
+- `Protect Training Jobs by Using an Amazon Virtual Private Cloud <https://docs.aws.amazon.com/sagemaker/latest/dg/train-vpc.html>`__
+- `Protect Endpoints by Using an Amazon Virtual Private Cloud <https://docs.aws.amazon.com/sagemaker/latest/dg/host-vpc.html>`__
+- `Protect Data in Batch Transform Jobs by Using an Amazon Virtual Private Cloud <https://docs.aws.amazon.com/sagemaker/latest/dg/batch-vpc.html>`__
+- `Working with VPCs and Subnets <https://docs.aws.amazon.com/vpc/latest/userguide/working-with-vpcs.html>`__
+
+You can also reference or reuse the example VPC created for integration tests: `tests/integ/vpc_test_utils.py <tests/integ/vpc_test_utils.py>`__
+
+To train a model using your own VPC, set the optional parameters ``subnets`` and ``security_group_ids`` on an ``Estimator``:
+
+.. code:: python
+
+    from sagemaker.mxnet import MXNet
+
+    # Configure an MXNet Estimator with subnets and security groups from your VPC
+    mxnet_vpc_estimator = MXNet('train.py',
+                            train_instance_type='ml.p2.xlarge',
+                            train_instance_count = 1,
+                            subnets=['subnet-1', 'subnet-2'],
+                            security_group_ids=['sg-1'])
+
+    # SageMaker Training Job will set VpcConfig and container instances will run in your VPC
+    mxnet_vpc_estimator.fit('s3://my_bucket/my_training_data/')
+
+When you create a ``Predictor`` from the ``Estimator`` using ``deploy()``, the same VPC configurations will be set on the SageMaker Model:
+
+.. code:: python
+
+    # Creates a SageMaker Model and Endpoint using the same VpcConfig
+    # Endpoint container instances will run in your VPC
+    mxnet_vpc_predictor = mxnet_vpc_estimator.deploy(initial_instance_count=1,
+                                                     instance_type='ml.p2.xlarge')
+
+    # You can also set ``vpc_config_override`` to use a different VpcConfig
+    other_vpc_config = {'Subnets': ['subnet-3', 'subnet-4'],
+                        'SecurityGroupIds': ['sg-2']}
+    mxnet_predictor_other_vpc = mxnet_vpc_estimator.deploy(initial_instance_count=1,
+                                                           instance_type='ml.p2.xlarge',
+                                                           vpc_config_override=other_vpc_config)
+
+    # Setting ``vpc_config_override=None`` will disable VpcConfig
+    mxnet_predictor_no_vpc = mxnet_vpc_estimator.deploy(initial_instance_count=1,
+                                                        instance_type='ml.p2.xlarge',
+                                                        vpc_config_override=None)
+
+Likewise, when you create ``Transformer`` from the ``Estimator`` using ``transformer()``, the same VPC configurations will be set on the SageMaker Model:
+
+.. code:: python
+
+    # Creates a SageMaker Model using the same VpcConfig
+    mxnet_vpc_transformer = mxnet_vpc_estimator.transformer(instance_count=1,
+                                                            instance_type='ml.p2.xlarge')
+
+    # Transform Job container instances will run in your VPC
+    mxnet_vpc_transformer.transform('s3://my-bucket/batch-transform-input')
+
+
 FAQ
 ---
 

src/sagemaker/amazon/factorization_machines.py

@@ -19,6 +19,7 @@
 from sagemaker.predictor import RealTimePredictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
 
 class FactorizationMachines(AmazonAlgorithmEstimatorBase):
@@ -163,11 +164,19 @@ def __init__(self, role, train_instance_count, train_instance_type,
         self.factors_init_sigma = factors_init_sigma
         self.factors_init_value = factors_init_value
 
-    def create_model(self):
+    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
         """Return a :class:`~sagemaker.amazon.FactorizationMachinesModel` referencing the latest
-        s3 model data produced by this Estimator."""
+        s3 model data produced by this Estimator.
 
-        return FactorizationMachinesModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session)
+        Args:
+            vpc_config_override (dict[str, list[str]]): Optional override for VpcConfig set on the model.
+                Default: use subnets and security groups from this Estimator.
+                * 'Subnets' (list[str]): List of subnet ids.
+                * 'SecurityGroupIds' (list[str]): List of security group ids.
+
+        """
+        return FactorizationMachinesModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session,
+                                          vpc_config=self.get_vpc_config(vpc_config_override))
 
 
 class FactorizationMachinesPredictor(RealTimePredictor):
@@ -195,12 +204,13 @@ class FactorizationMachinesModel(Model):
     """Reference S3 model data created by FactorizationMachines estimator. Calling :meth:`~sagemaker.model.Model.deploy`
     creates an Endpoint and returns :class:`FactorizationMachinesPredictor`."""
 
-    def __init__(self, model_data, role, sagemaker_session=None):
+    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
         sagemaker_session = sagemaker_session or Session()
         repo = '{}:{}'.format(FactorizationMachines.repo_name, FactorizationMachines.repo_version)
         image = '{}/{}'.format(registry(sagemaker_session.boto_session.region_name), repo)
         super(FactorizationMachinesModel, self).__init__(model_data,
                                                          image,
                                                          role,
                                                          predictor_cls=FactorizationMachinesPredictor,
-                                                         sagemaker_session=sagemaker_session)
+                                                         sagemaker_session=sagemaker_session,
+                                                         **kwargs)

src/sagemaker/amazon/kmeans.py

@@ -19,6 +19,7 @@
 from sagemaker.predictor import RealTimePredictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
 
 class KMeans(AmazonAlgorithmEstimatorBase):
@@ -102,10 +103,18 @@ def __init__(self, role, train_instance_count, train_instance_type, k, init_meth
         self.center_factor = center_factor
         self.eval_metrics = eval_metrics
 
-    def create_model(self):
+    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
         """Return a :class:`~sagemaker.amazon.kmeans.KMeansModel` referencing the latest
-        s3 model data produced by this Estimator."""
-        return KMeansModel(self.model_data, self.role, self.sagemaker_session)
+        s3 model data produced by this Estimator.
+
+        Args:
+            vpc_config_override (dict[str, list[str]]): Optional override for VpcConfig set on the model.
+                Default: use subnets and security groups from this Estimator.
+                * 'Subnets' (list[str]): List of subnet ids.
+                * 'SecurityGroupIds' (list[str]): List of security group ids.
+        """
+        return KMeansModel(self.model_data, self.role, self.sagemaker_session,
+                           vpc_config=self.get_vpc_config(vpc_config_override))
 
     def _prepare_for_training(self, records, mini_batch_size=5000, job_name=None):
         super(KMeans, self)._prepare_for_training(records, mini_batch_size=mini_batch_size, job_name=job_name)
@@ -138,9 +147,10 @@ class KMeansModel(Model):
     """Reference KMeans s3 model data. Calling :meth:`~sagemaker.model.Model.deploy` creates an Endpoint and return
     a Predictor to performs k-means cluster assignment."""
 
-    def __init__(self, model_data, role, sagemaker_session=None):
+    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
         sagemaker_session = sagemaker_session or Session()
         repo = '{}:{}'.format(KMeans.repo_name, KMeans.repo_version)
         image = '{}/{}'.format(registry(sagemaker_session.boto_session.region_name), repo)
         super(KMeansModel, self).__init__(model_data, image, role, predictor_cls=KMeansPredictor,
-                                          sagemaker_session=sagemaker_session)
+                                          sagemaker_session=sagemaker_session,
+                                          **kwargs)

src/sagemaker/amazon/knn.py

@@ -19,6 +19,7 @@
 from sagemaker.predictor import RealTimePredictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
 
 class KNN(AmazonAlgorithmEstimatorBase):
@@ -97,11 +98,18 @@ def __init__(self, role, train_instance_count, train_instance_type, k, sample_si
         if dimension_reduction_type and not dimension_reduction_target:
             raise ValueError('"dimension_reduction_target" is required when "dimension_reduction_type" is set.')
 
-    def create_model(self):
+    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
         """Return a :class:`~sagemaker.amazon.KNNModel` referencing the latest
-        s3 model data produced by this Estimator."""
+        s3 model data produced by this Estimator.
 
-        return KNNModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session)
+        Args:
+            vpc_config_override (dict[str, list[str]]): Optional override for VpcConfig set on the model.
+                Default: use subnets and security groups from this Estimator.
+                * 'Subnets' (list[str]): List of subnet ids.
+                * 'SecurityGroupIds' (list[str]): List of security group ids.
+        """
+        return KNNModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session,
+                        vpc_config=self.get_vpc_config(vpc_config_override))
 
     def _prepare_for_training(self, records, mini_batch_size=None, job_name=None):
         super(KNN, self)._prepare_for_training(records, mini_batch_size=mini_batch_size, job_name=job_name)
@@ -128,9 +136,9 @@ class KNNModel(Model):
     """Reference S3 model data created by KNN estimator. Calling :meth:`~sagemaker.model.Model.deploy`
     creates an Endpoint and returns :class:`KNNPredictor`."""
 
-    def __init__(self, model_data, role, sagemaker_session=None):
+    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
         sagemaker_session = sagemaker_session or Session()
         repo = '{}:{}'.format(KNN.repo_name, KNN.repo_version)
         image = '{}/{}'.format(registry(sagemaker_session.boto_session.region_name, KNN.repo_name), repo)
         super(KNNModel, self).__init__(model_data, image, role, predictor_cls=KNNPredictor,
-                                       sagemaker_session=sagemaker_session)
+                                       sagemaker_session=sagemaker_session, **kwargs)

src/sagemaker/amazon/lda.py

@@ -19,6 +19,7 @@
 from sagemaker.predictor import RealTimePredictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
 
 class LDA(AmazonAlgorithmEstimatorBase):
@@ -89,11 +90,18 @@ def __init__(self, role, train_instance_type, num_topics,
         self.max_iterations = max_iterations
         self.tol = tol
 
-    def create_model(self):
+    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
         """Return a :class:`~sagemaker.amazon.LDAModel` referencing the latest
-        s3 model data produced by this Estimator."""
+        s3 model data produced by this Estimator.
 
-        return LDAModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session)
+        Args:
+            vpc_config_override (dict[str, list[str]]): Optional override for VpcConfig set on the model.
+                Default: use subnets and security groups from this Estimator.
+                * 'Subnets' (list[str]): List of subnet ids.
+                * 'SecurityGroupIds' (list[str]): List of security group ids.
+        """
+        return LDAModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session,
+                        vpc_config=self.get_vpc_config(vpc_config_override))
 
     def _prepare_for_training(self, records, mini_batch_size, job_name=None):
         # mini_batch_size is required, prevent explicit calls with None
@@ -124,9 +132,9 @@ class LDAModel(Model):
     """Reference LDA s3 model data. Calling :meth:`~sagemaker.model.Model.deploy` creates an Endpoint and return
     a Predictor that transforms vectors to a lower-dimensional representation."""
 
-    def __init__(self, model_data, role, sagemaker_session=None):
+    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
         sagemaker_session = sagemaker_session or Session()
         repo = '{}:{}'.format(LDA.repo_name, LDA.repo_version)
         image = '{}/{}'.format(registry(sagemaker_session.boto_session.region_name, LDA.repo_name), repo)
         super(LDAModel, self).__init__(model_data, image, role, predictor_cls=LDAPredictor,
-                                       sagemaker_session=sagemaker_session)
+                                       sagemaker_session=sagemaker_session, **kwargs)

src/sagemaker/amazon/linear_learner.py

@@ -19,6 +19,7 @@
 from sagemaker.predictor import RealTimePredictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
 
 class LinearLearner(AmazonAlgorithmEstimatorBase):
@@ -246,11 +247,18 @@ def __init__(self, role, train_instance_count, train_instance_type, predictor_ty
             raise ValueError(
                 "For predictor_type 'multiclass_classifier', 'num_classes' should be set to a value greater than 2.")
 
-    def create_model(self):
+    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
         """Return a :class:`~sagemaker.amazon.LinearLearnerModel` referencing the latest
-        s3 model data produced by this Estimator."""
+        s3 model data produced by this Estimator.
 
-        return LinearLearnerModel(self.model_data, self.role, self.sagemaker_session)
+        Args:
+            vpc_config_override (dict[str, list[str]]): Optional override for VpcConfig set on the model.
+                Default: use subnets and security groups from this Estimator.
+                * 'Subnets' (list[str]): List of subnet ids.
+                * 'SecurityGroupIds' (list[str]): List of security group ids.
+        """
+        return LinearLearnerModel(self.model_data, self.role, self.sagemaker_session,
+                                  vpc_config=self.get_vpc_config(vpc_config_override))
 
     def _prepare_for_training(self, records, mini_batch_size=None, job_name=None):
         num_records = None
@@ -293,10 +301,11 @@ class LinearLearnerModel(Model):
     """Reference LinearLearner s3 model data. Calling :meth:`~sagemaker.model.Model.deploy` creates an Endpoint
     and returns a :class:`LinearLearnerPredictor`"""
 
-    def __init__(self, model_data, role, sagemaker_session=None):
+    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
         sagemaker_session = sagemaker_session or Session()
         repo = '{}:{}'.format(LinearLearner.repo_name, LinearLearner.repo_version)
         image = '{}/{}'.format(registry(sagemaker_session.boto_session.region_name), repo)
         super(LinearLearnerModel, self).__init__(model_data, image, role,
                                                  predictor_cls=LinearLearnerPredictor,
-                                                 sagemaker_session=sagemaker_session)
+                                                 sagemaker_session=sagemaker_session,
+                                                 **kwargs)

src/sagemaker/amazon/ntm.py

@@ -19,6 +19,7 @@
 from sagemaker.predictor import RealTimePredictor
 from sagemaker.model import Model
 from sagemaker.session import Session
+from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
 
 class NTM(AmazonAlgorithmEstimatorBase):
@@ -107,11 +108,18 @@ def __init__(self, role, train_instance_count, train_instance_type, num_topics,
         self.weight_decay = weight_decay
         self.learning_rate = learning_rate
 
-    def create_model(self):
+    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
         """Return a :class:`~sagemaker.amazon.NTMModel` referencing the latest
-        s3 model data produced by this Estimator."""
+        s3 model data produced by this Estimator.
 
-        return NTMModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session)
+        Args:
+            vpc_config_override (dict[str, list[str]]): Optional override for VpcConfig set on the model.
+                Default: use subnets and security groups from this Estimator.
+                * 'Subnets' (list[str]): List of subnet ids.
+                * 'SecurityGroupIds' (list[str]): List of security group ids.
+        """
+        return NTMModel(self.model_data, self.role, sagemaker_session=self.sagemaker_session,
+                        vpc_config=self.get_vpc_config(vpc_config_override))
 
     def _prepare_for_training(self, records, mini_batch_size, job_name=None):
         if mini_batch_size is not None and (mini_batch_size < 1 or mini_batch_size > 10000):
@@ -140,9 +148,10 @@ class NTMModel(Model):
     """Reference NTM s3 model data. Calling :meth:`~sagemaker.model.Model.deploy` creates an Endpoint and return
     a Predictor that transforms vectors to a lower-dimensional representation."""
 
-    def __init__(self, model_data, role, sagemaker_session=None):
+    def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
         sagemaker_session = sagemaker_session or Session()
         repo = '{}:{}'.format(NTM.repo_name, NTM.repo_version)
         image = '{}/{}'.format(registry(sagemaker_session.boto_session.region_name, NTM.repo_name), repo)
         super(NTMModel, self).__init__(model_data, image, role, predictor_cls=NTMPredictor,
-                                       sagemaker_session=sagemaker_session)
+                                       sagemaker_session=sagemaker_session,
+                                       **kwargs)