Skip to content

Commit f1a809c

Browse files
knikureknikure
authored
feature: SDK Defaults - Environment Variables (#3858)
1 parent 14b2af7 commit f1a809c

File tree

15 files changed

+453
-44
lines changed

15 files changed

+453
-44
lines changed

‎src/sagemaker/config/__init__.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,19 @@
2020
TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH,
2121
TRAINING_JOB_ROLE_ARN_PATH,
2222
TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH,
23+
TRAINING_JOB_ENVIRONMENT_PATH,
2324
TRAINING_JOB_VPC_CONFIG_PATH,
2425
TRAINING_JOB_OUTPUT_DATA_CONFIG_PATH,
2526
TRAINING_JOB_RESOURCE_CONFIG_PATH,
2627
PROCESSING_JOB_INPUTS_PATH,
2728
PROCESSING_JOB,
29+
PROCESSING_JOB_ENVIRONMENT_PATH,
2830
PROCESSING_JOB_INTER_CONTAINER_ENCRYPTION_PATH,
2931
PROCESSING_JOB_ROLE_ARN_PATH,
3032
PROCESSING_JOB_NETWORK_CONFIG_PATH,
3133
PROCESSING_OUTPUT_CONFIG_PATH,
3234
PROCESSING_JOB_PROCESSING_RESOURCES_PATH,
35+
MONITORING_JOB_ENVIRONMENT_PATH,
3336
MONITORING_JOB_ROLE_ARN_PATH,
3437
MONITORING_JOB_VOLUME_KMS_KEY_ID_PATH,
3538
MONITORING_JOB_NETWORK_CONFIG_PATH,
@@ -53,13 +56,17 @@
5356
VOLUME_KMS_KEY_ID,
5457
TRANSFORM_JOB_VOLUME_KMS_KEY_ID_PATH,
5558
MODEL,
59+
MODEL_CONTAINERS_PATH,
5660
MODEL_EXECUTION_ROLE_ARN_PATH,
5761
MODEL_ENABLE_NETWORK_ISOLATION_PATH,
5862
MODEL_VPC_CONFIG_PATH,
5963
MODEL_PACKAGE_VALIDATION_ROLE_PATH,
6064
VALIDATION_ROLE,
6165
VALIDATION_PROFILES,
66+
MODEL_PACKAGE_INFERENCE_SPECIFICATION_CONTAINERS_PATH,
6267
MODEL_PACKAGE_VALIDATION_PROFILES_PATH,
68+
MODEL_PRIMARY_CONTAINER_PATH,
69+
MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH,
6370
ENDPOINT_CONFIG_PRODUCTION_VARIANTS_PATH,
6471
KMS_KEY_ID,
6572
ENDPOINT_CONFIG_KMS_KEY_ID_PATH,
@@ -117,6 +124,7 @@
117124
MODEL_PACKAGE,
118125
VALIDATION_SPECIFICATION,
119126
TRANSFORM_JOB_DEFINITION,
127+
TRANSFORM_JOB_ENVIRONMENT_PATH,
120128
TRANSFORM_OUTPUT,
121129
TRANSFORM_RESOURCES,
122130
OFFLINE_STORE_CONFIG,
@@ -138,4 +146,8 @@
138146
DEFAULT_S3_BUCKET,
139147
DEFAULT_S3_OBJECT_KEY_PREFIX,
140148
SESSION,
149+
ENVIRONMENT,
150+
CONTAINERS,
151+
PRIMARY_CONTAINER,
152+
INFERENCE_SPECIFICATION,
141153
)

‎src/sagemaker/config/config_schema.py

Lines changed: 130 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,10 @@
9292
SESSION = "Session"
9393
DEFAULT_S3_BUCKET = "DefaultS3Bucket"
9494
DEFAULT_S3_OBJECT_KEY_PREFIX = "DefaultS3ObjectKeyPrefix"
95+
ENVIRONMENT = "Environment"
96+
CONTAINERS = "Containers"
97+
PRIMARY_CONTAINER = "PrimaryContainer"
98+
INFERENCE_SPECIFICATION = "InferenceSpecification"
9599

96100

97101
def _simple_path(*args: str):
@@ -106,6 +110,7 @@ def _simple_path(*args: str):
106110
)
107111
COMPILATION_JOB_OUTPUT_CONFIG_PATH = _simple_path(SAGEMAKER, COMPILATION_JOB, OUTPUT_CONFIG)
108112
COMPILATION_JOB_ROLE_ARN_PATH = _simple_path(SAGEMAKER, COMPILATION_JOB, ROLE_ARN)
113+
TRAINING_JOB_ENVIRONMENT_PATH = _simple_path(SAGEMAKER, TRAINING_JOB, ENVIRONMENT)
109114
TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH = _simple_path(
110115
SAGEMAKER, TRAINING_JOB, ENABLE_NETWORK_ISOLATION
111116
)
@@ -164,6 +169,7 @@ def _simple_path(*args: str):
164169
MONITORING_JOB_DEFINITION_PREFIX = _simple_path(
165170
SAGEMAKER, MONITORING_SCHEDULE, MONITORING_SCHEDULE_CONFIG, MONITORING_JOB_DEFINITION
166171
)
172+
MONITORING_JOB_ENVIRONMENT_PATH = _simple_path(MONITORING_JOB_DEFINITION_PREFIX, ENVIRONMENT)
167173
MONITORING_JOB_OUTPUT_KMS_KEY_ID_PATH = _simple_path(
168174
MONITORING_JOB_DEFINITION_PREFIX, MONITORING_OUTPUT_CONFIG, KMS_KEY_ID
169175
)
@@ -184,6 +190,7 @@ def _simple_path(*args: str):
184190
MONITORING_JOB_ROLE_ARN_PATH = _simple_path(MONITORING_JOB_DEFINITION_PREFIX, ROLE_ARN)
185191
PIPELINE_ROLE_ARN_PATH = _simple_path(SAGEMAKER, PIPELINE, ROLE_ARN)
186192
PIPELINE_TAGS_PATH = _simple_path(SAGEMAKER, PIPELINE, TAGS)
193+
TRANSFORM_JOB_ENVIRONMENT_PATH = _simple_path(SAGEMAKER, TRANSFORM_JOB, ENVIRONMENT)
187194
TRANSFORM_OUTPUT_KMS_KEY_ID_PATH = _simple_path(
188195
SAGEMAKER, TRANSFORM_JOB, TRANSFORM_OUTPUT, KMS_KEY_ID
189196
)
@@ -196,9 +203,12 @@ def _simple_path(*args: str):
196203
TRANSFORM_JOB_VOLUME_KMS_KEY_ID_PATH = _simple_path(
197204
SAGEMAKER, TRANSFORM_JOB, TRANSFORM_RESOURCES, VOLUME_KMS_KEY_ID
198205
)
206+
MODEL_CONTAINERS_PATH = _simple_path(SAGEMAKER, MODEL, CONTAINERS)
199207
MODEL_VPC_CONFIG_PATH = _simple_path(SAGEMAKER, MODEL, VPC_CONFIG)
200208
MODEL_ENABLE_NETWORK_ISOLATION_PATH = _simple_path(SAGEMAKER, MODEL, ENABLE_NETWORK_ISOLATION)
201209
MODEL_EXECUTION_ROLE_ARN_PATH = _simple_path(SAGEMAKER, MODEL, EXECUTION_ROLE_ARN)
210+
MODEL_PRIMARY_CONTAINER_PATH = _simple_path(SAGEMAKER, MODEL, PRIMARY_CONTAINER)
211+
MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH = _simple_path(MODEL_PRIMARY_CONTAINER_PATH, ENVIRONMENT)
202212
PROCESSING_JOB_ENABLE_NETWORK_ISOLATION_PATH = _simple_path(
203213
SAGEMAKER, PROCESSING_JOB, NETWORK_CONFIG, ENABLE_NETWORK_ISOLATION
204214
)
@@ -229,6 +239,9 @@ def _simple_path(*args: str):
229239
SAGEMAKER, PROCESSING_JOB, PROCESSING_RESOURCES, CLUSTER_CONFIG, VOLUME_KMS_KEY_ID
230240
)
231241
PROCESSING_JOB_ROLE_ARN_PATH = _simple_path(SAGEMAKER, PROCESSING_JOB, ROLE_ARN)
242+
MODEL_PACKAGE_INFERENCE_SPECIFICATION_CONTAINERS_PATH = _simple_path(
243+
SAGEMAKER, MODEL_PACKAGE, INFERENCE_SPECIFICATION, CONTAINERS
244+
)
232245
MODEL_PACKAGE_VALIDATION_ROLE_PATH = _simple_path(
233246
SAGEMAKER, MODEL_PACKAGE, VALIDATION_SPECIFICATION, VALIDATION_ROLE
234247
)
@@ -292,6 +305,7 @@ def _simple_path(*args: str):
292305
SECURITY_CONFIG,
293306
ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
294307
)
308+
PROCESSING_JOB_ENVIRONMENT_PATH = _simple_path(SAGEMAKER, PROCESSING_JOB, ENVIRONMENT)
295309
PROCESSING_JOB_INTER_CONTAINER_ENCRYPTION_PATH = _simple_path(
296310
SAGEMAKER, PROCESSING_JOB, NETWORK_CONFIG, ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION
297311
)
@@ -373,6 +387,9 @@ def _simple_path(*args: str):
373387
TYPE: OBJECT,
374388
ADDITIONAL_PROPERTIES: False,
375389
PROPERTIES: {
390+
ENVIRONMENT: {
391+
"$ref": "#/definitions/environment-Length10240-Properties16",
392+
},
376393
TRANSFORM_OUTPUT: {
377394
TYPE: OBJECT,
378395
ADDITIONAL_PROPERTIES: False,
@@ -463,6 +480,94 @@ def _simple_path(*args: str):
463480
"minLength": 3,
464481
"maxLength": 63,
465482
},
483+
# Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_MonitoringJobDefinition.html#sagemaker-Type-MonitoringJobDefinition-Environment
484+
"environment-Length256-Properties50": {
485+
TYPE: OBJECT,
486+
ADDITIONAL_PROPERTIES: False,
487+
PATTERN_PROPERTIES: {
488+
r"([a-zA-Z_][a-zA-Z0-9_]*){1,256}": {
489+
TYPE: "string",
490+
"pattern": r"[\S\s]*",
491+
"maxLength": 256,
492+
}
493+
},
494+
"maxProperties": 50,
495+
},
496+
# Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTransformJob.html#sagemaker-CreateTransformJob-request-Environment
497+
"environment-Length10240-Properties16": {
498+
TYPE: OBJECT,
499+
ADDITIONAL_PROPERTIES: False,
500+
PATTERN_PROPERTIES: {
501+
r"[a-zA-Z_][a-zA-Z0-9_]{0,1023}": {
502+
TYPE: "string",
503+
"pattern": r"[\S\s]*",
504+
"maxLength": 10240,
505+
}
506+
},
507+
"maxProperties": 16,
508+
},
509+
# Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ContainerDefinition.html#sagemaker-Type-ContainerDefinition-Environment
510+
"environment-Length1024-Properties16": {
511+
TYPE: OBJECT,
512+
ADDITIONAL_PROPERTIES: False,
513+
PATTERN_PROPERTIES: {
514+
r"([a-zA-Z_][a-zA-Z0-9_]*){1,1024}": {
515+
TYPE: "string",
516+
"pattern": r"[\S\s]*",
517+
"maxLength": 1024,
518+
}
519+
},
520+
"maxProperties": 16,
521+
},
522+
# Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateProcessingJob.html#sagemaker-CreateProcessingJob-request-Environment
523+
"environment-Length256-Properties100": {
524+
TYPE: OBJECT,
525+
ADDITIONAL_PROPERTIES: False,
526+
PATTERN_PROPERTIES: {
527+
r"([a-zA-Z_][a-zA-Z0-9_]*){1,256}": {
528+
TYPE: "string",
529+
"pattern": r"[\S\s]*",
530+
"maxLength": 256,
531+
}
532+
},
533+
"maxProperties": 100,
534+
},
535+
# Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTrainingJob.html#sagemaker-CreateTrainingJob-request-Environment
536+
"environment-Length512-Properties48": {
537+
TYPE: OBJECT,
538+
ADDITIONAL_PROPERTIES: False,
539+
PATTERN_PROPERTIES: {
540+
r"([a-zA-Z_][a-zA-Z0-9_]*){1,512}": {
541+
TYPE: "string",
542+
"pattern": r"[\S\s]*",
543+
"maxLength": 512,
544+
}
545+
},
546+
"maxProperties": 48,
547+
},
548+
"container": {
549+
TYPE: OBJECT,
550+
ADDITIONAL_PROPERTIES: False,
551+
PROPERTIES: {
552+
ENVIRONMENT: {"$ref": "#/definitions/environment-Length1024-Properties16"},
553+
},
554+
},
555+
"containers": {
556+
TYPE: "array",
557+
"items": {"$ref": "#/definitions/container"},
558+
# https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ContainerDefinition.html
559+
# According to the API docs, the array can have maximum of 15 items.
560+
"minItems": 0,
561+
"maxItems": 15,
562+
},
563+
"modelPackageContainers": {
564+
TYPE: "array",
565+
"items": {"$ref": "#/definitions/container"},
566+
# https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ModelPackageContainerDefinition.html
567+
# According to the API docs, the array can have minimum of 1 and maximum of 15 items.
568+
"minItems": 1,
569+
"maxItems": 15,
570+
},
466571
},
467572
PROPERTIES: {
468573
SCHEMA_VERSION: {
@@ -595,6 +700,11 @@ def _simple_path(*args: str):
595700
TYPE: OBJECT,
596701
ADDITIONAL_PROPERTIES: False,
597702
PROPERTIES: {
703+
ENVIRONMENT: {
704+
"$ref": (
705+
"#/definitions/environment-Length256-Properties50"
706+
),
707+
},
598708
MONITORING_OUTPUT_CONFIG: {
599709
TYPE: OBJECT,
600710
ADDITIONAL_PROPERTIES: False,
@@ -710,6 +820,7 @@ def _simple_path(*args: str):
710820
ADDITIONAL_PROPERTIES: False,
711821
PROPERTIES: {KMS_KEY_ID: {"$ref": "#/definitions/kmsKeyId"}},
712822
},
823+
ENVIRONMENT: {"$ref": "#/definitions/environment-Length10240-Properties16"},
713824
TRANSFORM_OUTPUT: {
714825
TYPE: OBJECT,
715826
ADDITIONAL_PROPERTIES: False,
@@ -752,22 +863,38 @@ def _simple_path(*args: str):
752863
},
753864
# Model
754865
# https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateModel.html
866+
# Note: It is recommended that customers should use either of
867+
# [Containers] or [PrimaryContainer] but not both
868+
# Since SDK will throw a ValueError as both can't be passed in.
755869
MODEL: {
756870
TYPE: OBJECT,
757871
ADDITIONAL_PROPERTIES: False,
758872
PROPERTIES: {
873+
CONTAINERS: {"$ref": "#/definitions/containers"},
759874
ENABLE_NETWORK_ISOLATION: {TYPE: "boolean"},
760875
EXECUTION_ROLE_ARN: {"$ref": "#/definitions/roleArn"},
876+
PRIMARY_CONTAINER: {"$ref": "#/definitions/container"},
761877
VPC_CONFIG: {"$ref": "#/definitions/vpcConfig"},
762878
TAGS: {"$ref": "#/definitions/tags"},
763879
},
764880
},
765881
# Model Package
766-
# https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateModelPackage.html
882+
# https://docs.aws.amazon.com/sagemaker/latest
883+
# /APIReference/API_CreateModelPackage.html
884+
# At present, we are not supporting [Environment] defaults for
885+
# [AdditionalInfereceSpecifcations] Since this parameter is not supported
886+
# in the SDK, we are not adding it to the config schema as well.
767887
MODEL_PACKAGE: {
768888
TYPE: OBJECT,
769889
ADDITIONAL_PROPERTIES: False,
770890
PROPERTIES: {
891+
INFERENCE_SPECIFICATION: {
892+
TYPE: OBJECT,
893+
ADDITIONAL_PROPERTIES: False,
894+
PROPERTIES: {
895+
CONTAINERS: {"$ref": "#/definitions/modelPackageContainers"}
896+
},
897+
},
771898
VALIDATION_SPECIFICATION: {
772899
TYPE: OBJECT,
773900
ADDITIONAL_PROPERTIES: False,
@@ -792,6 +919,7 @@ def _simple_path(*args: str):
792919
TYPE: OBJECT,
793920
ADDITIONAL_PROPERTIES: False,
794921
PROPERTIES: {
922+
ENVIRONMENT: {"$ref": "#/definitions/environment-Length256-Properties100"},
795923
NETWORK_CONFIG: {
796924
TYPE: OBJECT,
797925
ADDITIONAL_PROPERTIES: False,
@@ -837,6 +965,7 @@ def _simple_path(*args: str):
837965
PROPERTIES: {
838966
ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION: {TYPE: "boolean"},
839967
ENABLE_NETWORK_ISOLATION: {TYPE: "boolean"},
968+
ENVIRONMENT: {"$ref": "#/definitions/environment-Length512-Properties48"},
840969
OUTPUT_DATA_CONFIG: {
841970
TYPE: OBJECT,
842971
ADDITIONAL_PROPERTIES: False,

‎src/sagemaker/estimator.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
TRAINING_JOB_KMS_KEY_ID_PATH,
3737
TRAINING_JOB_ROLE_ARN_PATH,
3838
TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH,
39+
TRAINING_JOB_ENVIRONMENT_PATH,
3940
TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH,
4041
)
4142
from sagemaker.debugger import ( # noqa: F401 # pylint: disable=unused-import
@@ -677,7 +678,12 @@ def __init__(
677678
self.profiler_config = profiler_config
678679
self.disable_profiler = disable_profiler
679680

680-
self.environment = environment
681+
self.environment = resolve_value_from_config(
682+
direct_input=environment,
683+
config_path=TRAINING_JOB_ENVIRONMENT_PATH,
684+
default_value=None,
685+
sagemaker_session=self.sagemaker_session,
686+
)
681687

682688
self.max_retry_attempts = max_retry_attempts
683689

‎src/sagemaker/model.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,11 @@
3333
COMPILATION_JOB_ROLE_ARN_PATH,
3434
EDGE_PACKAGING_KMS_KEY_ID_PATH,
3535
EDGE_PACKAGING_ROLE_ARN_PATH,
36+
MODEL_CONTAINERS_PATH,
3637
MODEL_VPC_CONFIG_PATH,
3738
MODEL_ENABLE_NETWORK_ISOLATION_PATH,
3839
MODEL_EXECUTION_ROLE_ARN_PATH,
40+
MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH,
3941
ENDPOINT_CONFIG_ASYNC_KMS_KEY_ID_PATH,
4042
)
4143
from sagemaker.session import Session
@@ -53,6 +55,7 @@
5355
update_container_with_inference_params,
5456
to_string,
5557
resolve_value_from_config,
58+
resolve_nested_dict_value_from_config,
5659
)
5760
from sagemaker.async_inference import AsyncInferenceConfig
5861
from sagemaker.predictor_async import AsyncPredictor
@@ -292,7 +295,6 @@ def __init__(
292295
self.model_data = model_data
293296
self.image_uri = image_uri
294297
self.predictor_cls = predictor_cls
295-
self.env = env or {}
296298
self.name = name
297299
self._base_name = None
298300
self.sagemaker_session = sagemaker_session
@@ -313,6 +315,12 @@ def __init__(
313315
MODEL_ENABLE_NETWORK_ISOLATION_PATH,
314316
sagemaker_session=self.sagemaker_session,
315317
)
318+
self.env = resolve_value_from_config(
319+
env,
320+
MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH,
321+
default_value={},
322+
sagemaker_session=self.sagemaker_session,
323+
)
316324
if self._enable_network_isolation is None:
317325
self._enable_network_isolation = False
318326
self.model_kms_key = model_kms_key
@@ -735,6 +743,12 @@ def _create_sagemaker_model(
735743
MODEL_ENABLE_NETWORK_ISOLATION_PATH,
736744
sagemaker_session=self.sagemaker_session,
737745
)
746+
self.env = resolve_nested_dict_value_from_config(
747+
self.env,
748+
["Environment"],
749+
MODEL_CONTAINERS_PATH,
750+
sagemaker_session=self.sagemaker_session,
751+
)
738752
create_model_args = dict(
739753
name=self.name,
740754
role=self.role,

0 commit comments

Comments
 (0)