Skip to content

Commit 67fc282

Browse files
balajisankar15Balaji Sankar
authored and
Ruban Hussain
committed
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]>
1 parent 67e8d94 commit 67fc282

File tree

107 files changed

+3019
-383
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

107 files changed

+3019
-383
lines changed

‎src/sagemaker/algorithm.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ class AlgorithmEstimator(EstimatorBase):
4646
def __init__(
4747
self,
4848
algorithm_arn: str,
49-
role: str,
49+
role: str = None,
5050
instance_count: Optional[Union[int, PipelineVariable]] = None,
5151
instance_type: Optional[Union[str, PipelineVariable]] = None,
5252
volume_size: Union[int, PipelineVariable] = 30,

‎src/sagemaker/amazon/amazon_estimator.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ class AmazonAlgorithmEstimatorBase(EstimatorBase):
5050

5151
def __init__(
5252
self,
53-
role: str,
53+
role: Optional[Union[str, PipelineVariable]] = None,
5454
instance_count: Optional[Union[int, PipelineVariable]] = None,
5555
instance_type: Optional[Union[str, PipelineVariable]] = None,
5656
data_location: Optional[str] = None,

‎src/sagemaker/amazon/factorization_machines.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ class FactorizationMachines(AmazonAlgorithmEstimatorBase):
8787

8888
def __init__(
8989
self,
90-
role: str,
90+
role: Optional[Union[str, PipelineVariable]] = None,
9191
instance_count: Optional[Union[int, PipelineVariable]] = None,
9292
instance_type: Optional[Union[str, PipelineVariable]] = None,
9393
num_factors: Optional[int] = None,
@@ -326,7 +326,7 @@ class FactorizationMachinesModel(Model):
326326
def __init__(
327327
self,
328328
model_data: Union[str, PipelineVariable],
329-
role: str,
329+
role: Optional[str] = None,
330330
sagemaker_session: Optional[Session] = None,
331331
**kwargs
332332
):

‎src/sagemaker/amazon/ipinsights.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ class IPInsights(AmazonAlgorithmEstimatorBase):
6363

6464
def __init__(
6565
self,
66-
role: str,
66+
role: Optional[Union[str, PipelineVariable]] = None,
6767
instance_count: Optional[Union[int, PipelineVariable]] = None,
6868
instance_type: Optional[Union[str, PipelineVariable]] = None,
6969
num_entity_vectors: Optional[int] = None,
@@ -229,7 +229,7 @@ class IPInsightsModel(Model):
229229
def __init__(
230230
self,
231231
model_data: Union[str, PipelineVariable],
232-
role: str,
232+
role: Optional[str] = None,
233233
sagemaker_session: Optional[Session] = None,
234234
**kwargs
235235
):

‎src/sagemaker/amazon/kmeans.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ class KMeans(AmazonAlgorithmEstimatorBase):
6262

6363
def __init__(
6464
self,
65-
role: str,
65+
role: Optional[Union[str, PipelineVariable]] = None,
6666
instance_count: Optional[Union[int, PipelineVariable]] = None,
6767
instance_type: Optional[Union[str, PipelineVariable]] = None,
6868
k: Optional[int] = None,
@@ -255,7 +255,7 @@ class KMeansModel(Model):
255255
def __init__(
256256
self,
257257
model_data: Union[str, PipelineVariable],
258-
role: str,
258+
role: Optional[str] = None,
259259
sagemaker_session: Optional[Session] = None,
260260
**kwargs
261261
):

‎src/sagemaker/amazon/knn.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ class KNN(AmazonAlgorithmEstimatorBase):
7373

7474
def __init__(
7575
self,
76-
role: str,
76+
role: Optional[Union[str, PipelineVariable]] = None,
7777
instance_count: Optional[Union[int, PipelineVariable]] = None,
7878
instance_type: Optional[Union[str, PipelineVariable]] = None,
7979
k: Optional[int] = None,
@@ -246,7 +246,7 @@ class KNNModel(Model):
246246
def __init__(
247247
self,
248248
model_data: Union[str, PipelineVariable],
249-
role: str,
249+
role: Optional[str] = None,
250250
sagemaker_session: Optional[Session] = None,
251251
**kwargs
252252
):

‎src/sagemaker/amazon/lda.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ class LDA(AmazonAlgorithmEstimatorBase):
5252

5353
def __init__(
5454
self,
55-
role: str,
55+
role: Optional[Union[str, PipelineVariable]] = None,
5656
instance_type: Optional[Union[str, PipelineVariable]] = None,
5757
num_topics: Optional[int] = None,
5858
alpha0: Optional[float] = None,
@@ -230,7 +230,7 @@ class LDAModel(Model):
230230
def __init__(
231231
self,
232232
model_data: Union[str, PipelineVariable],
233-
role: str,
233+
role: Optional[str] = None,
234234
sagemaker_session: Optional[Session] = None,
235235
**kwargs
236236
):

‎src/sagemaker/amazon/linear_learner.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ class LinearLearner(AmazonAlgorithmEstimatorBase):
145145

146146
def __init__(
147147
self,
148-
role: str,
148+
role: Optional[Union[str, PipelineVariable]] = None,
149149
instance_count: Optional[Union[int, PipelineVariable]] = None,
150150
instance_type: Optional[Union[str, PipelineVariable]] = None,
151151
predictor_type: Optional[str] = None,
@@ -499,7 +499,7 @@ class LinearLearnerModel(Model):
499499
def __init__(
500500
self,
501501
model_data: Union[str, PipelineVariable],
502-
role: str,
502+
role: Optional[str] = None,
503503
sagemaker_session: Optional[Session] = None,
504504
**kwargs
505505
):

‎src/sagemaker/amazon/ntm.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ class NTM(AmazonAlgorithmEstimatorBase):
7474

7575
def __init__(
7676
self,
77-
role: str,
77+
role: Optional[Union[str, PipelineVariable]] = None,
7878
instance_count: Optional[Union[int, PipelineVariable]] = None,
7979
instance_type: Optional[Union[str, PipelineVariable]] = None,
8080
num_topics: Optional[int] = None,
@@ -263,7 +263,7 @@ class NTMModel(Model):
263263
def __init__(
264264
self,
265265
model_data: Union[str, PipelineVariable],
266-
role: str,
266+
role: Optional[str] = None,
267267
sagemaker_session: Optional[Session] = None,
268268
**kwargs
269269
):

‎src/sagemaker/amazon/object2vec.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ class Object2Vec(AmazonAlgorithmEstimatorBase):
153153

154154
def __init__(
155155
self,
156-
role: str,
156+
role: Optional[Union[str, PipelineVariable]] = None,
157157
instance_count: Optional[Union[int, PipelineVariable]] = None,
158158
instance_type: Optional[Union[str, PipelineVariable]] = None,
159159
epochs: Optional[int] = None,
@@ -361,7 +361,7 @@ class Object2VecModel(Model):
361361
def __init__(
362362
self,
363363
model_data: Union[str, PipelineVariable],
364-
role: str,
364+
role: Optional[str] = None,
365365
sagemaker_session: Optional[Session] = None,
366366
**kwargs
367367
):

‎src/sagemaker/amazon/pca.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ class PCA(AmazonAlgorithmEstimatorBase):
6060

6161
def __init__(
6262
self,
63-
role: str,
63+
role: Optional[Union[str, PipelineVariable]] = None,
6464
instance_count: Optional[Union[int, PipelineVariable]] = None,
6565
instance_type: Optional[Union[str, PipelineVariable]] = None,
6666
num_components: Optional[int] = None,
@@ -243,7 +243,7 @@ class PCAModel(Model):
243243
def __init__(
244244
self,
245245
model_data: Union[str, PipelineVariable],
246-
role: str,
246+
role: Optional[str] = None,
247247
sagemaker_session: Optional[Session] = None,
248248
**kwargs
249249
):

‎src/sagemaker/amazon/randomcutforest.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ class RandomCutForest(AmazonAlgorithmEstimatorBase):
5454

5555
def __init__(
5656
self,
57-
role: str,
57+
role: Optional[Union[str, PipelineVariable]] = None,
5858
instance_count: Optional[Union[int, PipelineVariable]] = None,
5959
instance_type: Optional[Union[str, PipelineVariable]] = None,
6060
num_samples_per_tree: Optional[int] = None,
@@ -216,7 +216,7 @@ class RandomCutForestModel(Model):
216216
def __init__(
217217
self,
218218
model_data: Union[str, PipelineVariable],
219-
role: str,
219+
role: Optional[str] = None,
220220
sagemaker_session: Optional[Session] = None,
221221
**kwargs
222222
):

‎src/sagemaker/async_inference/async_inference_config.py

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,9 +57,30 @@ def __init__(
5757
"""
5858
self.output_path = output_path
5959
self.max_concurrent_invocations_per_instance = max_concurrent_invocations_per_instance
60-
self.kms_key_id = kms_key_id
60+
self._kms_key_id = kms_key_id
6161
self.notification_config = notification_config
6262

63+
@property
64+
def kms_key_id(self):
65+
"""Getter for kms_key_id
66+
67+
Returns:
68+
str: The KMS Key ID.
69+
"""
70+
return self._kms_key_id
71+
72+
@kms_key_id.setter
73+
def kms_key_id(self, kms_key_id: str):
74+
"""Setter for kms_key_id
75+
76+
Args:
77+
kms_key_id: The new kms_key_id to replace the existing one.
78+
79+
Returns:
80+
81+
"""
82+
self._kms_key_id = kms_key_id
83+
6384
def _to_request_dict(self):
6485
"""Generates a request dictionary using the parameters provided to the class."""
6586
request_dict = {

‎src/sagemaker/automl/automl.py

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,15 @@
1919

2020
from sagemaker import Model, PipelineModel
2121
from sagemaker.automl.candidate_estimator import CandidateEstimator
22-
from sagemaker.config.config_schema import (
22+
from sagemaker.job import _Job
23+
from sagemaker.session import (
24+
Session,
25+
AUTO_ML_KMS_KEY_ID_PATH,
26+
AUTO_ML_ROLE_ARN_PATH,
27+
AUTO_ML_VPC_CONFIG_PATH,
28+
AUTO_ML_VOLUME_KMS_KEY_ID_PATH,
2329
PATH_V1_AUTO_ML_INTER_CONTAINER_ENCRYPTION,
2430
)
25-
from sagemaker.job import _Job
26-
from sagemaker.session import Session
2731
from sagemaker.utils import name_from_base
2832
from sagemaker.workflow.entities import PipelineVariable
2933
from sagemaker.workflow.pipeline_context import runnable_by_pipeline
@@ -101,8 +105,8 @@ class AutoML(object):
101105

102106
def __init__(
103107
self,
104-
role: str,
105-
target_attribute_name: str,
108+
role: Optional[str] = None,
109+
target_attribute_name: str = None,
106110
output_kms_key: Optional[str] = None,
107111
output_path: Optional[str] = None,
108112
base_job_name: Optional[str] = None,
@@ -179,13 +183,10 @@ def __init__(
179183
Returns:
180184
AutoML object.
181185
"""
182-
self.role = role
183-
self.output_kms_key = output_kms_key
184186
self.output_path = output_path
185187
self.base_job_name = base_job_name
186188
self.compression_type = compression_type
187-
self.volume_kms_key = volume_kms_key
188-
self.vpc_config = vpc_config
189+
self.encrypt_inter_container_traffic = encrypt_inter_container_traffic
189190
self.problem_type = problem_type
190191
self.max_candidate = max_candidates
191192
self.max_runtime_per_training_job_in_seconds = max_runtime_per_training_job_in_seconds
@@ -206,6 +207,24 @@ def __init__(
206207
self._auto_ml_job_desc = None
207208
self._best_candidate = None
208209
self.sagemaker_session = sagemaker_session or Session()
210+
self.vpc_config = self.sagemaker_session.get_sagemaker_config_override(
211+
AUTO_ML_VPC_CONFIG_PATH, default_value=vpc_config
212+
)
213+
self.volume_kms_key = self.sagemaker_session.get_sagemaker_config_override(
214+
AUTO_ML_VOLUME_KMS_KEY_ID_PATH, default_value=volume_kms_key
215+
)
216+
self.output_kms_key = self.sagemaker_session.get_sagemaker_config_override(
217+
AUTO_ML_KMS_KEY_ID_PATH, default_value=output_kms_key
218+
)
219+
self.role = self.sagemaker_session.get_sagemaker_config_override(
220+
AUTO_ML_ROLE_ARN_PATH, default_value=role
221+
)
222+
if not self.role:
223+
# Originally IAM role was a required parameter.
224+
# Now we marked that as Optional because we can fetch it from SageMakerConfig
225+
# Because of marking that parameter as optional, we should validate if it is None, even
226+
# after fetching the config.
227+
raise ValueError("IAM role should be provided for creating AutoML jobs.")
209228

210229
self.encrypt_inter_container_traffic = self.sagemaker_session.resolve_value_from_config(
211230
direct_input=encrypt_inter_container_traffic,

‎src/sagemaker/automl/candidate_estimator.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,12 @@
1515

1616
from six import string_types
1717

18-
from sagemaker import Session
19-
from sagemaker.config.config_schema import PATH_V1_TRAINING_JOB_INTER_CONTAINER_ENCRYPTION
18+
from sagemaker.session import (
19+
Session,
20+
TRAINING_JOB_VPC_CONFIG_PATH,
21+
TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH,
22+
PATH_V1_TRAINING_JOB_INTER_CONTAINER_ENCRYPTION,
23+
)
2024
from sagemaker.job import _Job
2125
from sagemaker.utils import name_from_base
2226

@@ -102,7 +106,12 @@ def fit(
102106
"""Logs can only be shown if wait is set to True.
103107
Please either set wait to True or set logs to False."""
104108
)
105-
109+
vpc_config = self.sagemaker_session.get_sagemaker_config_override(
110+
TRAINING_JOB_VPC_CONFIG_PATH, default_value=vpc_config
111+
)
112+
volume_kms_key = self.sagemaker_session.get_sagemaker_config_override(
113+
TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH, default_value=volume_kms_key
114+
)
106115
self.name = candidate_name or self.name
107116
running_jobs = {}
108117

‎src/sagemaker/chainer/model.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,8 +82,8 @@ class ChainerModel(FrameworkModel):
8282
def __init__(
8383
self,
8484
model_data: Union[str, PipelineVariable],
85-
role: str,
86-
entry_point: str,
85+
role: Optional[str] = None,
86+
entry_point: Optional[str] = None,
8787
image_uri: Optional[Union[str, PipelineVariable]] = None,
8888
framework_version: Optional[str] = None,
8989
py_version: Optional[str] = None,

‎src/sagemaker/clarify.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1205,9 +1205,9 @@ class SageMakerClarifyProcessor(Processor):
12051205

12061206
def __init__(
12071207
self,
1208-
role: str,
1209-
instance_count: int,
1210-
instance_type: str,
1208+
role: Optional[str] = None,
1209+
instance_count: int = None,
1210+
instance_type: str = None,
12111211
volume_size_in_gb: int = 30,
12121212
volume_kms_key: Optional[str] = None,
12131213
output_kms_key: Optional[str] = None,

0 commit comments

Comments
 (0)