From 4bf5363e1771cd7c692011c7ba59997679a2634e Mon Sep 17 00:00:00 2001 From: Yadan Wei Date: Tue, 1 Sep 2026 01:49:35 +0000 Subject: [PATCH] change: add ray/llama-cpp CPU images and move DLC serving frameworks to device-selectable configs The DLC serving-framework image_uri_configs added in #6218/#6220 exposed only GPU (cuda) images. DLC also publishes CPU images for ray-serve and llama-cpp; expose them, and prepare the remaining GPU-only frameworks so a CPU variant can be added later without changing how GPU callers resolve. All serving configs now use the image_uris processor schema (processors + processor_in_tag:false + a per-processor container_version tail) instead of a verbatim whole-tag: - ray-serve, llama-cpp: processors=[cpu, gpu]; instance_type selects the device. GPU tags unchanged; adds ray:serve-ml-sagemaker-cpu-v* and llama-cpp:server-sagemaker-cpu-v*. - vllm-server, vllm-omni, sglang-server, whisperx: processors=[gpu] only. Resolution is byte-identical to before (locked by literal-tag tests) and instance_type stays optional. Adding a CPU image later is a data-only change. Behavior change: for ray-serve and llama-cpp, instance_type is now required (previously defaulted to the GPU tag). For the GPU-only configs, a non-GPU instance type now raises instead of silently returning the GPU image. Both are safe: these configs shipped only in #6218/#6220. llama-cpp-arm64 (arm64 CPU, separate repo) is unchanged; select it by framework name. Tests restructured into whole-tag / gpu-only / multi-processor tiers with cpu+gpu coverage, required/optional instance_type checks, and literal repo:tag pins. --- .../core/image_uri_config/llama-cpp.json | 22 +- .../core/image_uri_config/ray-serve.json | 22 +- .../core/image_uri_config/sglang-server.json | 18 +- .../core/image_uri_config/vllm-omni.json | 18 +- .../core/image_uri_config/vllm-server.json | 18 +- .../core/image_uri_config/whisperx.json | 9 +- .../image_uris/test_dlc_serving_frameworks.py | 206 +++++++++++++++++- 7 files changed, 291 insertions(+), 22 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/image_uri_config/llama-cpp.json b/sagemaker-core/src/sagemaker/core/image_uri_config/llama-cpp.json index dd0e5241ef..2fdd19ed49 100644 --- a/sagemaker-core/src/sagemaker/core/image_uri_config/llama-cpp.json +++ b/sagemaker-core/src/sagemaker/core/image_uri_config/llama-cpp.json @@ -48,7 +48,16 @@ "us-west-2": "763104351884" }, "repository": "llama-cpp", - "tag_prefix": "server-sagemaker-cuda-v1" + "processors": [ + "cpu", + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "server-sagemaker", + "container_version": { + "cpu": "cpu-v1", + "gpu": "cuda-v1" + } }, "1.0": { "registries": { @@ -92,7 +101,16 @@ "us-west-2": "763104351884" }, "repository": "llama-cpp", - "tag_prefix": "server-sagemaker-cuda-v1.0" + "processors": [ + "cpu", + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "server-sagemaker", + "container_version": { + "cpu": "cpu-v1.0", + "gpu": "cuda-v1.0" + } } } } diff --git a/sagemaker-core/src/sagemaker/core/image_uri_config/ray-serve.json b/sagemaker-core/src/sagemaker/core/image_uri_config/ray-serve.json index 6fa72b8cb3..c410b5a29a 100644 --- a/sagemaker-core/src/sagemaker/core/image_uri_config/ray-serve.json +++ b/sagemaker-core/src/sagemaker/core/image_uri_config/ray-serve.json @@ -48,7 +48,16 @@ "us-west-2": "763104351884" }, "repository": "ray", - "tag_prefix": "serve-ml-sagemaker-cuda-v1" + "processors": [ + "cpu", + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "serve-ml-sagemaker", + "container_version": { + "cpu": "cpu-v1", + "gpu": "cuda-v1" + } }, "1.4": { "registries": { @@ -92,7 +101,16 @@ "us-west-2": "763104351884" }, "repository": "ray", - "tag_prefix": "serve-ml-sagemaker-cuda-v1.4" + "processors": [ + "cpu", + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "serve-ml-sagemaker", + "container_version": { + "cpu": "cpu-v1.4", + "gpu": "cuda-v1.4" + } } } } diff --git a/sagemaker-core/src/sagemaker/core/image_uri_config/sglang-server.json b/sagemaker-core/src/sagemaker/core/image_uri_config/sglang-server.json index 6e006a07a9..e0dec153a9 100644 --- a/sagemaker-core/src/sagemaker/core/image_uri_config/sglang-server.json +++ b/sagemaker-core/src/sagemaker/core/image_uri_config/sglang-server.json @@ -48,7 +48,14 @@ "us-west-2": "763104351884" }, "repository": "sglang", - "tag_prefix": "server-sagemaker-cuda-v1" + "processors": [ + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "server-sagemaker", + "container_version": { + "gpu": "cuda-v1" + } }, "1.3": { "registries": { @@ -92,7 +99,14 @@ "us-west-2": "763104351884" }, "repository": "sglang", - "tag_prefix": "server-sagemaker-cuda-v1.3" + "processors": [ + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "server-sagemaker", + "container_version": { + "gpu": "cuda-v1.3" + } } } } diff --git a/sagemaker-core/src/sagemaker/core/image_uri_config/vllm-omni.json b/sagemaker-core/src/sagemaker/core/image_uri_config/vllm-omni.json index 2d5f913f15..15ca82ab46 100644 --- a/sagemaker-core/src/sagemaker/core/image_uri_config/vllm-omni.json +++ b/sagemaker-core/src/sagemaker/core/image_uri_config/vllm-omni.json @@ -48,7 +48,14 @@ "us-west-2": "763104351884" }, "repository": "vllm", - "tag_prefix": "omni-sagemaker-cuda-v1" + "processors": [ + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "omni-sagemaker", + "container_version": { + "gpu": "cuda-v1" + } }, "1.6": { "registries": { @@ -92,7 +99,14 @@ "us-west-2": "763104351884" }, "repository": "vllm", - "tag_prefix": "omni-sagemaker-cuda-v1.6" + "processors": [ + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "omni-sagemaker", + "container_version": { + "gpu": "cuda-v1.6" + } } } } diff --git a/sagemaker-core/src/sagemaker/core/image_uri_config/vllm-server.json b/sagemaker-core/src/sagemaker/core/image_uri_config/vllm-server.json index 54a10f604b..d5e04cbb5c 100644 --- a/sagemaker-core/src/sagemaker/core/image_uri_config/vllm-server.json +++ b/sagemaker-core/src/sagemaker/core/image_uri_config/vllm-server.json @@ -48,7 +48,14 @@ "us-west-2": "763104351884" }, "repository": "vllm", - "tag_prefix": "server-sagemaker-cuda-v2" + "processors": [ + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "server-sagemaker", + "container_version": { + "gpu": "cuda-v2" + } }, "2.4": { "registries": { @@ -92,7 +99,14 @@ "us-west-2": "763104351884" }, "repository": "vllm", - "tag_prefix": "server-sagemaker-cuda-v2.4" + "processors": [ + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "server-sagemaker", + "container_version": { + "gpu": "cuda-v2.4" + } } } } diff --git a/sagemaker-core/src/sagemaker/core/image_uri_config/whisperx.json b/sagemaker-core/src/sagemaker/core/image_uri_config/whisperx.json index 62ea0640af..23b492da5c 100644 --- a/sagemaker-core/src/sagemaker/core/image_uri_config/whisperx.json +++ b/sagemaker-core/src/sagemaker/core/image_uri_config/whisperx.json @@ -48,7 +48,14 @@ "us-west-2": "763104351884" }, "repository": "whisperx", - "tag_prefix": "3.8-cu128-amzn2023-sagemaker" + "processors": [ + "gpu" + ], + "processor_in_tag": false, + "tag_prefix": "3.8", + "container_version": { + "gpu": "cu128-amzn2023-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 9f9909bec7..5004940daf 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 @@ -16,25 +16,40 @@ from sagemaker.core import image_uris -# GPU instance; these configs omit "processors", so the processor slot is dropped -# and the whole tag is taken verbatim from tag_prefix (instance type is ignored). +# GPU instance used wherever a gpu image is expected. INSTANCE_TYPE = "ml.g5.2xlarge" -# DLC serving-framework configs that use the whole-tag (djl-lmi style) pattern: -# each version's tag_prefix is the full image tag (channel or amzn2023), with no -# processor/py/container_version tokens appended. -SERVING_CONFIG_FILES = [ +# Instance types that resolve to each processor in image_uris._processor(). +# m5 is a general-purpose (CPU) family; g5 is a GPU family. +PROCESSOR_INSTANCE_TYPES = {"cpu": "ml.m5.xlarge", "gpu": "ml.g5.2xlarge"} + +# 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 = [ + "llama-cpp-arm64.json", +] + +# GPU-only configs on the processor schema: processors=["gpu"], processor_in_tag=false, +# and the tag tail in container_version["gpu"]. instance_type is optional today (single +# processor) and resolves to the gpu image; a cpu entry can be added later as pure data +# without changing how a gpu caller resolves. +GPU_ONLY_PROCESSOR_FILES = [ "vllm-server.json", "vllm-omni.json", "sglang-server.json", - "llama-cpp.json", - "llama-cpp-arm64.json", - "ray-serve.json", "whisperx.json", ] +# Configs shipping both a cpu and a gpu image under one repository: processors=["cpu","gpu"], +# processor_in_tag=false, per-processor container_version tail. instance_type selects the +# device and is therefore required. +MULTI_PROCESSOR_FILES = [ + "ray-serve.json", + "llama-cpp.json", +] -@pytest.mark.parametrize("load_config_and_file_name", SERVING_CONFIG_FILES, indirect=True) + +@pytest.mark.parametrize("load_config_and_file_name", WHOLE_TAG_CONFIG_FILES, indirect=True) def test_serving_framework_uris(load_config_and_file_name): """Every (version, region) resolves to the expected repo:tag verbatim.""" config, file_name = load_config_and_file_name @@ -56,7 +71,7 @@ def test_serving_framework_uris(load_config_and_file_name): assert uri.endswith(f"/{repo}:{tag}"), uri -@pytest.mark.parametrize("load_config_and_file_name", SERVING_CONFIG_FILES, indirect=True) +@pytest.mark.parametrize("load_config_and_file_name", WHOLE_TAG_CONFIG_FILES, indirect=True) def test_serving_framework_latest_alias(load_config_and_file_name): """The 'latest' alias resolves to its target version's tag.""" config, file_name = load_config_and_file_name @@ -71,3 +86,172 @@ def test_serving_framework_latest_alias(load_config_and_file_name): instance_type=INSTANCE_TYPE, ) assert uri.endswith(f"/{expected['repository']}:{expected['tag_prefix']}"), uri + + +@pytest.mark.parametrize("load_config_and_file_name", GPU_ONLY_PROCESSOR_FILES, indirect=True) +def test_gpu_only_processor_serving_framework_uris(load_config_and_file_name): + """GPU-only framework on the processor schema resolves to the gpu tail, and because it + has a single processor, omitting instance_type still yields the gpu 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"] + gpu_tail = version_config["container_version"]["gpu"] + expected_tag = f"{prefix}-{gpu_tail}" + for region, account in version_config["registries"].items(): + uri = image_uris.retrieve( + framework=framework, + region=region, + version=version, + image_scope="inference", + instance_type=INSTANCE_TYPE, + ) + 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 these configs 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 GPU_ONLY_PROCESSOR_FILES]) +def test_gpu_only_processor_rejects_cpu_instance(framework): + """Until a cpu image is added, a cpu instance type is rejected (not silently served gpu).""" + with pytest.raises(ValueError): + image_uris.retrieve( + framework=framework, + region="us-west-2", + version="latest", + image_scope="inference", + instance_type=PROCESSOR_INSTANCE_TYPES["cpu"], + ) + + +@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.""" + 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"] + for processor, tail in version_config["container_version"].items(): + expected_tag = f"{prefix}-{tail}" + instance_type = PROCESSOR_INSTANCE_TYPES[processor] + for region, account in version_config["registries"].items(): + uri = image_uris.retrieve( + framework=framework, + region=region, + version=version, + image_scope="inference", + instance_type=instance_type, + ) + assert uri.startswith(f"{account}.dkr.ecr.{region}."), uri + assert uri.endswith(f"/{repo}:{expected_tag}"), uri + + +@pytest.mark.parametrize("load_config_and_file_name", MULTI_PROCESSOR_FILES, indirect=True) +def test_processor_serving_framework_latest_alias(load_config_and_file_name): + """The 'latest' alias resolves to its target version's per-processor tag.""" + config, file_name = load_config_and_file_name + framework = file_name[: -len(".json")] + target = config["version_aliases"]["latest"] + expected = config["versions"][target] + repo = expected["repository"] + prefix = expected["tag_prefix"] + for processor, tail in expected["container_version"].items(): + uri = image_uris.retrieve( + framework=framework, + region="us-west-2", + version="latest", + image_scope="inference", + instance_type=PROCESSOR_INSTANCE_TYPES[processor], + ) + assert uri.endswith(f"/{repo}:{prefix}-{tail}"), uri + + +@pytest.mark.parametrize("framework", [f[: -len(".json")] for f in MULTI_PROCESSOR_FILES]) +def test_processor_serving_framework_requires_instance_type(framework): + """With both cpu and gpu offered, instance_type is required to disambiguate.""" + with pytest.raises(ValueError): + image_uris.retrieve( + framework=framework, + region="us-west-2", + version="latest", + image_scope="inference", + ) + + +# Exact repo:tag each (framework, version, processor) must resolve to. These pin the +# literal strings independent of the config dict: the gpu rows lock backward compatibility, +# the cpu rows lock the newly added tags. A self-consistent typo in tag_prefix/ +# container_version would fail here even though it passes the mechanism tests. +EXPECTED_REPO_TAGS = { + "ray-serve": { + ("1", "gpu"): "ray:serve-ml-sagemaker-cuda-v1", + ("1", "cpu"): "ray:serve-ml-sagemaker-cpu-v1", + ("1.4", "gpu"): "ray:serve-ml-sagemaker-cuda-v1.4", + ("1.4", "cpu"): "ray:serve-ml-sagemaker-cpu-v1.4", + }, + "llama-cpp": { + ("1", "gpu"): "llama-cpp:server-sagemaker-cuda-v1", + ("1", "cpu"): "llama-cpp:server-sagemaker-cpu-v1", + ("1.0", "gpu"): "llama-cpp:server-sagemaker-cuda-v1.0", + ("1.0", "cpu"): "llama-cpp:server-sagemaker-cpu-v1.0", + }, +} + +# gpu-only frameworks: the resolved gpu tag must be byte-identical to the pre-conversion +# whole-tag value, so the processor-schema change is a no-op for existing gpu callers. +GPU_ONLY_EXPECTED_REPO_TAGS = { + "vllm-server": { + "2": "vllm:server-sagemaker-cuda-v2", + "2.4": "vllm:server-sagemaker-cuda-v2.4", + }, + "vllm-omni": { + "1": "vllm:omni-sagemaker-cuda-v1", + "1.6": "vllm:omni-sagemaker-cuda-v1.6", + }, + "sglang-server": { + "1": "sglang:server-sagemaker-cuda-v1", + "1.3": "sglang:server-sagemaker-cuda-v1.3", + }, + "whisperx": { + "3.8": "whisperx:3.8-cu128-amzn2023-sagemaker", + }, +} + + +@pytest.mark.parametrize("framework", list(EXPECTED_REPO_TAGS)) +def test_processor_serving_framework_literal_tags(framework): + """Pin the exact repo:tag per (version, processor), not just the resolution mechanism.""" + for (version, processor), repo_tag in EXPECTED_REPO_TAGS[framework].items(): + uri = image_uris.retrieve( + framework=framework, + region="us-west-2", + version=version, + image_scope="inference", + instance_type=PROCESSOR_INSTANCE_TYPES[processor], + ) + assert uri.startswith("763104351884.dkr.ecr.us-west-2."), uri + assert uri.endswith(f"/{repo_tag}"), uri + + +@pytest.mark.parametrize("framework", list(GPU_ONLY_EXPECTED_REPO_TAGS)) +def test_gpu_only_processor_literal_tags(framework): + """The gpu tag is byte-identical to the pre-processor-schema (whole-tag) value.""" + for version, repo_tag in GPU_ONLY_EXPECTED_REPO_TAGS[framework].items(): + uri = image_uris.retrieve( + framework=framework, + region="us-west-2", + version=version, + image_scope="inference", + instance_type=INSTANCE_TYPE, + ) + assert uri.endswith(f"/{repo_tag}"), uri