fix: migrate llama-cpp-arm64 to cpu processor schema and align image_uri unit tests

sagemaker-core/src/sagemaker/core/image_uri_config/llama-cpp-arm64.json

@@ -7,6 +7,13 @@
     },
     "versions": {
         "1": {
+            "container_version": {
+                "cpu": "cpu-v1"
+            },
+            "processor_in_tag": false,
+            "processors": [
+                "cpu"
+            ],
             "registries": {
                 "af-south-1": "626614931356",
                 "ap-east-1": "871362719292",
@@ -48,9 +55,16 @@
                 "us-west-2": "763104351884"
             },
             "repository": "llama-cpp-arm64",
-            "tag_prefix": "server-sagemaker-cpu-v1"
+            "tag_prefix": "server-sagemaker"
         },
         "1.0": {
+            "container_version": {
+                "cpu": "cpu-v1.0"
+            },
+            "processor_in_tag": false,
+            "processors": [
+                "cpu"
+            ],
             "registries": {
                 "af-south-1": "626614931356",
                 "ap-east-1": "871362719292",
@@ -92,7 +106,7 @@
                 "us-west-2": "763104351884"
             },
             "repository": "llama-cpp-arm64",
-            "tag_prefix": "server-sagemaker-cpu-v1.0"
+            "tag_prefix": "server-sagemaker"
         }
     }
 }

sagemaker-core/tests/unit/image_uris/test_dlc_serving_frameworks.py

@@ -25,7 +25,14 @@
 
 # Single-variant configs whose tag_prefix is the full image tag, taken verbatim
 # (no processors / processor_in_tag / container_version). Instance type is ignored.
-WHOLE_TAG_CONFIG_FILES = [
+# (None currently: llama-cpp-arm64 migrated to the cpu processor schema below.)
+WHOLE_TAG_CONFIG_FILES = []
+
+# CPU-only configs on the processor schema: processors=["cpu"], processor_in_tag=false,
+# and the tag tail in container_version["cpu"]. Mirror of the gpu-only case: instance_type
+# is optional (single processor) and resolves to the cpu image; a gpu instance is rejected
+# until a gpu image is added.
+CPU_ONLY_PROCESSOR_FILES = [
     "llama-cpp-arm64.json",
 ]
 
@@ -133,6 +140,51 @@ def test_gpu_only_processor_rejects_cpu_instance(framework):
         )
 
 
+@pytest.mark.parametrize("load_config_and_file_name", CPU_ONLY_PROCESSOR_FILES, indirect=True)
+def test_cpu_only_processor_serving_framework_uris(load_config_and_file_name):
+    """CPU-only framework on the processor schema resolves to the cpu tail, and because it
+    has a single processor, omitting instance_type still yields the cpu image."""
+    config, file_name = load_config_and_file_name
+    framework = file_name[: -len(".json")]
+    for version, version_config in config["versions"].items():
+        repo = version_config["repository"]
+        prefix = version_config["tag_prefix"]
+        cpu_tail = version_config["container_version"]["cpu"]
+        expected_tag = f"{prefix}-{cpu_tail}"
+        for region, account in version_config["registries"].items():
+            uri = image_uris.retrieve(
+                framework=framework,
+                region=region,
+                version=version,
+                image_scope="inference",
+                instance_type=PROCESSOR_INSTANCE_TYPES["cpu"],
+            )
+            assert uri.startswith(f"{account}.dkr.ecr.{region}."), uri
+            assert uri.endswith(f"/{repo}:{expected_tag}"), uri
+        # instance_type is optional for a single-processor config (backward-compatible
+        # with the whole-tag form this config used before the processor-schema change).
+        uri_no_instance = image_uris.retrieve(
+            framework=framework,
+            region="us-west-2",
+            version=version,
+            image_scope="inference",
+        )
+        assert uri_no_instance.endswith(f"/{repo}:{expected_tag}"), uri_no_instance
+
+
+@pytest.mark.parametrize("framework", [f[: -len(".json")] for f in CPU_ONLY_PROCESSOR_FILES])
+def test_cpu_only_processor_rejects_gpu_instance(framework):
+    """Until a gpu image is added, a gpu instance type is rejected (not silently served cpu)."""
+    with pytest.raises(ValueError):
+        image_uris.retrieve(
+            framework=framework,
+            region="us-west-2",
+            version="latest",
+            image_scope="inference",
+            instance_type=PROCESSOR_INSTANCE_TYPES["gpu"],
+        )
+
+
 @pytest.mark.parametrize("load_config_and_file_name", MULTI_PROCESSOR_FILES, indirect=True)
 def test_processor_serving_framework_uris(load_config_and_file_name):
     """CPU/GPU share one config: the instance type selects the per-processor tag tail."""

sagemaker-core/tests/unit/image_uris/test_tensorflow.py

@@ -77,10 +77,23 @@ def test_tensorflow_latest_version_is_registered(load_config, scope):
 @pytest.mark.parametrize("scope", ["inference", "training"])
 @pytest.mark.parametrize("load_config", ["tensorflow.json"], indirect=True)
 def test_tensorflow_latest_version_registries_match_previous_release(load_config, scope):
-    """The newest version ships in the same regions and accounts as the previous release."""
+    """The newest version ships in at least the same regions and accounts as the previous
+    release. It may additionally carry a newly launched region that is still inside the
+    image-copy window and therefore not yet backfilled onto the previous release, so this
+    checks the reference registries are a subset (superset regressions are still caught)."""
     version, _ = LATEST[scope]
     versions = load_config[scope]["versions"]
-    assert versions[version]["registries"] == versions[REGISTRY_REFERENCE_VERSION]["registries"]
+    latest_registries = versions[version]["registries"]
+    reference_registries = versions[REGISTRY_REFERENCE_VERSION]["registries"]
+    missing = {
+        region: account
+        for region, account in reference_registries.items()
+        if latest_registries.get(region) != account
+    }
+    assert not missing, (
+        f"latest {scope} version {version} is missing regions/accounts present in the "
+        f"previous release {REGISTRY_REFERENCE_VERSION}: {missing}"
+    )
 
 
 @pytest.mark.parametrize("scope", ["inference", "training"])