aws
diff --git a/‎requirements/extras/test_requirements.txt
Lines changed: 2 additions & 0 deletions b/‎requirements/extras/test_requirements.txt
Lines changed: 2 additions & 0 deletions
diff --git a/‎setup.py
Lines changed: 2 additions & 0 deletions b/‎setup.py
Lines changed: 2 additions & 0 deletions
diff --git a/‎src/sagemaker/base_deserializers.py
Lines changed: 0 additions & 34 deletions b/‎src/sagemaker/base_deserializers.py
Lines changed: 0 additions & 34 deletions
diff --git a/‎src/sagemaker/base_serializers.py
Lines changed: 0 additions & 25 deletions b/‎src/sagemaker/base_serializers.py
Lines changed: 0 additions & 25 deletions
diff --git a/‎src/sagemaker/deserializers.py
Lines changed: 0 additions & 1 deletion b/‎src/sagemaker/deserializers.py
Lines changed: 0 additions & 1 deletion
diff --git a/‎src/sagemaker/djl_inference/model.py
Lines changed: 5 additions & 0 deletions b/‎src/sagemaker/djl_inference/model.py
Lines changed: 5 additions & 0 deletions
diff --git a/‎src/sagemaker/huggingface/model.py
Lines changed: 1 addition & 0 deletions b/‎src/sagemaker/huggingface/model.py
Lines changed: 1 addition & 0 deletions
diff --git a/‎src/sagemaker/serializers.py
Lines changed: 0 additions & 1 deletion b/‎src/sagemaker/serializers.py
Lines changed: 0 additions & 1 deletion
diff --git a/‎src/sagemaker/serve/builder/djl_builder.py
Lines changed: 21 additions & 22 deletions b/‎src/sagemaker/serve/builder/djl_builder.py
Lines changed: 21 additions & 22 deletions
diff --git a/‎src/sagemaker/serve/builder/jumpstart_builder.py
Lines changed: 19 additions & 10 deletions b/‎src/sagemaker/serve/builder/jumpstart_builder.py
Lines changed: 19 additions & 10 deletions
@@ -36,4 +36,6 @@ transformers==4.32.0
 sentencepiece==0.1.99
 # https://github.com/triton-inference-server/server/issues/6246
 tritonclient[http]<2.37.0
+onnx==1.14.1
+# tf2onnx==1.15.1
 nbformat>=5.9,<6
@@ -70,6 +70,7 @@ def read_requirements(filename):
     "requests",
     "docker",
     "tqdm",
+    "psutil",
 ]
 
 # Specific use case dependencies
@@ -106,6 +107,7 @@ def read_requirements(filename):
     description="Open source library for training and deploying models on Amazon SageMaker.",
     packages=find_packages("src"),
     package_dir={"": "src"},
+    package_data={"": ["*.whl"]},
     py_modules=[os.path.splitext(os.path.basename(path))[0] for path in glob("src/*.py")],
     include_package_data=True,
     long_description=read("README.rst"),
 
@@ -22,7 +22,6 @@
 
 import numpy as np
 from six import with_metaclass
-import cloudpickle
 
 from sagemaker.utils import DeferredError
 
@@ -378,36 +377,3 @@ def deserialize(self, stream, content_type="tensor/pt"):
                 "Unable to deserialize your data to torch.Tensor.\
                     Please provide custom deserializer in InferenceSpec."
             )
-
-
-class PickleDeserializer(SimpleBaseDeserializer):
-    """Deserialize stream to object using cloudpickle module.
-
-    Args:
-           stream (botocore.response.StreamingBody): Data to be deserialized.
-           content_type (str): The MIME type of the data.
-
-       Returns:
-           object: The data deserialized into a torch Tensor.
-    """
-
-    def __init__(self, accept="application/x-pkl"):
-        super(PickleDeserializer, self).__init__(accept)
-
-    def deserialize(self, stream, content_type="application/x-pkl"):
-        """Deserialize pickle data from an inference endpoint.
-
-        Args:
-            stream (botocore.response.StreamingBody): Data to be deserialized.
-            content_type (str): The MIME type of the data.
-
-        Returns:
-            list: A list of piclke serializable objects.
-        """
-        try:
-            return cloudpickle.loads(stream.read())
-        except Exception:
-            raise ValueError(
-                "Cannot deserialize bytes to object with cloudpickle.\
-                    Please provide custom deserializer."
-            )
@@ -21,7 +21,6 @@
 import numpy as np
 from pandas import DataFrame
 from six import with_metaclass
-import cloudpickle
 
 from sagemaker.utils import DeferredError
 
@@ -465,27 +464,3 @@ def serialize(self, data):
                 )
 
         raise ValueError("Object of type %s is not a torch.Tensor" % type(data))
-
-
-class PickleSerializer(SimpleBaseSerializer):
-    """Serialize an arbitrary object using cloudpickle module."""
-
-    def __init__(self, content_type="application/x-pkl"):
-        super(PickleSerializer, self).__init__(content_type)
-
-    def serialize(self, data: object) -> bytes:
-        """Serialize an arbitrary object using cloudpickle module.
-
-        Args:
-            data (object): Data to be serialized. The data must be of torch.Tensor type.
-        Returns:
-            raw-bytes: The data serialized as raw-bytes from the input.
-        """
-        try:
-            return cloudpickle.dumps(data)
-        except Exception:
-            raise ValueError(
-                "Cannot serialize your object of type %s into bytes with cloudpickle.\
-                    Please provide custom serializer."
-                % type(data)
-            )
@@ -31,7 +31,6 @@
     StreamDeserializer,
     StringDeserializer,
     TorchTensorDeserializer,
-    PickleDeserializer,
 )
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 
@@ -180,6 +180,11 @@ def _get_model_config_properties_from_hf(model_id: str, hf_hub_token: str = None
                 model_config = json.load(response)
                 break
         except (HTTPError, URLError, TimeoutError, JSONDecodeError) as e:
+            if "HTTP Error 401: Unauthorized" in str(e):
+                raise ValueError(
+                    "Trying to access a gated/private HuggingFace model without valid credentials. "
+                    "Please provide a HUGGING_FACE_HUB_TOKEN in env_vars"
+                )
             logger.warning(
                 "Exception encountered while trying to read config file %s. " "Details: %s",
                 config_file_url,
 
@@ -326,6 +326,7 @@ def deploy(
             container_startup_health_check_timeout=container_startup_health_check_timeout,
             inference_recommendation_id=inference_recommendation_id,
             explainer_config=explainer_config,
+            endpoint_logging=kwargs.get("endpoint_logging", False),
         )
 
     def register(
 
@@ -29,7 +29,6 @@
     SimpleBaseSerializer,
     SparseMatrixSerializer,
     TorchTensorSerializer,
-    PickleSerializer,
     StringSerializer,
 )
 
 
@@ -1,3 +1,15 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
 """Holds mixin logic to support deployment of Model ID"""
 from __future__ import absolute_import
 import logging
@@ -26,7 +38,10 @@
     _get_admissible_dtypes,
 )
 from sagemaker.serve.utils.local_hardware import _get_nb_instance, _get_ram_usage_mb
-from sagemaker.serve.model_server.djl_serving.prepare import prepare_for_djl_serving
+from sagemaker.serve.model_server.djl_serving.prepare import (
+    prepare_for_djl_serving,
+    _create_dir_structure,
+)
 from sagemaker.serve.utils.predictors import DjlLocalModePredictor
 from sagemaker.serve.utils.types import ModelServer, _DjlEngine
 from sagemaker.serve.mode.function_pointers import Mode
@@ -40,7 +55,6 @@
 
 logger = logging.getLogger(__name__)
 
-_JUMP_START_HUGGING_FACE_PREFIX = "huggingface"
 # Match JumpStart DJL entrypoint format
 _DJL_MODEL_BUILDER_ENTRY_POINT = "inference.py"
 _CODE_FOLDER = "code"
@@ -86,16 +100,9 @@ def _prepare_for_mode(self):
     def _get_client_translators(self):
         """Placeholder docstring"""
 
-    def _validate_model_server(self):
+    def _is_djl(self):
         """Placeholder docstring"""
-        if self.model_server != ModelServer.DJL_SERVING:
-            messaging = (
-                "HuggingFace Model ID support on model server: "
-                f"{self.model_server} is not currently supported. "
-                f"Defaulting to {ModelServer.DJL_SERVING}"
-            )
-            logger.warning(messaging)
-            self.model_server = ModelServer.DJL_SERVING
+        return self.model_server == ModelServer.DJL_SERVING
 
     def _validate_djl_serving_sample_data(self):
         """Placeholder docstring"""
@@ -112,12 +119,6 @@ def _validate_djl_serving_sample_data(self):
         ):
             raise ValueError(_INVALID_SAMPLE_DATA_EX)
 
-    def _is_jumpstart_model_id(self) -> bool:
-        """Placeholder docstring"""
-        # this will potentially extend in the future so leave like this
-        # for now, only hf jumpstart model ids will be considered
-        return self.model.startswith(_JUMP_START_HUGGING_FACE_PREFIX)
-
     def _create_djl_model(self) -> Type[Model]:
         """Placeholder docstring"""
         code_dir = str(Path(self.model_path).joinpath(_CODE_FOLDER))
@@ -211,9 +212,6 @@ def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
             ram_usage_after = _get_ram_usage_mb()
 
             self.ram_usage_model_load = max(ram_usage_after - ram_usage_before, 0)
-            logger.info(
-                "RAM used to load the %s locally was %s MB", self.model, self.ram_usage_model_load
-            )
 
             return predictor
 
@@ -237,7 +235,8 @@ def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
             self.pysdk_model.env["TRANSFORMERS_CACHE"] = "/tmp"
             self.pysdk_model.env["HUGGINGFACE_HUB_CACHE"] = "/tmp"
 
-        kwargs["endpoint_logging"] = True
+        if "endpoint_logging" not in kwargs:
+            kwargs["endpoint_logging"] = True
         if self.nb_instance_type and "instance_type" not in kwargs:
             kwargs.update({"instance_type": self.nb_instance_type})
 
@@ -253,6 +252,7 @@ def _build_for_hf_djl(self):
         """Placeholder docstring"""
         self.overwrite_props_from_file = True
         self.nb_instance_type = _get_nb_instance()
+        _create_dir_structure(self.model_path)
         self.engine, self.hf_model_config = _auto_detect_engine(
             self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
         )
@@ -463,7 +463,6 @@ def _tune_for_hf_djl(self, max_tuning_duration: int = 1800):
 
     def _build_for_djl(self):
         """Placeholder docstring"""
-        self._validate_model_server()
         self._validate_djl_serving_sample_data()
         self.secret_key = None
 
 
@@ -1,3 +1,15 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
 """Placeholder docstring"""
 from __future__ import absolute_import
 
@@ -8,7 +20,7 @@
 from sagemaker.model import Model
 from sagemaker import model_uris
 from sagemaker.serve.model_server.djl_serving.prepare import prepare_djl_js_resources
-from sagemaker.serve.model_server.tgi.prepare import prepare_tgi_js_resources
+from sagemaker.serve.model_server.tgi.prepare import prepare_tgi_js_resources, _create_dir_structure
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -20,7 +32,6 @@
 from sagemaker.base_predictor import PredictorBase
 from sagemaker.jumpstart.model import JumpStartModel
 
-_JUMP_START_HUGGING_FACE_PREFIX = "huggingface"
 _DJL_MODEL_BUILDER_ENTRY_POINT = "inference.py"
 _NO_JS_MODEL_EX = "HuggingFace JumpStart Model ID not detected. Building for HuggingFace Model ID."
 _JS_SCOPE = "inference"
@@ -60,6 +71,7 @@ def __init__(self):
         self.schema_builder = None
         self.nb_instance_type = None
         self.ram_usage_model_load = None
+        self.jumpstart = None
 
     @abstractmethod
     def _prepare_for_mode(self):
@@ -77,10 +89,6 @@ def _is_jumpstart_model_id(self) -> bool:
             logger.warning(_NO_JS_MODEL_EX)
             return False
 
-        if not self.model.startswith(_JUMP_START_HUGGING_FACE_PREFIX):
-            logger.warning(_NO_JS_MODEL_EX)
-            return False
-
         logger.info("JumpStart Model ID detected.")
         return True
 
@@ -156,17 +164,16 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
                 None,
                 predictor,
                 self.pysdk_model.env,
+                jumpstart=True,
             )
             ram_usage_after = _get_ram_usage_mb()
 
             self.ram_usage_model_load = max(ram_usage_after - ram_usage_before, 0)
-            logger.info(
-                "RAM used to load the %s locally was %s MB", self.model, self.ram_usage_model_load
-            )
 
             return predictor
 
-        kwargs["endpoint_logging"] = True
+        if "endpoint_logging" not in kwargs:
+            kwargs["endpoint_logging"] = True
         if hasattr(self, "nb_instance_type"):
             kwargs.update({"instance_type": self.nb_instance_type})
 
@@ -186,6 +193,7 @@ def _build_for_djl_jumpstart(self):
         """Placeholder docstring"""
 
         env = {}
+        _create_dir_structure(self.model_path)
         if self.mode == Mode.LOCAL_CONTAINER:
             if not hasattr(self, "prepared_for_djl"):
                 (
@@ -227,6 +235,7 @@ def _build_for_jumpstart(self):
         """Placeholder docstring"""
         # we do not pickle for jumpstart. set to none
         self.secret_key = None
+        self.jumpstart = True
 
         pysdk_model = self._create_pre_trained_js_model()
Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,6 @@`
`31`	`31`	`StreamDeserializer,`
`32`	`32`	`StringDeserializer,`
`33`	`33`	`TorchTensorDeserializer,`
`34`		`- PickleDeserializer,`
`35`	`34`	`)`
`36`	`35`
`37`	`36`	`from sagemaker.jumpstart import artifacts, utils as jumpstart_utils`
Original file line number	Diff line number	Diff line change
`@@ -326,6 +326,7 @@ def deploy(`
`326`	`326`	`container_startup_health_check_timeout=container_startup_health_check_timeout,`
`327`	`327`	`inference_recommendation_id=inference_recommendation_id,`
`328`	`328`	`explainer_config=explainer_config,`
	`329`	`+ endpoint_logging=kwargs.get("endpoint_logging", False),`
`329`	`330`	`)`
`330`	`331`
`331`	`332`	`def register(`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,6 @@`
`29`	`29`	`SimpleBaseSerializer,`
`30`	`30`	`SparseMatrixSerializer,`
`31`	`31`	`TorchTensorSerializer,`
`32`		`- PickleSerializer,`
`33`	`32`	`StringSerializer,`
`34`	`33`	`)`
`35`	`34`