feature: Added support for VPC Config, EnableNetworkIsolation, KMS Key ID, Volume KMS Key ID, IAM role to be fetched from Config (aws#846)

Co-authored-by: Balaji Sankar <[email protected]>

Author: 2 people

src/sagemaker/algorithm.py

@@ -46,7 +46,7 @@ class AlgorithmEstimator(EstimatorBase):
     def __init__(
         self,
         algorithm_arn: str,
-        role: str,
+        role: str = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         volume_size: Union[int, PipelineVariable] = 30,

src/sagemaker/amazon/amazon_estimator.py

@@ -50,7 +50,7 @@ class AmazonAlgorithmEstimatorBase(EstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         data_location: Optional[str] = None,

src/sagemaker/amazon/factorization_machines.py

@@ -87,7 +87,7 @@ class FactorizationMachines(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         num_factors: Optional[int] = None,
@@ -326,7 +326,7 @@ class FactorizationMachinesModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/amazon/ipinsights.py

@@ -63,7 +63,7 @@ class IPInsights(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         num_entity_vectors: Optional[int] = None,
@@ -229,7 +229,7 @@ class IPInsightsModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/amazon/kmeans.py

@@ -62,7 +62,7 @@ class KMeans(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         k: Optional[int] = None,
@@ -255,7 +255,7 @@ class KMeansModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/amazon/knn.py

@@ -73,7 +73,7 @@ class KNN(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         k: Optional[int] = None,
@@ -246,7 +246,7 @@ class KNNModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/amazon/lda.py

@@ -52,7 +52,7 @@ class LDA(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         num_topics: Optional[int] = None,
         alpha0: Optional[float] = None,
@@ -230,7 +230,7 @@ class LDAModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/amazon/linear_learner.py

@@ -145,7 +145,7 @@ class LinearLearner(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         predictor_type: Optional[str] = None,
@@ -499,7 +499,7 @@ class LinearLearnerModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/amazon/ntm.py

@@ -74,7 +74,7 @@ class NTM(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         num_topics: Optional[int] = None,
@@ -263,7 +263,7 @@ class NTMModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/amazon/object2vec.py

@@ -153,7 +153,7 @@ class Object2Vec(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         epochs: Optional[int] = None,
@@ -361,7 +361,7 @@ class Object2VecModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/amazon/pca.py

@@ -60,7 +60,7 @@ class PCA(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         num_components: Optional[int] = None,
@@ -243,7 +243,7 @@ class PCAModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/amazon/randomcutforest.py

@@ -54,7 +54,7 @@ class RandomCutForest(AmazonAlgorithmEstimatorBase):
 
     def __init__(
         self,
-        role: str,
+        role: Optional[Union[str, PipelineVariable]] = None,
         instance_count: Optional[Union[int, PipelineVariable]] = None,
         instance_type: Optional[Union[str, PipelineVariable]] = None,
         num_samples_per_tree: Optional[int] = None,
@@ -216,7 +216,7 @@ class RandomCutForestModel(Model):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
+        role: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         **kwargs
     ):

src/sagemaker/async_inference/async_inference_config.py

@@ -57,9 +57,30 @@ def __init__(
         """
         self.output_path = output_path
         self.max_concurrent_invocations_per_instance = max_concurrent_invocations_per_instance
-        self.kms_key_id = kms_key_id
+        self._kms_key_id = kms_key_id
         self.notification_config = notification_config
 
+    @property
+    def kms_key_id(self):
+        """Getter for kms_key_id
+
+        Returns:
+            str: The KMS Key ID.
+        """
+        return self._kms_key_id
+
+    @kms_key_id.setter
+    def kms_key_id(self, kms_key_id: str):
+        """Setter for kms_key_id
+
+        Args:
+            kms_key_id: The new kms_key_id to replace the existing one.
+
+        Returns:
+
+        """
+        self._kms_key_id = kms_key_id
+
     def _to_request_dict(self):
         """Generates a request dictionary using the parameters provided to the class."""
         request_dict = {

src/sagemaker/automl/automl.py

@@ -19,11 +19,15 @@
 
 from sagemaker import Model, PipelineModel
 from sagemaker.automl.candidate_estimator import CandidateEstimator
-from sagemaker.config.config_schema import (
+from sagemaker.job import _Job
+from sagemaker.session import (
+    Session,
+    AUTO_ML_KMS_KEY_ID_PATH,
+    AUTO_ML_ROLE_ARN_PATH,
+    AUTO_ML_VPC_CONFIG_PATH,
+    AUTO_ML_VOLUME_KMS_KEY_ID_PATH,
     PATH_V1_AUTO_ML_INTER_CONTAINER_ENCRYPTION,
 )
-from sagemaker.job import _Job
-from sagemaker.session import Session
 from sagemaker.utils import name_from_base
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline
@@ -101,8 +105,8 @@ class AutoML(object):
 
     def __init__(
         self,
-        role: str,
-        target_attribute_name: str,
+        role: Optional[str] = None,
+        target_attribute_name: str = None,
         output_kms_key: Optional[str] = None,
         output_path: Optional[str] = None,
         base_job_name: Optional[str] = None,
@@ -179,13 +183,10 @@ def __init__(
         Returns:
             AutoML object.
         """
-        self.role = role
-        self.output_kms_key = output_kms_key
         self.output_path = output_path
         self.base_job_name = base_job_name
         self.compression_type = compression_type
-        self.volume_kms_key = volume_kms_key
-        self.vpc_config = vpc_config
+        self.encrypt_inter_container_traffic = encrypt_inter_container_traffic
         self.problem_type = problem_type
         self.max_candidate = max_candidates
         self.max_runtime_per_training_job_in_seconds = max_runtime_per_training_job_in_seconds
@@ -206,6 +207,24 @@ def __init__(
         self._auto_ml_job_desc = None
         self._best_candidate = None
         self.sagemaker_session = sagemaker_session or Session()
+        self.vpc_config = self.sagemaker_session.get_sagemaker_config_override(
+            AUTO_ML_VPC_CONFIG_PATH, default_value=vpc_config
+        )
+        self.volume_kms_key = self.sagemaker_session.get_sagemaker_config_override(
+            AUTO_ML_VOLUME_KMS_KEY_ID_PATH, default_value=volume_kms_key
+        )
+        self.output_kms_key = self.sagemaker_session.get_sagemaker_config_override(
+            AUTO_ML_KMS_KEY_ID_PATH, default_value=output_kms_key
+        )
+        self.role = self.sagemaker_session.get_sagemaker_config_override(
+            AUTO_ML_ROLE_ARN_PATH, default_value=role
+        )
+        if not self.role:
+            # Originally IAM role was a required parameter.
+            # Now we marked that as Optional because we can fetch it from SageMakerConfig
+            # Because of marking that parameter as optional, we should validate if it is None, even
+            # after fetching the config.
+            raise ValueError("IAM role should be provided for creating AutoML jobs.")
 
         self.encrypt_inter_container_traffic = self.sagemaker_session.resolve_value_from_config(
             direct_input=encrypt_inter_container_traffic,

src/sagemaker/automl/candidate_estimator.py

@@ -15,8 +15,12 @@
 
 from six import string_types
 
-from sagemaker import Session
-from sagemaker.config.config_schema import PATH_V1_TRAINING_JOB_INTER_CONTAINER_ENCRYPTION
+from sagemaker.session import (
+    Session,
+    TRAINING_JOB_VPC_CONFIG_PATH,
+    TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH,
+    PATH_V1_TRAINING_JOB_INTER_CONTAINER_ENCRYPTION,
+)
 from sagemaker.job import _Job
 from sagemaker.utils import name_from_base
 
@@ -102,7 +106,12 @@ def fit(
                 """Logs can only be shown if wait is set to True.
                 Please either set wait to True or set logs to False."""
             )
-
+        vpc_config = self.sagemaker_session.get_sagemaker_config_override(
+            TRAINING_JOB_VPC_CONFIG_PATH, default_value=vpc_config
+        )
+        volume_kms_key = self.sagemaker_session.get_sagemaker_config_override(
+            TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH, default_value=volume_kms_key
+        )
         self.name = candidate_name or self.name
         running_jobs = {}
 

src/sagemaker/chainer/model.py

@@ -82,8 +82,8 @@ class ChainerModel(FrameworkModel):
     def __init__(
         self,
         model_data: Union[str, PipelineVariable],
-        role: str,
-        entry_point: str,
+        role: Optional[str] = None,
+        entry_point: Optional[str] = None,
         image_uri: Optional[Union[str, PipelineVariable]] = None,
         framework_version: Optional[str] = None,
         py_version: Optional[str] = None,

src/sagemaker/clarify.py

@@ -1205,9 +1205,9 @@ class SageMakerClarifyProcessor(Processor):
 
     def __init__(
         self,
-        role: str,
-        instance_count: int,
-        instance_type: str,
+        role: Optional[str] = None,
+        instance_count: int = None,
+        instance_type: str = None,
         volume_size_in_gb: int = 30,
         volume_kms_key: Optional[str] = None,
         output_kms_key: Optional[str] = None,