feature: Add DJL Inference support

Author: siddvenk

src/sagemaker/djl_inference/__init__.py

@@ -0,0 +1,19 @@
+# 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
+
+from sagemaker.djl_inference.model import DJLPredictor  # noqa: F401
+from sagemaker.djl_inference.model import DJLModel  # noqa: F401
+from sagemaker.djl_inference.model import DeepSpeedModel  # noqa: F401
+from sagemaker.djl_inference.model import HuggingFaceAccelerateModel  # noqa: F401

src/sagemaker/djl_inference/defaults.py

@@ -0,0 +1,50 @@
+# 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
+
+STABLE_DIFFUSION_MODEL_TYPE = "stable-diffusion"
+
+DEEPSPEED_RECOMMENDED_ARCHITECTURES = {
+    "bloom",
+    "opt",
+    "gpt_neox",
+    "gptj",
+    "gpt_neo",
+    "gpt2",
+    "xlm-roberta",
+    "roberta",
+    "bert",
+    STABLE_DIFFUSION_MODEL_TYPE,
+}
+
+DEEPSPEED_SUPPORTED_ARCHITECTURES = {
+    "bloom",
+    "opt",
+    "gpt_neox",
+    "gptj",
+    "gpt_neo",
+    "gpt2",
+    "xlm-roberta",
+    "roberta",
+    "bert",
+    STABLE_DIFFUSION_MODEL_TYPE,
+}
+
+ALLOWED_INSTANCE_FAMILIES = {
+    "ml.g4dn",
+    "ml.g5",
+    "ml.p3",
+    "ml.p4",
+    "local_gpu",
+}