fix: Update all model server prepare.py to use plain SHA-256

Remove generate_secret_key import and usage from TorchServe, MMS,
TF Serving, and SMD prepare functions. Switch compute_hash calls
from HMAC-SHA256 to plain SHA-256 (no secret_key parameter).

This is required because generate_secret_key was removed from
check_integrity.py in the previous commit. Without this change,
all model server imports fail with ImportError.

Author: Pravali Uppugunduri

sagemaker-serve/src/sagemaker/serve/model_builder_utils.py

@@ -26,7 +26,6 @@
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.detector.dependency_manager import capture_dependencies
 from sagemaker.serve.validations.check_integrity import (
-    generate_secret_key,
     compute_hash,
 )
 from sagemaker.core.remote_function.core.serialization import _MetaData
@@ -119,11 +118,8 @@ def prepare_for_mms(
 
     capture_dependencies(dependencies=dependencies, work_dir=code_dir)
 
-    secret_key = generate_secret_key()
     with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
         buffer = f.read()
-    hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    hash_value = compute_hash(buffer=buffer)
     with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
         metadata.write(_MetaData(hash_value).to_json())
-
-    return secret_key

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

@@ -35,7 +35,6 @@ def _start_serving(
         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:
@@ -47,7 +46,7 @@ def _start_serving(
             image,
             "serve",
             # network_mode="host",
-            ports={'8080/tcp': 8080},
+            ports={"8080/tcp": 8080},
             detach=True,
             auto_remove=True,
             volumes={
@@ -131,7 +130,6 @@ def _upload_server_artifacts(
             env_vars = {
                 "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
                 "SAGEMAKER_PROGRAM": "inference.py",
-                "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
                 "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
                 "SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
                 "LOCAL_PYTHON": platform.python_version(),

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

@@ -12,7 +12,6 @@
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.detector.dependency_manager import capture_dependencies
 from sagemaker.serve.validations.check_integrity import (
-    generate_secret_key,
     compute_hash,
 )
 from sagemaker.core.remote_function.core.serialization import _MetaData
@@ -64,11 +63,8 @@ def prepare_for_smd(
 
     capture_dependencies(dependencies=dependencies, work_dir=code_dir)
 
-    secret_key = generate_secret_key()
     with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
         buffer = f.read()
-    hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    hash_value = compute_hash(buffer=buffer)
     with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
         metadata.write(_MetaData(hash_value).to_json())
-
-    return secret_key

sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py

@@ -53,7 +53,6 @@ def _upload_smd_artifacts(
             "SAGEMAKER_INFERENCE_CODE_DIRECTORY": "/opt/ml/model/code",
             "SAGEMAKER_INFERENCE_CODE": "inference.handler",
             "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
-            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
             "LOCAL_PYTHON": platform.python_version(),
         }
         return s3_upload_path, env_vars

sagemaker-serve/src/sagemaker/serve/model_server/smd/server.py

@@ -38,8 +38,6 @@ def _start_tei_serving(
             secret_key: Secret key to use for authentication
             env_vars: Environment variables to set
         """
-        if env_vars and secret_key:
-            env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = secret_key
 
         self.container = client.containers.run(
             image,

sagemaker-serve/src/sagemaker/serve/model_server/tei/server.py

@@ -11,7 +11,6 @@
 )
 from sagemaker.serve.detector.dependency_manager import capture_dependencies
 from sagemaker.serve.validations.check_integrity import (
-    generate_secret_key,
     compute_hash,
 )
 from sagemaker.core.remote_function.core.serialization import _MetaData
@@ -56,12 +55,9 @@ def prepare_for_tf_serving(
     if not mlflow_saved_model_dir:
         raise ValueError("SavedModel is not found for Tensorflow or Keras flavor.")
     _move_contents(src_dir=mlflow_saved_model_dir, dest_dir=saved_model_bundle_dir)
-    
-    secret_key = generate_secret_key()
+
     with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
         buffer = f.read()
-    hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    hash_value = compute_hash(buffer=buffer)
     with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
         metadata.write(_MetaData(hash_value).to_json())
-
-    return secret_key

sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py

@@ -37,7 +37,7 @@ def _start_tensorflow_serving(
             detach=True,
             auto_remove=False,  # Temporarily disabled to see crash logs
             # network_mode="host",
-            ports={'8501/tcp': 8501},
+            ports={"8501/tcp": 8501},
             volumes={
                 Path(model_path): {
                     "bind": "/opt/ml/model",
@@ -47,7 +47,6 @@ def _start_tensorflow_serving(
             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,
             },
@@ -124,7 +123,6 @@ def _upload_tensorflow_serving_artifacts(
             "SAGEMAKER_PROGRAM": "inference.py",
             "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
             "SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
-            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
             "LOCAL_PYTHON": platform.python_version(),
         }
         return s3_upload_path, env_vars

sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/server.py

@@ -13,7 +13,6 @@
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.detector.dependency_manager import capture_dependencies
 from sagemaker.serve.validations.check_integrity import (
-    generate_secret_key,
     compute_hash,
 )
 from sagemaker.serve.validations.check_image_uri import is_1p_image_uri
@@ -56,7 +55,9 @@ def prepare_for_torchserve(
     # https://github.com/aws/sagemaker-python-sdk/issues/4288
     if is_1p_image_uri(image_uri=image_uri) and "xgboost" in image_uri:
         shutil.copy2(Path(__file__).parent.joinpath("xgboost_inference.py"), code_dir)
-        os.rename(str(code_dir.joinpath("xgboost_inference.py")), str(code_dir.joinpath("inference.py")))
+        os.rename(
+            str(code_dir.joinpath("xgboost_inference.py")), str(code_dir.joinpath("inference.py"))
+        )
     else:
         shutil.copy2(Path(__file__).parent.joinpath("inference.py"), code_dir)
 
@@ -67,11 +68,8 @@ def prepare_for_torchserve(
 
     capture_dependencies(dependencies=dependencies, work_dir=code_dir)
 
-    secret_key = generate_secret_key()
     with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
         buffer = f.read()
-    hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    hash_value = compute_hash(buffer=buffer)
     with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
         metadata.write(_MetaData(hash_value).to_json())
-
-    return secret_key

sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py

@@ -29,7 +29,7 @@ def _start_torch_serve(
             detach=True,
             auto_remove=True,
             # network_mode="host",
-            ports={'8080/tcp': 8080},
+            ports={"8080/tcp": 8080},
             volumes={
                 Path(model_path): {
                     "bind": "/opt/ml/model",
@@ -39,7 +39,6 @@ def _start_torch_serve(
             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,
             },
@@ -103,7 +102,6 @@ def _upload_torchserve_artifacts(
             "SAGEMAKER_PROGRAM": "inference.py",
             "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
             "SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
-            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
             "LOCAL_PYTHON": platform.python_version(),
         }
         return s3_upload_path, env_vars