fix: HMAC signing for ModelBuilder Triton python backend (#1282)

Author: SSRraymond

doc/overview.rst

@@ -801,7 +801,6 @@ def _create_sagemaker_model(
                 Specifies configuration related to serverless endpoint. Instance type is
                 not provided in serverless inference. So this is used to find image URIs.
         """
-
         if self.model_package_arn is not None or self.algorithm_arn is not None:
             model_package = ModelPackage(
                 role=self.role,

src/sagemaker/model.py

@@ -56,7 +56,6 @@ class ModelBuilder(Triton, DJL, JumpStart):
 
               * 1/Launch on SageMaker endpoint
               * 2/Launch locally with a container
-              * 3/Launch in process
 
         shared_libs (List[str]): Any shared libraries you want to bring into
             the model.
@@ -336,6 +335,7 @@ def _create_model(self):
         self.pysdk_model.deploy = self._model_builder_deploy_wrapper
         self._original_register = self.pysdk_model.register
         self.pysdk_model.register = self._model_builder_register_wrapper
+        self.model_package = None
         return self.pysdk_model
 
     @_capture_telemetry("register")
@@ -347,9 +347,23 @@ def _model_builder_register_wrapper(self, *args, **kwargs):
         if "response_types" not in kwargs:
             self.pysdk_model.response_types = deserializer.ACCEPT.split()
         new_model_package = self._original_register(*args, **kwargs)
-        new_model_package.deploy = self._model_builder_deploy_wrapper
+        self.pysdk_model.model_package_arn = new_model_package.model_package_arn
+        new_model_package.deploy = self._model_builder_deploy_model_package_wrapper
+        self.model_package = new_model_package
         return new_model_package
 
+    def _model_builder_deploy_model_package_wrapper(self, *args, **kwargs):
+        """Placeholder docstring"""
+        if self.pysdk_model.model_package_arn is not None:
+            return self._model_builder_register_wrapper(*args, **kwargs)
+
+        # need to set the model_package_arn
+        # so that the model is created using the model_package's configs
+        self.pysdk_model.model_package_arn = self.model_package.model_package_arn
+        predictor = self._model_builder_register_wrapper(*args, **kwargs)
+        self.pysdk_model.model_package_arn = None
+        return predictor
+
     @_capture_telemetry("torchserve.deploy")
     def _model_builder_deploy_wrapper(
         self, *args, container_timeout_in_second: int = 300, **kwargs

src/sagemaker/serve/builder/model_builder.py

@@ -19,10 +19,9 @@
 
 from pathlib import Path
 import logging
-import shutil
 import subprocess
 import sys
-
+import re
 
 _SUPPORTED_SUFFIXES = [".txt"]
 # TODO : Move PKL_FILE_NAME to common location
@@ -31,7 +30,7 @@
 logger = logging.getLogger(__name__)
 
 
-def capture_dependencies(dependencies: str, work_dir: Path, capture_all: bool = False):
+def capture_dependencies(dependencies: dict, work_dir: Path, capture_all: bool = False):
     """Placeholder docstring"""
     path = work_dir.joinpath("requirements.txt")
     if "auto" in dependencies and dependencies["auto"]:
@@ -53,24 +52,44 @@ def capture_dependencies(dependencies: str, work_dir: Path, capture_all: bool =
             check=True,
         )
 
-    if "requirements" in dependencies:
-        _capture_from_customer_provided_requirements(dependencies["requirements"], path)
+        with open(path, "r") as f:
+            autodetect_depedencies = f.read().splitlines()
+    else:
+        autodetect_depedencies = []
 
+    module_version_dict = _parse_dependency_list(autodetect_depedencies)
+
+    if "requirements" in dependencies:
+        module_version_dict = _process_customer_provided_requirements(
+            requirements_file=dependencies["requirements"], module_version_dict=module_version_dict
+        )
     if "custom" in dependencies:
+        module_version_dict = _process_custom_dependencies(
+            custom_dependencies=dependencies.get("custom"), module_version_dict=module_version_dict
+        )
+    with open(path, "w") as f:
+        for module, version in module_version_dict.items():
+            f.write(f"{module}{version}\n")
 
-        with open(path, "a+") as f:
-            for package in dependencies["custom"]:
-                f.write(f"{package}\n")
 
+def _process_custom_dependencies(custom_dependencies: list, module_version_dict: dict):
+    """Placeholder docstring"""
+    custom_module_version_dict = _parse_dependency_list(custom_dependencies)
+    module_version_dict.update(custom_module_version_dict)
+    return module_version_dict
 
-def _capture_from_customer_provided_requirements(requirements_file: str, output_path: Path):
+
+def _process_customer_provided_requirements(requirements_file: str, module_version_dict: dict):
     """Placeholder docstring"""
-    input_path = Path(requirements_file)
-    if not input_path.is_file() or not _is_valid_requirement_file(input_path):
+    requirements_file = Path(requirements_file)
+    if not requirements_file.is_file() or not _is_valid_requirement_file(requirements_file):
         raise Exception(f"Path: {requirements_file} to requirements.txt doesn't exist")
     logger.debug("Packaging provided requirements.txt from %s", requirements_file)
-    with open(output_path, "a+") as f:
-        shutil.copyfileobj(open(input_path, "r"), f)
+    with open(requirements_file, "r") as f:
+        custom_dependencies = f.read().splitlines()
+
+    module_version_dict.update(_parse_dependency_list(custom_dependencies))
+    return module_version_dict
 
 
 def _is_valid_requirement_file(path):
@@ -82,6 +101,32 @@ def _is_valid_requirement_file(path):
     return False
 
 
+def _parse_dependency_list(depedency_list: list) -> dict:
+    """Placeholder docstring"""
+
+    # Divide a string into 2 part, first part is the module name
+    # and second part is its version constraint or the url
+    # checkout tests/unit/sagemaker/serve/detector/test_dependency_manager.py
+    # for examples
+    pattern = r"^([\w.-]+)(@[^,\n]+|((?:[<>=!~]=?[\w.*-]+,?)+)?)$"
+
+    module_version_dict = {}
+
+    for dependency in depedency_list:
+        if dependency.startswith("#"):
+            continue
+        match = re.match(pattern, dependency)
+        if match:
+            package = match.group(1)
+            # Group 2 is either a URL or version constraint, if present
+            url_or_version = match.group(2) if match.group(2) else ""
+            module_version_dict.update({package: url_or_version})
+        else:
+            module_version_dict.update({dependency: ""})
+
+    return module_version_dict
+
+
 # only required for dev testing
 def prepare_wheel(code_artifact_client, whl_dir: str):
     """Placeholder docstring"""

src/sagemaker/serve/detector/dependency_manager.py

@@ -108,13 +108,14 @@ def get_requirements_for_pkl_file(pkl_path: Path, dest: Path):
     with open(dest, mode="w+") as out:
         for x in get_all_installed_packages():
             name = x["name"]
+            version = x["version"]
             # skip only for dev
             if name == "sagemaker":
-                out.write("/opt/ml/model/whl/sagemaker-2.185.1.dev0-py2.py3-none-any.whl\n")
+                out.write("/opt/ml/model/whl/sagemaker-2.195.2.dev0-py2.py3-none-any.whl\n")
             elif name == "boto3":
                 out.write("boto3==1.26.131\n")
             elif name in currently_used_packages:
-                out.write(f"{name}\n")
+                out.write(f"{name}=={version}\n")
 
 
 def get_all_requirements(dest: Path):
@@ -128,10 +129,8 @@ def get_all_requirements(dest: Path):
             # skip only for dev
             if name == "sagemaker":
                 continue
-            if name == "boto3":
-                out.write("boto3==1.26.131\n")
-            else:
-                out.write(f"{name}=={version}\n")
+
+            out.write(f"{name}=={version}\n")
 
 
 def parse_args():

src/sagemaker/serve/detector/pickle_dependencies.py

@@ -93,6 +93,7 @@ def create_server(
                 docker_client=self.client,
                 model_path=model_path if model_path else self.model_path,
                 image_uri=image,
+                secret_key=secret_key,
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             self._ping_container = self._triton_deep_ping

src/sagemaker/serve/mode/local_container_mode.py

@@ -70,6 +70,7 @@ def prepare(
             return self._upload_triton_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
+                secret_key=secret_key,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
             )

src/sagemaker/serve/mode/sagemaker_endpoint_mode.py

@@ -97,7 +97,9 @@ def _py_vs_parity_check():
 def _pickle_file_integrity_check():
     with open("/opt/ml/model/code/serve.pkl", "rb") as f:
         buffer = f.read()
-    perform_integrity_check(buffer=buffer)
+
+    metadeata_path = Path("/opt/ml/model/code/metadata.json")
+    perform_integrity_check(buffer=buffer, metadata_path=metadeata_path)
 
 
 # on import, execute

src/sagemaker/serve/model_server/torchserve/inference.py

@@ -4,9 +4,11 @@
 import logging
 import ssl
 from pathlib import Path
+import platform
 
 import triton_python_backend_utils as pb_utils
 import cloudpickle
+from sagemaker.serve.validations.check_integrity import perform_integrity_check
 
 logger = logging.getLogger(__name__)
 
@@ -52,3 +54,31 @@ def execute(self, requests):
             responses.append(response)
 
         return responses
+
+
+def _run_preflight_diagnostics():
+    _py_vs_parity_check()
+    _pickle_file_integrity_check()
+
+
+def _py_vs_parity_check():
+    container_py_vs = platform.python_version()
+    local_py_vs = os.getenv("LOCAL_PYTHON")
+
+    if not local_py_vs or container_py_vs.split(".")[1] != local_py_vs.split(".")[1]:
+        logger.warning(
+            f"The local python version {local_py_vs} differs from the python version "
+            f"{container_py_vs} on the container. Please align the two to avoid unexpected behavior"
+        )
+
+
+def _pickle_file_integrity_check():
+    serve_path = Path(TRITON_MODEL_DIR).joinpath("serve.pkl")
+    metadata_path = Path(TRITON_MODEL_DIR).joinpath("metadata.json")
+    with open(str(serve_path), "rb") as f:
+        buffer = f.read()
+    perform_integrity_check(buffer=buffer, metadata_path=metadata_path)
+
+
+# on import, execute
+_run_preflight_diagnostics()

src/sagemaker/serve/model_server/triton/model.py

@@ -3,6 +3,7 @@
 import uuid
 import logging
 import importlib
+import platform
 
 from sagemaker import fw_utils
 from sagemaker import Session
@@ -29,13 +30,20 @@ def _start_triton_server(
         self,
         docker_client: docker.DockerClient,
         model_path: str,
+        secret_key: str,
         image_uri: str,
         env_vars: dict,
     ):
         """Placeholder docstring"""
         self.container_name = "triton" + uuid.uuid1().hex
         model_repository = model_path + "/model_repository"
-        env_vars.update({"TRITON_MODEL_DIR": "/models/model"})
+        env_vars.update(
+            {
+                "TRITON_MODEL_DIR": "/models/model",
+                "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
+                "LOCAL_PYTHON": platform.python_version(),
+            }
+        )
 
         if "cpu" not in image_uri:
             self.container = docker_client.containers.run(
@@ -102,6 +110,7 @@ def _upload_triton_artifacts(
         self,
         model_path: str,
         sagemaker_session: Session,
+        secret_key: str,
         s3_model_data_url: str = None,
         image: str = None,
     ):
@@ -127,5 +136,7 @@ def _upload_triton_artifacts(
         env_vars = {
             "SAGEMAKER_TRITON_DEFAULT_MODEL_NAME": "model",
             "TRITON_MODEL_DIR": "/opt/ml/model/model",
+            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
+            "LOCAL_PYTHON": platform.python_version(),
         }
         return s3_upload_path, env_vars

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

@@ -21,6 +21,12 @@
 from sagemaker.base_deserializers import JSONDeserializer
 from sagemaker.serve.detector.pickler import save_pkl
 from sagemaker.serve.model_server.triton.config_template import CONFIG_TEMPLATE
+from sagemaker.serve.validations.check_integrity import (
+    generate_secret_key,
+    compute_hash,
+)
+
+from sagemaker.remote_function.core.serialization import _MetaData
 
 
 logger = logging.getLogger(__name__)
@@ -206,6 +212,8 @@ def _prepare_for_triton(self):
             export_path.mkdir(parents=True)
 
         if self.model:
+            self.secret_key = "dummy secret key for onnx backend"
+
             if self.framework == "pytorch":
                 self._export_pytorch_to_onnx(
                     export_path=export_path, model=self.model, schema_builder=self.schema_builder
@@ -228,10 +236,26 @@ def _prepare_for_triton(self):
 
             self._pack_conda_env(pkl_path=pkl_path)
 
+            self._hmac_signing()
+
             return
 
         raise ValueError("Either model or inference_spec should be provided to ModelBuilder.")
 
+    def _hmac_signing(self):
+        """Perform HMAC signing on picke file for integrity check"""
+        secret_key = generate_secret_key()
+        pkl_path = Path(self.model_path).joinpath("model_repository").joinpath("model")
+
+        with open(str(pkl_path.joinpath("serve.pkl")), "rb") as f:
+            buffer = f.read()
+        hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+
+        with open(str(pkl_path.joinpath("metadata.json")), "wb") as metadata:
+            metadata.write(_MetaData(hash_value).to_json())
+
+        self.secret_key = secret_key
+
     def _generate_config_pbtxt(self, pkl_path: Path):
         config_path = pkl_path.joinpath("config.pbtxt")
 
@@ -436,8 +460,6 @@ def _build_for_triton(self):
 
         self._auto_detect_image_for_triton()
 
-        self.secret_key = "dummy secret key"
-
         self._save_inference_spec()
 
         self._prepare_for_triton()

src/sagemaker/serve/model_server/triton/triton_builder.py

@@ -19,18 +19,16 @@ def compute_hash(buffer: bytes, secret_key: str) -> str:
     return hmac.new(secret_key.encode(), msg=buffer, digestmod=hashlib.sha256).hexdigest()
 
 
-def perform_integrity_check(buffer: bytes):
+def perform_integrity_check(buffer: bytes, metadata_path: Path):
     """Validates the integrity of bytes by comparing the hash value"""
     secret_key = os.environ.get("SAGEMAKER_SERVE_SECRET_KEY")
     actual_hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
 
-    metadata_path = Path("/opt/ml/model/code/metadata.json").resolve()
-
     if not Path.exists(metadata_path):
-        raise Exception("Path to metadata.json does not exist")
+        raise ValueError("Path to metadata.json does not exist")
 
-    with open(metadata_path, "rb") as md:
+    with open(str(metadata_path), "rb") as md:
         expected_hash_value = _MetaData.from_json(md.read()).sha256_hash
 
     if not hmac.compare_digest(expected_hash_value, actual_hash_value):
-        raise Exception("Integrity check for the serialized function or data failed.")
+        raise ValueError("Integrity check for the serialized function or data failed.")