Changes to environment, avoid duplicates

Author: Bryannah Hernandez

src/sagemaker/serve/builder/transformers_builder.py

@@ -16,9 +16,8 @@
 import os
 from abc import ABC, abstractmethod
 from typing import Type
-from packaging.version import Version
-
 from pathlib import Path
+from packaging.version import Version
 
 from sagemaker.model import Model
 from sagemaker import image_uris
@@ -356,7 +355,7 @@ def _build_for_transformers(self):
         # save the model or inference spec in cloud pickle format
         if self.inference_spec:
             save_pkl(code_path, (self.inference_spec, self.schema_builder))
-            logger.info("PKL file saved to file: {}".format(code_path))
+            logger.info("PKL file saved to file: %s", code_path)
 
         self._auto_detect_container()
 

src/sagemaker/serve/model_server/multi_model_server/prepare.py

@@ -85,6 +85,7 @@ def prepare_for_mms(
     inference_spec: InferenceSpec = None,
 ) -> str:
     """Prepares for InferenceSpec using model_path, writes inference.py, and captures dependencies to generate secret_key.
+
     Args:to
         model_path (str) : Argument
         shared_libs (List[]) : Argument

src/sagemaker/serve/model_server/multi_model_server/server.py

@@ -31,6 +31,17 @@ def _start_serving(
         env_vars: dict,
     ):
         """Placeholder docstring"""
+        env = {
+            "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+            "SAGEMAKER_PROGRAM": "inference.py",
+            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
+            "LOCAL_PYTHON": platform.python_version(),
+        }
+        if env_vars:
+            env_vars.update(env)
+        else:
+            env_vars = env
+
         self.container = client.containers.run(
             image,
             "serve",
@@ -43,13 +54,7 @@ def _start_serving(
                     "mode": "rw",
                 },
             },
-            environment={
-                "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
-                "SAGEMAKER_PROGRAM": "inference.py",
-                "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
-                "LOCAL_PYTHON": platform.python_version(),
-                **env_vars,
-            },
+            environment=env_vars,
         )
 
     def _invoke_multi_model_server_serving(self, request: object, content_type: str, accept: str):

tests/unit/sagemaker/serve/model_server/multi_model_server/test_multi_model_server_prepare.py

@@ -25,7 +25,10 @@
     LocalMultiModelServer,
 )
 
-CPU_TF_IMAGE = "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04"
+CPU_TF_IMAGE = (
+    "763104351884.dkr.ecr.us-east-1.amazonaws.com/"
+    "huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04"
+)
 MODEL_PATH = "model_path"
 MODEL_REPO = f"{MODEL_PATH}/1"
 ENV_VAR = {"KEY": "VALUE"}
@@ -57,16 +60,16 @@ def test_start_invoke_destroy_local_multi_model_server(self):
         mock_docker_client.containers.run.assert_called_once_with(
             CPU_TF_IMAGE,
             "serve",
+            network_mode="host",
             detach=True,
             auto_remove=True,
-            network_mode="host",
-            volumes={PosixPath("model_path"): {"bind": "/opt/ml/model", "mode": "rw"}},
+            volumes={PosixPath("model_path/code"): {"bind": "/opt/ml/model/", "mode": "rw"}},
             environment={
+                "KEY": "VALUE",
                 "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
                 "SAGEMAKER_PROGRAM": "inference.py",
                 "SAGEMAKER_SERVE_SECRET_KEY": "secret_key",
                 "LOCAL_PYTHON": platform.python_version(),
-                "KEY": "VALUE",
             },
         )