Refactoring

Author: Jonathan Makunga

src/sagemaker/serve/mode/sagemaker_endpoint_mode.py

@@ -68,8 +68,9 @@ def prepare(
                 + "session to be created or supply `sagemaker_session` into @serve.invoke."
             ) from e
 
+        upload_artifacts = None
         if self.model_server == ModelServer.TORCHSERVE:
-            return self._upload_torchserve_artifacts(
+            upload_artifacts = self._upload_torchserve_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 secret_key=secret_key,
@@ -78,7 +79,7 @@ def prepare(
             )
 
         if self.model_server == ModelServer.TRITON:
-            return self._upload_triton_artifacts(
+            upload_artifacts = self._upload_triton_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 secret_key=secret_key,
@@ -87,15 +88,15 @@ def prepare(
             )
 
         if self.model_server == ModelServer.DJL_SERVING:
-            return self._upload_djl_artifacts(
+            upload_artifacts = self._upload_djl_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
             )
 
         if self.model_server == ModelServer.TGI:
-            return self._upload_tgi_artifacts(
+            upload_artifacts = self._upload_tgi_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
@@ -104,15 +105,15 @@ def prepare(
             )
 
         if self.model_server == ModelServer.MMS:
-            return self._upload_server_artifacts(
+            upload_artifacts = self._upload_server_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
             )
 
         if self.model_server == ModelServer.TENSORFLOW_SERVING:
-            return self._upload_tensorflow_serving_artifacts(
+            upload_artifacts = self._upload_tensorflow_serving_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 secret_key=secret_key,
@@ -121,11 +122,14 @@ def prepare(
             )
 
         if self.model_server == ModelServer.TEI:
-            return self._upload_tei_artifacts(
+            upload_artifacts = self._upload_tei_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
             )
 
+        if isinstance(self.model_server, ModelServer) and upload_artifacts:
+            return upload_artifacts
+
         raise ValueError("%s model server is not supported" % self.model_server)

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

@@ -40,7 +40,7 @@ def _start_tei_serving(
             env_vars: Environment variables to set
         """
         if env_vars and secret_key:
-            env_vars['SAGEMAKER_SERVE_SECRET_KEY'] = secret_key
+            env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = secret_key
 
         self.container = client.containers.run(
             image,