Merge branch 'master' into master

src/sagemaker/image_uri_config/autogluon.json

@@ -12,7 +12,8 @@
             "0.7": "0.7.0",
             "0.8": "0.8.2",
             "1.0": "1.0.0",
-            "1.1": "1.1.1"
+            "1.1": "1.1.1",
+            "1.2": "1.2.0"
         },
         "versions": {
             "0.3.1": {
@@ -563,6 +564,47 @@
                 "py_versions": [
                     "py311"
                 ]
+            },
+            "1.2.0": {
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "il-central-1": "780543022126",
+                    "ap-east-1": "871362719292",
+                    "ap-northeast-1": "763104351884",
+                    "ap-northeast-2": "763104351884",
+                    "ap-northeast-3": "364406365360",
+                    "ap-south-1": "763104351884",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-3": "907027046896",
+                    "ap-southeast-4": "457447274322",
+                    "ca-central-1": "763104351884",
+                    "eu-central-1": "763104351884",
+                    "eu-north-1": "763104351884",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "me-south-1": "217643126080",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
+                    "us-iso-east-1": "886529160074",
+                    "us-isob-east-1": "094389454867",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884",
+                    "ca-west-1": "204538143572"
+                },
+                "repository": "autogluon-training",
+                "processors": [
+                    "cpu",
+                    "gpu"
+                ],
+                "py_versions": [
+                    "py311"
+                ]
             }
         }
     },
@@ -575,7 +617,8 @@
             "0.7": "0.7.0",
             "0.8": "0.8.2",
             "1.0": "1.0.0",
-            "1.1": "1.1.1"
+            "1.1": "1.1.1",
+            "1.2": "1.2.0"
         },
         "versions": {
             "0.3.1": {
@@ -1157,6 +1200,49 @@
                 "py_versions": [
                     "py311"
                 ]
+            },
+            "1.2.0": {
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "il-central-1": "780543022126",
+                    "ap-east-1": "871362719292",
+                    "ap-northeast-1": "763104351884",
+                    "ap-northeast-2": "763104351884",
+                    "ap-northeast-3": "364406365360",
+                    "ap-south-1": "763104351884",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-3": "907027046896",
+                    "ap-southeast-4": "457447274322",
+                    "ca-central-1": "763104351884",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-north-1": "763104351884",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "me-south-1": "217643126080",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
+                    "us-iso-east-1": "886529160074",
+                    "us-isob-east-1": "094389454867",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884",
+                    "ca-west-1": "204538143572"
+                },
+                "repository": "autogluon-inference",
+                "processors": [
+                    "cpu",
+                    "gpu"
+                ],
+                "py_versions": [
+                    "py311"
+                ]
             }
         }
     }

src/sagemaker/serve/builder/model_builder.py

@@ -1433,15 +1433,15 @@ def _model_builder_optimize_wrapper(
 
             # HF Model ID format = "meta-llama/Meta-Llama-3.1-8B"
             # JS Model ID format = "meta-textgeneration-llama-3-1-8b"
-            llama_3_1_keywords = ["llama-3.1", "llama-3-1"]
-            is_llama_3_1 = self.model and any(
-                keyword in self.model.lower() for keyword in llama_3_1_keywords
+            is_llama_3_plus = self.model and bool(
+                re.search(r"llama-3[\.\-][1-9]\d*", self.model.lower())
             )
 
             if is_gpu_instance and self.model and self.is_compiled:
-                if is_llama_3_1:
+                if is_llama_3_plus:
                     raise ValueError(
-                        "Compilation is not supported for Llama-3.1 with a GPU instance."
+                        "Compilation is not supported for models greater "
+                        "than Llama-3.0 with a GPU instance."
                     )
                 if speculative_decoding_config:
                     raise ValueError(

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

@@ -21,7 +21,7 @@
 METADATA_PATH = Path(__file__).parent.joinpath("metadata.json")
 
 
-def model_fn(model_dir):
+def model_fn(model_dir, context=None):
     """Overrides default method for loading a model"""
     shared_libs_path = Path(model_dir + "/shared_libs")
 
@@ -40,16 +40,26 @@ def model_fn(model_dir):
         return partial(inference_spec.invoke, model=inference_spec.load(model_dir))
 
 
-def input_fn(input_data, content_type):
+def input_fn(input_data, content_type, context=None):
     """Deserializes the bytes that were received from the model server"""
     try:
         if hasattr(schema_builder, "custom_input_translator"):
             deserialized_data = schema_builder.custom_input_translator.deserialize(
-                io.BytesIO(input_data), content_type
+                (
+                    io.BytesIO(input_data)
+                    if type(input_data) == bytes
+                    else io.BytesIO(input_data.encode("utf-8"))
+                ),
+                content_type,
             )
         else:
             deserialized_data = schema_builder.input_deserializer.deserialize(
-                io.BytesIO(input_data), content_type[0]
+                (
+                    io.BytesIO(input_data)
+                    if type(input_data) == bytes
+                    else io.BytesIO(input_data.encode("utf-8"))
+                ),
+                content_type[0],
             )
 
         # Check if preprocess method is defined and call it
@@ -62,12 +72,12 @@ def input_fn(input_data, content_type):
         raise Exception("Encountered error in deserialize_request.") from e
 
 
-def predict_fn(input_data, predict_callable):
+def predict_fn(input_data, predict_callable, context=None):
     """Invokes the model that is taken in by model server"""
     return predict_callable(input_data)
 
 
-def output_fn(predictions, accept_type):
+def output_fn(predictions, accept_type, context=None):
     """Prediction is serialized to bytes and sent back to the customer"""
     try:
         if hasattr(inference_spec, "postprocess"):

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

@@ -67,11 +67,21 @@ def input_fn(input_data, content_type):
     try:
         if hasattr(schema_builder, "custom_input_translator"):
             deserialized_data = schema_builder.custom_input_translator.deserialize(
-                io.BytesIO(input_data), content_type
+                (
+                    io.BytesIO(input_data)
+                    if type(input_data) == bytes
+                    else io.BytesIO(input_data.encode("utf-8"))
+                ),
+                content_type,
             )
         else:
             deserialized_data = schema_builder.input_deserializer.deserialize(
-                io.BytesIO(input_data), content_type[0]
+                (
+                    io.BytesIO(input_data)
+                    if type(input_data) == bytes
+                    else io.BytesIO(input_data.encode("utf-8"))
+                ),
+                content_type[0],
             )
 
         # Check if preprocess method is defined and call it

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

@@ -70,11 +70,21 @@ def input_fn(input_data, content_type):
     try:
         if hasattr(schema_builder, "custom_input_translator"):
             return schema_builder.custom_input_translator.deserialize(
-                io.BytesIO(input_data), content_type
+                (
+                    io.BytesIO(input_data)
+                    if type(input_data) == bytes
+                    else io.BytesIO(input_data.encode("utf-8"))
+                ),
+                content_type,
             )
         else:
             return schema_builder.input_deserializer.deserialize(
-                io.BytesIO(input_data), content_type[0]
+                (
+                    io.BytesIO(input_data)
+                    if type(input_data) == bytes
+                    else io.BytesIO(input_data.encode("utf-8"))
+                ),
+                content_type[0],
             )
     except Exception as e:
         raise Exception("Encountered error in deserialize_request.") from e

src/sagemaker/workflow/utilities.py

@@ -268,29 +268,29 @@ def get_config_hash(step: Entity):
 
 
 def hash_object(obj) -> str:
-    """Get the MD5 hash of an object.
+    """Get the SHA256 hash of an object.
 
     Args:
         obj (dict): The object
     Returns:
-        str: The MD5 hash of the object
+        str: The SHA256 hash of the object
     """
-    return hashlib.md5(str(obj).encode()).hexdigest()
+    return hashlib.sha256(str(obj).encode()).hexdigest()
 
 
 def hash_file(path: str) -> str:
-    """Get the MD5 hash of a file.
+    """Get the SHA256 hash of a file.
 
     Args:
         path (str): The local path for the file.
     Returns:
-        str: The MD5 hash of the file.
+        str: The SHA256 hash of the file.
     """
-    return _hash_file(path, hashlib.md5()).hexdigest()
+    return _hash_file(path, hashlib.sha256()).hexdigest()
 
 
 def hash_files_or_dirs(paths: List[str]) -> str:
-    """Get the MD5 hash of the contents of a list of files or directories.
+    """Get the SHA256 hash of the contents of a list of files or directories.
 
     Hash is changed if:
        * input list is changed
@@ -301,58 +301,58 @@ def hash_files_or_dirs(paths: List[str]) -> str:
     Args:
         paths: List of file or directory paths
     Returns:
-        str: The MD5 hash of the list of files or directories.
+        str: The SHA256 hash of the list of files or directories.
     """
-    md5 = hashlib.md5()
+    sha256 = hashlib.sha256()
     for path in sorted(paths):
-        md5 = _hash_file_or_dir(path, md5)
-    return md5.hexdigest()
+        sha256 = _hash_file_or_dir(path, sha256)
+    return sha256.hexdigest()
 
 
-def _hash_file_or_dir(path: str, md5: Hash) -> Hash:
+def _hash_file_or_dir(path: str, sha256: Hash) -> Hash:
     """Updates the inputted Hash with the contents of the current path.
 
     Args:
         path: path of file or directory
     Returns:
-        str: The MD5 hash of the file or directory
+        str: The SHA256 hash of the file or directory
     """
     if isinstance(path, str) and path.lower().startswith("file://"):
         path = unquote(urlparse(path).path)
-    md5.update(path.encode())
+    sha256.update(path.encode())
     if Path(path).is_dir():
-        md5 = _hash_dir(path, md5)
+        sha256 = _hash_dir(path, sha256)
     elif Path(path).is_file():
-        md5 = _hash_file(path, md5)
-    return md5
+        sha256 = _hash_file(path, sha256)
+    return sha256
 
 
-def _hash_dir(directory: Union[str, Path], md5: Hash) -> Hash:
+def _hash_dir(directory: Union[str, Path], sha256: Hash) -> Hash:
     """Updates the inputted Hash with the contents of the current path.
 
     Args:
         directory: path of the directory
     Returns:
-        str: The MD5 hash of the directory
+        str: The SHA256 hash of the directory
     """
     if not Path(directory).is_dir():
         raise ValueError(str(directory) + " is not a valid directory")
     for path in sorted(Path(directory).iterdir()):
-        md5.update(path.name.encode())
+        sha256.update(path.name.encode())
         if path.is_file():
-            md5 = _hash_file(path, md5)
+            sha256 = _hash_file(path, sha256)
         elif path.is_dir():
-            md5 = _hash_dir(path, md5)
-    return md5
+            sha256 = _hash_dir(path, sha256)
+    return sha256
 
 
-def _hash_file(file: Union[str, Path], md5: Hash) -> Hash:
+def _hash_file(file: Union[str, Path], sha256: Hash) -> Hash:
     """Updates the inputted Hash with the contents of the current path.
 
     Args:
         file: path of the file
     Returns:
-        str: The MD5 hash of the file
+        str: The SHA256 hash of the file
     """
     if isinstance(file, str) and file.lower().startswith("file://"):
         file = unquote(urlparse(file).path)
@@ -363,8 +363,8 @@ def _hash_file(file: Union[str, Path], md5: Hash) -> Hash:
             data = f.read(BUF_SIZE)
             if not data:
                 break
-            md5.update(data)
-    return md5
+            sha256.update(data)
+    return sha256
 
 
 def validate_step_args_input(

tests/unit/sagemaker/serve/builder/test_model_builder.py

@@ -3270,7 +3270,7 @@ def test_optimize_with_gpu_instance_and_llama_3_1_and_compilation(
 
         mock_pysdk_model = Mock()
         mock_pysdk_model.model_data = None
-        mock_pysdk_model.env = {"HF_MODEL_ID": "meta-llama/Meta-Llama-3-1-8B-Instruct"}
+        mock_pysdk_model.env = {"HF_MODEL_ID": "meta-llama/Meta-Llama-3-2-8B-Instruct"}
 
         sample_input = {"inputs": "dummy prompt", "parameters": {}}
 
@@ -3279,7 +3279,7 @@ def test_optimize_with_gpu_instance_and_llama_3_1_and_compilation(
         dummy_schema_builder = SchemaBuilder(sample_input, sample_output)
 
         model_builder = ModelBuilder(
-            model="meta-llama/Meta-Llama-3-1-8B-Instruct",
+            model="meta-llama/Meta-Llama-3-2-8B-Instruct",
             schema_builder=dummy_schema_builder,
             env_vars={"HF_TOKEN": "token"},
             model_metadata={
@@ -3293,7 +3293,7 @@ def test_optimize_with_gpu_instance_and_llama_3_1_and_compilation(
 
         self.assertRaisesRegex(
             ValueError,
-            "Compilation is not supported for Llama-3.1 with a GPU instance.",
+            "Compilation is not supported for models greater than Llama-3.0 with a GPU instance.",
             lambda: model_builder.optimize(
                 job_name="job_name-123",
                 instance_type="ml.g5.24xlarge",