test: adds regions in which djl images do not exist

Author: malte-aws

tests/unit/sagemaker/image_uris/test_djl.py

@@ -55,15 +55,69 @@ def _test_djl_uris(account, region, version, tag, djl_framework):
     "me-south-1", "me-central-1", "sa-east-1", "cn-north-1", "cn-northwest-1"
 }
 
+# Known missing framework:version:region combinations that don't exist in ECR
+KNOWN_MISSING_COMBINATIONS = {
+    "djl-lmi": {
+        "0.30.0-lmi12.0.0-cu124": {"ap-east-2"},
+        "0.29.0-lmi11.0.0-cu124": {"ap-east-2"},
+        "0.28.0-lmi10.0.0-cu124": {"ap-east-2"},
+    },
+    "djl-neuronx": {
+        "0.29.0-neuronx-sdk2.19.1": {"ap-east-1", "me-central-1", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
+        "0.28.0-neuronx-sdk2.18.2": {"ap-east-1", "me-central-1", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
+        "0.27.0-neuronx-sdk2.18.1": {"ap-east-1", "me-central-1", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
+        "0.26.0-neuronx-sdk2.16.0": {"ap-east-1", "me-central-1", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
+        "0.25.0-neuronx-sdk2.15.0": {"eu-north-1", "ap-east-1", "me-central-1", "eu-west-2", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
+        "0.24.0-neuronx-sdk2.14.1": {"eu-north-1", "ap-east-1", "me-central-1", "eu-west-2", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
+        "0.23.0-neuronx-sdk2.12.0": {"eu-north-1", "ap-east-1", "me-central-1", "eu-west-2", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
+        "0.22.1-neuronx-sdk2.10.0": {"eu-north-1", "ap-east-1", "me-central-1", "eu-west-2", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
+    },
+    "djl-tensorrtllm": {
+        "0.30.0-tensorrtllm0.12.0-cu125": {"ap-east-2"},
+        "0.29.0-tensorrtllm0.11.0-cu124": {"ap-east-2"},
+        "0.28.0-tensorrtllm0.9.0-cu122": {"ap-east-2"},
+        "0.27.0-tensorrtllm0.8.0-cu122": {"ap-east-2"},
+        "0.26.0-tensorrtllm0.7.1-cu122": {"ap-east-2"},
+        "0.25.0-tensorrtllm0.5.0-cu122": {"ap-east-2"},
+    },
+    "djl-fastertransformer": {
+        "0.24.0-fastertransformer5.3.0-cu118": {"ap-east-2"},
+        "0.23.0-fastertransformer5.3.0-cu118": {"ap-east-2"},
+        "0.22.1-fastertransformer5.3.0-cu118": {"ap-east-2"},
+        "0.21.0-fastertransformer5.3.0-cu117": {"ap-east-2"},
+    },
+    "djl-deepspeed": {
+        "0.27.0-deepspeed0.12.6-cu121": {"ap-east-2"},
+        "0.26.0-deepspeed0.12.6-cu121": {"ap-east-2"},
+        "0.25.0-deepspeed0.11.0-cu118": {"ap-east-2"},
+        "0.24.0-deepspeed0.10.0-cu118": {"ap-east-2"},
+        "0.23.0-deepspeed0.9.5-cu118": {"ap-east-2"},
+        "0.22.1-deepspeed0.9.2-cu118": {"ap-east-2"},
+        "0.21.0-deepspeed0.8.3-cu117": {"ap-east-2"},
+        "0.20.0-deepspeed0.7.5-cu116": {"ap-east-2"},
+        "0.19.0-deepspeed0.7.3-cu113": {"ap-east-2"},
+    }
+}
+
 
-def test_djl_lmi_config_for_framework_has_all_regions():
-    """Test that config_for_framework('djl-lmi') returns all expected regions for each version."""
-    config = image_uris.config_for_framework("djl-lmi")
+@pytest.mark.parametrize("framework", ["djl-deepspeed", "djl-fastertransformer", "djl-lmi", "djl-neuronx", "djl-tensorrtllm"])
+def test_djl_lmi_config_for_framework_has_all_regions(framework):
+    """Test that config_for_framework returns all expected regions for each version."""
+    config = image_uris.config_for_framework(framework)
 
-    # Check that each version has all expected regions
+    # Check that each version has all expected regions, excluding known missing combinations
     for version, version_config in config["versions"].items():
         actual_regions = set(version_config["registries"].keys())
-        missing_regions = EXPECTED_DJL_LMI_REGIONS - actual_regions
+        expected_regions_for_version = EXPECTED_DJL_LMI_REGIONS.copy()
+        
+        # Use tag_prefix for lookup if available, otherwise fall back to version
+        lookup_key = version_config.get("tag_prefix", version)
+        
+        # Remove regions that are known to be missing for this framework:version combination
+        missing_regions_for_version = KNOWN_MISSING_COMBINATIONS.get(framework, {}).get(lookup_key, set())
+        expected_regions_for_version -= missing_regions_for_version
+        
+        missing_regions = expected_regions_for_version - actual_regions
 
-        assert not missing_regions, f"Version {version} missing regions: {missing_regions}"
+        assert not missing_regions, f"Framework {framework} version {version} missing regions: {missing_regions}"