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_server/multi_model_server/prepare.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,10 @@ 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
+    return ""

sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.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,10 @@ 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
+    return ""

sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.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
@@ -57,11 +56,10 @@ def prepare_for_tf_serving(
         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
+    return ""

sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.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
@@ -67,11 +66,10 @@ 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
+    return ""

sagemaker-serve/tests/unit/test_model_builder_utils_triton.py

@@ -265,7 +265,7 @@ class TestHMACSignin(unittest.TestCase):
     """Test _compute_integrity_hash method."""
 
     def test_compute_integrity_hash(self):
-        """Test HMAC signing."""
+        """Test SHA-256 integrity hash computation."""
         utils = _ModelBuilderUtils()
 
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -278,8 +278,7 @@ def test_compute_integrity_hash(self):
 
             utils._compute_integrity_hash()
 
-            # Secret key is generated, not mocked
-            self.assertIsNotNone(utils.secret_key)
+            # metadata.json should be created with the SHA-256 hash
             self.assertTrue((pkl_path / "metadata.json").exists())
 
 

sagemaker-serve/tests/unit/validations/test_check_integrity.py

@@ -1,50 +1,45 @@
 import unittest
-import tempfile
 from pathlib import Path
 from unittest.mock import patch, mock_open
 from sagemaker.serve.validations.check_integrity import (
-    generate_secret_key,
     compute_hash,
     perform_integrity_check
 )
 
 
 class TestCheckIntegrity(unittest.TestCase):
-    def test_generate_secret_key(self):
-        key = generate_secret_key()
-        self.assertIsInstance(key, str)
-        self.assertEqual(len(key), 64)
-
-    def test_generate_secret_key_custom_bytes(self):
-        key = generate_secret_key(nbytes=16)
-        self.assertEqual(len(key), 32)
-
     def test_compute_hash(self):
         buffer = b"test data"
-        secret_key = "test_secret"
-        hash_value = compute_hash(buffer, secret_key)
+        hash_value = compute_hash(buffer)
         self.assertIsInstance(hash_value, str)
         self.assertEqual(len(hash_value), 64)
 
     def test_compute_hash_consistency(self):
         buffer = b"test data"
-        secret_key = "test_secret"
-        hash1 = compute_hash(buffer, secret_key)
-        hash2 = compute_hash(buffer, secret_key)
+        hash1 = compute_hash(buffer)
+        hash2 = compute_hash(buffer)
         self.assertEqual(hash1, hash2)
 
-    @patch.dict("os.environ", {"SAGEMAKER_SERVE_SECRET_KEY": "test_key"})
+    def test_compute_hash_different_data(self):
+        hash1 = compute_hash(b"data1")
+        hash2 = compute_hash(b"data2")
+        self.assertNotEqual(hash1, hash2)
+
     @patch("pathlib.Path.exists")
     @patch("builtins.open", new_callable=mock_open, read_data=b'{"sha256_hash": "test_hash"}')
     @patch("sagemaker.serve.validations.check_integrity._MetaData.from_json")
     def test_perform_integrity_check_failure(self, mock_metadata, mock_file, mock_exists):
         mock_exists.return_value = True
         mock_meta = type("obj", (object,), {"sha256_hash": "wrong_hash"})()
         mock_metadata.return_value = mock_meta
-        
+
         with self.assertRaises(ValueError):
             perform_integrity_check(b"test", Path("/tmp/metadata.json"))
 
+    def test_perform_integrity_check_missing_metadata(self):
+        with self.assertRaises(ValueError, msg="Path to metadata.json does not exist"):
+            perform_integrity_check(b"test", Path("/nonexistent/metadata.json"))
+
 
 if __name__ == "__main__":
     unittest.main()