fix: add environment key to model container only if it is not empty

CHANGELOG.md

@@ -977,7 +977,7 @@
  * default repack encryption
  * support large pipeline
  * add support for pytorch 1.10.0
- 
+
 ### Documentation Changes
 
  * SageMaker model parallel library 1.6.0 API doc

src/sagemaker/image_uri_config/autogluon.json

@@ -489,12 +489,15 @@
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
                     "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "me-south-1": "217643126080",
+                    "me-central-1": "914824155844",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",

src/sagemaker/local/entities.py

@@ -390,7 +390,7 @@ def _get_container_environment(self, **kwargs):
             container
         """
         environment = {}
-        environment.update(self.primary_container["Environment"])
+        environment.update(self.primary_container.get("Environment", {}))
         environment["SAGEMAKER_BATCH"] = "True"
         if "MaxPayloadInMB" in kwargs:
             environment["SAGEMAKER_MAX_PAYLOAD_IN_MB"] = str(kwargs["MaxPayloadInMB"])
@@ -591,18 +591,15 @@ def serve(self):
         instance_count = self.production_variant["InitialInstanceCount"]
 
         accelerator_type = self.production_variant.get("AcceleratorType")
+        environment = self.primary_container.get("Environment", {})
         if accelerator_type == "local_sagemaker_notebook":
-            self.primary_container["Environment"][
-                "SAGEMAKER_INFERENCE_ACCELERATOR_PRESENT"
-            ] = "true"
+            environment["SAGEMAKER_INFERENCE_ACCELERATOR_PRESENT"] = "true"
 
         self.create_time = datetime.datetime.now()
         self.container = _SageMakerContainer(
             instance_type, instance_count, image, self.local_session
         )
-        self.container.serve(
-            self.primary_container["ModelDataUrl"], self.primary_container["Environment"]
-        )
+        self.container.serve(self.primary_container["ModelDataUrl"], environment)
 
         serving_port = get_config_value("local.serving_port", self.local_session.config) or 8080
         _wait_for_serving_container(serving_port)

src/sagemaker/multidatamodel.py

@@ -139,7 +139,7 @@ def prepare_container_def(
         if self.model:
             container_definition = self.model.prepare_container_def(instance_type, accelerator_type)
             image_uri = container_definition["Image"]
-            environment = container_definition["Environment"]
+            environment = container_definition.get("Environment", {})
         else:
             image_uri = self.image_uri
             environment = self.env

src/sagemaker/session.py

@@ -5153,9 +5153,9 @@ def container_def(image_uri, model_data_url=None, env=None, container_mode=None,
         dict[str, str]: A complete container definition object usable with the CreateModel API if
         passed via `PrimaryContainers` field.
     """
-    if env is None:
-        env = {}
-    c_def = {"Image": image_uri, "Environment": env}
+    c_def = {"Image": image_uri}
+    if env:
+        c_def["Environment"] = env
     if model_data_url:
         c_def["ModelDataUrl"] = model_data_url
     if container_mode:

tests/unit/sagemaker/model/test_model.py

@@ -134,7 +134,7 @@ def test_prepare_container_def_with_model_data():
     model = Model(MODEL_IMAGE)
     container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")
 
-    expected = {"Image": MODEL_IMAGE, "Environment": {}}
+    expected = {"Image": MODEL_IMAGE}
     assert expected == container_def
 
 
@@ -158,7 +158,6 @@ def test_prepare_container_def_with_image_config():
     expected = {
         "Image": MODEL_IMAGE,
         "ImageConfig": {"RepositoryAccessMode": "Vpc"},
-        "Environment": {},
     }
 
     container_def = model.prepare_container_def()

tests/unit/sagemaker/tensorflow/test_tfs.py

@@ -89,7 +89,7 @@ def test_tfs_model(retrieve_image_uri, sagemaker_session, tensorflow_inference_v
         serverless_inference_config=None,
     )
     assert IMAGE == cdef["Image"]
-    assert {} == cdef["Environment"]
+    assert cdef.get("Environment") is None
 
     predictor = model.deploy(INSTANCE_COUNT, INSTANCE_TYPE)
     assert isinstance(predictor, TensorFlowPredictor)
@@ -485,7 +485,6 @@ def test_register_tfs_model_auto_infer_framework(sagemaker_session, tensorflow_i
         "containers": [
             {
                 "Image": image_uri,
-                "Environment": ANY,
                 "ModelDataUrl": ANY,
                 "Framework": "TENSORFLOW",
                 "FrameworkVersion": tensorflow_inference_version,

tests/unit/sagemaker/workflow/test_airflow.py

@@ -1014,7 +1014,6 @@ def test_amazon_alg_model_config(sagemaker_session):
         "ModelName": "pca-%s" % TIME_STAMP,
         "PrimaryContainer": {
             "Image": "174872318107.dkr.ecr.us-west-2.amazonaws.com/pca:1",
-            "Environment": {},
             "ModelDataUrl": "{{ model_data }}",
         },
         "ExecutionRoleArn": "{{ role }}",
@@ -1108,7 +1107,6 @@ def test_model_config_from_amazon_alg_estimator(sagemaker_session):
         "ModelName": "knn-%s" % TIME_STAMP,
         "PrimaryContainer": {
             "Image": "174872318107.dkr.ecr.us-west-2.amazonaws.com/knn:1",
-            "Environment": {},
             "ModelDataUrl": "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Tuning']['BestTrainingJob']"
             "['TrainingJobName'] }}/output/model.tar.gz",
         },
@@ -1309,7 +1307,6 @@ def test_transform_config_from_amazon_alg_estimator(sagemaker_session):
             "ModelName": "knn-%s" % TIME_STAMP,
             "PrimaryContainer": {
                 "Image": "174872318107.dkr.ecr.us-west-2.amazonaws.com/knn:1",
-                "Environment": {},
                 "ModelDataUrl": "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']['TrainingJobName'] }}"
                 "/output/model.tar.gz",
             },
@@ -1413,7 +1410,6 @@ def test_deploy_amazon_alg_model_config(sagemaker_session):
             "ModelName": "pca-%s" % TIME_STAMP,
             "PrimaryContainer": {
                 "Image": "174872318107.dkr.ecr.us-west-2.amazonaws.com/pca:1",
-                "Environment": {},
                 "ModelDataUrl": "{{ model_data }}",
             },
             "ExecutionRoleArn": "{{ role }}",
@@ -1549,7 +1545,6 @@ def test_deploy_config_from_amazon_alg_estimator(sagemaker_session):
             "ModelName": "knn-%s" % TIME_STAMP,
             "PrimaryContainer": {
                 "Image": "174872318107.dkr.ecr.us-west-2.amazonaws.com/knn:1",
-                "Environment": {},
                 "ModelDataUrl": "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Tuning']['BestTrainingJob']"
                 "['TrainingJobName'] }}/output/model.tar.gz",
             },

tests/unit/sagemaker/workflow/test_model_step.py

@@ -621,7 +621,7 @@ def test_conditional_model_create_and_regis(
             container = arguments["PrimaryContainer"]
             assert container["Image"] == _IMAGE_URI
             assert container["ModelDataUrl"] == {"Get": "Parameters.ModelData"}
-            assert not container.get("Environment", {})
+            assert container.get("Environment") is None
         else:
             raise Exception("A step exists in the collection of an invalid type.")
     adjacency_list = PipelineGraph.from_pipeline(pipeline).adjacency_list
@@ -913,7 +913,7 @@ def _verify_register_model_container_definition(
     containers = request["InferenceSpecification"]["Containers"]
     assert len(containers) == 1
     isinstance(containers[0].pop("ModelDataUrl"), expected_model_data_type)
-    container_env = containers[0]["Environment"]
+    container_env = containers[0].get("Environment", {})
     assert container_env.pop(_SAGEMAKER_PROGRAM, None) == expected_program
     submit_dir = container_env.pop(_SAGEMAKER_SUBMIT_DIRECTORY, None)
     if submit_dir and not submit_dir.startswith("s3://"):

tests/unit/sagemaker/workflow/test_step_collections.py

@@ -1176,7 +1176,6 @@ def test_estimator_transformer(estimator):
                 "Arguments": {
                     "ExecutionRoleArn": "DummyRole",
                     "PrimaryContainer": {
-                        "Environment": {},
                         "Image": "fakeimage",
                         "ModelDataUrl": "s3://my-bucket/model.tar.gz",
                     },
@@ -1292,7 +1291,6 @@ def test_estimator_transformer_with_model_repack_with_estimator(estimator):
             assert arguments == {
                 "ExecutionRoleArn": "DummyRole",
                 "PrimaryContainer": {
-                    "Environment": {},
                     "Image": "fakeimage",
                 },
             }

tests/unit/sagemaker/workflow/test_steps.py

@@ -803,7 +803,7 @@ def test_create_model_step(sagemaker_session):
         "DependsOn": ["TestStep", "SecondTestStep"],
         "Arguments": {
             "ExecutionRoleArn": "DummyRole",
-            "PrimaryContainer": {"Environment": {}, "Image": "fakeimage"},
+            "PrimaryContainer": {"Image": "fakeimage"},
         },
     }
     assert step.properties.ModelName.expr == {"Get": "Steps.MyCreateModelStep.ModelName"}

tests/unit/sagemaker/workflow/test_training_step.py

@@ -778,6 +778,7 @@ def test_training_step_with_algorithm_base_local_code(
     # test idempotency
     step_def2 = json.loads(pipeline.definition())["Steps"][0]
     del step_def2["Arguments"]["InputDataConfig"][0]["DataSource"]["S3DataSource"]["S3Uri"]
+
     assert step_def == step_def2
 
 

tests/unit/test_create_deploy_entities.py

@@ -23,7 +23,7 @@
 ROLE = "myimrole"
 EXPANDED_ROLE = "arn:aws:iam::111111111111:role/ExpandedRole"
 IMAGE = "myimage"
-FULL_CONTAINER_DEF = {"Environment": {}, "Image": IMAGE, "ModelDataUrl": "s3://mybucket/mymodel"}
+FULL_CONTAINER_DEF = {"Image": IMAGE, "ModelDataUrl": "s3://mybucket/mymodel"}
 VPC_CONFIG = {"Subnets": ["subnet-foo"], "SecurityGroups": ["sg-foo"]}
 INITIAL_INSTANCE_COUNT = 1
 INSTANCE_TYPE = "ml.c4.xlarge"
@@ -57,7 +57,7 @@ def test_create_model_expand_primary_container(sagemaker_session):
     sagemaker_session.create_model(name=MODEL_NAME, role=ROLE, container_defs=IMAGE)
 
     _1, _2, create_model_kwargs = sagemaker_session.sagemaker_client.create_model.mock_calls[0]
-    assert create_model_kwargs["PrimaryContainer"] == {"Environment": {}, "Image": IMAGE}
+    assert create_model_kwargs["PrimaryContainer"] == {"Image": IMAGE}
 
 
 def test_create_endpoint_config(sagemaker_session):

tests/unit/test_endpoint_from_model_data.py

@@ -25,7 +25,7 @@
 ACCELERATOR_TYPE = "ml.eia.medium"
 S3_MODEL_ARTIFACTS = "s3://mybucket/mymodel"
 DEPLOY_IMAGE = "mydeployimage"
-CONTAINER_DEF = {"Environment": {}, "Image": DEPLOY_IMAGE, "ModelDataUrl": S3_MODEL_ARTIFACTS}
+CONTAINER_DEF = {"Image": DEPLOY_IMAGE, "ModelDataUrl": S3_MODEL_ARTIFACTS}
 VPC_CONFIG = {"Subnets": ["foo"], "SecurityGroupIds": ["bar"]}
 DEPLOY_ROLE = "mydeployrole"
 ENV_VARS = {"PYTHONUNBUFFERED": "TRUE", "some": "nonsense"}

tests/unit/test_estimator.py

@@ -2780,7 +2780,6 @@ def test_fit_deploy_tags_in_estimator(name_from_base, sagemaker_session):
         role="DummyRole",
         container_defs={
             "ModelDataUrl": "s3://bucket/model.tar.gz",
-            "Environment": {},
             "Image": "fakeimage",
         },
         enable_network_isolation=False,
@@ -2834,7 +2833,6 @@ def test_fit_deploy_tags(name_from_base, sagemaker_session):
         role="DummyRole",
         container_defs={
             "ModelDataUrl": "s3://bucket/model.tar.gz",
-            "Environment": {},
             "Image": "fakeimage",
         },
         enable_network_isolation=False,

tests/unit/test_pipeline_model.py

@@ -339,7 +339,6 @@ def test_network_isolation(tfo, time, sagemaker_session):
             },
             {
                 "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:3.3",
-                "Environment": {},
                 "ModelDataUrl": "s3://bucket/model_2.tar.gz",
             },
         ],

tests/unit/test_session.py

@@ -1753,7 +1753,6 @@ def test_logs_for_transform_job_full_lifecycle(time, cw, sagemaker_session_full_
 
 MODEL_NAME = "some-model"
 PRIMARY_CONTAINER = {
-    "Environment": {},
     "Image": IMAGE,
     "ModelDataUrl": "s3://sagemaker-123/output/jobname/model/model.tar.gz",
 }