diff --git a/sagemaker-core/src/sagemaker/core/image_uri_config/llama-cpp-arm64.json b/sagemaker-core/src/sagemaker/core/image_uri_config/llama-cpp-arm64.json index d1d3d314ca..10fe51cd0c 100644 --- a/sagemaker-core/src/sagemaker/core/image_uri_config/llama-cpp-arm64.json +++ b/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" } } } diff --git a/sagemaker-core/tests/unit/image_uris/test_dlc_serving_frameworks.py b/sagemaker-core/tests/unit/image_uris/test_dlc_serving_frameworks.py index 5004940daf..aa94f05281 100644 --- a/sagemaker-core/tests/unit/image_uris/test_dlc_serving_frameworks.py +++ b/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.""" diff --git a/sagemaker-core/tests/unit/image_uris/test_tensorflow.py b/sagemaker-core/tests/unit/image_uris/test_tensorflow.py index e5fc733df2..68c8656402 100644 --- a/sagemaker-core/tests/unit/image_uris/test_tensorflow.py +++ b/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"])