Fix: Triton inference spec hmac exposure (aws#5654)

* fix: Remove hardcoded secret key from Triton ONNX export path

The ONNX export path in _prepare_for_triton() set self.secret_key to a
hardcoded value 'dummy secret key for onnx backend'. This key was then
passed as SAGEMAKER_SERVE_SECRET_KEY into container environment variables
and exposed in plaintext via DescribeModel/DescribeEndpointConfig APIs.

The ONNX path does not use pickle serialization — models are exported to
.onnx format and loaded natively by Triton's ONNX Runtime backend. There
is no serve.pkl, no metadata.json, and no integrity check to perform.
The secret key was dead code that also constituted a hardcoded credential
(CWE-798).

With this change, self.secret_key remains empty string (set by
_build_for_triton), and the existing cleanup in _build_for_transformers
removes empty SAGEMAKER_SERVE_SECRET_KEY from env_vars before CreateModel.

Addresses: P400136088 (Bug 2 - Hardcoded secret key)

* fix: Add HMAC integrity verification for Triton inference handler

Addresses P400136088 Bug 1 and V2146375387 (Triton path).

Three changes:

1. check_integrity.py: Switch from HMAC-SHA256 to plain SHA-256.
   - Remove generate_secret_key() — no longer needed
   - compute_hash() now uses hashlib.sha256() instead of hmac.new()
   - perform_integrity_check() no longer reads SAGEMAKER_SERVE_SECRET_KEY
     from environment

2. triton/model.py: Add integrity check in initialize() BEFORE
   cloudpickle deserialization. Previously the handler called
   cloudpickle.load() with no verification (acknowledged by a TODO
   comment). Now reads the file into a buffer, runs
   perform_integrity_check(), then deserializes with cloudpickle.loads().

3. triton/server.py: Remove SAGEMAKER_SERVE_SECRET_KEY from container
   environment variables in both local and SageMaker deployment modes.
   The key is no longer needed since integrity checking uses plain
   SHA-256.

4. model_builder_utils.py: Update _hmac_signing() to use plain SHA-256
   and stop generating/storing a secret key. Remove generate_secret_key
   import.

The integrity check still detects accidental corruption of model
artifacts in S3. The HMAC was providing a false sense of security since
the key was exposed via DescribeModel/DescribeEndpointConfig APIs.

* 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.

---------

Co-authored-by: Pravali Uppugunduri <upravali@amazon.com>

Author: pravali96

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