Merge branch 'aws:master' into asyncWaiterTimeoutHandle

Author: haozhx23

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/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) if type(input_data)== bytes else io.BytesIO(input_data.encode('utf-8')), 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) if type(input_data)== bytes else io.BytesIO(input_data.encode('utf-8')), 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) if type(input_data)== bytes else io.BytesIO(input_data.encode('utf-8')), 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) if type(input_data)== bytes else io.BytesIO(input_data.encode('utf-8')), 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) if type(input_data)== bytes else io.BytesIO(input_data.encode('utf-8')), 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) if type(input_data)== bytes else io.BytesIO(input_data.encode('utf-8')), 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/workflow/test_steps.py

@@ -671,7 +671,7 @@ def test_processing_step_normalizes_args_with_local_code(mock_normalize_args, sc
     mock_normalize_args.return_value = [step.inputs, step.outputs]
     step.to_request()
     mock_normalize_args.assert_called_with(
-        job_name="MyProcessingStep-3e89f0c7e101c356cbedf27d9d27e9db",
+        job_name="MyProcessingStep-a22fc59b38f13da26f6a40b18687ba598cf669f74104b793cefd9c63eddf4ac7",
         arguments=step.job_arguments,
         inputs=step.inputs,
         outputs=step.outputs,

tests/unit/sagemaker/workflow/test_utilities.py

@@ -31,14 +31,14 @@ def test_hash_file():
     with tempfile.NamedTemporaryFile() as tmp:
         tmp.write("hashme".encode())
         hash = hash_file(tmp.name)
-        assert hash == "d41d8cd98f00b204e9800998ecf8427e"
+        assert hash == "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
 
 
 def test_hash_file_uri():
     with tempfile.NamedTemporaryFile() as tmp:
         tmp.write("hashme".encode())
         hash = hash_file(f"file:///{tmp.name}")
-        assert hash == "d41d8cd98f00b204e9800998ecf8427e"
+        assert hash == "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
 
 
 def test_hash_files_or_dirs_with_file():

tests/unit/sagemaker/workflow/test_utils.py

@@ -82,7 +82,8 @@ def test_repack_model_step(estimator):
     assert hyperparameters["sagemaker_program"] == f'"{REPACK_SCRIPT_LAUNCHER}"'
     assert (
         hyperparameters["sagemaker_submit_directory"]
-        == '"s3://my-bucket/MyRepackModelStep-b5ea77f701b47a8d075605497462ccc2/source/sourcedir.tar.gz"'
+        == '"s3://my-bucket/MyRepackModelStep-717d7bdd388168c27e9ad2938ff0314e35be50b3157cf2498688c7525ea27e1e\
+/source/sourcedir.tar.gz"'
     )
 
     del request_dict["Arguments"]["HyperParameters"]