FruitStand support (aws#892)

Co-authored-by: Rohan Gujarathi <[email protected]>
Co-authored-by: Rohan Gujarathi <[email protected]>
Co-authored-by: Ao Guo <[email protected]>

Author: 4 people

src/sagemaker/config/config_schema.py

@@ -52,6 +52,7 @@
 S3_KMS_KEY_ID = "S3KmsKeyId"
 S3_ROOT_URI = "S3RootUri"
 JOB_CONDA_ENV = "JobCondaEnvironment"
+ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = "EnableInterContainerTrafficEncryption"
 OFFLINE_STORE_CONFIG = "OfflineStoreConfig"
 ONLINE_STORE_CONFIG = "OnlineStoreConfig"
 S3_STORAGE_CONFIG = "S3StorageConfig"

src/sagemaker/remote_function/client.py

@@ -78,6 +78,8 @@ def remote(
     tags: List[Tuple[str, str]] = None,
     volume_kms_key: str = None,
     volume_size: int = 30,
+    encrypt_inter_container_traffic: bool = None,
+    enable_network_isolation: bool = None,
 ):
     """Function that starts a new SageMaker job synchronously with overridden runtime settings.
 
@@ -114,6 +116,13 @@ def remote(
             instance.
         volume_size (int): Size in GB of the storage volume to use for storing input and output
             data. Default is 30.
+        encrypt_inter_container_traffic (bool): Specifies whether traffic between training
+            containers is encrypted for the training job. (default: ``False``).
+        enable_network_isolation (bool): Specifies whether container will
+            run in network isolation mode (default: ``False``). 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. Also known as Internet-free mode.
     """
 
     def _remote(func):
@@ -143,6 +152,8 @@ def wrapper(*args, **kwargs):
                 tags=tags,
                 volume_kms_key=volume_kms_key,
                 volume_size=volume_size,
+                encrypt_inter_container_traffic=encrypt_inter_container_traffic,
+                enable_network_isolation=enable_network_isolation,
             )
             job = _Job.start(job_settings, func, args, kwargs)
 
@@ -331,6 +342,8 @@ def __init__(
         tags: List[Tuple[str, str]] = None,
         volume_kms_key: str = None,
         volume_size: int = 30,
+        encrypt_inter_container_traffic: bool = None,
+        enable_network_isolation: bool = None,
     ):
         """Initiates a ``RemoteExecutor`` instance.
 
@@ -373,6 +386,13 @@ def __init__(
                 instance.
             volume_size (int): Size in GB of the storage volume to use for storing input and output
                 data. Defaults to 30.
+            encrypt_inter_container_traffic (bool): Specifies whether traffic between training
+                containers is encrypted for the training job. (default: ``False``).
+            enable_network_isolation (bool): Specifies whether container will
+                run in network isolation mode (default: ``False``). 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. Also known as Internet-free mode.
         """
         self.max_parallel_jobs = max_parallel_jobs
 
@@ -400,6 +420,8 @@ def __init__(
             tags=tags,
             volume_kms_key=volume_kms_key,
             volume_size=volume_size,
+            encrypt_inter_container_traffic=encrypt_inter_container_traffic,
+            enable_network_isolation=enable_network_isolation,
         )
 
         self._state_condition = threading.Condition()

src/sagemaker/remote_function/job.py

@@ -124,6 +124,8 @@ def __init__(
         tags: List[Tuple[str, str]] = None,
         volume_kms_key: str = None,
         volume_size: int = 30,
+        encrypt_inter_container_traffic: bool = None,
+        enable_network_isolation: bool = None,
     ):
 
         self.sagemaker_config = SageMakerConfigFactory.build_sagemaker_config(
@@ -157,6 +159,14 @@ def __init__(
         self.keep_alive_period_in_seconds = keep_alive_period_in_seconds
         self.job_conda_env = self._get_from_config(job_conda_env, config_schema.JOB_CONDA_ENV)
         self.job_name_prefix = job_name_prefix
+        self.encrypt_inter_container_traffic = self._get_from_config(
+            encrypt_inter_container_traffic,
+            config_schema.ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
+            default=False,
+        )
+        self.enable_network_isolation = self._get_from_config(
+            enable_network_isolation, config_schema.ENABLE_NETWORK_ISOLATION, default=False
+        )
 
         _role = self._get_from_config(role, config_schema.ROLE_ARN)
         if _role:
@@ -195,7 +205,7 @@ def _get_from_config(
         required=False,
     ):
         """Get value from sagemaker config."""
-        if override_value:
+        if override_value is not None:
             return override_value
         config_value = self.sagemaker_config.get_config_value(
             "{}.{}.{}.{}.{}".format(
@@ -206,9 +216,9 @@ def _get_from_config(
                 sagemaker_config_key,
             )
         )
-        if config_value:
+        if config_value is not None:
             return transform(config_value)
-        if not default and required:
+        if default is None and required:
             raise ValueError(f"{sagemaker_config_key} is a required parameter!")
         return default
 
@@ -391,6 +401,14 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info=Non
 
         request_dict["ResourceConfig"] = resource_config
 
+        if job_settings.enable_network_isolation is not None:
+            request_dict["EnableNetworkIsolation"] = job_settings.enable_network_isolation
+
+        if job_settings.encrypt_inter_container_traffic is not None:
+            request_dict[
+                "EnableInterContainerTrafficEncryption"
+            ] = job_settings.encrypt_inter_container_traffic
+
         if job_settings.vpc_config:
             request_dict["VpcConfig"] = job_settings.vpc_config
 

tests/data/remote_function/config.yaml

@@ -11,4 +11,5 @@ SageMaker:
         SecurityGroupIds: ["sg123"]
         Subnets: ["subnet-1234"]
         Tags: [{"someTagKey": "someTagValue"}, {"someTagKey2": "someTagValue2"}]
-        VolumeKmsKeyId: "someVolumeKmsKey"
+        VolumeKmsKeyId: "someVolumeKmsKey"
+        EnableNetworkIsolation: true

tests/unit/sagemaker/remote_function/test_job.py

@@ -149,6 +149,8 @@ def test_sagemaker_config_job_settings_with_configuration_file(
     assert job_settings.volume_kms_key == "someVolumeKmsKey"
     assert job_settings.s3_kms_key == "someS3KmsKey"
     assert job_settings.instance_type == "ml.m5.large"
+    assert job_settings.enable_network_isolation is True
+    assert job_settings.encrypt_inter_container_traffic is False
 
     monkeypatch.delenv("SAGEMAKER_DEFAULT_CONFIG_OVERRIDE")
 
@@ -210,6 +212,7 @@ def test_start(
         role=ROLE_ARN,
         include_local_workdir=True,
         instance_type="ml.m5.large",
+        encrypt_inter_container_traffic=True,
     )
 
     job = _Job.start(job_settings, job_function, func_args=(1, 2), func_kwargs={"c": 3, "d": 4})
@@ -297,6 +300,8 @@ def test_start(
             InstanceType="ml.m5.large",
             KeepAlivePeriodInSeconds=0,
         ),
+        EnableNetworkIsolation=False,
+        EnableInterContainerTrafficEncryption=True,
         Environment={"AWS_DEFAULT_REGION": "us-west-2"},
     )
 
@@ -416,6 +421,8 @@ def test_start_with_complete_job_settings(
             VolumeKmsKeyId=KMS_KEY_ARN,
             KeepAlivePeriodInSeconds=120,
         ),
+        EnableNetworkIsolation=False,
+        EnableInterContainerTrafficEncryption=False,
         VpcConfig=dict(Subnets=["subnet"], SecurityGroupIds=["sg"]),
         Environment={"AWS_DEFAULT_REGION": "us-east-2"},
     )