diff --git a/.cargo/config.toml b/.cargo/config.toml index 02ef6843e995..589d6e2a5ae8 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -19,3 +19,6 @@ rustflags = ["-C", "target-cpu=neoverse-n1", "-C", "force-frame-pointers=yes", " [env] PCRE2_SYS_STATIC = "1" +[registries.buf] +index = "sparse+https://buf.build/gen/cargo/" +credential-provider = "cargo:token" diff --git a/.github/actions/compliance-extract/action.yml b/.github/actions/compliance-extract/action.yml index a66a35f79915..a9b25a8f10a8 100644 --- a/.github/actions/compliance-extract/action.yml +++ b/.github/actions/compliance-extract/action.yml @@ -61,6 +61,9 @@ inputs: description: 'SCCACHE region (= aws_default_region the build passed). See sccache_bucket.' required: false default: '' + buf_token: + description: 'Buf Schema Registry token for cache-missed Cargo rebuilds.' + required: true diff_base_sha: description: | Commit SHA of the baseline build to diff this build's OSRB CSV against @@ -126,6 +129,7 @@ runs: env: SCCACHE_BUCKET: ${{ inputs.sccache_bucket }} SCCACHE_REGION: ${{ inputs.sccache_region }} + BUF_TOKEN: ${{ inputs.buf_token }} GIT_SHA: ${{ inputs.git_sha }} EPP_IMAGE: ${{ inputs.epp_image }} run: | @@ -149,6 +153,12 @@ runs: # Only forward the credentials when S3 sccache is actually configured # (bucket set) — a bucket-less caller's build has no use for them. SECRET_ARGS="" + if [ -n "${BUF_TOKEN:-}" ]; then + SECRET_ARGS+=" --secret id=buf_token,env=BUF_TOKEN" + else + echo "::error::BUF_TOKEN is required for Cargo access to the Buf registry" + exit 1 + fi if [ -n "${SCCACHE_BUCKET:-}" ]; then TOKEN_FILE="${AWS_WEB_IDENTITY_TOKEN_FILE:-}" if [ -n "$TOKEN_FILE" ] && [ -f "$TOKEN_FILE" ] && [ -n "${AWS_ROLE_ARN:-}" ]; then @@ -344,6 +354,7 @@ runs: env: SCCACHE_BUCKET: ${{ inputs.sccache_bucket }} SCCACHE_REGION: ${{ inputs.sccache_region }} + BUF_TOKEN: ${{ inputs.buf_token }} GIT_SHA: ${{ inputs.git_sha }} EPP_IMAGE: ${{ inputs.epp_image }} run: | @@ -353,6 +364,12 @@ runs: # as the compliance_artifact extract above). Same IRSA secrets too, so a # cache-missed stage rebuild can still authenticate sccache (see above). SECRET_ARGS="" + if [ -n "${BUF_TOKEN:-}" ]; then + SECRET_ARGS+=" --secret id=buf_token,env=BUF_TOKEN" + else + echo "::error::BUF_TOKEN is required for Cargo access to the Buf registry" + exit 1 + fi if [ -n "${SCCACHE_BUCKET:-}" ]; then TOKEN_FILE="${AWS_WEB_IDENTITY_TOKEN_FILE:-}" if [ -n "$TOKEN_FILE" ] && [ -f "$TOKEN_FILE" ] && [ -n "${AWS_ROLE_ARN:-}" ]; then diff --git a/.github/actions/docker-remote-build/action.yml b/.github/actions/docker-remote-build/action.yml index 47dfae7e97fa..41722566d257 100644 --- a/.github/actions/docker-remote-build/action.yml +++ b/.github/actions/docker-remote-build/action.yml @@ -25,6 +25,9 @@ inputs: sccache_s3_bucket: description: 'SCCache S3 Bucket' required: false + buf_token: + description: 'Buf Schema Registry token' + required: true no_cache: description: 'Disable Docker build cache' required: false @@ -66,6 +69,7 @@ runs: env: AWS_DEFAULT_REGION: ${{ inputs.aws_default_region }} SCCACHE_S3_BUCKET: ${{ inputs.sccache_s3_bucket }} + BUF_TOKEN: ${{ inputs.buf_token }} PLATFORM: ${{ inputs.platform }} GITHUB_RUN_ID: ${{ github.run_id }} GITHUB_JOB: ${{ github.job }} @@ -137,6 +141,12 @@ runs: # AWS_ROLE_ARN. We pass the token file and role ARN to BuildKit so sccache # can authenticate via STS AssumeRoleWithWebIdentity -- no static keys needed. SECRET_ARGS="" + if [ -n "${BUF_TOKEN:-}" ]; then + SECRET_ARGS+=" --secret id=buf_token,env=BUF_TOKEN" + else + echo "::error::BUF_TOKEN is required for Cargo access to the Buf registry" + exit 1 + fi if [ "${{ inputs.use_sccache }}" == "true" ]; then TOKEN_FILE="${AWS_WEB_IDENTITY_TOKEN_FILE:-}" if [ -n "$TOKEN_FILE" ] && [ -f "$TOKEN_FILE" ] && [ -n "${AWS_ROLE_ARN:-}" ]; then diff --git a/.github/filters.yaml b/.github/filters.yaml index 3a56722c349a..cd472fdb8def 100644 --- a/.github/filters.yaml +++ b/.github/filters.yaml @@ -310,9 +310,6 @@ rust: - '**/Cargo.toml' - '**/Cargo.lock' - 'deny.toml' - # Sidecar protobuf contracts are compiled by crate build scripts and need the - # same workspace checks as sidecar Rust sources. - - 'lib/sidecar/**/*.proto' benchmarks: - 'benchmarks/**' diff --git a/.github/scripts/test-filters.js b/.github/scripts/test-filters.js index 0669371fd115..a91f409b5baf 100755 --- a/.github/scripts/test-filters.js +++ b/.github/scripts/test-filters.js @@ -94,16 +94,16 @@ const testCases = [ desc: 'vllm component triggers only vllm' }, - // Sidecar Rust and proto files should trigger Rust checks without unrelated E2E + // Sidecar Rust files should trigger Rust checks without unrelated E2E { file: 'lib/sidecar/common/src/lib.rs', expect: { sidecar: true, rust: true, core: false, frontend: false, vllm: false, sglang: false, trtllm: false }, desc: 'common sidecar source avoids unrelated build and E2E filters' }, { - file: 'lib/sidecar/vllm/proto/vllm_grpc.proto', + file: 'lib/sidecar/vllm/build.rs', expect: { sidecar: true, rust: true, core: false, frontend: false, vllm: false, sglang: false, trtllm: false }, - desc: 'vllm sidecar proto triggers Rust checks without backend E2E' + desc: 'vllm sidecar build script triggers Rust checks without backend E2E' }, { file: 'lib/sidecar/sglang/src/lib.rs', diff --git a/.github/workflows/copyright-check.ps1 b/.github/workflows/copyright-check.ps1 index eeb0d84692f7..b22f6bcc35b9 100644 --- a/.github/workflows/copyright-check.ps1 +++ b/.github/workflows/copyright-check.ps1 @@ -84,7 +84,7 @@ $global:copyright_results = @{ $ignored_files = @('.clang-format', '.gitattributes', '.gitignore', '.gitkeep', '.patch', 'Cargo.lock', 'LICENSE', 'uv.lock', 'rust-toolchain.toml', 'codespell.txt', 'exclusions.txt') write-debug " ignored_files = ['$($ignored_files -join "','")']." -$ignored_paths = @('.github', '.mypy_cache', '.pytest_cache', 'lib/llm/tests/data/sample-models', 'lib/llm/tests/data/deepseek-v3.2', 'lib/llm/tests/data/deepseek-v4', 'container/compliance/spdx_licenses', 'lib/sidecar/vllm/proto/control.proto', 'lib/sidecar/vllm/proto/inference.proto') +$ignored_paths = @('.github', '.mypy_cache', '.pytest_cache', 'lib/llm/tests/data/sample-models', 'lib/llm/tests/data/deepseek-v3.2', 'lib/llm/tests/data/deepseek-v4', 'container/compliance/spdx_licenses') write-debug " ignored_paths = ['$($ignored_paths -join "','")']." $ignored_types = @('.bat', '.gif', '.ico', '.ipynb', '.jpg', '.jpeg', '.patch', '.png', '.pyc', '.pyi', '.rst', '.zip', '.md', '.json') write-debug " ignored_types = ['$($ignored_types -join "', '")']." diff --git a/.github/workflows/dynamo-pipeline.yml b/.github/workflows/dynamo-pipeline.yml index 6a1524da146f..11f4f70d63e3 100644 --- a/.github/workflows/dynamo-pipeline.yml +++ b/.github/workflows/dynamo-pipeline.yml @@ -73,6 +73,8 @@ on: required: false HF_TOKEN: required: false + BUF_TOKEN: + required: true jobs: @@ -113,6 +115,7 @@ jobs: # "Permission denied (os error 13)" while downloading crates. Redirect # CARGO_HOME to the runner-writable workspace so cargo owns its cache. CARGO_HOME: /__w/dynamo/dynamo/.cargo + CARGO_REGISTRIES_BUF_TOKEN: "Bearer ${{ secrets.BUF_TOKEN }}" CONTAINER_ID: test_${{ github.run_id }}_${{ github.run_attempt }}_rust_dynamo timeout-minutes: 30 steps: diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index 049d620d63f3..f30af50abbdb 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -965,6 +965,8 @@ jobs: dir: ['.', 'lib/bindings/python', 'lib/bindings/kvbm'] permissions: contents: read + env: + CARGO_REGISTRIES_BUF_TOKEN: "Bearer ${{ secrets.BUF_TOKEN }}" steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: diff --git a/.github/workflows/pre-merge.yml b/.github/workflows/pre-merge.yml index 15291e642b4b..ccdddaf8c07b 100644 --- a/.github/workflows/pre-merge.yml +++ b/.github/workflows/pre-merge.yml @@ -237,6 +237,8 @@ jobs: matrix: { dir: ['.', 'lib/bindings/python', 'lib/runtime/examples', 'lib/bindings/kvbm'] } permissions: contents: read + env: + CARGO_REGISTRIES_BUF_TOKEN: "Bearer ${{ secrets.BUF_TOKEN }}" steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: @@ -341,6 +343,8 @@ jobs: matrix: { dir: ['.', 'lib/bindings/python', 'lib/runtime/examples', 'lib/bindings/kvbm'] } permissions: contents: read + env: + CARGO_REGISTRIES_BUF_TOKEN: "Bearer ${{ secrets.BUF_TOKEN }}" steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: diff --git a/.github/workflows/shared-build-image.yml b/.github/workflows/shared-build-image.yml index c2cbb2642c3b..690ab9304ce8 100644 --- a/.github/workflows/shared-build-image.yml +++ b/.github/workflows/shared-build-image.yml @@ -186,6 +186,8 @@ on: required: false HF_TOKEN: required: false + BUF_TOKEN: + required: true outputs: target_tag_plain: description: 'Plain runtime image tag prefix' @@ -370,6 +372,7 @@ jobs: IMAGE_REPOSITORY: ${{ vars.ECR_REPOSITORY }} AWS_DEFAULT_REGION: ${{ vars.AWS_DEFAULT_REGION }} SCCACHE_S3_BUCKET: ${{ secrets.SCCACHE_S3_BUCKET }} + BUF_TOKEN: ${{ secrets.BUF_TOKEN }} timeout-minutes: 60 run: | set -x @@ -500,6 +503,7 @@ jobs: cuda_version: ${{ matrix.cuda_version }} aws_default_region: ${{ vars.AWS_DEFAULT_REGION }} sccache_s3_bucket: ${{ secrets.SCCACHE_S3_BUCKET }} + buf_token: ${{ secrets.BUF_TOKEN }} no_cache: ${{ inputs.no_cache }} extra_tags: ${{ steps.extra-tags.outputs.tags }} push_image: ${{ inputs.push_image }} @@ -606,6 +610,7 @@ jobs: # wheel_builder/pre_runtime cache instead of a cold rebuild. sccache_bucket: ${{ secrets.SCCACHE_S3_BUCKET }} sccache_region: ${{ vars.AWS_DEFAULT_REGION }} + buf_token: ${{ secrets.BUF_TOKEN }} # Must match the build's render: EFA images attribute libfabric / # aws-ofi-nccl via --make-efa, and a mismatch would cold-miss the cache. make_efa: ${{ inputs.make_efa }} diff --git a/Cargo.lock b/Cargo.lock index 0be7008ef562..c0dbbb9157fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2513,7 +2513,7 @@ dependencies = [ "tokio-util", "tonic 0.13.1", "tonic-build 0.13.1", - "tonic-health", + "tonic-health 0.13.1", "tracing", "tracing-subscriber", "uuid", @@ -2990,11 +2990,13 @@ dependencies = [ name = "dynamo-sidecar-common" version = "1.4.0" dependencies = [ + "async-trait", "clap", "dynamo-backend-common", "futures", "tokio", "tonic 0.13.1", + "tonic 0.14.6", "tracing", "url", ] @@ -3080,11 +3082,11 @@ dependencies = [ "dynamo-mocker", "dynamo-vllm-sidecar", "futures", - "prost-types 0.13.5", + "prost-types 0.14.3", "tokio", "tokio-stream", - "tonic 0.13.1", - "tonic-health", + "tonic 0.14.6", + "tonic-health 0.14.6", "tracing", "tracing-subscriber", "uuid", @@ -3101,16 +3103,16 @@ dependencies = [ "dynamo-backend-common", "dynamo-sidecar-common", "futures", - "prost 0.13.5", - "prost-types 0.13.5", + "prost-types 0.14.3", "serde_json", "tokio", "tokio-stream", "tokio-util", - "tonic 0.13.1", - "tonic-build 0.13.1", - "tonic-health", + "tonic 0.14.6", + "tonic-health 0.14.6", "tracing", + "vllm-project_vllm_community_neoeinstein-prost", + "vllm-project_vllm_community_neoeinstein-tonic", ] [[package]] @@ -9869,6 +9871,19 @@ dependencies = [ "tonic 0.13.1", ] +[[package]] +name = "tonic-health" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082" +dependencies = [ + "prost 0.14.3", + "tokio", + "tokio-stream", + "tonic 0.14.6", + "tonic-prost", +] + [[package]] name = "tonic-prost" version = "0.14.6" @@ -10701,6 +10716,27 @@ version = "0.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1" +[[package]] +name = "vllm-project_vllm_community_neoeinstein-prost" +version = "0.5.0-00000000000000-7726adbdafb3.2" +source = "registry+sparse+https://buf.build/gen/cargo/" +checksum = "0f83f6ba9c6750bc44f1bfac07dde88cb781385b1bd3919baba244d0dbd4089d" +dependencies = [ + "prost 0.14.3", + "prost-types 0.14.3", +] + +[[package]] +name = "vllm-project_vllm_community_neoeinstein-tonic" +version = "0.5.0-00000000000000-7726adbdafb3.4" +source = "registry+sparse+https://buf.build/gen/cargo/" +checksum = "350ca64b22b0397d015661ee1958f16af5f6ec7c54a984b0a9b653fb4cdd44f3" +dependencies = [ + "tonic 0.14.6", + "tonic-prost", + "vllm-project_vllm_community_neoeinstein-prost", +] + [[package]] name = "vsimd" version = "0.8.0" diff --git a/README.md b/README.md index 274366dd854b..789f33adc2c9 100644 --- a/README.md +++ b/README.md @@ -204,6 +204,10 @@ sudo apt install -y build-essential libhwloc-dev libudev-dev pkg-config libclang # Install Rust curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh && source $HOME/.cargo/env +# Authenticate Cargo to the Buf Schema Registry +export BUF_TOKEN="your-buf-token" +cargo login --registry buf "Bearer ${BUF_TOKEN}" + # Create venv and build uv venv dynamo && source dynamo/bin/activate uv pip install pip 'maturin[patchelf]' diff --git a/container/templates/wheel_builder.Dockerfile b/container/templates/wheel_builder.Dockerfile index ef21df68f540..e0706e22e005 100644 --- a/container/templates/wheel_builder.Dockerfile +++ b/container/templates/wheel_builder.Dockerfile @@ -571,11 +571,13 @@ ARG USE_SCCACHE {% if framework != "sglang" %} ARG ENABLE_MEDIA_FFMPEG {% endif %} -RUN --mount=type=secret,id=aws-web-identity-token,target=/run/secrets/aws-token \ +RUN --mount=type=secret,id=buf_token,required=true \ + --mount=type=secret,id=aws-web-identity-token,target=/run/secrets/aws-token \ --mount=type=secret,id=aws-role-arn,env=AWS_ROLE_ARN \ --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ --mount=type=cache,target=/root/.cargo/git,sharing=shared \ --mount=type=cache,id=uv-root-{{ context.dynamo.uv_version }},target=/root/.cache/uv,sharing=shared \ + export CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" && \ export AWS_WEB_IDENTITY_TOKEN_FILE=/run/secrets/aws-token && \ export UV_CACHE_DIR=/root/.cache/uv && \ export SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX:-${TARGETARCH}} && \ @@ -655,8 +657,10 @@ ARG ENABLE_SOURCE_ARCHIVAL=false # Mount cargo registry + git caches so re-runs don't re-download the # ~750 crates from crates.io every build. `sharing=shared` lets parallel # builds (e.g. multiple frameworks in CI) read the same cache concurrently. -RUN --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ +RUN --mount=type=secret,id=buf_token,required=true \ + --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ --mount=type=cache,target=/root/.cargo/git,sharing=shared \ + export CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" && \ if [ "$ENABLE_SOURCE_ARCHIVAL" = "true" ]; then \ mkdir -p /tmp/dynamo-vendor-full && \ cd /opt/dynamo && \ @@ -796,11 +800,13 @@ COPY components/ /opt/dynamo/components/ # Build kvbm wheel (with nixl linkage via auditwheel repair) ARG ENABLE_KVBM -RUN --mount=type=secret,id=aws-web-identity-token,target=/run/secrets/aws-token \ +RUN --mount=type=secret,id=buf_token,required=true \ + --mount=type=secret,id=aws-web-identity-token,target=/run/secrets/aws-token \ --mount=type=secret,id=aws-role-arn,env=AWS_ROLE_ARN \ --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ --mount=type=cache,target=/root/.cargo/git,sharing=shared \ --mount=type=cache,id=uv-root-{{ context.dynamo.uv_version }},target=/root/.cache/uv,sharing=shared \ + export CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" && \ export AWS_WEB_IDENTITY_TOKEN_FILE=/run/secrets/aws-token && \ export UV_CACHE_DIR=/root/.cache/uv && \ export SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX:-${TARGETARCH}} && \ diff --git a/deploy/inference-gateway/epp/Dockerfile b/deploy/inference-gateway/epp/Dockerfile index f0e2e694e359..4eedb8165589 100644 --- a/deploy/inference-gateway/epp/Dockerfile +++ b/deploy/inference-gateway/epp/Dockerfile @@ -87,10 +87,12 @@ COPY --from=dynamo deploy/inference-gateway/ext-proc/ deploy/inference-gateway/e # git caches are content-addressed (safe to persist); no target/ mount -- # sccache caches compilations in S3 where stale artifacts can't be linked # against newer source. -RUN --mount=type=cache,target=/usr/local/cargo/registry,id=cargo-registry-${TARGETARCH},sharing=shared \ +RUN --mount=type=secret,id=buf_token,required=true \ + --mount=type=cache,target=/usr/local/cargo/registry,id=cargo-registry-${TARGETARCH},sharing=shared \ --mount=type=cache,target=/usr/local/cargo/git,id=cargo-git-${TARGETARCH},sharing=shared \ --mount=type=secret,id=aws-web-identity-token,target=/run/secrets/aws-token \ --mount=type=secret,id=aws-role-arn,env=AWS_ROLE_ARN \ + export CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" && \ export AWS_WEB_IDENTITY_TOKEN_FILE=/run/secrets/aws-token && \ export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-epp-${TARGETARCH}}" && \ if [ "$USE_SCCACHE" = "true" ]; then \ diff --git a/deploy/inference-gateway/epp/Makefile b/deploy/inference-gateway/epp/Makefile index eaf2133a5623..d42bc68cb91b 100644 --- a/deploy/inference-gateway/epp/Makefile +++ b/deploy/inference-gateway/epp/Makefile @@ -26,6 +26,7 @@ MULTIARCH_PLATFORMS ?= linux/amd64,linux/arm64 # Docker proxy for avoiding rate limits (e.g., ECR mirror) DOCKER_PROXY ?= EXTRA_BUILD_ARGS ?= +BUF_SECRET_ARGS = $(if $(BUF_TOKEN),--secret id=buf_token,env=BUF_TOKEN,$(error BUF_TOKEN is required for Cargo access to the Buf registry)) # sccache configuration for Rust compilation caching (CI only). # Leave USE_SCCACHE unset locally to build without S3 cache. @@ -96,7 +97,7 @@ image-build: ## Build the Docker image (self-contained, no host prerequisites) --build-arg BUILDER_IMAGE=$(BUILDER_IMAGE) \ --build-arg COMMIT_SHA=$(GIT_COMMIT_SHA) \ --build-arg BUILD_REF=$(GIT_TAG) \ - $(EXTRA_BUILD_ARGS) $(PUSH) $(LOAD) . + $(BUF_SECRET_ARGS) $(EXTRA_BUILD_ARGS) $(PUSH) $(LOAD) . .PHONY: image-push image-push: PUSH=--push ## Build and push the Docker image @@ -138,7 +139,7 @@ image-multiarch: ## Build multi-arch image (requires --push ; --load not support --build-arg BUILDER_IMAGE=$(BUILDER_IMAGE) \ --build-arg COMMIT_SHA=$(GIT_COMMIT_SHA) \ --build-arg BUILD_REF=$(GIT_TAG) \ - $(SCCACHE_ARGS) $(EXTRA_BUILD_ARGS) $(PUSH) . + $(BUF_SECRET_ARGS) $(SCCACHE_ARGS) $(EXTRA_BUILD_ARGS) $(PUSH) . .PHONY: image-multiarch-push image-multiarch-push: PUSH=--push ## Build and push multi-arch image to registry @@ -164,7 +165,7 @@ sbom-export: ## Export Go SBOM + license texts to SBOM_DEST (reuses build cache) --build-arg BUILDER_IMAGE=$(BUILDER_IMAGE) \ --build-arg COMMIT_SHA=$(GIT_COMMIT_SHA) \ --build-arg BUILD_REF=$(GIT_TAG) \ - $(SCCACHE_ARGS) $(EXTRA_BUILD_ARGS) \ + $(BUF_SECRET_ARGS) $(SCCACHE_ARGS) $(EXTRA_BUILD_ARGS) \ --output type=local,dest=$(SBOM_DEST) . diff --git a/deploy/inference-gateway/ext-proc/Dockerfile b/deploy/inference-gateway/ext-proc/Dockerfile index 0134980d01d7..a3b91ff04e5d 100644 --- a/deploy/inference-gateway/ext-proc/Dockerfile +++ b/deploy/inference-gateway/ext-proc/Dockerfile @@ -71,10 +71,12 @@ COPY --from=dynamo lib/ lib/ COPY --from=dynamo deploy/inference-gateway/ext-proc/ deploy/inference-gateway/ext-proc/ # Build the binary -RUN --mount=type=cache,target=/usr/local/cargo/registry,id=cargo-registry-${TARGETARCH} \ +RUN --mount=type=secret,id=buf_token,required=true \ + --mount=type=cache,target=/usr/local/cargo/registry,id=cargo-registry-${TARGETARCH} \ --mount=type=cache,target=/usr/local/cargo/git,id=cargo-git-${TARGETARCH} \ --mount=type=secret,id=aws-web-identity-token,target=/run/secrets/aws-token \ --mount=type=secret,id=aws-role-arn,env=AWS_ROLE_ARN \ + export CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" && \ export AWS_WEB_IDENTITY_TOKEN_FILE=/run/secrets/aws-token && \ export SCCACHE_S3_KEY_PREFIX="${SCCACHE_S3_KEY_PREFIX:-rust-epp-${TARGETARCH}}" && \ if [ "$USE_SCCACHE" = "true" ]; then \ diff --git a/deploy/inference-gateway/ext-proc/Makefile b/deploy/inference-gateway/ext-proc/Makefile index 0e5df9e2d751..19eaa6ea8ffc 100644 --- a/deploy/inference-gateway/ext-proc/Makefile +++ b/deploy/inference-gateway/ext-proc/Makefile @@ -26,6 +26,7 @@ endif MULTIARCH_PLATFORMS ?= linux/amd64,linux/arm64 DOCKER_PROXY ?= EXTRA_BUILD_ARGS ?= +BUF_SECRET_ARGS = $(if $(BUF_TOKEN),--secret id=buf_token,env=BUF_TOKEN,$(error BUF_TOKEN is required for Cargo access to the Buf registry)) # sccache configuration (CI only) USE_SCCACHE ?= @@ -88,7 +89,7 @@ image-build: ## Build the Docker image --build-context dynamo=$(DYNAMO_DIR) \ --build-arg RUST_IMAGE=$(RUST_IMAGE) \ --build-arg BASE_IMAGE=$(BASE_IMAGE) \ - $(EXTRA_BUILD_ARGS) $(PUSH) $(LOAD) . + $(BUF_SECRET_ARGS) $(EXTRA_BUILD_ARGS) $(PUSH) $(LOAD) . .PHONY: image-push image-push: PUSH=--push ## Build and push the Docker image @@ -124,7 +125,7 @@ image-multiarch: ## Build multi-arch image (requires --push) --build-context dynamo=$(DYNAMO_DIR) \ --build-arg RUST_IMAGE=$(RUST_IMAGE) \ --build-arg BASE_IMAGE=$(BASE_IMAGE) \ - $(SCCACHE_ARGS) $(EXTRA_BUILD_ARGS) $(PUSH) . + $(BUF_SECRET_ARGS) $(SCCACHE_ARGS) $(EXTRA_BUILD_ARGS) $(PUSH) . .PHONY: image-multiarch-push image-multiarch-push: PUSH=--push ## Build and push multi-arch image diff --git a/docs/fern/pages/developer-guide/advanced-customizations/building-from-source.md b/docs/fern/pages/developer-guide/advanced-customizations/building-from-source.md index 8657f448cafb..28aca239886f 100644 --- a/docs/fern/pages/developer-guide/advanced-customizations/building-from-source.md +++ b/docs/fern/pages/developer-guide/advanced-customizations/building-from-source.md @@ -34,6 +34,13 @@ curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh source $HOME/.cargo/env ``` +The workspace lockfile includes generated vLLM SDKs from the Buf Schema Registry, so Cargo requires [BSR authentication](https://buf.build/docs/bsr/generated-sdks/cargo/) for any workspace build: + +```bash +export BUF_TOKEN="your-buf-token" +cargo login --registry buf "Bearer ${BUF_TOKEN}" +``` + ## 3. Create a Python Virtual Environment Install [uv](https://docs.astral.sh/uv/#installation) if you don't have it: diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/mocker/rust-backend.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/mocker/rust-backend.md index 8d30ab754e3f..bb9b22043260 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/mocker/rust-backend.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/mocker/rust-backend.md @@ -25,6 +25,7 @@ backend — all built from source in this repo: ```bash cd lib/backend-common/examples/mocker +export BUF_TOKEN="your-buf-token" docker compose up --build ``` diff --git a/lib/backend-common/examples/mocker/Dockerfile b/lib/backend-common/examples/mocker/Dockerfile index 67929b08300f..022c63ab2ece 100644 --- a/lib/backend-common/examples/mocker/Dockerfile +++ b/lib/backend-common/examples/mocker/Dockerfile @@ -31,9 +31,11 @@ COPY . . # /build/target — compiled artifacts, lock while writing # Cache mounts are NOT part of the resulting image, so copy the binary out # of the cache to /out/ within the same RUN so later stages can COPY it. -RUN --mount=type=cache,target=/usr/local/cargo/registry,sharing=shared \ +RUN --mount=type=secret,id=buf_token,required=true \ + --mount=type=cache,target=/usr/local/cargo/registry,sharing=shared \ --mount=type=cache,target=/build/target,sharing=locked \ - mkdir -p /out \ + export CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" \ + && mkdir -p /out \ && cargo build -p dynamo-mocker-backend --release \ && cp /build/target/release/dynamo-mocker-backend /out/ diff --git a/lib/backend-common/examples/mocker/Dockerfile.frontend b/lib/backend-common/examples/mocker/Dockerfile.frontend index 740b3b066444..8d51f173f493 100644 --- a/lib/backend-common/examples/mocker/Dockerfile.frontend +++ b/lib/backend-common/examples/mocker/Dockerfile.frontend @@ -38,9 +38,11 @@ COPY . . # written so incremental rebuilds don't trample each other. The wheel # output goes to /tmp/wheels which is NOT a cache mount, so it persists into # the next stage. -RUN --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ +RUN --mount=type=secret,id=buf_token,required=true \ + --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ --mount=type=cache,target=/build/lib/bindings/python/target,sharing=locked \ - cd lib/bindings/python \ + export CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" \ + && cd lib/bindings/python \ && maturin build --release --out /tmp/wheels FROM python:3.12-slim-bookworm diff --git a/lib/backend-common/examples/mocker/docker-compose.yml b/lib/backend-common/examples/mocker/docker-compose.yml index cac9606a917c..907c6e85dff0 100644 --- a/lib/backend-common/examples/mocker/docker-compose.yml +++ b/lib/backend-common/examples/mocker/docker-compose.yml @@ -5,6 +5,7 @@ # scheduler in the `LLMEngine` contract from `dynamo-backend-common`. # # docker compose up --build +# (requires BUF_TOKEN in the shell environment) # curl http://localhost:8000/v1/chat/completions \ # -H 'Content-Type: application/json' \ # -d '{"model": "mocker-model", @@ -64,6 +65,8 @@ services: build: context: ../../../.. dockerfile: lib/backend-common/examples/mocker/Dockerfile.frontend + secrets: + - buf_token command: ["--http-port", "8000"] environment: - NATS_SERVER=nats://nats-server:4222 @@ -88,6 +91,8 @@ services: # Build context is the workspace root so cargo can see all crates. context: ../../../.. dockerfile: lib/backend-common/examples/mocker/Dockerfile + secrets: + - buf_token # --model-path points at a real HF repo so the frontend can load a # tokenizer + chat template. The engine still emits mocked token # IDs (no weights needed); Qwen3-0.6B is just a small, openly @@ -123,3 +128,7 @@ services: volumes: huggingface-cache: + +secrets: + buf_token: + environment: BUF_TOKEN diff --git a/lib/backend-common/src/lib.rs b/lib/backend-common/src/lib.rs index bf9a6ed42c1d..d84ee749eb80 100644 --- a/lib/backend-common/src/lib.rs +++ b/lib/backend-common/src/lib.rs @@ -35,9 +35,9 @@ pub use engine::{ AsyncEngineContext, BootstrapInfo, CompletionUsage, ComponentSnapshot, EngineConfig, FinishReason, GenerateContext, GuidedDecodingOptions, HEALTH_CHECK_KEY, KvEventPublisher, KvEventSource, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, LlmRegistration, LogProbs, - Metrics, MetricsBindings, MetricsCtx, OnPublisherReady, OnSnapshotPublisherReady, - OutputOptions, PrefillResult, PreprocessedRequest, RawEngine, SamplingOptions, StopConditions, - StopReason, TopLogprob, TopLogprobs, chunk, usage, + Metrics, MetricsBindings, MetricsCtx, MultimodalData, OnPublisherReady, + OnSnapshotPublisherReady, OutputOptions, PrefillResult, PreprocessedRequest, RawEngine, + SamplingOptions, StopConditions, StopReason, TopLogprob, TopLogprobs, chunk, usage, }; pub use error::{BackendError, DynamoError, ErrorType}; pub use metrics::{ComponentGauges, EngineMetrics, LifecycleGauges}; diff --git a/lib/mocker/servers/vllm/Cargo.toml b/lib/mocker/servers/vllm/Cargo.toml index 17bcd78bf655..88e592aef5fc 100644 --- a/lib/mocker/servers/vllm/Cargo.toml +++ b/lib/mocker/servers/vllm/Cargo.toml @@ -25,10 +25,10 @@ async-stream = { workspace = true } blake3 = { workspace = true } clap = { version = "4", features = ["derive", "env"] } futures = { workspace = true } -prost-types = { workspace = true } +prost-types = "0.14.1" tokio = { workspace = true } -tonic = { workspace = true } -tonic-health = { workspace = true } +tonic = "0.14.1" +tonic-health = "0.14.1" tracing = { workspace = true } tracing-subscriber = { workspace = true } uuid = { workspace = true } diff --git a/lib/mocker/servers/vllm/src/server.rs b/lib/mocker/servers/vllm/src/server.rs index c41f8bd87a01..e08b3774a706 100644 --- a/lib/mocker/servers/vllm/src/server.rs +++ b/lib/mocker/servers/vllm/src/server.rs @@ -131,7 +131,6 @@ impl VllmMockerService { anyhow::anyhow!("max_num_batched_tokens exceeds the Control API range") })? .unwrap_or_default(), - supports_explicit_data_parallel_rank: true, }; Ok(Self { config: Arc::new(config), @@ -156,14 +155,38 @@ impl VllmMockerService { async fn start_generation( &self, - request: pb::GenerateRequest, + request: Request, ) -> Result<(PreparedRequest, LiveRequest, OwnedSemaphorePermit), Status> { + let data_parallel_rank = request + .metadata() + .get("x-data-parallel-rank") + .map(|value| { + value + .to_str() + .ok() + .and_then(|value| value.parse::().ok()) + .ok_or_else(|| { + Box::new(Status::invalid_argument( + "x-data-parallel-rank metadata must be an unsigned 32-bit integer", + )) + }) + }) + .transpose() + .map_err(|status| *status)?; + if let Some(rank) = data_parallel_rank + && rank != DP_RANK + { + return Err(Status::invalid_argument(format!( + "data_parallel_rank {rank} is not served; expected {DP_RANK}" + ))); + } let permit = self .request_permits .clone() .try_acquire_owned() .map_err(|_| Status::resource_exhausted("Mocker concurrent request limit reached"))?; - let prepared = PreparedRequest::new(request, &self.config).map_err(|status| *status)?; + let prepared = + PreparedRequest::new(request.into_inner(), &self.config).map_err(|status| *status)?; let live = self .engine .submit(prepared.direct_request()) @@ -184,7 +207,7 @@ impl pb::inference_server::Inference for VllmMockerService { &self, request: Request, ) -> Result, Status> { - let (prepared, mut live, _permit) = self.start_generation(request.into_inner()).await?; + let (prepared, mut live, _permit) = self.start_generation(request).await?; let mut output_ids = Vec::with_capacity(prepared.max_output_tokens); while let Some(signal) = live.recv().await { let token_id = checked_token(&signal).map_err(|status| *status)?; @@ -205,7 +228,7 @@ impl pb::inference_server::Inference for VllmMockerService { &self, request: Request, ) -> Result, Status> { - let (prepared, mut live, permit) = self.start_generation(request.into_inner()).await?; + let (prepared, mut live, permit) = self.start_generation(request).await?; // Decouple LiveEngine's small fixed per-request buffer from client and // transport pacing. A pump drains the engine promptly into a buffer // bounded by this request's own token budget, so a bursty producer diff --git a/lib/mocker/servers/vllm/src/server_request.rs b/lib/mocker/servers/vllm/src/server_request.rs index 61e5736b72a9..d736bc048f39 100644 --- a/lib/mocker/servers/vllm/src/server_request.rs +++ b/lib/mocker/servers/vllm/src/server_request.rs @@ -75,14 +75,6 @@ impl PreparedRequest { )) .into()); } - if let Some(rank) = request.data_parallel_rank - && rank != DP_RANK - { - return Err(Status::invalid_argument(format!( - "data_parallel_rank {rank} is not served; expected {DP_RANK}" - )) - .into()); - } let mut prompt_tokens = match request.prompt.take() { Some(pb::generate_request::Prompt::TokenIds(tokens)) => tokens.ids, Some(pb::generate_request::Prompt::Text(_)) => { diff --git a/lib/mocker/servers/vllm/src/server_tests.rs b/lib/mocker/servers/vllm/src/server_tests.rs index 0260360e3917..41e7822439f4 100644 --- a/lib/mocker/servers/vllm/src/server_tests.rs +++ b/lib/mocker/servers/vllm/src/server_tests.rs @@ -311,11 +311,15 @@ fn decode_rejects_a_handoff_missing_the_opacity_sentinel() { async fn unary_generate_accumulates_output_and_terminal_metadata() { let service = VllmMockerService::new(MockerServerConfig::default(), admitting_args()).unwrap(); - let response = - pb::inference_server::Inference::generate(&service, Request::new(request("unary"))) - .await - .unwrap() - .into_inner(); + let mut routed_request = Request::new(request("unary")); + routed_request.metadata_mut().insert( + "x-data-parallel-rank", + tonic::metadata::MetadataValue::from(DP_RANK), + ); + let response = pb::inference_server::Inference::generate(&service, routed_request) + .await + .unwrap() + .into_inner(); assert!(response.prompt_info.is_some()); let outputs = response @@ -333,6 +337,16 @@ async fn unary_generate_accumulates_output_and_terminal_metadata() { ); assert_eq!(finish.num_output_tokens, 2); assert_eq!(service.active_request_count(), 0); + + let mut wrong_rank_request = Request::new(request("wrong-rank")); + wrong_rank_request.metadata_mut().insert( + "x-data-parallel-rank", + tonic::metadata::MetadataValue::from(DP_RANK + 1), + ); + let error = pb::inference_server::Inference::generate(&service, wrong_rank_request) + .await + .unwrap_err(); + assert_eq!(error.code(), tonic::Code::InvalidArgument); } #[tokio::test] diff --git a/lib/sidecar/common/Cargo.toml b/lib/sidecar/common/Cargo.toml index b2c8d900594f..e53e1bed73d3 100644 --- a/lib/sidecar/common/Cargo.toml +++ b/lib/sidecar/common/Cargo.toml @@ -16,7 +16,9 @@ dynamo-backend-common = { workspace = true } clap = { version = "4", features = ["derive", "env"] } futures = { workspace = true } +async-trait = { workspace = true } tokio = { workspace = true } tonic = { workspace = true } +tonic-v14 = { package = "tonic", version = "0.14.1" } tracing = { workspace = true } url = { workspace = true } diff --git a/lib/sidecar/common/src/endpoint.rs b/lib/sidecar/common/src/endpoint.rs index c25957eb3add..1e9fe07fd583 100644 --- a/lib/sidecar/common/src/endpoint.rs +++ b/lib/sidecar/common/src/endpoint.rs @@ -9,7 +9,10 @@ use crate::invalid_argument; /// Validated plaintext gRPC endpoint containing only a scheme and authority. #[derive(Clone, Debug, Eq, PartialEq)] -pub struct GrpcEndpoint(String); +pub struct GrpcEndpoint { + endpoint: String, + authority_host: String, +} impl GrpcEndpoint { pub fn parse(raw: &str, argument: &str) -> Result { @@ -39,11 +42,16 @@ impl GrpcEndpoint { let parsed = url::Url::parse(&normalized).map_err(|error| { invalid_argument(format!("invalid gRPC endpoint for `{argument}`: {error}")) })?; - if parsed.host().is_none() { - return Err(invalid_argument(format!( - "`{argument}` must include a host" - ))); - } + let authority_host = match parsed.host() { + Some(url::Host::Domain(host)) => host.to_string(), + Some(url::Host::Ipv4(host)) => host.to_string(), + Some(url::Host::Ipv6(host)) => format!("[{host}]"), + None => { + return Err(invalid_argument(format!( + "`{argument}` must include a host" + ))); + } + }; if !parsed.username().is_empty() || parsed.password().is_some() { return Err(invalid_argument(format!( "`{argument}` must not include user information" @@ -56,17 +64,25 @@ impl GrpcEndpoint { } let authority = &parsed[url::Position::BeforeHost..url::Position::AfterPort]; - Ok(Self(format!("http://{authority}"))) + Ok(Self { + endpoint: format!("http://{authority}"), + authority_host, + }) } pub fn as_str(&self) -> &str { - &self.0 + &self.endpoint + } + + /// Host formatted for use in a URI authority, including IPv6 brackets. + pub fn authority_host(&self) -> &str { + &self.authority_host } } impl fmt::Display for GrpcEndpoint { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str(&self.0) + formatter.write_str(&self.endpoint) } } @@ -96,6 +112,9 @@ mod tests { .as_str(), "http://server:50051" ); + let ipv6 = GrpcEndpoint::parse("http://[2001:db8::1]:50051", ARGUMENT).unwrap(); + assert_eq!(ipv6.as_str(), "http://[2001:db8::1]:50051"); + assert_eq!(ipv6.authority_host(), "[2001:db8::1]"); } #[test] diff --git a/lib/sidecar/common/src/error.rs b/lib/sidecar/common/src/error.rs index e5b224864ac6..a0bf3fbefe85 100644 --- a/lib/sidecar/common/src/error.rs +++ b/lib/sidecar/common/src/error.rs @@ -34,7 +34,19 @@ pub fn connection_timeout(message: impl Into) -> DynamoError { } pub fn status_to_dynamo(rpc: &str, status: tonic::Status) -> DynamoError { - let kind = match status.code() { + status_to_dynamo_parts(rpc, status.message(), status.code()) +} + +pub fn status_to_dynamo_v14(rpc: &str, status: tonic_v14::Status) -> DynamoError { + status_to_dynamo_parts( + rpc, + status.message(), + tonic::Code::from_i32(status.code() as i32), + ) +} + +fn status_to_dynamo_parts(rpc: &str, message: &str, code: tonic::Code) -> DynamoError { + let kind = match code { tonic::Code::InvalidArgument | tonic::Code::NotFound | tonic::Code::OutOfRange @@ -45,10 +57,7 @@ pub fn status_to_dynamo(rpc: &str, status: tonic::Status) -> DynamoError { tonic::Code::DeadlineExceeded => BackendError::ConnectionTimeout, _ => BackendError::Unknown, }; - backend( - kind, - format!("{rpc}: {} ({:?})", status.message(), status.code()), - ) + backend(kind, format!("{rpc}: {message} ({code:?})")) } #[cfg(test)] diff --git a/lib/sidecar/common/src/lib.rs b/lib/sidecar/common/src/lib.rs index 222657bc95e1..f6643538319c 100644 --- a/lib/sidecar/common/src/lib.rs +++ b/lib/sidecar/common/src/lib.rs @@ -12,6 +12,6 @@ pub use args::{GrpcTransportArgs, GrpcTransportConfig, SidecarArgs}; pub use endpoint::GrpcEndpoint; pub use error::{ cannot_connect, connection_timeout, engine_shutdown, invalid_argument, protocol_error, - status_to_dynamo, + status_to_dynamo, status_to_dynamo_v14, }; -pub use transport::{DEFAULT_MAX_GRPC_MESSAGE_SIZE, GrpcChannelPool}; +pub use transport::{DEFAULT_MAX_GRPC_MESSAGE_SIZE, GrpcChannelPool, GrpcChannelPoolV14}; diff --git a/lib/sidecar/common/src/transport.rs b/lib/sidecar/common/src/transport.rs index f1e547556151..85b7bb2da8f9 100644 --- a/lib/sidecar/common/src/transport.rs +++ b/lib/sidecar/common/src/transport.rs @@ -5,10 +5,12 @@ use std::fmt::Write as _; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; +use async_trait::async_trait; use dynamo_backend_common::DynamoError; use futures::future::try_join_all; use tokio::time::{Instant, sleep_until, timeout_at}; use tonic::transport::{Channel, Endpoint}; +use tonic_v14::transport::{Channel as ChannelV14, Endpoint as EndpointV14}; use crate::{GrpcEndpoint, GrpcTransportConfig, cannot_connect, invalid_argument}; @@ -21,76 +23,143 @@ pub struct GrpcChannelPool { next: AtomicUsize, } -impl GrpcChannelPool { - pub async fn connect( - peer: &str, - endpoint: &GrpcEndpoint, - transport: GrpcTransportConfig, - ) -> Result { - let endpoint_label = endpoint.to_string(); - let tonic_endpoint = Endpoint::from_shared(endpoint_label.clone()).map_err(|error| { - invalid_argument(format!("invalid {peer} endpoint after validation: {error}")) - })?; - let deadline = checked_instant_add( - Instant::now(), - transport.startup_deadline, - "gRPC startup deadline", - )?; - let first = connect_until_ready( - peer, - tonic_endpoint.clone(), - endpoint_label.clone(), - 1, - transport, - deadline, - ) - .await?; - let mut channels = vec![first]; - let remaining = try_join_all((1..transport.connections.get()).map(|index| { - let endpoint = tonic_endpoint.clone(); - let endpoint_label = endpoint_label.clone(); - async move { - connect_until_ready( - peer, - endpoint, - endpoint_label, - index + 1, - transport, - deadline, - ) - .await +/// Connected tonic 0.14 channels distributed in round-robin order. +pub struct GrpcChannelPoolV14 { + channels: Vec, + next: AtomicUsize, +} + +macro_rules! impl_channel_pool { + ($pool:ident, $endpoint:ty, $channel:ty) => { + impl $pool { + pub async fn connect( + peer: &str, + endpoint: &GrpcEndpoint, + transport: GrpcTransportConfig, + ) -> Result { + Ok(Self { + channels: connect_channels::<$endpoint>(peer, endpoint, transport).await?, + next: AtomicUsize::new(0), + }) + } + + pub fn len(&self) -> usize { + self.channels.len() + } + + pub fn is_empty(&self) -> bool { + self.channels.is_empty() + } + + pub fn next_channel(&self) -> $channel { + let index = self.next.fetch_add(1, Ordering::Relaxed) % self.channels.len(); + self.channels[index].clone() } - })) - .await?; - channels.extend(remaining); - Ok(Self { - channels, - next: AtomicUsize::new(0), - }) + } + }; +} + +impl_channel_pool!(GrpcChannelPool, Endpoint, Channel); +impl_channel_pool!(GrpcChannelPoolV14, EndpointV14, ChannelV14); + +#[async_trait] +trait ConnectEndpoint: Clone + Send + Sync { + type Channel: Clone + Send + Sync; + type Error: std::error::Error + Send + Sync + 'static; + + fn from_shared(uri: String) -> Result; + fn connect_timeout(self, timeout: Duration) -> Self; + async fn connect(&self) -> Result; +} + +#[async_trait] +impl ConnectEndpoint for Endpoint { + type Channel = Channel; + type Error = tonic::transport::Error; + + fn from_shared(uri: String) -> Result { + Endpoint::from_shared(uri) } - pub fn len(&self) -> usize { - self.channels.len() + fn connect_timeout(self, timeout: Duration) -> Self { + Endpoint::connect_timeout(self, timeout) + } + + async fn connect(&self) -> Result { + Endpoint::connect(self).await + } +} + +#[async_trait] +impl ConnectEndpoint for EndpointV14 { + type Channel = ChannelV14; + type Error = tonic_v14::transport::Error; + + fn from_shared(uri: String) -> Result { + EndpointV14::from_shared(uri) } - pub fn is_empty(&self) -> bool { - self.channels.is_empty() + fn connect_timeout(self, timeout: Duration) -> Self { + EndpointV14::connect_timeout(self, timeout) } - pub fn next_channel(&self) -> Channel { - let index = self.next.fetch_add(1, Ordering::Relaxed) % self.channels.len(); - self.channels[index].clone() + async fn connect(&self) -> Result { + EndpointV14::connect(self).await } } -async fn connect_until_ready( +async fn connect_channels( + peer: &str, + endpoint: &GrpcEndpoint, + transport: GrpcTransportConfig, +) -> Result, DynamoError> { + let endpoint_label = endpoint.to_string(); + let tonic_endpoint = E::from_shared(endpoint_label.clone()).map_err(|error| { + invalid_argument(format!("invalid {peer} endpoint after validation: {error}")) + })?; + let deadline = checked_instant_add( + Instant::now(), + transport.startup_deadline, + "gRPC startup deadline", + )?; + let first = connect_until_ready( + peer, + tonic_endpoint.clone(), + endpoint_label.clone(), + 1, + transport, + deadline, + ) + .await?; + let mut channels = vec![first]; + let remaining = try_join_all((1..transport.connections.get()).map(|index| { + let endpoint = tonic_endpoint.clone(); + let endpoint_label = endpoint_label.clone(); + async move { + connect_until_ready( + peer, + endpoint, + endpoint_label, + index + 1, + transport, + deadline, + ) + .await + } + })) + .await?; + channels.extend(remaining); + Ok(channels) +} + +async fn connect_until_ready( peer: &str, - endpoint: Endpoint, + endpoint: E, endpoint_label: String, pool_slot: usize, transport: GrpcTransportConfig, deadline: Instant, -) -> Result { +) -> Result { let started = Instant::now(); let mut attempt = 0_u64; let mut last_error = None; diff --git a/lib/sidecar/sglang/Dockerfile b/lib/sidecar/sglang/Dockerfile index 6b3343dbc916..a73960d82a9d 100644 --- a/lib/sidecar/sglang/Dockerfile +++ b/lib/sidecar/sglang/Dockerfile @@ -5,7 +5,8 @@ # A pure gRPC connector — no GPU, no engine runtime — that runs beside the # engine container over loopback (see deploy/agg.yaml). # -# docker build -f lib/sidecar/sglang/Dockerfile -t dynamo-sglang-sidecar:1.3.0 . +# docker build --secret id=buf_token,env=BUF_TOKEN \ +# -f lib/sidecar/sglang/Dockerfile -t dynamo-sglang-sidecar:1.3.0 . FROM rust:1.96.1-bookworm AS builder RUN apt-get update && apt-get install -y --no-install-recommends \ @@ -20,7 +21,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ WORKDIR /src COPY . . -RUN cargo build --release --locked -p dynamo-sglang-sidecar \ +RUN --mount=type=secret,id=buf_token,required=true \ + CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" \ + cargo build --release --locked -p dynamo-sglang-sidecar \ && strip target/release/dynamo-sglang-sidecar FROM debian:bookworm-slim AS runtime diff --git a/lib/sidecar/trtllm/Dockerfile b/lib/sidecar/trtllm/Dockerfile index 864f647e5d22..e6ccdc44a47f 100644 --- a/lib/sidecar/trtllm/Dockerfile +++ b/lib/sidecar/trtllm/Dockerfile @@ -5,7 +5,8 @@ # A pure gRPC connector — no GPU, no engine runtime — that runs beside the # engine container over loopback (see deploy/agg.yaml). # -# docker build -f lib/sidecar/trtllm/Dockerfile -t dynamo-trtllm-sidecar:1.3.0 . +# docker build --secret id=buf_token,env=BUF_TOKEN \ +# -f lib/sidecar/trtllm/Dockerfile -t dynamo-trtllm-sidecar:1.3.0 . FROM rust:1.96.1-bookworm AS builder RUN apt-get update && apt-get install -y --no-install-recommends \ @@ -20,7 +21,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ WORKDIR /src COPY . . -RUN cargo build --release --locked -p dynamo-trtllm-sidecar \ +RUN --mount=type=secret,id=buf_token,required=true \ + CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" \ + cargo build --release --locked -p dynamo-trtllm-sidecar \ && strip target/release/dynamo-trtllm-sidecar FROM debian:bookworm-slim AS runtime diff --git a/lib/sidecar/vllm/Cargo.toml b/lib/sidecar/vllm/Cargo.toml index ea6603f2714f..bf2f4f2c4492 100644 --- a/lib/sidecar/vllm/Cargo.toml +++ b/lib/sidecar/vllm/Cargo.toml @@ -11,10 +11,6 @@ homepage.workspace = true repository.workspace = true description = "Rust sidecar for vLLM's native gRPC server" -[package.metadata.cargo-machete] -# Referenced by protobuf code generated into OUT_DIR. -ignored = ["prost"] - [[bin]] name = "dynamo-vllm-sidecar" path = "src/main.rs" @@ -33,13 +29,11 @@ tokio = { workspace = true } tokio-util = { workspace = true } tracing = { workspace = true } -prost = { workspace = true } -prost-types = { workspace = true } -tonic = { workspace = true } -tonic-health = { workspace = true } - -[build-dependencies] -tonic-build = { workspace = true } +prost-types = "0.14.1" +tonic = "0.14.1" +tonic-health = "0.14.1" +vllm-grpc = { package = "vllm-project_vllm_community_neoeinstein-tonic", version = "=0.5.0-00000000000000-7726adbdafb3.4", registry = "buf" } +vllm-proto = { package = "vllm-project_vllm_community_neoeinstein-prost", version = "=0.5.0-00000000000000-7726adbdafb3.2", registry = "buf" } [dev-dependencies] dynamo-backend-common = { workspace = true, features = ["testing"] } diff --git a/lib/sidecar/vllm/Dockerfile b/lib/sidecar/vllm/Dockerfile index 3abdc9e35133..d4d016668bb2 100644 --- a/lib/sidecar/vllm/Dockerfile +++ b/lib/sidecar/vllm/Dockerfile @@ -9,7 +9,8 @@ # deploy/agg.yaml). # # Build from the repository root (the workspace is needed to compile the crate): -# docker build -f lib/sidecar/vllm/Dockerfile -t dynamo-vllm-sidecar:1.3.0 . +# docker build --secret id=buf_token,env=BUF_TOKEN \ +# -f lib/sidecar/vllm/Dockerfile -t dynamo-vllm-sidecar:1.3.0 . # ---- builder ---------------------------------------------------------------- # Pinned to the workspace toolchain (rust-toolchain.toml: 1.96.1). @@ -35,7 +36,9 @@ COPY . . # Build only the sidecar binary; --locked pins the checked-in Cargo.lock. There # is no dependency-cache layer, so a source change recompiles deps — acceptable # for this occasional image build. -RUN cargo build --release --locked -p dynamo-vllm-sidecar \ +RUN --mount=type=secret,id=buf_token,required=true \ + CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" \ + cargo build --release --locked -p dynamo-vllm-sidecar \ && strip target/release/dynamo-vllm-sidecar # ---- runtime (minimal) ------------------------------------------------------ diff --git a/lib/sidecar/vllm/README.md b/lib/sidecar/vllm/README.md index b25604b15c02..6ff103ec817c 100644 --- a/lib/sidecar/vllm/README.md +++ b/lib/sidecar/vllm/README.md @@ -25,18 +25,31 @@ It is a standalone Rust executable. - Token and text requests through Dynamo preprocessing - Sampling, stop conditions, structured output, logprobs, cache options, and priority - Opaque `kv_transfer_params` handoff +- Data-parallel rank routing and KV-event source discovery +- Image URL and data-URI inputs, including media UUIDs -The initial protocol does not support multimodal input, LoRA, KV-aware data -parallel routing, encode workers, beam search, or `n > 1`. +The protocol does not support LoRA, encode workers, beam search, `n > 1`, +preprocessed multimodal features, audio/video media, or Dynamo tool-call and +reasoning parsers. Parser defaults returned by Control are intentionally not +advertised to the Dynamo frontend because the current inference protocol does +not preserve all parser-related request semantics. ## Run -Start vLLM with its released gRPC listener: +Start a vLLM build with the split Inference and Control services. Data-parallel routing requires a build containing [vLLM PR #51178](https://github.com/vllm-project/vllm/pull/51178) or a release that includes it: ```bash vllm-rs serve Qwen/Qwen3-0.6B --host 127.0.0.1 --grpc-port 50051 ``` +The sidecar uses the generated Rust SDKs from the pinned [`vllm-project/vllm`](https://buf.build/vllm-project/vllm/docs/nightly) BSR commit `7726adbdafb34bda85e25c8fc5e192f4`. The BSR Cargo registry requires authentication. Create a BSR token and export it for Cargo before building: + +```bash +export CARGO_REGISTRIES_BUF_TOKEN="Bearer ${BUF_TOKEN}" +``` + +The repository's `.cargo/config.toml` configures the registry and credential provider. Because the SDKs are in the workspace lockfile, authenticate before running any workspace Cargo command. CI reads `BUF_TOKEN` from the repository's Actions secrets. + This listener is unauthenticated and plaintext. Keep colocated deployments on loopback or a private interface. Remote access requires network controls or a secure proxy. @@ -51,9 +64,9 @@ dynamo-vllm-sidecar \ Use `VLLM_GRPC_ENDPOINT` instead of `--vllm-endpoint` when the endpoint is provided through the environment. -The sidecar discovers `model_id`, the served name, context length, KV capacity, and scheduler limits through `vllm.Control`. `model_id` must be readable locally or fetchable by Dynamo for tokenization and chat templates. Parser defaults are not advertised because the current inference protocol cannot preserve all parser-related request semantics. +The sidecar discovers `model_id`, the served name, context length, KV capacity, scheduler limits, data-parallel topology, and KV-event sources through `vllm.Control`. `model_id` must be readable locally or fetchable by Dynamo for tokenization and chat templates. -Data-parallel registration is omitted because Control reports global topology, not the rank range hosted by the connected frontend. +The sidecar currently supports one vLLM frontend hosting the complete data-parallel group starting at rank 0. Control reports the global size; Dynamo forwards the selected rank as `x-data-parallel-rank` gRPC metadata on each generation request. Partial and hybrid rank ownership are unsupported because the protocol does not report the locally hosted rank count, and a nonzero starting rank is rejected. When KV routing is enabled, Control must return one unique ZMQ event source for every rank in the group. Aggregated serving is the default. Set the existing `--disaggregation-mode` to `prefill` or `decode` only for non-aggregated deployments; the current Control API does not report engine role. @@ -67,6 +80,8 @@ pool. Override them with `--grpc-connect-attempt-timeout-secs`, `--grpc-retry-interval-secs`, and `--grpc-startup-deadline-secs`, or with the corresponding `DYN_SIDECAR_GRPC_*` environment variables. +Each request owns its response stream but borrows a channel from the shared pool. Aggregate and prefill cancellation drops only that request's stream. Decode cancellation first submits the decode request and retains its stream until the first output token or a response containing `finish_info`, so a NIXL receiver can complete and release the transferred KV; it then drops the stream. If the stream ends early, returns a gRPC error, or produces an invalid response after cancellation, the sidecar logs the failure and reports the request as cancelled. vLLM automatically aborts the corresponding engine request while the pooled HTTP/2 connection remains available to other requests. The sidecar does not call the Control `Abort` RPC. + ## Test without vLLM or a GPU Use the CPU-only `dynamo-vllm-mocker-server` to exercise the same Inference, Control, and health contracts: @@ -94,7 +109,7 @@ disaggregated prefill/decode with NIXL KV transfer. There is no published vLLM sidecar image yet, so you build and push your own from `Dockerfile` — the same pattern as the TensorRT-LLM and SGLang sidecars. -The sidecar waits for both the Control and Inference services through the standard gRPC health API before registering the worker. The deployment manifests retain lightweight socket probes for container lifecycle monitoring. The engine image must include a `vllm-rs` build compatible with the vendored protocol. +The sidecar waits for both the Control and Inference services through the standard gRPC health API before registering the worker. The deployment manifests retain lightweight socket probes for container lifecycle monitoring. The engine image must include a `vllm-rs` build compatible with the pinned BSR protocol. ### Prerequisites @@ -104,6 +119,7 @@ The sidecar waits for both the Control and Inference services through the standa `restartPolicy: Always`), which requires that version. - `kubectl` set to that cluster, and a namespace to deploy into. - A Hugging Face token for the model. +- A BSR token in `BUF_TOKEN` for building the sidecar. - A container registry you can push to and the cluster can pull from. ### 1. Build and push the sidecar image @@ -113,6 +129,7 @@ Build a multi-arch image so it runs on any node — `amd64` (x86) or `arm64` ```bash docker buildx build --platform linux/amd64,linux/arm64 \ + --secret id=buf_token,env=BUF_TOKEN \ -f lib/sidecar/vllm/Dockerfile \ -t /dynamo-vllm-sidecar:1.3.0 --push . ``` diff --git a/lib/sidecar/vllm/build.rs b/lib/sidecar/vllm/build.rs deleted file mode 100644 index 7f01c4e896f2..000000000000 --- a/lib/sidecar/vllm/build.rs +++ /dev/null @@ -1,14 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -fn main() -> Result<(), Box> { - tonic_build::configure() - .protoc_arg("--experimental_allow_proto3_optional") - .compile_protos( - &["proto/inference.proto", "proto/control.proto"], - &["proto"], - )?; - println!("cargo:rerun-if-changed=proto/inference.proto"); - println!("cargo:rerun-if-changed=proto/control.proto"); - Ok(()) -} diff --git a/lib/sidecar/vllm/proto/README.md b/lib/sidecar/vllm/proto/README.md deleted file mode 100644 index bab4a03cc097..000000000000 --- a/lib/sidecar/vllm/proto/README.md +++ /dev/null @@ -1,13 +0,0 @@ - - -# Vendored vLLM protocol - -- Source: [`rust/proto/inference.proto`](https://github.com/connorcarpenter15/vllm/blob/2d2c3af18c52e8e4efa4b0b4903843b15c0dba0e/rust/proto/inference.proto) and [`rust/proto/control.proto`](https://github.com/connorcarpenter15/vllm/blob/2d2c3af18c52e8e4efa4b0b4903843b15c0dba0e/rust/proto/control.proto) -- Commit: `2d2c3af18c52e8e4efa4b0b4903843b15c0dba0e` -- `inference.proto` SHA-256: `a0d196dc240683e1c09abb54f324d4428d0c122a6802b44916ad2d96b491b06c` -- `control.proto` SHA-256: `cd4e7a8043f19d05929a2f59f5a5442894a037ef2d65832d3f7099992b1f1dbd` - -The files are copied without modification. Update the revision and checksums together. `dynamo-vllm-sidecar` generates and temporarily exports these types for `dynamo-vllm-mocker-server`. diff --git a/lib/sidecar/vllm/proto/control.proto b/lib/sidecar/vllm/proto/control.proto deleted file mode 100644 index d2ec9da4e7cc..000000000000 --- a/lib/sidecar/vllm/proto/control.proto +++ /dev/null @@ -1,78 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -syntax = "proto3"; -package vllm; - -service Control { - rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {} - rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} - rpc Abort (AbortRequest) returns (AbortResponse) {} - rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {} -} - -message GetServerInfoRequest {} - -message ServerInfo { - string engine_version = 1; - string api_version = 2; - string instance_id = 3; - ParallelismInfo parallelism = 4; - uint32 max_model_len = 5; - uint32 kv_block_size = 6; - uint64 total_kv_blocks = 7; - uint64 max_running_requests = 8; - uint64 max_batched_tokens = 9; - // GenerateRequest.data_parallel_rank is honored by this server. Clients - // that require deterministic rank routing must fail closed when this is - // false, because older servers accept and silently discard the field. - bool supports_explicit_data_parallel_rank = 10; -} - -message ParallelismInfo { - uint32 tensor_parallel_size = 1; - uint32 pipeline_parallel_size = 2; - uint32 data_parallel_size = 3; - uint32 data_parallel_rank = 4; - uint32 decode_context_parallel_size = 5; -} - -message GetModelInfoRequest {} - -message ModelInfo { - string model_id = 1; - string served_model_name = 2; - repeated string served_model_aliases = 3; - - bool supports_text_input = 20; - bool supports_token_ids_input = 21; - bool supports_multimodal = 23; - string reasoning_parser = 24; - string tool_call_parser = 25; -} - -message AbortRequest { - repeated string request_ids = 1; -} - -message AbortResponse {} - -// ====================================================================================== -// KV discovery -// ====================================================================================== - -message GetKvEventSourcesRequest {} -message GetKvEventSourcesResponse { repeated KvEventSource sources = 1; } - -message KvEventSource { - string transport = 1; - string endpoint = 2; - string topic = 3; - string replay_endpoint = 4; - optional uint32 data_parallel_rank = 5; - string encoding = 6; - uint32 schema_version = 7; - uint32 buffer_steps = 8; - uint32 hwm = 9; - uint32 max_queue_size = 10; -} diff --git a/lib/sidecar/vllm/proto/inference.proto b/lib/sidecar/vllm/proto/inference.proto deleted file mode 100644 index 4acb6826504a..000000000000 --- a/lib/sidecar/vllm/proto/inference.proto +++ /dev/null @@ -1,232 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -syntax = "proto3"; -package vllm; - -import "google/protobuf/struct.proto"; - - -service Inference { - // Generates text given a prompt - rpc Generate (GenerateRequest) returns (GenerateResponse) {} - // Generates text given a prompt, streaming the outputs - rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {} -} - -// ====================================================================================== -// Generate Request -// ====================================================================================== - -message GenerateRequest { - string request_id = 1; - string model = 2; - - oneof prompt { - string text = 3; - TokenIds token_ids = 4; - } - - // Temperature, defaults to model-specific default or 0 - optional float temperature = 5; - // Parameters controlling random sampling, not applicable if temperature == 0 - RandomSampling sampling = 6; - // Parameters for conditionally penalizing/boosting - // candidate tokens during decoding - DecodingParameters decoding = 7; - // Parameters controlling when generation should stop - StoppingCriteria stopping = 8; - // Flags to control what is returned in the response - ResponseOptions response = 9; - // Parameters controlling KV cache/distribution - KVCacheParameters kv = 10; - - // Truncate prompt tokens; default (0) means no truncation - uint32 truncate_prompt_tokens = 11; - - int32 priority = 12; - - optional string session_id = 13; - - // Multimodal inputs aligned with placeholder markers in token_ids. - repeated MediaItem media = 14; - - // Global data-parallel rank advertised by the Control service. - optional uint32 data_parallel_rank = 15; -} - -message RandomSampling { - uint32 num_sequences = 1; // "n", default (0) means 1 - uint32 top_k = 2; // 0 means default - float top_p = 3; // 0 means default - float min_p = 4; // 0 means default - optional int64 seed = 5; -} - -message DecodingParameters { - // Penalties - float presence_penalty = 1; // Default (0.0) means no penalty - float frequency_penalty = 2; // Default (0.0) means no penalty - float repetition_penalty = 3; // Default (0.0) means no penalty - map logit_bias = 4; - repeated uint32 allowed_token_ids = 5; - - message StringChoices { - repeated string choices = 1; - } - - // Control structured outputs - oneof structured_output { - string json = 6; - string regex = 7; - StringChoices choice = 8; - string grammar = 9; - bool json_object = 10; - string structural_tag = 11; - } -} - -message StoppingCriteria { - // Default (0) is currently 20 - uint32 max_new_tokens = 1; - // Default (0) means no minimum - uint32 min_new_tokens = 2; - - repeated uint32 stop_token_ids = 3; - repeated string stop_strings = 4; - bool include_stop_strings = 5; - - bool ignore_eos = 6; -} - -message ResponseOptions { - // Prompt options - bool prompt_token_ids = 1; - bool prompt_logprobs = 2; - optional CandidateTokens prompt_candidates = 3; - - // Output options; output_text defaults to true - optional bool output_text = 4; - bool output_token_ids = 5; - bool output_logprobs = 6; - optional CandidateTokens output_candidates = 7; -} - -message KVCacheParameters { - bool bypass_prefix_cache = 1; - string cache_salt = 2; - - // KV Connector transfer parameters - google.protobuf.Struct kv_transfer_params = 3; - - // Encoder cache connector transfer parameters - google.protobuf.Struct ec_transfer_params = 4; -} - -// Controls which extra candidate tokens at each position should be returned -message CandidateTokens { - oneof select { - uint32 top_n = 1; - TokenIds token_ids = 2; - bool all = 3; - } -} - -// ====================================================================================== -// Generate Response -// ====================================================================================== - -message GenerateResponse { - // Only present in first response - optional PromptInfo prompt_info = 1; - SequenceOutput outputs = 2; -} - -message SequenceOutput { - // Index of output sequence for num_sequences > 1. - uint32 index = 1; - - string text = 2; - uint32 num_tokens = 3; // Number of tokens in this chunk - repeated uint32 token_ids = 4; // If requested - repeated float logprobs = 5; // If requested - repeated uint32 ranks = 6; // If logprobs were requested - repeated CandidateTokenInfo candidate_tokens = 7; // If requested - - // Only present in final output for this sequence - optional FinishInfo finish_info = 8; -} - -// Prompt info, returned in the first response -message PromptInfo { - uint32 num_prompt_tokens = 1; - repeated uint32 token_ids = 2; // If requested - repeated float logprobs = 3; // If requested - repeated uint32 ranks = 4; // If logprobs were requested - repeated CandidateTokenInfo candidate_tokens = 5; -} - -// Finish info, returned in the final response -message FinishInfo { - uint32 num_output_tokens = 1; - - enum FinishReason { - NOT_FINISHED = 0; // Possibly more tokens to be streamed - LENGTH = 1; // Finished due to length constraint - STOP = 2; // Stop string/token or EOS encountered - ABORTED = 3; // Request aborted/cancelled - } - - FinishReason finish_reason = 2; - // One of these will be set when finish_reason == STOP - oneof stop_reason { - uint32 stop_token_id = 3; - uint32 eos_token_id = 4; - string stop_string = 5; - } - - google.protobuf.Struct kv_transfer_params = 6; - //uint64 seed = 7; - google.protobuf.Struct ec_transfer_params = 8; -} - -// Info for candidate tokens other than the input/sampled -// token at a given position -message CandidateTokenInfo { - message TokenInfo { - uint32 id = 1; - float logprob = 2; - uint32 rank = 3; - // string text = 4; - // bytes token_bytes = 5; - } - // Candidate token infos at this position - repeated TokenInfo tokens = 1; -} - -// Token ids used for prompt -message TokenIds { - repeated uint32 ids = 1; -} - -// ====================================================================================== -// Media -// ====================================================================================== - -enum Modality { - MODALITY_UNSPECIFIED = 0; - MODALITY_IMAGE = 1; - MODALITY_VIDEO = 2; - MODALITY_AUDIO = 3; -} - -message MediaItem { - Modality modality = 1; - oneof source { - string url = 2; // http:// or https:// - string data_uri = 3; // data: - bytes raw_bytes = 4; - } - string mime_type = 5; - string uuid = 6; -} diff --git a/lib/sidecar/vllm/src/client.rs b/lib/sidecar/vllm/src/client.rs index e3cf9c19e902..1596b643e5eb 100644 --- a/lib/sidecar/vllm/src/client.rs +++ b/lib/sidecar/vllm/src/client.rs @@ -5,21 +5,25 @@ use std::time::Duration; use dynamo_backend_common::DynamoError; use dynamo_sidecar_common::{ - DEFAULT_MAX_GRPC_MESSAGE_SIZE, GrpcChannelPool, GrpcEndpoint, GrpcTransportConfig, + DEFAULT_MAX_GRPC_MESSAGE_SIZE, GrpcChannelPoolV14, GrpcEndpoint, GrpcTransportConfig, }; use tokio::time::{Instant, sleep_until, timeout_at}; +use tonic::metadata::MetadataValue; use tonic_health::pb::health_check_response::ServingStatus; use tonic_health::pb::{HealthCheckRequest, health_client::HealthClient}; -pub(crate) use dynamo_sidecar_common::{engine_shutdown, invalid_argument, status_to_dynamo}; +pub(crate) use dynamo_sidecar_common::{ + engine_shutdown, invalid_argument, status_to_dynamo_v14 as status_to_dynamo, +}; use crate::proto as pb; pub(crate) const CONTROL_SERVICE: &str = "vllm.Control"; pub(crate) const INFERENCE_SERVICE: &str = "vllm.Inference"; +const DATA_PARALLEL_RANK_METADATA_KEY: &str = "x-data-parallel-rank"; pub(crate) struct VllmClient { - pool: GrpcChannelPool, + pool: GrpcChannelPoolV14, } impl VllmClient { @@ -30,7 +34,7 @@ impl VllmClient { ) -> Result { let pool = timeout_at( startup_deadline, - GrpcChannelPool::connect("vLLM", endpoint, transport), + GrpcChannelPoolV14::connect("vLLM", endpoint, transport), ) .await .map_err(|_| { @@ -147,16 +151,35 @@ impl VllmClient { pub(crate) async fn generate_stream( &self, request: pb::GenerateRequest, + data_parallel_rank: Option, ) -> Result, DynamoError> { let mut client = pb::inference_client::InferenceClient::new(self.pool.next_channel()) .max_encoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE) .max_decoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE); + let mut request = tonic::Request::new(request); + if let Some(rank) = data_parallel_rank { + request + .metadata_mut() + .insert(DATA_PARALLEL_RANK_METADATA_KEY, MetadataValue::from(rank)); + } client .generate_stream(request) .await .map(tonic::Response::into_inner) .map_err(|status| status_to_dynamo("GenerateStream", status)) } + + pub(crate) async fn kv_event_sources(&self) -> Result, DynamoError> { + let mut client = pb::control_client::ControlClient::new(self.pool.next_channel()) + .max_encoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE) + .max_decoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE); + client + .get_kv_event_sources(pb::GetKvEventSourcesRequest {}) + .await + .map(tonic::Response::into_inner) + .map(|response| response.sources) + .map_err(|status| status_to_dynamo("GetKvEventSources", status)) + } } pub(crate) fn startup_deadline(duration: Duration) -> Result { diff --git a/lib/sidecar/vllm/src/convert.rs b/lib/sidecar/vllm/src/convert.rs index 193f5c21c1dc..9f93ca4e2abe 100644 --- a/lib/sidecar/vllm/src/convert.rs +++ b/lib/sidecar/vllm/src/convert.rs @@ -2,8 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 use dynamo_backend_common::{ - DisaggregationMode, DynamoError, GuidedDecodingOptions, LLMEngineOutput, PrefillResult, - PreprocessedRequest, StopReason, TopLogprob, usage, + DisaggregationMode, DynamoError, GuidedDecodingOptions, LLMEngineOutput, MultimodalData, + PrefillResult, PreprocessedRequest, StopReason, TopLogprob, usage, }; use crate::client; @@ -11,6 +11,10 @@ use crate::json::{json_to_struct, struct_to_json}; use crate::proto as pb; const VLLM_LOGPROB_FLOOR: f64 = -9999.0; +const MULTIMODAL_PROMPT_TOKEN_IDS_KEY: &str = "_dynamo_sidecar_multimodal_prompt_token_ids"; +const MM_HASHES_KEY: &str = "mm_hashes"; +// Must match DYNAMO_CACHE_SALT_PREFIX in lib/kv-router/src/zmq_wire/extra_keys.rs. +const DYNAMO_CACHE_SALT_PREFIX: &str = "dynamo-cache-salt:"; pub(crate) fn build_generate_request( request: PreprocessedRequest, @@ -19,6 +23,26 @@ pub(crate) fn build_generate_request( ) -> Result { validate_request(&request, mode)?; + let has_media = request + .multi_modal_data + .as_ref() + .is_some_and(|media| media.values().any(|items| !items.is_empty())); + // Decode reuses the prefill-expanded tokens without reprocessing media. + let forwarded_mm_uuids = if has_media && !mode.is_decode() { + forwarded_mm_uuids(&request)? + } else { + None + }; + let media = if mode.is_decode() { + Vec::new() + } else { + build_media(&request, forwarded_mm_uuids.as_deref())? + }; + let mut prefill_result = request.prefill_result; + let mut token_ids = request.token_ids; + if mode.is_decode() && has_media { + token_ids = take_multimodal_prompt_token_ids(&mut prefill_result)?; + } let prompt_logprobs = request.output_options.prompt_logprobs; let output_logprobs = request.output_options.logprobs; let max_new_tokens = if mode.is_prefill() { @@ -38,18 +62,25 @@ pub(crate) fn build_generate_request( .unwrap_or(0); let cache_salt = routing .as_mut() - .and_then(|routing| routing.cache_namespace.take()) - .or(request.mdc_sum); + .and_then(|routing| routing.cache_namespace.take()); let sampling = request.sampling_options; let stop_conditions = request.stop_conditions; - let kv = build_kv_parameters(request.extra_args, request.prefill_result, cache_salt, mode)?; + let mut extra_args = request.extra_args; + consume_redundant_nvext(&mut extra_args, cache_salt.as_deref())?; + if has_media && let Some(serde_json::Value::Object(extra)) = extra_args.as_mut() { + // These fields are already represented by token_ids and media. + extra.remove("messages"); + extra.remove("formatted_prompt"); + extra.remove(MM_HASHES_KEY); + } + let kv = build_kv_parameters(extra_args, prefill_result, cache_salt, mode)?; Ok(pb::GenerateRequest { request_id, model: String::new(), prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { - ids: request.token_ids, + ids: token_ids, })), temperature: sampling.temperature, sampling: Some(pb::RandomSampling { @@ -79,7 +110,7 @@ pub(crate) fn build_generate_request( ignore_eos: stop_conditions.ignore_eos.unwrap_or(false), }), response: Some(pb::ResponseOptions { - prompt_token_ids: prompt_logprobs.is_some(), + prompt_token_ids: prompt_logprobs.is_some() || (has_media && mode.is_prefill()), prompt_logprobs: prompt_logprobs.is_some(), prompt_candidates: prompt_logprobs.map(top_n_candidates).transpose()?, output_text: Some(true), @@ -91,11 +122,242 @@ pub(crate) fn build_generate_request( truncate_prompt_tokens: 0, priority, session_id: None, - media: Vec::new(), - data_parallel_rank: None, + media, }) } +pub(crate) fn data_parallel_rank( + request: &PreprocessedRequest, + mode: DisaggregationMode, +) -> Option { + request.routing.as_ref().and_then(|routing| { + if mode.is_prefill() { + routing.prefill_dp_rank.or(routing.dp_rank) + } else { + routing.dp_rank + } + }) +} + +fn consume_redundant_nvext( + extra_args: &mut Option, + cache_namespace: Option<&str>, +) -> Result<(), DynamoError> { + let Some(serde_json::Value::Object(extra)) = extra_args.as_mut() else { + return Ok(()); + }; + let remove_nvext = { + let Some(serde_json::Value::Object(nvext)) = extra.get_mut("nvext") else { + return Ok(()); + }; + if let Some(value) = nvext.remove("cache_salt") { + let value = value + .as_str() + .filter(|value| !value.is_empty()) + .ok_or_else(|| { + client::invalid_argument( + "extra_args.nvext.cache_salt must be a non-empty string", + ) + })?; + match cache_namespace { + Some(expected) if value == expected => {} + Some(expected) => { + return Err(client::invalid_argument(format!( + "extra_args.nvext.cache_salt `{value}` does not match routing.cache_namespace `{expected}`" + ))); + } + None => { + return Err(client::invalid_argument( + "extra_args.nvext.cache_salt requires routing.cache_namespace", + )); + } + } + } + if let Some(token_in) = nvext.remove("token_in") + && token_in != serde_json::Value::Bool(true) + { + return Err(client::invalid_argument( + "extra_args.nvext.token_in must be true when present", + )); + } + nvext.is_empty() + }; + if remove_nvext { + extra.remove("nvext"); + } + Ok(()) +} + +fn take_multimodal_prompt_token_ids( + prefill_result: &mut Option, +) -> Result, DynamoError> { + let params = &mut prefill_result + .as_mut() + .ok_or_else(|| { + client::invalid_argument("multimodal decode request is missing the prefill result") + })? + .disaggregated_params; + let value = params + .as_object_mut() + .and_then(|params| params.remove(MULTIMODAL_PROMPT_TOKEN_IDS_KEY)) + .ok_or_else(|| { + client::invalid_argument( + "multimodal decode request is missing expanded prefill token IDs", + ) + })?; + let token_ids: Vec = serde_json::from_value(value).map_err(|error| { + client::invalid_argument(format!("multimodal prefill token IDs are invalid: {error}")) + })?; + if token_ids.is_empty() { + return Err(client::invalid_argument( + "multimodal prefill token IDs must not be empty", + )); + } + Ok(token_ids) +} + +fn media_source(source: &str) -> Result { + if source.starts_with("data:") { + Ok(pb::media_item::Source::DataUri(source.to_string())) + } else if source.starts_with("http://") || source.starts_with("https://") { + Ok(pb::media_item::Source::Url(source.to_string())) + } else { + Err(client::invalid_argument( + "vLLM gRPC image input must use an http://, https://, or data: URI", + )) + } +} + +fn forwarded_mm_uuids(request: &PreprocessedRequest) -> Result>, DynamoError> { + let has_user_uuid = request + .multi_modal_uuids + .as_ref() + .is_some_and(|by_modality| { + by_modality + .values() + .flatten() + .any(|uuid| uuid.as_ref().is_some_and(|uuid| !uuid.is_empty())) + }); + if has_user_uuid { + return Ok(None); + } + + let hashes = match request.extra_args.as_ref() { + Some(serde_json::Value::Object(extra)) => extra.get(MM_HASHES_KEY), + _ => None, + }; + let Some(hashes) = hashes else { + return Ok(None); + }; + let hashes = hashes.as_array().ok_or_else(|| { + client::invalid_argument("extra_args.mm_hashes must be an array of strings") + })?; + if hashes.is_empty() { + return Ok(None); + } + hashes + .iter() + .enumerate() + .map(|(index, hash)| { + let hash = hash + .as_str() + .filter(|hash| !hash.is_empty()) + .ok_or_else(|| { + client::invalid_argument(format!( + "extra_args.mm_hashes[{index}] must be a non-empty string" + )) + })?; + let mut uuid = hash.to_string(); + if uuid.len() < 64 { + uuid.extend(std::iter::repeat_n('0', 64 - uuid.len())); + } + Ok(uuid) + }) + .collect::, _>>() + .map(Some) +} + +fn build_media( + request: &PreprocessedRequest, + forwarded_uuids: Option<&[String]>, +) -> Result, DynamoError> { + let Some(media_by_modality) = request.multi_modal_data.as_ref() else { + if request + .multi_modal_uuids + .as_ref() + .is_some_and(|uuids| !uuids.is_empty()) + { + return Err(client::invalid_argument( + "multi_modal_uuids were provided without multi_modal_data", + )); + } + return Ok(Vec::new()); + }; + + let mut media = Vec::new(); + for (key, items) in media_by_modality { + if items.is_empty() { + continue; + } + if key != "image_url" { + return Err(client::invalid_argument(format!( + "vLLM gRPC currently supports image_url media only; got `{key}`" + ))); + } + let uuids = request + .multi_modal_uuids + .as_ref() + .and_then(|by_modality| by_modality.get(key)); + if let Some(uuids) = uuids + && uuids.len() != items.len() + { + return Err(client::invalid_argument(format!( + "multi_modal_uuids.{key} has {} entries for {} media items", + uuids.len(), + items.len() + ))); + } + if let Some(uuids) = forwarded_uuids + && uuids.len() != items.len() + { + return Err(client::invalid_argument(format!( + "extra_args.mm_hashes has {} entries for {} media items", + uuids.len(), + items.len() + ))); + } + + for (index, item) in items.iter().enumerate() { + let source = match item { + MultimodalData::Url(url) => media_source(url.as_str())?, + MultimodalData::RawUrl(source) => media_source(source)?, + MultimodalData::Decoded(_) => { + return Err(client::invalid_argument( + "vLLM sidecar cannot dereference pre-decoded RDMA media; configure URL passthrough", + )); + } + MultimodalData::UuidOnly(_) => { + return Err(client::invalid_argument( + "vLLM gRPC requires a media source and cannot resolve UUID-only media", + )); + } + }; + let uuid = uuids + .and_then(|uuids| uuids.get(index)) + .and_then(Clone::clone) + .or_else(|| forwarded_uuids.and_then(|uuids| uuids.get(index)).cloned()) + .unwrap_or_default(); + media.push(pb::MediaItem { + modality: pb::Modality::Image as i32, + source: Some(source), + mime_type: String::new(), + uuid, + }); + } + } + Ok(media) +} + fn top_n_candidates(count: u32) -> Result { i32::try_from(count).map_err(|_| { client::invalid_argument(format!( @@ -242,7 +504,9 @@ fn build_kv_parameters( Ok(pb::KvCacheParameters { bypass_prefix_cache, - cache_salt: cache_salt.unwrap_or_default(), + cache_salt: cache_salt + .map(|cache_salt| format!("{DYNAMO_CACHE_SALT_PREFIX}{cache_salt}")) + .unwrap_or_default(), kv_transfer_params: kv_transfer_params.map(json_to_struct).transpose()?, ec_transfer_params: None, }) @@ -301,13 +565,9 @@ fn validate_request( "prompt embeddings are not supported by vLLM gRPC v0.25.1", )); } - if request.multi_modal_data.is_some() - || request.mm_routing_info.is_some() - || request.mm_processor_kwargs.is_some() - || request.encoder_result.is_some() - { + if request.mm_processor_kwargs.is_some() || request.encoder_result.is_some() { return Err(client::invalid_argument( - "multimodal requests are not supported by vLLM gRPC v0.25.1", + "preprocessed multimodal features are not supported by vLLM gRPC", )); } if mode.is_encode() { @@ -325,15 +585,6 @@ fn validate_request( "LoRA request selection is not supported by vLLM gRPC v0.25.1", )); } - if request - .routing - .as_ref() - .is_some_and(|routing| routing.dp_rank.is_some() || routing.prefill_dp_rank.is_some()) - { - return Err(client::invalid_argument( - "KV-aware data-parallel routing is not supported by vLLM gRPC v0.25.1", - )); - } if request.bootstrap_info.is_some() { return Err(client::invalid_argument( "Dynamo bootstrap handoff is not supported by the vLLM sidecar", @@ -381,6 +632,8 @@ fn validate_request( pub(crate) struct ResponseState { prompt_tokens: u32, + has_media: bool, + multimodal_prompt_token_ids: Option>, completion_tokens: u32, is_prefill: bool, output_logprobs: Option, @@ -392,6 +645,11 @@ impl ResponseState { pub(crate) fn new(request: &PreprocessedRequest, mode: DisaggregationMode) -> Self { Self { prompt_tokens: request.token_ids.len() as u32, + has_media: request + .multi_modal_data + .as_ref() + .is_some_and(|media| media.values().any(|items| !items.is_empty())), + multimodal_prompt_token_ids: None, completion_tokens: 0, is_prefill: mode.is_prefill(), output_logprobs: request.output_options.logprobs, @@ -514,6 +772,28 @@ impl ResponseState { "prefill terminal is missing kv_transfer_params", )); } + if self.is_prefill && self.has_media { + let token_ids = self.multimodal_prompt_token_ids.take().ok_or_else(|| { + client::protocol_error( + "multimodal prefill did not return expanded prompt token IDs", + ) + })?; + let params = mapped + .disaggregated_params + .as_mut() + .and_then(serde_json::Value::as_object_mut) + .ok_or_else(|| { + client::protocol_error("prefill kv_transfer_params is not a JSON object") + })?; + params.insert( + MULTIMODAL_PROMPT_TOKEN_IDS_KEY.to_string(), + serde_json::to_value(token_ids).map_err(|error| { + client::protocol_error(format!( + "failed to encode multimodal prefill token IDs: {error}" + )) + })?, + ); + } self.attach_prompt_data(&mut mapped); Ok(Some(mapped)) } @@ -526,10 +806,24 @@ impl ResponseState { fn consume_prompt_info(&mut self, prompt: pb::PromptInfo) -> Result<(), DynamoError> { if prompt.num_prompt_tokens != self.prompt_tokens { - return Err(client::protocol_error(format!( - "prompt token count {} does not match request count {}", - prompt.num_prompt_tokens, self.prompt_tokens - ))); + if !self.has_media { + return Err(client::protocol_error(format!( + "prompt token count {} does not match request count {}", + prompt.num_prompt_tokens, self.prompt_tokens + ))); + } + // vLLM's count includes expanded media tokens. + self.prompt_tokens = prompt.num_prompt_tokens; + } + if self.is_prefill && self.has_media { + if prompt.token_ids.len() != prompt.num_prompt_tokens as usize { + return Err(client::protocol_error(format!( + "multimodal prefill returned {} prompt token IDs for {} prompt tokens", + prompt.token_ids.len(), + prompt.num_prompt_tokens + ))); + } + self.multimodal_prompt_token_ids = Some(prompt.token_ids.clone()); } if !self.expect_prompt_logprobs { return Ok(()); diff --git a/lib/sidecar/vllm/src/engine.rs b/lib/sidecar/vllm/src/engine.rs index 347e93e09db4..6a246d16564a 100644 --- a/lib/sidecar/vllm/src/engine.rs +++ b/lib/sidecar/vllm/src/engine.rs @@ -1,9 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashSet; + use async_trait::async_trait; use dynamo_backend_common::{ - DisaggregationMode, DynamoError, GenerateContext, LLMEngine, LLMEngineOutput, + DisaggregationMode, DynamoError, GenerateContext, KvEventSource, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, WorkerConfig, usage, }; use dynamo_sidecar_common::{GrpcEndpoint, GrpcTransportConfig}; @@ -14,7 +16,7 @@ use tokio_util::sync::CancellationToken; use crate::args::Args; use crate::client::{self, CONTROL_SERVICE, INFERENCE_SERVICE, VllmClient}; -use crate::convert::{ResponseState, build_generate_request}; +use crate::convert::{ResponseState, build_generate_request, data_parallel_rank}; use crate::model::DiscoveredModel; pub struct VllmSidecarEngine { @@ -121,13 +123,14 @@ impl VllmSidecarEngine { custom_jinja_template: args.sidecar.common.custom_jinja_template, model_name: model.source.clone(), served_model_name: Some(model.served_name.clone()), + // gRPC cannot yet preserve the parser request semantics. tool_call_parser: None, reasoning_parser: None, exclude_tools_when_tool_choice_none: args .sidecar .common .exclude_tools_when_tool_choice_none, - enable_kv_routing: false, + enable_kv_routing: true, disaggregation_mode: mode, route_to_encoder: false, ..Default::default() @@ -183,26 +186,45 @@ impl LLMEngine for VllmSidecarEngine { request: dynamo_backend_common::PreprocessedRequest, ctx: GenerateContext, ) -> Result>, DynamoError> { + if request + .multi_modal_data + .as_ref() + .is_some_and(|media| media.values().any(|items| !items.is_empty())) + && !self.model.supports_multimodal + { + return Err(client::invalid_argument(format!( + "model `{}` does not advertise multimodal support", + self.model.served_name + ))); + } let client = self .client .get() .ok_or_else(|| client::engine_shutdown("vLLM sidecar is not started"))?; let request_id = ctx.id().to_string(); let mut state = ResponseState::new(&request, self.mode); + let data_parallel_rank = data_parallel_rank(&request, self.mode); let mut proto_request = build_generate_request(request, request_id, self.mode)?; proto_request.model.clone_from(&self.model.served_name); + let defer_request_cancellation = self.mode.is_decode(); let stopped_ctx = ctx.inner_arc(); let shutdown = self.cancel.clone(); - let mut cancellation = Box::pin(async move { + let mut request_cancellation = Box::pin(async move { stopped_ctx.stopped().await }); + let mut shutdown_cancellation = Box::pin(async move { shutdown.cancelled().await }); + let stream = if defer_request_cancellation { + // Decode must reach vLLM so NIXL can release transferred KV. + tokio::select! { + biased; + _ = shutdown_cancellation.as_mut() => None, + result = client.generate_stream(proto_request, data_parallel_rank) => Some(result?), + } + } else { tokio::select! { - _ = stopped_ctx.stopped() => {} - _ = shutdown.cancelled() => {} + biased; + _ = shutdown_cancellation.as_mut() => None, + _ = request_cancellation.as_mut() => None, + result = client.generate_stream(proto_request, data_parallel_rank) => Some(result?), } - }); - let stream = tokio::select! { - biased; - _ = cancellation.as_mut() => None, - result = client.generate_stream(proto_request) => Some(result?), }; let Some(mut stream) = stream else { let output = cancelled(&state); @@ -210,41 +232,100 @@ impl LLMEngine for VllmSidecarEngine { }; Ok(Box::pin(async_stream::stream! { + let mut request_cancelled = false; + let mut first_token_observed = false; loop { - tokio::select! { - biased; - _ = cancellation.as_mut() => { - yield Ok(cancelled(&state)); - break; + let message = if request_cancelled { + tokio::select! { + biased; + _ = shutdown_cancellation.as_mut() => None, + message = stream.message() => Some(message), + } + } else { + tokio::select! { + biased; + _ = shutdown_cancellation.as_mut() => None, + _ = request_cancellation.as_mut() => { + if defer_request_cancellation && !first_token_observed { + request_cancelled = true; + continue; + } + None + } + message = stream.message() => Some(message), } - message = stream.message() => { - match message { - Ok(Some(response)) => match state.convert(response) { - Ok(Some(output)) => { - let terminal = output.finish_reason.is_some(); - yield Ok(output); - if terminal { - break; + }; + + let Some(message) = message else { + yield Ok(cancelled(&state)); + break; + }; + match message { + Ok(Some(response)) => { + let response_has_token = response + .outputs + .as_ref() + .is_some_and(|output| output.num_tokens > 0); + let transfer_completed = response.outputs.as_ref().is_some_and(|output| { + output.num_tokens > 0 || output.finish_info.is_some() + }); + match state.convert(response) { + Ok(Some(output)) => { + first_token_observed |= response_has_token; + if request_cancelled && transfer_completed { + // Dropping this stream aborts only this request. + if first_token_observed { + ctx.notify_first_token(); } + yield Ok(cancelled(&state)); + break; } - Ok(None) => {} - Err(error) => { - yield Err(error); + let terminal = output.finish_reason.is_some(); + yield Ok(output); + if terminal { break; } - }, - Ok(None) => { - yield Err(client::protocol_error( - "GenerateStream ended before a terminal response", - )); + } + Ok(None) => {} + Err(error) if request_cancelled => { + tracing::warn!( + %error, + "vLLM response conversion failed after request cancellation" + ); + yield Ok(cancelled(&state)); break; } - Err(status) => { - yield Err(client::status_to_dynamo("GenerateStream", status)); + Err(error) => { + yield Err(error); break; } } } + Ok(None) if request_cancelled => { + tracing::warn!( + "vLLM GenerateStream ended before transfer completion after request cancellation" + ); + yield Ok(cancelled(&state)); + break; + } + Ok(None) => { + yield Err(client::protocol_error( + "GenerateStream ended before a terminal response", + )); + break; + } + Err(status) if request_cancelled => { + tracing::warn!( + %status, + "vLLM GenerateStream failed before transfer completion after request cancellation" + ); + yield Ok(cancelled(&state)); + break; + } + Err(status) => { + yield Err(client::status_to_dynamo("GenerateStream", status)); + break; + } } } })) @@ -254,6 +335,70 @@ impl LLMEngine for VllmSidecarEngine { self.cancel.cancel(); Ok(()) } + + async fn kv_event_sources(&self) -> Result, DynamoError> { + let client = self + .client + .get() + .ok_or_else(|| client::engine_shutdown("vLLM sidecar is not started"))?; + let expected_dp_size = self.model.data_parallel_size(); + let mut ranks = HashSet::new(); + let mut sources = Vec::new(); + for source in client.kv_event_sources().await? { + if source.transport != "zmq" { + tracing::warn!( + transport = %source.transport, + endpoint = %source.endpoint, + "Skipping unsupported vLLM KV-event transport" + ); + continue; + } + let dp_rank = source.data_parallel_rank.ok_or_else(|| { + client::protocol_error( + "GetKvEventSources returned a ZMQ source without data_parallel_rank", + ) + })?; + if dp_rank >= expected_dp_size { + return Err(client::protocol_error(format!( + "GetKvEventSources returned rank {dp_rank}, outside the expected range 0..{expected_dp_size}", + ))); + } + if !ranks.insert(dp_rank) { + return Err(client::protocol_error(format!( + "GetKvEventSources returned duplicate rank {dp_rank}", + ))); + } + if source.endpoint.trim().is_empty() { + return Err(client::protocol_error( + "GetKvEventSources returned a ZMQ source without an endpoint", + )); + } + sources.push(KvEventSource::Zmq { + endpoint: zmq_connect_endpoint(&source.endpoint, &self.endpoint), + topic: source.topic, + dp_rank, + }); + } + if ranks.len() != expected_dp_size as usize { + return Err(client::protocol_error(format!( + "GetKvEventSources returned ZMQ sources for {} of {expected_dp_size} data-parallel ranks; KV routing requires one source for every rank", + ranks.len() + ))); + } + Ok(sources) + } +} + +fn zmq_connect_endpoint(endpoint: &str, grpc_endpoint: &GrpcEndpoint) -> String { + let port = endpoint + .strip_prefix("tcp://*:") + .or_else(|| endpoint.strip_prefix("tcp://0.0.0.0:")) + .or_else(|| endpoint.strip_prefix("tcp://[::]:")); + let Some(port) = port else { + return endpoint.to_string(); + }; + + format!("tcp://{}:{port}", grpc_endpoint.authority_host()) } fn bootstrap_discover( diff --git a/lib/sidecar/vllm/src/lib.rs b/lib/sidecar/vllm/src/lib.rs index b0bfff6ecd56..6e911f86b06d 100644 --- a/lib/sidecar/vllm/src/lib.rs +++ b/lib/sidecar/vllm/src/lib.rs @@ -10,8 +10,7 @@ mod engine; mod json; mod model; -/// Generated vLLM gRPC types, temporarily exposed for the Mocker server until -/// vLLM publishes its upstream protocol package. +/// vLLM gRPC types published through the Buf Schema Registry. #[doc(hidden)] pub mod proto; diff --git a/lib/sidecar/vllm/src/model.rs b/lib/sidecar/vllm/src/model.rs index b40c3191191c..03e82919e775 100644 --- a/lib/sidecar/vllm/src/model.rs +++ b/lib/sidecar/vllm/src/model.rs @@ -21,6 +21,7 @@ struct ModelIdentity { pub(crate) struct DiscoveredModel { pub source: String, pub served_name: String, + pub supports_multimodal: bool, identity: ModelIdentity, server: pb::ServerInfo, } @@ -36,6 +37,19 @@ impl DiscoveredModel { server.api_version ))); } + if let Some(parallelism) = server.parallelism.as_ref() { + if parallelism.data_parallel_size == 0 { + return Err(client::protocol_error( + "vLLM reports a data-parallel size of zero", + )); + } + if parallelism.data_parallel_rank != 0 { + return Err(client::protocol_error(format!( + "vLLM reports data_parallel_rank {}; the sidecar currently requires one frontend hosting the complete data-parallel group starting at rank 0", + parallelism.data_parallel_rank + ))); + } + } let source = required("model_id", model.model_id)?; let served_name = required("served_model_name", model.served_model_name)?; if !model.supports_token_ids_input { @@ -55,6 +69,7 @@ impl DiscoveredModel { Ok(Self { source, served_name, + supports_multimodal: model.supports_multimodal, identity, server, }) @@ -71,6 +86,7 @@ impl DiscoveredModel { } pub(crate) fn engine_config(&self) -> EngineConfig { + let parallelism = self.server.parallelism.as_ref(); EngineConfig { model: self.source.clone(), served_model_name: Some(self.served_name.clone()), @@ -82,10 +98,20 @@ impl DiscoveredModel { total_kv_blocks: nonzero(self.server.total_kv_blocks), max_num_seqs: nonzero(self.server.max_running_requests), max_num_batched_tokens: nonzero(self.server.max_batched_tokens), + data_parallel_size: parallelism + .and_then(|parallelism| nonzero(parallelism.data_parallel_size)), + data_parallel_start_rank: parallelism.map(|_| 0), ..Default::default() }), } } + + pub(crate) fn data_parallel_size(&self) -> u32 { + self.server + .parallelism + .as_ref() + .map_or(1, |parallelism| parallelism.data_parallel_size) + } } fn required(field: &str, value: String) -> Result { diff --git a/lib/sidecar/vllm/src/proto.rs b/lib/sidecar/vllm/src/proto.rs index de5928a0cbdd..8e80140a4ed2 100644 --- a/lib/sidecar/vllm/src/proto.rs +++ b/lib/sidecar/vllm/src/proto.rs @@ -1,7 +1,5 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -#![allow(clippy::all)] -#![allow(missing_docs)] - -tonic::include_proto!("vllm"); +pub use vllm_grpc::vllm::tonic::*; +pub use vllm_proto::vllm::*; diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index 05043e15cc10..e7cfda2e2253 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -8,9 +8,10 @@ use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; +use dynamo_backend_common::engine::RoutingHints; use dynamo_backend_common::{ - DisaggregationMode, FinishReason, GenerateContext, LLMEngine, OutputOptions, PrefillResult, - PreprocessedRequest, SamplingOptions, StopConditions, + DisaggregationMode, FinishReason, GenerateContext, LLMEngine, MultimodalData, OutputOptions, + PrefillResult, PreprocessedRequest, SamplingOptions, StopConditions, }; use dynamo_sidecar_common::{GrpcEndpoint, GrpcTransportConfig}; use futures::{Stream, StreamExt}; @@ -31,6 +32,7 @@ use crate::proto as pb; #[derive(Clone, Default)] struct FakeVllm { requests: Arc>>, + data_parallel_rank_metadata: Arc>>>, peers: Arc>>, model_info_override: Arc>>, reject: Arc, @@ -38,6 +40,10 @@ struct FakeVllm { hang_before_headers: Arc, headers_pending: Arc, release_headers: Arc, + hold_before_first_token: Arc, + close_before_first_token: Arc, + first_token_pending: Arc, + release_first_token: Arc, server_stream_dropped: Arc, } @@ -68,6 +74,16 @@ impl pb::inference_server::Inference for FakeVllm { if let Some(peer) = request.remote_addr() { self.peers.lock().await.push(peer); } + let data_parallel_rank = request + .metadata() + .get("x-data-parallel-rank") + .map(|value| value.to_str().map(str::to_owned)) + .transpose() + .map_err(|error| Status::invalid_argument(error.to_string()))?; + self.data_parallel_rank_metadata + .lock() + .await + .push(data_parallel_rank); let request = request.into_inner(); self.requests.lock().await.push(request.clone()); if self.hang_before_headers.load(Ordering::SeqCst) { @@ -86,10 +102,19 @@ impl pb::inference_server::Inference for FakeVllm { } None => return Err(Status::invalid_argument("prompt required")), }; + let prompt_tokens = if request.media.is_empty() { + prompt_tokens + } else { + 601 + }; let wants_logprobs = request .response .as_ref() .is_some_and(|response| response.output_logprobs); + let wants_prompt_token_ids = request + .response + .as_ref() + .is_some_and(|response| response.prompt_token_ids); let wants_prompt_logprobs = request .response .as_ref() @@ -120,33 +145,51 @@ impl pb::inference_server::Inference for FakeVllm { "nested": {"flags": [true, null, "opaque"]}, }); let hang = self.hang.load(Ordering::SeqCst); + let hold_before_first_token = self.hold_before_first_token.load(Ordering::SeqCst); + let close_before_first_token = self.close_before_first_token.load(Ordering::SeqCst); + let first_token_pending = self.first_token_pending.clone(); + let release_first_token = self.release_first_token.clone(); let dropped = self.server_stream_dropped.clone(); let stream = async_stream::try_stream! { let _drop_signal = DropSignal(dropped); - let prompt_info = if wants_prompt_logprobs { - pb::PromptInfo { - num_prompt_tokens: prompt_tokens, - token_ids: vec![11, 22, 33], - logprobs: vec![0.0, -0.2, -0.3], - ranks: vec![0, 1, 2], - candidate_tokens: vec![ - pb::CandidateTokenInfo { tokens: vec![] }, - pb::CandidateTokenInfo { tokens: vec![] }, - pb::CandidateTokenInfo { tokens: vec![] }, - ], - } - } else { - pb::PromptInfo { - num_prompt_tokens: prompt_tokens, - ..Default::default() - } + let prompt_info = pb::PromptInfo { + num_prompt_tokens: prompt_tokens, + token_ids: if wants_prompt_token_ids { + (0..prompt_tokens).collect() + } else { + Vec::new() + }, + logprobs: if wants_prompt_logprobs { + vec![-0.2; prompt_tokens as usize] + } else { + Vec::new() + }, + ranks: if wants_prompt_logprobs { + vec![1; prompt_tokens as usize] + } else { + Vec::new() + }, + candidate_tokens: if wants_prompt_logprobs { + vec![pb::CandidateTokenInfo::default(); prompt_tokens as usize] + } else { + Vec::new() + }, }; yield pb::GenerateResponse { prompt_info: Some(prompt_info), outputs: None, }; + if hold_before_first_token { + first_token_pending.store(true, Ordering::SeqCst); + release_first_token.notified().await; + first_token_pending.store(false, Ordering::SeqCst); + } + if close_before_first_token { + return; + } + if hang { loop { yield sequence_response(false, wants_logprobs, None); @@ -197,7 +240,20 @@ impl pb::control_server::Control for FakeVllm { _request: Request, ) -> Result, Status> { Ok(Response::new(pb::GetKvEventSourcesResponse { - sources: Vec::new(), + sources: (0..2) + .map(|rank| pb::KvEventSource { + transport: "zmq".to_string(), + endpoint: format!("tcp://*:{}", 20081 + rank), + topic: String::new(), + replay_endpoint: String::new(), + data_parallel_rank: Some(rank), + encoding: "msgpack".to_string(), + schema_version: 1, + buffer_steps: 0, + hwm: 0, + max_queue_size: 0, + }) + .collect(), })) } } @@ -223,8 +279,8 @@ fn server_info() -> pb::ServerInfo { parallelism: Some(pb::ParallelismInfo { tensor_parallel_size: 2, pipeline_parallel_size: 1, - data_parallel_size: 4, - data_parallel_rank: 2, + data_parallel_size: 2, + data_parallel_rank: 0, decode_context_parallel_size: 1, }), max_model_len: 8192, @@ -232,7 +288,6 @@ fn server_info() -> pb::ServerInfo { total_kv_blocks: 4096, max_running_requests: 128, max_batched_tokens: 2048, - supports_explicit_data_parallel_rank: false, } } @@ -471,8 +526,13 @@ fn request() -> PreprocessedRequest { prompt_logprobs: Some(1), ..Default::default() }) - .mdc_sum(Some("cache-salt".to_string())) + .mdc_sum(Some("model-checksum".to_string())) + .routing(Some(RoutingHints { + cache_namespace: Some("cache-salt".to_string()), + ..Default::default() + })) .extra_args(Some(json!({ + "nvext": {"cache_salt": "cache-salt", "token_in": true}, "bypass_prefix_cache": true, "kv_transfer_params": { "connector_data": {"values": [1, true, null]} @@ -482,14 +542,35 @@ fn request() -> PreprocessedRequest { .expect("request") } -fn engine(endpoint: &str, mode: DisaggregationMode, connections: usize) -> VllmSidecarEngine { +fn decode_request() -> PreprocessedRequest { + let mut request = request(); + request.prefill_result = Some(PrefillResult { + disaggregated_params: json!({ + "do_remote_decode": false, + "do_remote_prefill": true, + "remote_engine_id": "prefill-0", + "remote_host": "127.0.0.1", + "remote_port": 20097, + "remote_block_ids": [7, 8], + }), + prompt_tokens_details: None, + }); + request +} + +fn engine( + endpoint: &str, + mode: DisaggregationMode, + connections: usize, + model: pb::ModelInfo, +) -> VllmSidecarEngine { let transport = GrpcTransportConfig { connections: NonZeroUsize::new(connections).expect("non-zero connection count"), ..Default::default() }; VllmSidecarEngine::new( GrpcEndpoint::parse(endpoint, "--vllm-endpoint").expect("valid test endpoint"), - DiscoveredModel::from_proto(model_info(), server_info()).expect("valid discovery"), + DiscoveredModel::from_proto(model, server_info()).expect("valid discovery"), mode, transport, ) @@ -586,10 +667,40 @@ async fn aggregated_generation_converts_request_stream_and_usage() { assert_eq!(registration.total_kv_blocks, Some(4096)); assert_eq!(registration.max_num_seqs, Some(128)); assert_eq!(registration.max_num_batched_tokens, Some(2048)); - assert_eq!(registration.data_parallel_size, None); - assert_eq!(registration.data_parallel_start_rank, None); + assert_eq!(registration.data_parallel_size, Some(2)); + assert_eq!(registration.data_parallel_start_rank, Some(0)); + + let sources = engine.kv_event_sources().await.expect("KV event sources"); + assert_eq!(sources.len(), 2); + assert_eq!( + sources + .iter() + .map(|source| source.dp_rank()) + .collect::>(), + BTreeSet::from([0, 1]) + ); + assert!(sources.iter().all(|source| matches!( + source, + dynamo_backend_common::KvEventSource::Zmq { topic, .. } if topic.is_empty() + ))); + assert_eq!( + sources + .iter() + .map(|source| match source { + dynamo_backend_common::KvEventSource::Zmq { endpoint, .. } => endpoint.as_str(), + dynamo_backend_common::KvEventSource::Push { .. } => unreachable!(), + }) + .collect::>(), + ["tcp://127.0.0.1:20081", "tcp://127.0.0.1:20082"] + ); - let outputs = collect(&engine, request()).await; + let mut routed_request = serde_json::to_value(request()).expect("serialize request"); + routed_request["routing"] = json!({"dp_rank": 1, "cache_salt": "cache-salt"}); + let outputs = collect( + &engine, + serde_json::from_value(routed_request).expect("deserialize routed request"), + ) + .await; assert_eq!(outputs.len(), 1); let terminal = &outputs[0]; assert_eq!(terminal.token_ids, [42]); @@ -605,6 +716,10 @@ async fn aggregated_generation_converts_request_stream_and_usage() { let sent = requests.first().expect("recorded request"); assert_eq!(sent.model, "served-model"); assert_eq!(sent.priority, 0); + assert_eq!( + server.service.data_parallel_rank_metadata.lock().await[0], + Some("1".to_string()) + ); let sampling = sent.sampling.as_ref().unwrap(); assert_eq!( (sampling.top_k, sampling.top_p, sampling.min_p), @@ -631,19 +746,151 @@ async fn aggregated_generation_converts_request_stream_and_usage() { assert!(stopping.ignore_eos); let kv = sent.kv.as_ref().unwrap(); assert!(kv.bypass_prefix_cache); - assert_eq!(kv.cache_salt, "cache-salt"); + assert_eq!(kv.cache_salt, "dynamo-cache-salt:cache-salt"); assert_eq!( struct_to_json(kv.kv_transfer_params.clone().unwrap()).unwrap(), json!({"connector_data": {"values": [1, true, null]}}) ); } +#[tokio::test] +async fn multimodal_image_is_forwarded_with_uuid() { + let service = FakeVllm::default(); + let mut discovered = model_info(); + discovered.supports_multimodal = true; + *service.model_info_override.lock().await = Some(discovered.clone()); + let server = FakeServer::start(service).await; + let (aggregate, _) = engine_from_args(&server.endpoint).await; + aggregate.start(0).await.expect("start"); + + let mut image_request = request(); + image_request.multi_modal_data = Some(std::collections::HashMap::from([( + "image_url".to_string(), + vec![MultimodalData::RawUrl( + "data:image/png;base64,iVBORw0KGgo=".to_string(), + )], + )])); + image_request.output_options.prompt_logprobs = None; + image_request + .extra_args + .as_mut() + .and_then(serde_json::Value::as_object_mut) + .expect("object extra_args") + .extend([ + ( + "messages".to_string(), + json!([{"role": "user", "content": [{"type": "image_url"}]}]), + ), + ("formatted_prompt".to_string(), json!("\nDescribe.")), + ("mm_hashes".to_string(), json!(["0123456789abcdef"])), + ]); + + let outputs = collect(&aggregate, image_request.clone()).await; + assert_eq!(outputs[0].finish_reason, Some(FinishReason::Stop)); + assert_eq!( + outputs[0] + .completion_usage + .as_ref() + .expect("usage") + .prompt_tokens, + 601 + ); + + let requests = server.service.requests.lock().await; + let media = &requests.last().expect("recorded request").media; + assert_eq!(media.len(), 1); + assert_eq!(media[0].modality(), pb::Modality::Image); + assert_eq!( + media[0].uuid, + "0123456789abcdef000000000000000000000000000000000000000000000000" + ); + assert!(matches!( + media[0].source.as_ref(), + Some(pb::media_item::Source::DataUri(_)) + )); + drop(requests); + + let prefill = engine( + &server.endpoint, + DisaggregationMode::Prefill, + 1, + discovered.clone(), + ); + let decode = engine(&server.endpoint, DisaggregationMode::Decode, 1, discovered); + prefill.start(1).await.expect("start prefill"); + decode.start(2).await.expect("start decode"); + + let prefill_outputs = collect(&prefill, image_request.clone()).await; + let handoff = prefill_outputs[0] + .disaggregated_params + .clone() + .expect("multimodal handoff"); + assert_eq!( + handoff["_dynamo_sidecar_multimodal_prompt_token_ids"] + .as_array() + .expect("expanded prompt token IDs") + .len(), + 601 + ); + + let mut decode_request = image_request; + decode_request.prefill_result = Some(PrefillResult { + disaggregated_params: handoff, + prompt_tokens_details: None, + }); + let decode_outputs = collect(&decode, decode_request).await; + assert_eq!( + decode_outputs[0] + .completion_usage + .as_ref() + .expect("decode usage") + .prompt_tokens, + 601 + ); + + let requests = server.service.requests.lock().await; + let prefill_wire = &requests[requests.len() - 2]; + let decode_wire = &requests[requests.len() - 1]; + assert_eq!(prefill_wire.media.len(), 1); + assert!( + prefill_wire + .response + .as_ref() + .expect("prefill response options") + .prompt_token_ids + ); + assert!(decode_wire.media.is_empty()); + assert_eq!( + decode_wire.prompt.as_ref(), + Some(&pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: (0..601).collect(), + })) + ); + let decode_kv = struct_to_json( + decode_wire + .kv + .as_ref() + .and_then(|kv| kv.kv_transfer_params.clone()) + .expect("decode KV handoff"), + ) + .expect("decode KV JSON"); + assert!( + decode_kv["_dynamo_sidecar_multimodal_prompt_token_ids"].is_null(), + "sidecar metadata must not reach vLLM" + ); +} + #[tokio::test] async fn grpc_request_errors_are_propagated() { let service = FakeVllm::default(); service.reject.store(true, Ordering::SeqCst); let server = FakeServer::start(service).await; - let engine = engine(&server.endpoint, DisaggregationMode::Aggregated, 1); + let engine = engine( + &server.endpoint, + DisaggregationMode::Aggregated, + 1, + model_info(), + ); engine.start(0).await.expect("start"); let context = dynamo_backend_common::testing::mock_context(); @@ -657,8 +904,18 @@ async fn grpc_request_errors_are_propagated() { #[tokio::test] async fn prefill_decode_handoff_is_opaque_and_repeatable() { let server = FakeServer::start(FakeVllm::default()).await; - let prefill = engine(&server.endpoint, DisaggregationMode::Prefill, 1); - let decode = engine(&server.endpoint, DisaggregationMode::Decode, 1); + let prefill = engine( + &server.endpoint, + DisaggregationMode::Prefill, + 1, + model_info(), + ); + let decode = engine( + &server.endpoint, + DisaggregationMode::Decode, + 1, + model_info(), + ); prefill.start(0).await.expect("start prefill"); decode.start(1).await.expect("start decode"); @@ -734,11 +991,14 @@ async fn pool_uses_each_configured_connection() { for index in 0..4 { let mut stream = client - .generate_stream(pb::GenerateRequest { - request_id: format!("request-{index}"), - prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), - ..Default::default() - }) + .generate_stream( + pb::GenerateRequest { + request_id: format!("request-{index}"), + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), + ..Default::default() + }, + None, + ) .await .expect("start stream"); while stream.message().await.expect("message").is_some() {} @@ -753,6 +1013,15 @@ async fn pool_uses_each_configured_connection() { .map(SocketAddr::port) .collect(); assert_eq!(ports.len(), 2); + assert!( + server + .service + .data_parallel_rank_metadata + .lock() + .await + .iter() + .all(Option::is_none) + ); } #[tokio::test] @@ -760,7 +1029,12 @@ async fn cancellation_drops_the_remote_stream() { let service = FakeVllm::default(); service.hang.store(true, Ordering::SeqCst); let server = FakeServer::start(service).await; - let engine = engine(&server.endpoint, DisaggregationMode::Aggregated, 1); + let engine = engine( + &server.endpoint, + DisaggregationMode::Aggregated, + 1, + model_info(), + ); engine.start(0).await.expect("start"); let context = dynamo_backend_common::testing::mock_context(); @@ -789,7 +1063,12 @@ async fn cancellation_interrupts_pending_response_headers() { let service = FakeVllm::default(); service.hang_before_headers.store(true, Ordering::SeqCst); let server = FakeServer::start(service).await; - let engine = engine(&server.endpoint, DisaggregationMode::Aggregated, 1); + let engine = engine( + &server.endpoint, + DisaggregationMode::Aggregated, + 1, + model_info(), + ); engine.start(0).await.expect("start"); let context = dynamo_backend_common::testing::mock_context(); @@ -815,10 +1094,127 @@ async fn cancellation_interrupts_pending_response_headers() { server.service.release_headers.notify_waiters(); } +#[tokio::test] +async fn decode_cancellation_waits_for_submission_and_first_token() { + let service = FakeVllm::default(); + service.hang_before_headers.store(true, Ordering::SeqCst); + service + .hold_before_first_token + .store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine( + &server.endpoint, + DisaggregationMode::Decode, + 1, + model_info(), + ); + engine.start(0).await.expect("start"); + + let context = dynamo_backend_common::testing::mock_context(); + let generate = engine.generate( + decode_request(), + GenerateContext::new(context.clone(), None), + ); + tokio::pin!(generate); + + tokio::select! { + _ = &mut generate => panic!("decode returned before response headers were gated"), + _ = async { + while !server.service.headers_pending.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + } => {} + } + assert_eq!(server.service.requests.lock().await.len(), 1); + context.stop_generating(); + tokio::select! { + _ = &mut generate => panic!("decode cancellation returned before response headers"), + _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => {} + } + + server.service.release_headers.notify_one(); + let mut stream = tokio::time::timeout(std::time::Duration::from_secs(2), &mut generate) + .await + .expect("decode response headers") + .expect("decode stream"); + let next = stream.next(); + tokio::pin!(next); + tokio::select! { + _ = &mut next => panic!("decode returned before the first token was gated"), + _ = async { + while !server.service.first_token_pending.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + } => {} + } + assert!( + !server.service.server_stream_dropped.load(Ordering::SeqCst), + "decode stream dropped before the first token" + ); + tokio::select! { + _ = &mut next => panic!("decode cancellation completed before the first token"), + _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => {} + } + + server.service.release_first_token.notify_one(); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), &mut next) + .await + .expect("first token did not release decode cancellation") + .expect("cancelled terminal") + .expect("cancelled output"); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); + drop(stream); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while !server.service.server_stream_dropped.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("server stream dropped after first token"); +} + +#[tokio::test] +async fn decode_cancellation_maps_premature_eof_to_cancelled() { + let service = FakeVllm::default(); + service + .close_before_first_token + .store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine( + &server.endpoint, + DisaggregationMode::Decode, + 1, + model_info(), + ); + engine.start(0).await.expect("start"); + + let context = dynamo_backend_common::testing::mock_context(); + let mut stream = engine + .generate( + decode_request(), + GenerateContext::new(context.clone(), None), + ) + .await + .expect("decode stream"); + context.stop_generating(); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) + .await + .expect("premature EOF did not release decode cancellation") + .expect("cancelled terminal") + .expect("cancelled output"); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); +} + #[tokio::test] async fn unsupported_features_fail_before_rpc_submission() { let server = FakeServer::start(FakeVllm::default()).await; - let engine = engine(&server.endpoint, DisaggregationMode::Aggregated, 1); + let engine = engine( + &server.endpoint, + DisaggregationMode::Aggregated, + 1, + model_info(), + ); engine.start(0).await.expect("start"); let mut requests = Vec::new(); @@ -835,11 +1231,14 @@ async fn unsupported_features_fail_before_rpc_submission() { multimodal.mm_processor_kwargs = Some(json!({"use_audio_in_video": true})); requests.push(multimodal); - for routing in [json!({"lora_name": "adapter"}), json!({"dp_rank": 1})] { - let mut value = serde_json::to_value(request()).expect("serialize request"); - value["routing"] = routing; - requests.push(serde_json::from_value(value).expect("deserialize request")); - } + let mut lora_request = serde_json::to_value(request()).expect("serialize request"); + lora_request["routing"] = json!({"lora_name": "adapter"}); + requests.push(serde_json::from_value(lora_request).expect("deserialize request")); + + let mut mismatched_cache_salt = request(); + mismatched_cache_salt.extra_args.as_mut().unwrap()["nvext"]["cache_salt"] = + json!("different-cache-salt"); + requests.push(mismatched_cache_salt); for unsupported in requests { let context = dynamo_backend_common::testing::mock_context(); diff --git a/recipes/glm-5-nvfp4/sglang/disagg/efa/Dockerfile.efa b/recipes/glm-5-nvfp4/sglang/disagg/efa/Dockerfile.efa index a820215539b6..7c26dc5b7a74 100644 --- a/recipes/glm-5-nvfp4/sglang/disagg/efa/Dockerfile.efa +++ b/recipes/glm-5-nvfp4/sglang/disagg/efa/Dockerfile.efa @@ -52,8 +52,10 @@ RUN cargo install maturin --locked RUN git clone https://github.com/ai-dynamo/dynamo.git /build/dynamo && \ cd /build/dynamo && git checkout ${DYNAMO_COMMIT} -RUN --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ +RUN --mount=type=secret,id=buf_token,required=true \ + --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ --mount=type=cache,target=/root/.cargo/git,sharing=shared \ + export CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)" && \ cd /build/dynamo/lib/bindings/python && \ maturin build --release && \ mkdir -p /build/dist && \ diff --git a/recipes/glm-5-nvfp4/sglang/disagg/efa/README.md b/recipes/glm-5-nvfp4/sglang/disagg/efa/README.md index 58abc2a13ced..0c40769f345b 100644 --- a/recipes/glm-5-nvfp4/sglang/disagg/efa/README.md +++ b/recipes/glm-5-nvfp4/sglang/disagg/efa/README.md @@ -30,6 +30,7 @@ Identical to the non-EFA recipe. - A Kubernetes cluster with the Dynamo Operator installed. - The NVIDIA `ComputeDomain` operator (for the MNNVL ResourceClaim used here). - Shared NFS PVC for model weights (same as the non-EFA recipe). +- A Buf Schema Registry token exported as `BUF_TOKEN`. The libfabric is built into the image — no cluster-side DaemonSet is required. @@ -37,6 +38,7 @@ The libfabric is built into the image — no cluster-side DaemonSet is required. ```bash docker buildx build \ + --secret id=buf_token,env=BUF_TOKEN \ --platform linux/arm64 \ --build-arg ARCH=arm64 \ -t /sglang-dynamo-glm5-efa:latest \ diff --git a/recipes/kimi-k2.5/tokenspeed/agg/nvidia/Dockerfile b/recipes/kimi-k2.5/tokenspeed/agg/nvidia/Dockerfile index 8936678b2358..317187ea48ae 100644 --- a/recipes/kimi-k2.5/tokenspeed/agg/nvidia/Dockerfile +++ b/recipes/kimi-k2.5/tokenspeed/agg/nvidia/Dockerfile @@ -105,9 +105,11 @@ RUN --mount=type=cache,target=/root/.cargo/registry \ FROM build_tools AS wheel_builder WORKDIR /workspace/dynamo COPY . /workspace/dynamo -RUN --mount=type=cache,target=/root/.cargo/registry \ +RUN --mount=type=secret,id=buf_token,required=true \ + --mount=type=cache,target=/root/.cargo/registry \ --mount=type=cache,target=/root/.cargo/git \ set -eux; \ + export CARGO_REGISTRIES_BUF_TOKEN="Bearer $(cat /run/secrets/buf_token)"; \ cd /workspace/dynamo/lib/bindings/python; \ maturin build --release -o /tmp/dynamo-dist diff --git a/recipes/kimi-k2.5/tokenspeed/agg/nvidia/README.md b/recipes/kimi-k2.5/tokenspeed/agg/nvidia/README.md index acc4b4ba1c56..9dd1598f875c 100644 --- a/recipes/kimi-k2.5/tokenspeed/agg/nvidia/README.md +++ b/recipes/kimi-k2.5/tokenspeed/agg/nvidia/README.md @@ -57,12 +57,12 @@ update the `image:` fields in [`deploy.yaml`](deploy.yaml). ### 1. Build the Dynamo+TokenSpeed image -The build context must be the **Dynamo repo root** (the Dockerfile `COPY`s the source -tree in to build the Dynamo Python wheel via `maturin`). +The build context must be the **Dynamo repo root** (the Dockerfile `COPY`s the source tree in to build the Dynamo Python wheel via `maturin`). Export a Buf Schema Registry token as `BUF_TOKEN` before building. ```bash # From the repo root. docker build \ + --secret id=buf_token,env=BUF_TOKEN \ -f recipes/kimi-k2.5/tokenspeed/agg/nvidia/Dockerfile \ --target dev \ -t /dynamo-tokenspeed:dev \