support latest container version in image_uris and DJLModel for lmi containers

Author: siddvenk

src/sagemaker/djl_inference/model.py

@@ -43,7 +43,7 @@ def __init__(
         self,
         model_id: Optional[str] = None,
         engine: Optional[str] = None,
-        djl_version: str = "0.28.0",
+        djl_version: str = "latest",
         djl_framework: Optional[str] = None,
         task: Optional[str] = None,
         dtype: Optional[str] = None,

src/sagemaker/image_uri_config/djl-lmi.json

@@ -2,6 +2,9 @@
     "scope": [
         "inference"
     ],
+    "version_aliases": {
+        "latest": "0.29.0"
+    },
     "versions": {
         "0.29.0": {
             "registries": {

src/sagemaker/image_uri_config/djl-neuronx.json

@@ -2,6 +2,9 @@
     "scope": [
         "inference"
     ],
+    "version_aliases": {
+        "latest": "0.29.0"
+    },
     "versions": {
         "0.29.0": {
             "registries": {

src/sagemaker/image_uri_config/djl-tensorrtllm.json

@@ -2,6 +2,9 @@
     "scope": [
         "inference"
     ],
+    "version_aliases": {
+        "latest": "0.29.0"
+    },
     "versions": {
         "0.29.0": {
             "registries": {

tests/unit/sagemaker/image_uris/test_djl.py

@@ -18,12 +18,7 @@
 
 @pytest.mark.parametrize(
     "load_config_and_file_name",
-    [
-        "djl-neuronx.json",
-        "djl-fastertransformer.json",
-        "djl-deepspeed.json",
-        "djl-tensorrtllm.json",
-    ],
+    ["djl-neuronx.json", "djl-tensorrtllm.json", "djl-lmi.json"],
     indirect=True,
 )
 def test_djl_uris(load_config_and_file_name):

tests/unit/test_djl_inference.py

@@ -26,7 +26,7 @@
 HF_MODEL_ID = "hf_hub_model_id"
 ROLE = "dummy_role"
 REGION = "us-west-2"
-VERSION = "0.28.0"
+VERSION = "latest"
 
 LMI_IMAGE_URI = image_uris.retrieve(framework="djl-lmi", version=VERSION, region=REGION)
 TRT_IMAGE_URI = image_uris.retrieve(framework="djl-tensorrtllm", version=VERSION, region=REGION)