From ffae72b7d27a3cedcab140ee31b32d77c13f9ff3 Mon Sep 17 00:00:00 2001 From: Harrison Saturley-Hall <454891+saturley-hall@users.noreply.github.com> Date: Fri, 15 Aug 2025 16:43:55 -0400 Subject: [PATCH] fix: remove kvmanager feature from python 3.12 ai-dynamo-runtime wheel (#2456) --- Cargo.lock | 18 +++++++++--------- Cargo.toml | 2 +- container/Dockerfile | 9 ++++++++- container/Dockerfile.vllm | 9 ++++++++- container/build.sh | 8 ++++++++ lib/bindings/python/Cargo.toml | 2 +- lib/bindings/python/pyproject.toml | 2 +- .../python/tests/test_block_manager.py | 8 +++++++- lib/runtime/examples/Cargo.toml | 2 +- pyproject.toml | 4 ++-- 10 files changed, 46 insertions(+), 18 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 21b8a2468..00bbc3538 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1814,7 +1814,7 @@ dependencies = [ [[package]] name = "dynamo-engine-llamacpp" -version = "0.4.0" +version = "0.4.0+post0" dependencies = [ "async-stream", "dynamo-llm", @@ -1826,7 +1826,7 @@ dependencies = [ [[package]] name = "dynamo-engine-mistralrs" -version = "0.4.0" +version = "0.4.0+post0" dependencies = [ "anyhow", "async-openai", @@ -1844,7 +1844,7 @@ dependencies = [ [[package]] name = "dynamo-llm" -version = "0.4.0" +version = "0.4.0+post0" dependencies = [ "ahash", "akin", @@ -1923,7 +1923,7 @@ dependencies = [ [[package]] name = "dynamo-run" -version = "0.4.0" +version = "0.4.0+post0" dependencies = [ "anyhow", "async-openai", @@ -1952,7 +1952,7 @@ dependencies = [ [[package]] name = "dynamo-runtime" -version = "0.4.0" +version = "0.4.0+post0" dependencies = [ "anyhow", "arc-swap", @@ -2009,7 +2009,7 @@ dependencies = [ [[package]] name = "dynamo-tokens" -version = "0.4.0" +version = "0.4.0+post0" dependencies = [ "bytemuck", "derive-getters", @@ -3776,7 +3776,7 @@ checksum = "d750af042f7ef4f724306de029d18836c26c1765a54a6a3f094cbd23a7267ffa" [[package]] name = "libdynamo_llm" -version = "0.4.0" +version = "0.4.0+post0" dependencies = [ "anyhow", "async-once-cell", @@ -4088,7 +4088,7 @@ dependencies = [ [[package]] name = "metrics" -version = "0.4.0" +version = "0.4.0+post0" dependencies = [ "axum 0.8.3", "clap 4.5.40", @@ -5899,7 +5899,7 @@ dependencies = [ [[package]] name = "router" -version = "0.4.0" +version = "0.4.0+post0" dependencies = [ "clap 4.5.40", "dynamo-llm", diff --git a/Cargo.toml b/Cargo.toml index ecb1965f5..3a9b0ecb1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,7 @@ members = [ resolver = "3" [workspace.package] -version = "0.4.0" +version = "0.4.0+post0" edition = "2021" description = "Dynamo Inference Framework" authors = ["NVIDIA Inc. "] diff --git a/container/Dockerfile b/container/Dockerfile index 2f5fbc336..612d53b93 100644 --- a/container/Dockerfile +++ b/container/Dockerfile @@ -8,6 +8,7 @@ ARG BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base" # can be updated to later versions. ARG BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04" ARG RELEASE_BUILD=false +ARG ENABLE_KVBM=false # Define general architecture ARGs for supporting both x86 and aarch64 builds. # ARCH: Used for package suffixes (e.g., amd64, arm64) @@ -197,6 +198,8 @@ ARG CARGO_BUILD_JOBS ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16} # Use build arg RELEASE_BUILD = true to generate wheels for Python 3.10, 3.11 and 3.12. ARG RELEASE_BUILD +# Use arg ENABLE_KVBM = true to turn on the block-manager feature +ARG ENABLE_KVBM WORKDIR /opt/dynamo @@ -228,7 +231,11 @@ COPY components/ /opt/dynamo/components/ RUN uv build --wheel --out-dir /opt/dynamo/dist && \ cd /opt/dynamo/lib/bindings/python && \ uv pip install maturin[patchelf] && \ - maturin build --release --features block-manager --out /opt/dynamo/dist && \ + if [ "$ENABLE_KVBM" = "true" ]; then \ + maturin build --release --features block-manager --out /opt/dynamo/dist; \ + else \ + maturin build --release --out /opt/dynamo/dist; \ + fi && \ if [ "$RELEASE_BUILD" = "true" ]; then \ # do not enable KVBM feature, ensure compatibility with lower glibc uv run --python 3.11 maturin build --release --out /opt/dynamo/dist && \ diff --git a/container/Dockerfile.vllm b/container/Dockerfile.vllm index 3bc6d0a09..71f23a144 100644 --- a/container/Dockerfile.vllm +++ b/container/Dockerfile.vllm @@ -8,6 +8,7 @@ ARG BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base" # can be updated to later versions. ARG BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04" ARG RELEASE_BUILD +ARG ENABLE_KVBM=false ARG RUNTIME_IMAGE="nvcr.io/nvidia/cuda" ARG RUNTIME_IMAGE_TAG="12.8.1-runtime-ubuntu24.04" @@ -320,6 +321,8 @@ ARG CARGO_BUILD_JOBS ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16} # Use build arg RELEASE_BUILD = true to generate wheels for Python 3.10, 3.11 and 3.12. ARG RELEASE_BUILD +# Use arg ENABLE_KVBM = true to turn on the block-manager feature +ARG ENABLE_KVBM # Keep in sync with the base image. ENV NIXL_PREFIX=/opt/nvidia/nvda_nixl @@ -369,7 +372,11 @@ RUN cargo build \ RUN uv build --wheel --out-dir /workspace/dist && \ cd /workspace/lib/bindings/python && \ uv pip install maturin[patchelf] && \ - maturin build --release --features block-manager --out /workspace/dist && \ + if [ "$ENABLE_KVBM" = "true" ]; then \ + maturin build --release --features block-manager --out /workspace/dist; \ + else \ + maturin build --release --out /workspace/dist; \ + fi && \ if [ "$RELEASE_BUILD" = "true" ]; then \ # do not enable KVBM feature, ensure compatibility with lower glibc uv run --python 3.11 maturin build --release --out /workspace/dist && \ diff --git a/container/build.sh b/container/build.sh index c1159f73b..3300c7b73 100755 --- a/container/build.sh +++ b/container/build.sh @@ -274,6 +274,9 @@ get_options() { --release-build) RELEASE_BUILD=true ;; + --enable-kvbm) + ENABLE_KVBM=true + ;; --make-efa) NIXL_UCX_REF=$NIXL_UCX_EFA_REF ;; @@ -530,6 +533,11 @@ if [ ! -z ${RELEASE_BUILD} ]; then BUILD_ARGS+=" --build-arg RELEASE_BUILD=${RELEASE_BUILD} " fi +if [ ! -z ${ENABLE_KVBM} ]; then + echo "Enabling the KVBM in the ai-dynamo-runtime" + BUILD_ARGS+=" --build-arg ENABLE_KVBM=${ENABLE_KVBM} " +fi + if [ -n "${NIXL_UCX_REF}" ]; then BUILD_ARGS+=" --build-arg NIXL_UCX_REF=${NIXL_UCX_REF} " fi diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index 3f631a5b2..58a84cf91 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -19,7 +19,7 @@ [package] name = "dynamo-py3" -version = "0.4.0" +version = "0.4.0+post0" edition = "2021" authors = ["NVIDIA"] license = "Apache-2.0" diff --git a/lib/bindings/python/pyproject.toml b/lib/bindings/python/pyproject.toml index 26647d20a..91ca9554f 100644 --- a/lib/bindings/python/pyproject.toml +++ b/lib/bindings/python/pyproject.toml @@ -16,7 +16,7 @@ [project] name = "ai-dynamo-runtime" -dynamic = ["version"] +version = "0.4.0.post0" description = "Dynamo Inference Framework Runtime" readme = "README.md" authors = [ diff --git a/lib/bindings/python/tests/test_block_manager.py b/lib/bindings/python/tests/test_block_manager.py index 94c7b455d..ea2c21c7e 100644 --- a/lib/bindings/python/tests/test_block_manager.py +++ b/lib/bindings/python/tests/test_block_manager.py @@ -19,7 +19,13 @@ import asyncio import pytest import torch -from dynamo.llm import BlockManager +# Attempt to import the optional module +try: + from dynamo.llm import BlockManager +except ImportError: + pytest.importorskip( + "optional_module", reason="block-manager feature is not enabled" + ) pytestmark = pytest.mark.pre_merge diff --git a/lib/runtime/examples/Cargo.toml b/lib/runtime/examples/Cargo.toml index 855f20385..cb9a7cb6b 100644 --- a/lib/runtime/examples/Cargo.toml +++ b/lib/runtime/examples/Cargo.toml @@ -22,7 +22,7 @@ members = [ resolver = "3" [workspace.package] -version = "0.4.0" +version = "0.4.0+post0" edition = "2021" authors = ["NVIDIA"] license = "Apache-2.0" diff --git a/pyproject.toml b/pyproject.toml index 69783cb5a..2719e6a9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ [project] name = "ai-dynamo" -version = "0.4.0" +version = "0.4.0.post0" description = "Distributed Inference Framework" readme = "README.md" authors = [ @@ -13,7 +13,7 @@ license = { text = "Apache-2.0" } license-files = ["LICENSE"] requires-python = ">=3.10" dependencies = [ - "ai-dynamo-runtime==0.4.0", + "ai-dynamo-runtime==0.4.0.post0", "pytest>=8.3.4", "types-psutil>=7.0.0.20250218", "kubernetes>=32.0.1,<33.0.0",