diff --git a/.github/actions/compliance-extract/action.yml b/.github/actions/compliance-extract/action.yml index b470a2fa6935..076deaf5420c 100644 --- a/.github/actions/compliance-extract/action.yml +++ b/.github/actions/compliance-extract/action.yml @@ -130,7 +130,6 @@ runs: GIT_SHA: ${{ inputs.git_sha }} EPP_IMAGE: ${{ inputs.epp_image }} run: | - mkdir -p /tmp/compliance # Reuse the warm cache the build just populated by matching ALL the # cache-key-affecting build-args the build passed. Otherwise the shared # stages rebuild here instead of cache-hitting: @@ -174,7 +173,8 @@ runs: # --build-arg (it would clear the ARG and reintroduce the failure). EPP_IMAGE_ARG="" [ -n "${EPP_IMAGE:-}" ] && EPP_IMAGE_ARG="--build-arg EPP_IMAGE=${EPP_IMAGE}" - docker buildx build \ + bash "${{ github.action_path }}/extract-with-retry.sh" /tmp/compliance \ + docker buildx build \ --builder ${{ inputs.builder_name }} \ --platform ${{ inputs.platform }} \ --target compliance_artifact \ @@ -183,7 +183,6 @@ runs: --build-arg SCCACHE_BUCKET="${SCCACHE_BUCKET}" \ --build-arg SCCACHE_REGION="${SCCACHE_REGION}" \ $EPP_IMAGE_ARG \ - --output type=local,dest=/tmp/compliance \ --progress=plain \ -f "${DOCKERFILE_PATH}" \ $SECRET_ARGS \ @@ -357,7 +356,6 @@ runs: GIT_SHA: ${{ inputs.git_sha }} EPP_IMAGE: ${{ inputs.epp_image }} run: | - mkdir -p /tmp/sources # ENABLE_SOURCE_ARCHIVAL=true to actually run sources_collect; the other # build-args match the build so this shares the warm cache (same reason # as the compliance_artifact extract above). Same IRSA secrets too, so a @@ -385,7 +383,8 @@ runs: # resolves the GC'd default staging tag. Empty for non-frontend targets. EPP_IMAGE_ARG="" [ -n "${EPP_IMAGE:-}" ] && EPP_IMAGE_ARG="--build-arg EPP_IMAGE=${EPP_IMAGE}" - docker buildx build \ + bash "${{ github.action_path }}/extract-with-retry.sh" /tmp/sources \ + docker buildx build \ --builder ${{ inputs.builder_name }} \ --platform ${{ inputs.platform }} \ --target sources_archive \ @@ -395,7 +394,6 @@ runs: --build-arg SCCACHE_REGION="${SCCACHE_REGION}" \ --build-arg ENABLE_SOURCE_ARCHIVAL=true \ $EPP_IMAGE_ARG \ - --output type=local,dest=/tmp/sources \ --progress=plain \ -f "${DOCKERFILE_PATH}" \ $SECRET_ARGS \ diff --git a/.github/actions/compliance-extract/extract-with-retry.sh b/.github/actions/compliance-extract/extract-with-retry.sh new file mode 100755 index 000000000000..4feed8b62576 --- /dev/null +++ b/.github/actions/compliance-extract/extract-with-retry.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Retry only transient BuildKit transport failures, preserving each attempt's log. +set -euo pipefail + +destination=$1 +shift +attempt_root=$(mktemp -d "${RUNNER_TEMP:-${TMPDIR:-/tmp}}/compliance-extract.XXXXXX") +echo "Compliance extraction logs: ${attempt_root}" + +for attempt in 1 2 3; do + output="${attempt_root}/output-${attempt}" + log="${attempt_root}/attempt-${attempt}.log" + mkdir -p "$output" + status=0 + "$@" --output "type=local,dest=${output}" 2>&1 | tee "$log" || status=$? + if [ "$status" -eq 0 ]; then + # Publish only a complete export; failed attempts never reach consumers. + rm -rf -- "$destination" + mv -- "$output" "$destination" + exit 0 + fi + rm -rf -- "$output" + + # Inspect the terminal build error, not earlier warnings or command output. + error=$(grep '^ERROR:' "$log" | tail -n 1 || true) + if printf '%s\n' "$error" | grep -Eiq \ + '(process .*did not complete successfully|401 Unauthorized|403 Forbidden|denied:|manifest unknown)'; then + exit "$status" + fi + if [ "$attempt" -eq 3 ] || ! printf '%s\n' "$error" | grep -Eiq \ + '(Unavailable:.*error reading from server: (unexpected )?EOF|connection reset by peer|transport is closing|unexpected status.*(502 Bad Gateway|503 Service Unavailable|504 Gateway Timeout))'; then + exit "$status" + fi + delay=10 + [ "$attempt" -eq 2 ] && delay=30 + echo "::warning::Compliance extraction attempt ${attempt}/3 failed with a transient transport error; retrying in ${delay}s" + sleep "$delay" +done diff --git a/.github/actions/compliance-extract/test_extract_with_retry.py b/.github/actions/compliance-extract/test_extract_with_retry.py new file mode 100644 index 000000000000..616ae9342061 --- /dev/null +++ b/.github/actions/compliance-extract/test_extract_with_retry.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exercise the real retry wrapper with a fake BuildKit command and sleep.""" + +import os +import subprocess +import tempfile +import unittest +from pathlib import Path + +SCRIPT = Path(__file__).with_name("extract-with-retry.sh") +TRANSIENT = "ERROR: failed to build: failed to solve: Unavailable: error reading from server: EOF" + + +class ExtractRetryTest(unittest.TestCase): + def run_extract(self, errors): + with tempfile.TemporaryDirectory() as temp: + root = Path(temp) + for index, error in enumerate(errors, 1): + (root / f"error-{index}").write_text(error) + fake = root / "build" + fake.write_text( + """#!/usr/bin/env bash +set -eu +count=0 +[ ! -f "$STATE/count" ] || count=$(cat "$STATE/count") +count=$((count + 1)) +echo "$count" > "$STATE/count" +for arg in "$@"; do + case "$arg" in type=local,dest=*) output=${arg#type=local,dest=} ;; esac +done +mkdir -p "$output" +echo "$count" > "$output/attempt-$count" +error="$STATE/error-$count" +if [ -s "$error" ]; then + cat "$error" + exit 7 +fi +""" + ) + fake.chmod(0o755) + sleep = root / "sleep" + sleep.write_text('#!/usr/bin/env bash\necho "$1" >> "$STATE/delays"\n') + sleep.chmod(0o755) + destination = root / "result" + env = dict( + os.environ, + STATE=temp, + RUNNER_TEMP=temp, + PATH=temp + os.pathsep + os.environ["PATH"], + ) + result = subprocess.run( + ["bash", str(SCRIPT), str(destination), str(fake)], + env=env, + capture_output=True, + text=True, + ) + count = int((root / "count").read_text()) + delays = ( + (root / "delays").read_text().splitlines() + if (root / "delays").exists() + else [] + ) + files = sorted(p.name for p in destination.glob("*")) + logs = list(root.glob("compliance-extract.*/attempt-*.log")) + self.assertEqual(len(logs), count) + return result.returncode, count, delays, files + + def test_success(self): + self.assertEqual(self.run_extract([""]), (0, 1, [], ["attempt-1"])) + + def test_transient_recovers_without_partial_files(self): + self.assertEqual( + self.run_extract([TRANSIENT, ""]), (0, 2, ["10"], ["attempt-2"]) + ) + + def test_exhaustion_preserves_failure(self): + self.assertEqual(self.run_extract([TRANSIENT] * 3), (7, 3, ["10", "30"], [])) + + def test_deterministic_errors_do_not_retry(self): + for error in [ + "ERROR: failed to solve: process did not complete: exit code: 1", + "ERROR: failed to authorize: 401 Unauthorized", + "ERROR: manifest unknown", + "No solution found when resolving dependencies", + ]: + with self.subTest(error=error): + self.assertEqual(self.run_extract([error]), (7, 1, [], [])) + + def test_earlier_transient_does_not_mask_terminal_failure(self): + self.assertEqual( + self.run_extract([TRANSIENT + "\nERROR: manifest unknown"]), (7, 1, [], []) + ) + + def test_other_transport_errors_recover(self): + for error in [ + "ERROR: failed to solve: connection reset by peer", + "ERROR: failed to solve: rpc error: transport is closing", + "ERROR: failed to solve: unexpected status from HEAD request: 503 Service Unavailable", + ]: + with self.subTest(error=error): + self.assertEqual( + self.run_extract([error, ""]), (0, 2, ["10"], ["attempt-2"]) + ) + + def test_build_command_containing_transport_text_does_not_retry(self): + error = 'ERROR: process "echo connection reset by peer" did not complete successfully: exit code: 1' + self.assertEqual(self.run_extract([error]), (7, 1, [], [])) + + def test_permanent_error_after_retry_stops(self): + self.assertEqual( + self.run_extract([TRANSIENT, "ERROR: manifest unknown"]), (7, 2, ["10"], []) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/actions/connect-vcluster/action.yml b/.github/actions/connect-vcluster/action.yml index eafc9792a41b..77889ae55bf8 100644 --- a/.github/actions/connect-vcluster/action.yml +++ b/.github/actions/connect-vcluster/action.yml @@ -167,5 +167,8 @@ runs: fi KUBECONFIG_B64=$(base64 < "${KUBECONFIG_PATH}" | tr -d '\n') + # Masked before export: the value carries a vCluster client certificate, + # and every consumer interpolates it into a shell script. + echo "::add-mask::${KUBECONFIG_B64}" echo "kubeconfig_base64=${KUBECONFIG_B64}" >> $GITHUB_OUTPUT echo "::endgroup::" diff --git a/.github/actions/pytest-local/action.yml b/.github/actions/pytest-local/action.yml index e7bae3c105ca..d3ce10dead83 100644 --- a/.github/actions/pytest-local/action.yml +++ b/.github/actions/pytest-local/action.yml @@ -176,10 +176,19 @@ runs: # Artifacts go to the GitHub Actions checkout so the per-job upload in # shared-test.yml can find them. TEST_RESULTS_DIR="${GITHUB_WORKSPACE}/test-results" - mkdir -p "${TEST_RESULTS_DIR}" + mkdir -p "${TEST_RESULTS_DIR}/status" chmod 777 "${TEST_RESULTS_DIR}" echo "๐Ÿ“ Test results will be saved to: ${TEST_RESULTS_DIR}" + # Emptied up front so the file read at the end of the job can only have + # come from this run of this stage. Truncated rather than removed: the + # distroless planner test image ships no rm (see the tool list in + # container/Dockerfile.test), and an empty file is the better signal + # anyway - absent means the stage never started, empty means it started + # and was cut off before recording an exit code. + STATUS_FILE="${TEST_RESULTS_DIR}/status/${STR_TEST_TYPE}.exitcode" + : > "${STATUS_FILE}" + # NVMe guard: GPU runners mount ${NVME_DIR} (nvme-scratch emptyDir, NVMe-backed via # instanceStorePolicy: RAID0). CPU/ARM runners have no ${NVME_DIR} mount. if [[ -d "${NVME_DIR}" ]]; then @@ -303,8 +312,40 @@ runs: echo "โš ๏ธ JUnit XML file not found - test results may not be available for upload" fi + # Checked by "Verify test stages completed" in shared-test.yml, which + # cannot use the exit below or $GITHUB_ENV: a truncated script is still + # reported as a successful step and its $GITHUB_ENV writes are dropped. + # Keep this last - its absence is what proves the script was truncated. + echo "${TEST_EXIT_CODE}" > "${STATUS_FILE}" + exit ${TEST_EXIT_CODE} + # conftest.py's autouse `logger` fixture writes every test's full root-logger + # output (ManagedProcess/EngineProcess health-check status, retry reasons, + # subprocess log tails, ...) to /tmp/dynamo_tests//test.log.txt -- + # a much richer trace than the JUnit XML's truncated captured-output, but + # nothing ever uploaded it. failure()-gated (not always()) to keep this + # opt-in-by-outcome: pytest-local is shared by every framework's CI, and + # most runs pass. + - name: Collect per-test log files + if: failure() + shell: bash + run: | + set -euo pipefail + mkdir -p test-results/dynamo-tests-logs + if [ -d /tmp/dynamo_tests ]; then + cp -r /tmp/dynamo_tests/. test-results/dynamo-tests-logs/ + fi + + - name: Upload per-test log files + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + if: failure() + with: + name: dynamo-tests-logs-${{ inputs.test_suite_name }}-${{ inputs.test_type }}-${{ inputs.platform_arch }}-${{ github.run_id }}-${{ job.check_run_id }} + path: test-results/dynamo-tests-logs/ + retention-days: 7 + if-no-files-found: ignore + # Allure results and coverage data are uploaded per stage, under the same # artifact names as before the step collapse, because an external CI-metrics # pipeline consumes them by name. `always()` because the step above exits diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 552e7bd495ab..ea6c2beedc4d 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -234,7 +234,6 @@ areas: - components/src/dynamo/common/utils/ - components/src/dynamo/experimental/ - examples/backends/sample/ - - examples/backends/tritonserver/ - examples/common/ - examples/custom_backend/ - lib/backend-common/ @@ -459,6 +458,7 @@ areas: github_team: '@ai-dynamo/dynamo-rl-codeowners' path_globs: - components/src/dynamo/common/rl/ + - examples/rl/ - lib/rl/ - tests/rl/ - label: performance @@ -1063,7 +1063,7 @@ shared: owners: - router - kv-memory -- glob: lib/llm/src/preprocessor/lightseek_mm.rs +- glob: lib/llm/src/preprocessor/mm_routing/ owners: - frontend - multimodal @@ -1275,10 +1275,6 @@ shared: owners: - ops - backend-vllm -- glob: examples/backends/tritonserver/Dockerfile - owners: - - runtime - - ops - glob: examples/diffusers/Dockerfile owners: - diffusion diff --git a/.github/filters.yaml b/.github/filters.yaml index 15864f15c6d0..1287320c1f20 100644 --- a/.github/filters.yaml +++ b/.github/filters.yaml @@ -352,6 +352,7 @@ sidecar: # acceptable CI fan-out is agreed. Keep the remaining scope narrow for now. - 'lib/sidecar/**' - '.github/workflows/shared-build-sidecar.yml' + - 'tests/serve/test_sidecar.py' - 'Cargo.toml' - 'Cargo.lock' # Compliance policy and baseline changes can fail the sidecar image build. diff --git a/.github/scripts/dep_upgrade/baseline_stem.py b/.github/scripts/dep_upgrade/baseline_stem.py new file mode 100755 index 000000000000..fe2d47f39e02 --- /dev/null +++ b/.github/scripts/dep_upgrade/baseline_stem.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Print the baseline_sbom stem a capture just recorded for an image:tag. + +Reads the stem back out of the manifest rather than scraping capture logs. +container/context.yaml holds a single stem per framework and the licenses stage +appends -${TARGETARCH}, so per-arch stems that disagree are unrepresentable +downstream; this fails instead of shipping a half-right baseline. + +Stdout: the bare stem, e.g. release@2366e4b4. +""" +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +DEFAULT_MANIFEST = Path("container/compliance/base_sboms/manifest.json") + + +def stem_for(manifest: dict, image: str, tag: str, arches: list[str]) -> str: + """Return the shared baseline_sbom stem for image:tag across `arches`. + + Raises SystemExit if any arch is missing, matches more than one row, has an + unexpected filename, or if the arches disagree on the stem. + """ + if not arches: + raise SystemExit("no architectures given; --arches must name at least one") + entries = manifest.get("entries", []) + stems: dict[str, str] = {} + for arch in arches: + rows = [ + e + for e in entries + if e.get("from_image") == image + and e.get("from_tag") == tag + and e.get("platform") == f"linux/{arch}" + ] + if len(rows) != 1: + raise SystemExit( + f"expected exactly 1 {arch} row for {image}:{tag}, found {len(rows)}" + ) + suffix = f"-{arch}.cdx.json" + name = rows[0].get("baseline_sbom", "") + if not name.endswith(suffix): + raise SystemExit(f"{arch} baseline_sbom {name!r} does not end in {suffix}") + stems[arch] = name[: -len(suffix)] + if len(set(stems.values())) != 1: + raise SystemExit(f"per-arch baseline stems disagree: {stems}") + return next(iter(stems.values())) + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--image", required=True, help="from_image, e.g. nvcr.io/nvidia/x") + p.add_argument("--tag", required=True, help="from_tag, e.g. 1.3.0rc25") + p.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST) + p.add_argument( + "--arches", + default="amd64,arm64", + help="comma-separated arches that must agree (default: %(default)s)", + ) + args = p.parse_args() + + manifest = json.loads(args.manifest.read_text(encoding="utf-8")) + arches = [a.strip() for a in args.arches.split(",") if a.strip()] + print(stem_for(manifest, args.image, args.tag, arches)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/scripts/dep_upgrade/bump_dependency.py b/.github/scripts/dep_upgrade/bump_dependency.py index b9d7693185d8..8efa7a11a8d8 100755 --- a/.github/scripts/dep_upgrade/bump_dependency.py +++ b/.github/scripts/dep_upgrade/bump_dependency.py @@ -27,7 +27,7 @@ # Match runtime_image_tag inside the trtllm: block (first sub-block, # e.g. cuda13.1) without colliding with vllm/sglang earlier in the file. re.compile( - r"(?m)(^trtllm:\s*?\n(?:[ \t]+[^\n]*\n)*?[ \t]+runtime_image_tag:\s+)\S+", + r"(?m)(^trtllm:\s*?\n(?:[ \t][^\n]*\n)*?[ \t]+runtime_image_tag:\s+)\S+", ), r"\g<1>{ver}", ), @@ -51,6 +51,41 @@ "trtllm": TRTLLM_TARGETS, } +# Sibling of FRAMEWORK_TARGETS, kept separate because it is driven by a value +# only the compliance capture can produce (the new baseline's digest stem), not +# by the release version. Same block anchoring as runtime_image_tag above, so it +# cannot reach the vllm/sglang baseline_sbom lines earlier in the file. +FRAMEWORK_BASELINES: dict[str, tuple[str, Pattern[str], str]] = { + "trtllm": ( + "container/context.yaml", + re.compile( + r"(?m)(^trtllm:\s*?\n(?:[ \t][^\n]*\n)*?[ \t]+baseline_sbom:\s+)\S+", + ), + r"\g<1>{stem}", + ), +} + + +def set_baseline_stem(framework: str, stem: str, repo_root: Path) -> bool: + """Point the framework's baseline_sbom at a freshly captured stem. + + The runtime image's licenses stage subtracts this baseline from NOTICES, so + leaving it behind a runtime_image_tag bump makes every moved package version + read as net-new and fails the license gate. Returns True if the file changed. + """ + rel, pat, tmpl = FRAMEWORK_BASELINES[framework] + path = repo_root / rel + text = path.read_text() + new_text, n = pat.subn(tmpl.replace("{stem}", stem), text) + if n != 1: + raise SystemExit( + f"{rel}: matched {n} {framework} baseline_sbom pins; expected 1" + ) + if new_text == text: + return False + path.write_text(new_text) + return True + def apply(framework: str, version: str, repo_root: Path) -> int: """Apply all bumps for the framework. Returns count of files written.""" @@ -81,10 +116,24 @@ def main() -> int: required=True, help="version without 'v' prefix, e.g. 1.3.0rc12", ) + p.add_argument( + "--baseline-stem", + help=( + "Compliance baseline stem from capture_baseline_sbom.py, e.g. " + "release@2366e4b4. When given, also repoints the framework's " + "baseline_sbom in container/context.yaml. Omit to bump versions only." + ), + ) p.add_argument("--repo-root", type=Path, default=Path(".")) args = p.parse_args() n = apply(args.framework, args.version, args.repo_root) print(f"changed {n} files") + if args.baseline_stem: + changed = set_baseline_stem(args.framework, args.baseline_stem, args.repo_root) + print( + f"baseline_sbom -> {args.baseline_stem}" + f"{'' if changed else ' (already set)'}" + ) return 0 diff --git a/.github/scripts/dep_upgrade/check_bump_dependency.py b/.github/scripts/dep_upgrade/check_bump_dependency.py index 920cea352ece..f6aa5556664b 100755 --- a/.github/scripts/dep_upgrade/check_bump_dependency.py +++ b/.github/scripts/dep_upgrade/check_bump_dependency.py @@ -6,7 +6,12 @@ version and asserts it rewrites exactly MAIN_TOT.trtllm in docs/fern/components/releases.data.ts โ€” one line, inside the MAIN_TOT block, no RELEASES pin touched, and the result still parses through -gen_llms_tables.py. Run from anywhere; exit 0 == safe to bump. +gen_llms_tables.py. + +Then runs set_baseline_stem with a synthetic stem and asserts it rewrites +exactly the baseline_sbom inside the trtllm: block of container/context.yaml, +leaving runtime_image_tag and the vllm/sglang blocks alone. Run from anywhere; +exit 0 == safe to bump. """ from __future__ import annotations @@ -14,6 +19,7 @@ import shutil import sys import tempfile +import traceback from pathlib import Path REPO = Path(__file__).resolve().parents[3] @@ -25,6 +31,8 @@ SEED_FILES = ["container/context.yaml", "pyproject.toml", DATA_REL] SYNTH = "9.9.9rc99" # synthetic upstream TRT-LLM version, absent from real data +SYNTH_STEM = "release@0bad0bad" # synthetic baseline stem, absent from real data +CTX_REL = "container/context.yaml" def load(name: str, path: Path): @@ -35,10 +43,21 @@ def load(name: str, path: Path): def fail(msg: str) -> None: + """Report a failed assertion and exit non-zero (exit 0 == safe to bump).""" print(f"FAIL: {msg}") sys.exit(1) +def fail_exc(what: str, exc: BaseException) -> None: + """Same, for an unexpected crash: keep the traceback for debugging. + + The traceback is printed rather than propagated so the script keeps its + single contract -- exit 0 means safe to bump, anything else means stop. + """ + traceback.print_exc() + fail(f"{what} raised {type(exc).__name__}: {exc}") + + def main() -> int: bump = load("bump_dependency", BUMP) gen = load("gen_llms_tables", GEN) @@ -62,7 +81,7 @@ def main() -> int: except SystemExit as exc: fail(f"bump apply() raised SystemExit: {exc}") except Exception as exc: # noqa: BLE001 - any crash is a red result - fail(f"bump apply() raised {type(exc).__name__}: {exc}") + fail_exc("bump apply()", exc) after = data_path.read_text() if after == before: @@ -116,9 +135,69 @@ def main() -> int: if not rel_pins: fail("no RELEASES trtllm pins parsed; cross-check is meaningless") + # --- baseline stem rewrite ------------------------------------------ + # Runs against the already-bumped tree, matching the workflow's order. + ctx_path = root / CTX_REL + ctx_before = ctx_path.read_text() + + try: + if not bump.set_baseline_stem("trtllm", SYNTH_STEM, root): + fail("set_baseline_stem reported no change") + except SystemExit as exc: + fail(f"set_baseline_stem raised SystemExit: {exc}") + except Exception as exc: # noqa: BLE001 - any crash is a red result + fail_exc("set_baseline_stem", exc) + + ctx_after = ctx_path.read_text() + cb, ca = ctx_before.splitlines(), ctx_after.splitlines() + if len(cb) != len(ca): + fail("context.yaml line count changed; expected an in-place stem edit") + ctx_changed = [i for i, (x, y) in enumerate(zip(cb, ca)) if x != y] + if len(ctx_changed) != 1: + fail( + f"expected exactly 1 changed context.yaml line, " + f"got {len(ctx_changed)}: {ctx_changed}" + ) + stem_idx = ctx_changed[0] + if ca[stem_idx].strip() != f"baseline_sbom: {SYNTH_STEM}": + fail(f"changed line is not the baseline_sbom pin: {ca[stem_idx]!r}") + + # It must sit inside the trtllm: block, not vllm/sglang earlier in the file. + fw_start = next( + (i for i, ln in enumerate(ca) if ln.rstrip() == "trtllm:"), None + ) + if fw_start is None: + fail("no top-level 'trtllm:' block found in context.yaml") + fw_end = next( + ( + i + for i in range(fw_start + 1, len(ca)) + if ca[i][:1].strip() and not ca[i].startswith("#") + ), + len(ca), + ) + if not (fw_start < stem_idx < fw_end): + fail( + f"baseline_sbom line {stem_idx} is outside the trtllm: block " + f"({fw_start}..{fw_end})" + ) + + # The version bump and the stem must not clobber each other. + if f"runtime_image_tag: {SYNTH}" not in ctx_after: + fail("runtime_image_tag lost its bumped value after the stem rewrite") + if ctx_after.count(SYNTH_STEM) != 1: + fail( + f"{SYNTH_STEM} appears {ctx_after.count(SYNTH_STEM)} times; expected 1" + ) + + # Idempotent: re-applying the same stem is a no-op, not a second edit. + if bump.set_baseline_stem("trtllm", SYNTH_STEM, root): + fail("set_baseline_stem is not idempotent") + print("PASS: bump rewrote only MAIN_TOT.trtllm; RELEASES pins intact; file parses") print(f" MAIN_TOT.trtllm -> {SYNTH}") print(f" {len(rel_pins)} RELEASES trtllm pins left untouched") + print(f"PASS: baseline_sbom -> {SYNTH_STEM} (trtllm block only, idempotent)") return 0 diff --git a/.github/scripts/dep_upgrade/detect_latest_versions.py b/.github/scripts/dep_upgrade/detect_latest_versions.py index 2be0d669eb8c..41b5372ff367 100755 --- a/.github/scripts/dep_upgrade/detect_latest_versions.py +++ b/.github/scripts/dep_upgrade/detect_latest_versions.py @@ -3,7 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 """Detect latest upstream version for a Dynamo framework dependency. -Stdout: one-line JSON: {"framework", "current", "latest", "upgrade_needed"}. +Stdout: one-line JSON: +{"framework", "current", "latest", "upgrade_needed", "runtime_image"}. +runtime_image is the registry repo the tag belongs to, so callers can build the +full ref (:) without re-parsing context.yaml. Exit 0 on success regardless of upgrade_needed; non-zero only on real errors. """ from __future__ import annotations @@ -24,7 +27,11 @@ # Anchored inside the trtllm: block so we don't pick up # vllm/sglang runtime_image_tag lines earlier in the file. "current_regex": re.compile( - r"(?m)^trtllm:\s*?\n(?:[ \t]+[^\n]*\n)*?[ \t]+runtime_image_tag:\s*(\S+)\s*$", + r"(?m)^trtllm:\s*?\n(?:[ \t][^\n]*\n)*?[ \t]+runtime_image_tag:\s*(\S+)\s*$", + ), + # Same anchoring for the sibling runtime_image the tag hangs off. + "runtime_image_regex": re.compile( + r"(?m)^trtllm:\s*?\n(?:[ \t][^\n]*\n)*?[ \t]+runtime_image:\s*(\S+)\s*$", ), }, } @@ -75,15 +82,34 @@ def latest_release(repo: str, include_prereleases: bool) -> str: return max(candidates, key=parse_version) -def current_pin(framework: str, repo_root: Path) -> str: +def _context_field(framework: str, key: str, repo_root: Path) -> str: + """Return the single capture of FRAMEWORK_SOURCES[framework][key] in context.yaml. + + Raises SystemExit if the pattern does not match, so a drifted pin format + fails loudly rather than silently reporting a stale value. + """ src = FRAMEWORK_SOURCES[framework] text = (repo_root / "container" / "context.yaml").read_text() - m = src["current_regex"].search(text) # type: ignore[union-attr] + m = src[key].search(text) # type: ignore[union-attr] if not m: - raise SystemExit(f"no current pin found for {framework}") + raise SystemExit(f"no {key} match for {framework} in container/context.yaml") return m.group(1) +def current_pin(framework: str, repo_root: Path) -> str: + """Return the framework's currently pinned version, e.g. 1.3.0rc25.""" + return _context_field(framework, "current_regex", repo_root) + + +def runtime_image(framework: str, repo_root: Path) -> str: + """Return the registry repo the pin belongs to, without a tag. + + e.g. nvcr.io/nvidia/tensorrt-llm/release. Callers join this with the + detected version to form the ref to capture a baseline against. + """ + return _context_field(framework, "runtime_image_regex", repo_root) + + def main() -> int: p = argparse.ArgumentParser(description=__doc__) p.add_argument("--framework", required=True, choices=sorted(FRAMEWORK_SOURCES)) @@ -105,6 +131,7 @@ def main() -> int: "current": cur, "latest": latest, "upgrade_needed": parse_version(latest) > parse_version(cur), + "runtime_image": runtime_image(args.framework, args.repo_root), } ) ) diff --git a/.github/workflows/auto-dep-upgrade-trigger.yml b/.github/workflows/auto-dep-upgrade-trigger.yml index 4c5ead538553..13101beeb513 100644 --- a/.github/workflows/auto-dep-upgrade-trigger.yml +++ b/.github/workflows/auto-dep-upgrade-trigger.yml @@ -1,5 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Bumping runtime_image_tag moves the runtime image out from under the +# baseline_sbom that the licenses stage subtracts from NOTICES. Subtraction +# keys on (name, version), so every package whose version moved in the new tag +# stops subtracting, reads as net-new, resolves to UNKNOWN, and fails the +# license gate. The baseline is therefore re-captured here, in the same commit +# as the version bump, and a capture that cannot complete pushes nothing. name: Auto Dependency Upgrade - Trigger @@ -13,26 +20,60 @@ on: type: choice options: [trtllm] default: trtllm + dry_run: + description: 'Rehearse the baseline capture only; push nothing' + type: boolean + default: false permissions: contents: write actions: write +# A capture holds a scarce runner for tens of minutes, so never run two at once +# for the same ref. Queue rather than cancel: a half-finished capture wastes the +# scan it already paid for, and the branch idempotency check makes the second +# run cheap anyway. +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + jobs: - upgrade: - name: ${{ matrix.framework }} + # The matrix is single-entry, which is what makes the job outputs below + # well-defined: a matrix job's outputs are last-writer-wins across instances. + # Adding a second framework means extracting these three jobs into a + # workflow_call file and running the matrix over that instead. + detect: + name: ${{ matrix.framework }} / detect runs-on: ubuntu-latest + permissions: + contents: read strategy: fail-fast: false matrix: framework: [trtllm] + outputs: + framework: ${{ matrix.framework }} + current: ${{ steps.detect.outputs.current }} + latest: ${{ steps.detect.outputs.latest }} + runtime_image: ${{ steps.detect.outputs.runtime_image }} + branch: ${{ steps.branch.outputs.name }} + proceed: ${{ steps.branch.outputs.proceed }} steps: - - name: Checkout main + # No explicit ref: checkout defaults to the ref that triggered the run. + # A scheduled run is always the default branch, so cron behaviour is + # unchanged, while a workflow_dispatch runs the pipeline code from the ref + # you dispatched -- which is what makes a change to it testable pre-merge. + # + # Because that ref is caller-chosen, this job must hold no write-capable + # credential: it runs repository Python, so persisting the bot PAT here + # would let anyone able to dispatch escalate to whatever that PAT can + # reach. Nothing here needs it -- the release lookup uses GITHUB_TOKEN and + # ls-remote only reads. + - name: Checkout uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - ref: main fetch-depth: 0 - token: ${{ secrets.OPS_BOT_PAT }} + persist-credentials: false - name: Set up Python uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 @@ -51,55 +92,197 @@ jobs: echo "current=$(echo "$OUT" | jq -r .current)" echo "latest=$(echo "$OUT" | jq -r .latest)" echo "upgrade_needed=$(echo "$OUT" | jq -r .upgrade_needed)" + echo "runtime_image=$(echo "$OUT" | jq -r .runtime_image)" } >> "$GITHUB_OUTPUT" - - name: Already up to date - if: steps.detect.outputs.upgrade_needed != 'true' - run: echo "No upgrade needed (current=${{ steps.detect.outputs.current }}, latest=${{ steps.detect.outputs.latest }})." - - name: Compute branch name + idempotency check - if: steps.detect.outputs.upgrade_needed == 'true' id: branch + env: + UPGRADE_NEEDED: ${{ steps.detect.outputs.upgrade_needed }} run: | BRANCH="deps/upgrade-${{ matrix.framework }}-v${{ steps.detect.outputs.latest }}" echo "name=$BRANCH" >> "$GITHUB_OUTPUT" - if git ls-remote --heads origin "$BRANCH" | grep -q "$BRANCH"; then + if [ "$UPGRADE_NEEDED" != "true" ]; then + echo "No upgrade needed (current=${{ steps.detect.outputs.current }}, latest=${{ steps.detect.outputs.latest }})." + echo "proceed=false" >> "$GITHUB_OUTPUT" + elif git ls-remote --heads origin "$BRANCH" | grep -q "$BRANCH"; then echo "Branch $BRANCH already exists on origin; skipping." - echo "exists=true" >> "$GITHUB_OUTPUT" + echo "proceed=false" >> "$GITHUB_OUTPUT" else - echo "exists=false" >> "$GITHUB_OUTPUT" + echo "proceed=true" >> "$GITHUB_OUTPUT" + fi + + # syft catalogs the full multi-GB runtime image. Runner history is recorded in + # shared-compliance-audit.yml: small-v2 and prod-default-v2 run out of disk and + # prod-builder-v3 OOMs, so this pool -- which runs these images for tests -- is + # the one with the headroom. Both architectures are captured from here: syft + # selects the platform via SYFT_PLATFORM and scans the registry, so the arm64 + # pass needs neither an arm64 runner nor emulation. + capture: + name: ${{ needs.detect.outputs.framework }} / baseline capture + needs: detect + if: needs.detect.outputs.proceed == 'true' + # Fastchecker rather than prod-tester-amd-v2: that pool forbids containerless + # jobs and its admission rejects a push on a personal branch, so the pipeline + # could never be exercised before merge. Fastchecker runs the rust-tests / + # rust-clippy jobs containerless and admits this event, which keeps a change + # to this workflow testable. + runs-on: + group: Fastchecker + timeout-minutes: 90 + permissions: + contents: read + outputs: + stem: ${{ steps.capture.outputs.stem }} + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Install tooling + run: | + set -euo pipefail + docker buildx version + # container/compliance/overrides.py imports yaml. + python3 -c 'import yaml' 2>/dev/null \ + || python3 -m pip install --quiet --user --break-system-packages pyyaml \ + || python3 -m pip install --quiet --user pyyaml + mkdir -p "${RUNNER_TEMP}/bin" + curl -sSfL https://raw.githubusercontent.com/anchore/syft/main/install.sh \ + | sh -s -- -b "${RUNNER_TEMP}/bin" + echo "${RUNNER_TEMP}/bin" >> "$GITHUB_PATH" + + - name: Capture baseline SBOMs + id: capture + env: + REF: ${{ needs.detect.outputs.runtime_image }}:${{ needs.detect.outputs.latest }} + DRY_RUN: ${{ inputs.dry_run }} + run: | + set -euo pipefail + EXTRA=() + if [ "${DRY_RUN:-false}" = "true" ]; then EXTRA=(--dry-run); fi + # Sequential, not parallel: both passes read-modify-write the single + # manifest.json, so concurrent runs would drop one arch's entry. + for PLATFORM in linux/amd64 linux/arm64; do + echo "::group::capture ${PLATFORM}" + python3 container/compliance/base_sboms/capture_baseline_sbom.py \ + --from "$REF" --baseline "$REF" \ + --platform "$PLATFORM" --prune-superseded -v "${EXTRA[@]}" + echo "::endgroup::" + done + if [ "${DRY_RUN:-false}" = "true" ]; then + echo "Dry run: manifest not written, no stem to publish." + exit 0 fi + STEM=$(python3 .github/scripts/dep_upgrade/baseline_stem.py \ + --image "${REF%:*}" --tag "${REF##*:}") + echo "Captured baseline stem: $STEM" + echo "stem=$STEM" >> "$GITHUB_OUTPUT" + + - name: Upload refreshed baseline corpus + if: inputs.dry_run != true + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + with: + name: baseline-refresh-${{ needs.detect.outputs.framework }}-${{ needs.detect.outputs.latest }} + path: | + container/compliance/base_sboms/manifest.json + container/compliance/base_sboms/*.cdx.json + if-no-files-found: error + retention-days: 7 + + - name: Notify on capture failure + # A failed capture pushes no branch and opens no PR, so without this the + # upgrade stalls silently -- a genuinely disallowed license in a new tag + # would go unnoticed until someone wonders why bumps stopped arriving. + if: failure() + uses: slackapi/slack-github-action@91efab103c0de0a537f72a35f6b8cda0ee76bf0a # v2.1.1 + with: + webhook: ${{ secrets.SLACK_NOTIFY_NIGHTLY_WEBHOOK_URL }} + webhook-type: incoming-webhook + payload: | + blocks: + - type: "section" + text: + type: mrkdwn + text: ":package: *Auto Dep Upgrade* :x: Baseline capture failed โ€” ${{ needs.detect.outputs.framework }} v${{ needs.detect.outputs.latest }}" + - type: "section" + text: + type: mrkdwn + text: "<${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|Workflow run> ยท image `${{ needs.detect.outputs.runtime_image }}:${{ needs.detect.outputs.latest }}`" + - type: "section" + text: + type: mrkdwn + text: "No branch was pushed. Common causes: the container tag is not published yet (the cron will retry), or the new image introduces a license the policy denies (needs a human)." + + commit: + name: ${{ needs.detect.outputs.framework }} / commit + needs: [detect, capture] + if: needs.detect.outputs.proceed == 'true' && inputs.dry_run != true + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + fetch-depth: 0 + token: ${{ secrets.OPS_BOT_PAT }} + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: '3.12' + + - name: Clear stale baseline SBOMs + # Artifacts carry additions, not deletions, so without this the SBOMs + # that --prune-superseded removed would survive into the commit. + run: rm -f container/compliance/base_sboms/*.cdx.json + + - name: Download refreshed baseline corpus + uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + with: + name: baseline-refresh-${{ needs.detect.outputs.framework }}-${{ needs.detect.outputs.latest }} + path: container/compliance/base_sboms - name: Apply bump - if: steps.detect.outputs.upgrade_needed == 'true' && steps.branch.outputs.exists == 'false' + env: + FW: ${{ needs.detect.outputs.framework }} + VER: ${{ needs.detect.outputs.latest }} + STEM: ${{ needs.capture.outputs.stem }} run: | - # Pre-flight: prove the bump tooling still edits exactly the - # MAIN_TOT pin in releases.data.ts before touching the real tree. + set -euo pipefail + test -n "$STEM" + # Pre-flight: prove the bump tooling still edits exactly the MAIN_TOT + # pin in releases.data.ts and the trtllm baseline_sbom in context.yaml + # before touching the real tree. python3 .github/scripts/dep_upgrade/check_bump_dependency.py python3 .github/scripts/dep_upgrade/bump_dependency.py \ - --framework ${{ matrix.framework }} \ - --version ${{ steps.detect.outputs.latest }} + --framework "$FW" --version "$VER" --baseline-stem "$STEM" git status --short - name: Commit + push upgrade branch - if: steps.detect.outputs.upgrade_needed == 'true' && steps.branch.outputs.exists == 'false' + env: + FW: ${{ needs.detect.outputs.framework }} + VER: ${{ needs.detect.outputs.latest }} + BRANCH: ${{ needs.detect.outputs.branch }} run: | + set -euo pipefail git config user.name "dynamo-ops" git config user.email "170655669+dynamo-ops@users.noreply.github.com" - git checkout -b "${{ steps.branch.outputs.name }}" + git checkout -b "$BRANCH" # releases.data.ts feeds generated llms-tables twins and JSON/Atom assets; # regenerate them in the same commit or check_reference.sh flags the branch stale. python3 docs/fern/scripts/gen_llms_tables.py git add container/context.yaml pyproject.toml docs/fern/components/releases.data.ts git add docs/fern/ - git commit -s -m "chore: bump ${{ matrix.framework }} to v${{ steps.detect.outputs.latest }}" - git push origin "${{ steps.branch.outputs.name }}" + # -A so the pruned SBOMs are staged as deletions, not left behind. + git add -A container/compliance/base_sboms/ + git commit -s -m "chore: bump $FW to v$VER" + git push origin "$BRANCH" - name: Dispatch post-merge-ci on upgrade branch - if: steps.detect.outputs.upgrade_needed == 'true' && steps.branch.outputs.exists == 'false' env: GH_TOKEN: ${{ secrets.OPS_BOT_PAT }} + FW: ${{ needs.detect.outputs.framework }} + BRANCH: ${{ needs.detect.outputs.branch }} run: | - gh workflow run post-merge-ci.yml \ - --ref "${{ steps.branch.outputs.name }}" \ - -f framework=${{ matrix.framework }} + gh workflow run post-merge-ci.yml --ref "$BRANCH" -f framework="$FW" diff --git a/.github/workflows/docs-link-check.yml b/.github/workflows/docs-link-check.yml index a998ef723b17..93233292cfe1 100644 --- a/.github/workflows/docs-link-check.yml +++ b/.github/workflows/docs-link-check.yml @@ -39,6 +39,7 @@ jobs: --exclude-path ".*ATTRIBUTIONS.*" --exclude-path "./lib/llm/tests/data/.*" --exclude-path "./container/compliance/spdx_licenses/.*" + --exclude-path "(^|/)lychee/" --accept "200..=299, 403, 429" --exclude-all-private --exclude 0.0.0.0 @@ -107,71 +108,6 @@ jobs: with: python-version: '3.11' - - name: Check generated release docs determinism - run: | - set -euo pipefail - - python3 docs/fern/scripts/gen_llms_tables.py --check - - tmp_dir=$(mktemp -d) - trap 'rm -rf "$tmp_dir"' EXIT - mkdir -p "$tmp_dir/first" "$tmp_dir/second" - cp -R docs "$tmp_dir/first/docs" - cp -R docs "$tmp_dir/second/docs" - - outputs=( - docs/fern/pages/reference/general/compatibility.mdx - docs/fern/pages/reference/general/release-artifacts.mdx - docs/fern/pages/reference/general/model-early-access-builds.mdx - docs/fern/pages/reference/general/releases-machine-readable.mdx - docs/fern/pages/reference/general/releases/release-history.mdx - docs/fern/assets/releases.json - docs/fern/assets/releases-atom.xml - ) - - # Force both runs to actually WRITE every output: drop the whole-file - # assets and CRLF-corrupt the spliced pages. A copied-through current - # tree would skip every write on the unchanged fast-path, leaving cmp - # comparing the originals with themselves and proving nothing about - # how the generator writes. - for tree in first second; do - rm -f "$tmp_dir/$tree/docs/fern/assets/releases.json" \ - "$tmp_dir/$tree/docs/fern/assets/releases-atom.xml" - for output in "${outputs[@]}"; do - case "$output" in - *.mdx) sed -i 's/$/\r/' "$tmp_dir/$tree/$output" ;; - esac - done - done - - TZ=UTC LC_ALL=C PYTHONHASHSEED=1 SOURCE_DATE_EPOCH=0 \ - python3 "$tmp_dir/first/docs/fern/scripts/gen_llms_tables.py" - TZ=Pacific/Honolulu LC_ALL=C.UTF-8 PYTHONHASHSEED=2 SOURCE_DATE_EPOCH=4102444800 \ - python3 "$tmp_dir/second/docs/fern/scripts/gen_llms_tables.py" - - failed=0 - for output in "${outputs[@]}"; do - if ! cmp -s "$tmp_dir/first/$output" "$tmp_dir/second/$output"; then - echo "::error file=$output::Generator output is not byte-identical across runs" - diff -u "$tmp_dir/first/$output" "$tmp_dir/second/$output" || true - failed=1 - fi - # Committed-bytes comparison only where the output is committed - # (the two assets stop being tracked once publish-time generation - # lands; temp-vs-temp and the CR scan still cover them). - if [ -f "$output" ] && ! cmp -s "$tmp_dir/first/$output" "$output"; then - echo "::error file=$output::Regenerated output does not match the committed bytes" - failed=1 - fi - if LC_ALL=C grep -q $'\r' "$tmp_dir/first/$output"; then - echo "::error file=$output::Generator output contains CR bytes; outputs must be LF-only" - failed=1 - fi - done - if (( failed )); then - exit 1 - fi - - name: Install Python dependencies run: | python -m pip install --upgrade pip diff --git a/.github/workflows/fern-docs.yml b/.github/workflows/fern-docs.yml index b71494499c37..592020588e56 100644 --- a/.github/workflows/fern-docs.yml +++ b/.github/workflows/fern-docs.yml @@ -137,18 +137,58 @@ jobs: python-version: '3.13' - name: Generate API references - # The Python/Rust reference pages are publish-time artifacts: - # gitignored in the source tree and generated here against the exact - # ref being published, so source PRs never carry (or review) the - # regenerated MDX. Generation must precede the pages-dev sync below, - # which rsyncs whatever sits under docs/fern/pages/. The Kubernetes - # reference regenerates from a docs-owned source file and stays - # committed, so it keeps the freshness check. + id: api_refs + # Python/Rust references are publish-time artifacts: generated here + # against the ref being published, never committed. Generation must + # precede the pages-dev sync below, which rsyncs docs/fern/pages/. + # + # The Kubernetes reference stays committed, but this path REGENERATES + # it rather than gating on it. As `--check` the step failed, which + # skipped the sync and publish steps entirely, so a main left stale by + # merge skew silently froze the live site. Drift is still fatal, just + # after the site is current: the probe records it and the gate at the + # end of this job fails the run. run: | + set -euo pipefail python3 -m pip install 'griffe==2.1.0' python3 source-checkout/docs/fern/scripts/gen_python_api.py python3 source-checkout/docs/fern/scripts/gen_rust_api.py - python3 source-checkout/docs/fern/scripts/gen_kubernetes_api.py --check + if python3 source-checkout/docs/fern/scripts/gen_kubernetes_api.py --check >/dev/null 2>&1; then + echo "kubernetes_drift=false" >> "$GITHUB_OUTPUT" + else + echo "kubernetes_drift=true" >> "$GITHUB_OUTPUT" + fi + k8s_page=source-checkout/docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx + committed_bytes=$(wc -c < "$k8s_page") + python3 source-checkout/docs/fern/scripts/gen_kubernetes_api.py + regenerated_bytes=$(wc -c < "$k8s_page") + # Publishing through drift is the point of this step, but the + # regenerated page is what the sync below rsyncs, so a generator that + # exits 0 while emitting a gutted reference would reach the live site + # and only be reported afterwards. Ordinary skew moves this page by a + # few hundred bytes; losing a fifth of it means generation itself is + # what is wrong, so stop before anything is synced or published. + if [ "$regenerated_bytes" -lt "$(( committed_bytes * 4 / 5 ))" ]; then + echo "::error file=docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx::Regenerated Kubernetes API reference is $regenerated_bytes bytes against $committed_bytes committed, a loss of more than a fifth. Refusing to sync or publish. Check gen_kubernetes_api.py and its source, docs/fern/pages/reference/kubernetes-api/api-reference-k8s.md." + exit 1 + fi + + - name: Generate install-selector nightly data + # Writes components/nightly-selector-data.generated.ts, which the + # component sync below carries to docs-website. Needs full history and + # tags: each NGC nightly tag names the commit that build came from, and + # the generator reads container/context.yaml and pyproject.toml there. + run: | + git -C source-checkout fetch --unshallow --tags 2>/dev/null \ + || git -C source-checkout fetch --tags + # Fail rather than silently emit a collapsed nightly history from a + # still-shallow clone (e.g. a transient --unshallow failure). + if [ "$(git -C source-checkout rev-parse --is-shallow-repository)" != "false" ]; then + echo "::error::source-checkout is still shallow; nightly history would be incomplete" + exit 1 + fi + python3 -m pip install --quiet pyyaml + python3 source-checkout/docs/fern/scripts/gen_nightly_selector.py - name: Checkout docs-website branch uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -664,6 +704,25 @@ jobs: if: steps.ctx.outputs.is_main == 'true' run: python3 source-checkout/docs/fern/scripts/check_published_styles.py + - name: Fail if the committed Kubernetes API reference was stale + # Deliberately last, and deliberately not gated on is_main: the preview + # path should report drift too. Whether the site actually published is + # a separate question - Publish Docs needs both is_main and has_changes + # - so the message reports which of the two happened rather than + # asserting a publish. `!cancelled()` keeps the drift annotation on runs + # where an earlier step already failed; a bare `if:` carries an implicit + # success() and would drop the report on exactly those runs. + if: ${{ !cancelled() && steps.api_refs.outputs.kubernetes_drift == 'true' }} + run: | + set -euo pipefail + if [ "${{ steps.ctx.outputs.is_main }}" = "true" ] && [ "${{ steps.changes.outputs.has_changes }}" = "true" ]; then + outcome="The site published the regenerated output, so docs are current, but the committed page differs from the reviewed bytes." + else + outcome="Nothing was published on this run, so the live site is unchanged." + fi + echo "::error file=docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx::Committed Kubernetes API reference does not match its source. $outcome Fix with: python3 docs/fern/scripts/gen_kubernetes_api.py" + exit 1 + ############################################################################# # VERSION RELEASE - Run on new version tags (vX.Y.Z) ############################################################################# @@ -844,7 +903,16 @@ jobs: echo "::error::$SHARED not found โ€” the shared-reference exclusion would be a no-op and $TAG would freeze its own reference copy. If reference/general/ was renamed, update the path key in this workflow and in docs/fern/scripts/simulate_docs_website.sh." exit 1 fi - rm -rf "$SHARED" + # A page removed or moved on main after this tag was cut has no shared + # copy in pages-dev, so the tag keeps its own frozen copy of it. + while IFS= read -r f; do + rel="${f#docs-checkout/fern/pages-$TAG/}" + if [ -e "docs-checkout/fern/pages-dev/$rel" ]; then + rm -f "$f" + else + echo "::warning::$rel has no shared copy in pages-dev; $TAG keeps its own." + fi + done < <(find "$SHARED" -type f) find "docs-checkout/fern/pages-$TAG/reference" -type d -empty -delete 2>/dev/null || true - name: Snapshot translation mirrors @@ -975,6 +1043,14 @@ jobs: # which invalidates the whole navigation against Fern's schema. perl -pi -e "s|path: \.\./pages-$TAG/reference/general/|path: ../pages-dev/reference/general/|g" "$VERSION_FILE" + # Pages the snapshot kept (no shared copy) stay on the frozen snapshot. + if [ -d "docs-checkout/fern/pages-$TAG/reference/general" ]; then + while IFS= read -r f; do + rel="${f#docs-checkout/fern/pages-$TAG/}" + perl -pi -e "s|path: \.\./pages-dev/\Q$rel\E|path: ../pages-$TAG/$rel|g" "$VERSION_FILE" + done < <(find "docs-checkout/fern/pages-$TAG/reference/general" -type f) + fi + if ! grep -q "path: \.\./pages-dev/reference/general/" "$VERSION_FILE"; then echo "::error::$VERSION_FILE has no shared reference/general/ paths after the revert; $TAG would freeze its own reference copy." exit 1 diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index 5e14495deecf..fff53ec81f8a 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -641,11 +641,17 @@ jobs: # ============================================================================ # EFA DEPLOY TEST # ============================================================================ - # The only job here that runs against real EFA hardware. Everything else - # validates the -efa images statically; this one deploys a 1P1D vLLM stack on - # aws-dev-02 (p5 / H100, 32 EFA devices per node) and proves the prefill->decode - # KV transfer rode NIXL -> LIBFABRIC -> EFA instead of silently falling back to - # UCX. + # The only jobs here that run against real EFA hardware. Everything else + # validates the -efa images statically; these deploy a 1P1D stack on aws-dev-02 + # (p5 / H100, 32 EFA devices per node) -- vLLM first, then SGLang -- and prove + # the prefill->decode KV transfer rode NIXL -> LIBFABRIC -> EFA instead of + # silently falling back to UCX. + # + # The two framework tests share one vCluster and run one after the other rather + # than together, to keep peak GPU use on the shared CI pool at one 1P1D stack. + # SGLang is gated on the vLLM job having *finished*, not on it having passed + # (`!cancelled()` rather than `success()`), so a vLLM regression does not also + # delete the SGLang result. # # It cannot reuse shared-deploy-test.yml: that workflow hardcodes its runner, # timeout, ACR image reference and `-m framework_only` selector, and forwards no @@ -670,7 +676,7 @@ jobs: # nscleanup/ttl label for dynamo-deploy-test, so a leaked vCluster holds p5 # GPUs and EFA interfaces on a shared cluster with no automated reaper. efa-deploy-operator: - name: vLLM EFA Deploy Operator + name: EFA Deploy Operator needs: [compute-release-mode, operator-build, resolve-source-sha] if: ${{ needs.compute-release-mode.outputs.run_tests == 'true' }} runs-on: prod-deploy-tester-v1 @@ -754,10 +760,95 @@ jobs: image: ${{ vars.ECR_REGISTRY }}/${{ vars.ECR_REPOSITORY }}:${{ needs.resolve-source-sha.outputs.source_sha }}-${{ needs.vllm-efa-build.outputs.target_tag_plain }}-nightly platform_arch: amd64 + efa-deploy-test-sglang: + name: SGLang EFA Deploy Test + needs: [compute-release-mode, operator-build, resolve-source-sha, sglang-efa-build, efa-deploy-operator, efa-deploy-test-vllm] + # Runs after the vLLM test whatever that test concluded, so the two never hold + # GPUs on the shared CI pool at the same time, but a vLLM failure still leaves + # SGLang measured. Cancellation stops the lane; failure does not. + # + # `!cancelled()` replaces the default "every need succeeded" rule, so the two + # needs this job genuinely cannot run without -- the vCluster and the image -- + # have to be spelled out again. Only efa-deploy-test-vllm is allowed to fail. + if: ${{ !cancelled() && needs.compute-release-mode.outputs.run_tests == 'true' && needs.efa-deploy-operator.result == 'success' && needs.sglang-efa-build.result == 'success' }} + runs-on: prod-deploy-tester-v1 + # Same reasoning as the vLLM job. Bring-up is dominated by the cold image + # pull rather than by engine startup (measured on aws-dev-02: ~10 min to all + # pods ready cold, ~1.5 min once the image is on the nodes), which the test's + # own 900s readiness wait covers. + timeout-minutes: 30 + permissions: + contents: read + env: + KUBE_CONTEXT: nv-prd-dgxc.teleport.sh-dynamo-aws-dev-02 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + ref: ${{ needs.resolve-source-sha.outputs.source_sha }} + + - name: Set kubectl context to aws-dev-02 + uses: ./.github/actions/set-kube-context + + - name: Connect to vCluster + id: connect-vcluster + uses: ./.github/actions/connect-vcluster + with: + vcluster_name: ${{ needs.efa-deploy-operator.outputs.vcluster_name }} + vcluster_namespace: ${{ needs.efa-deploy-operator.outputs.namespace }} + + # The vLLM stack normally tears itself down in the test's own fixture, but + # that does not survive a job timeout or a lost runner, and + # dynamo-deploy-test's fallback cleanup cannot help: its GRAPH_NAME reads + # steps.deploy.outputs.graph_name, which the pytest step never sets, so it + # always reports "already cleaned up". Without this, a stranded vLLM stack + # would still hold GPUs while SGLang runs -- the one thing running the two + # lanes sequentially is meant to prevent. + - name: Delete any stranded vLLM EFA deployment + shell: bash + env: + KUBECONFIG_B64: ${{ steps.connect-vcluster.outputs.kubeconfig_base64 }} + PREFLIGHT_KUBECONFIG: ${{ github.workspace }}/.kubeconfig-preflight + run: | + set -euo pipefail + printf '%s' "${KUBECONFIG_B64}" | base64 -d > "${PREFLIGHT_KUBECONFIG}" + chmod 600 "${PREFLIGHT_KUBECONFIG}" + export KUBECONFIG="${PREFLIGHT_KUBECONFIG}" + kubectl delete dynamographdeployment vllm-disagg-efa \ + -n default --ignore-not-found --timeout=180s + # Deletion returns once the DGD object is gone; the pods it owned can + # outlive it, and those are what hold the GPUs and EFA devices. + SEL="nvidia.com/dynamo-graph-deployment-name=vllm-disagg-efa" + if ! kubectl wait --for=delete pod -l "${SEL}" -n default --timeout=180s; then + # wait also exits non-zero when the selector matched nothing, which + # is the normal path. Only pods that outlived the timeout matter, and + # they are worth a warning rather than a failure: SGLang may still + # schedule, and a hard failure here would lose that run entirely. + REMAINING=$(kubectl get pod -l "${SEL}" -n default -o name | wc -l | tr -d ' ') + [ "${REMAINING}" -eq 0 ] || echo "::warning::${REMAINING} stranded vllm-disagg-efa pod(s) still hold GPUs/EFA devices" + fi + rm -f "${PREFLIGHT_KUBECONFIG}" + + - name: Run EFA deploy test + uses: ./.github/actions/dynamo-deploy-test + with: + kubeconfig_base64: ${{ steps.connect-vcluster.outputs.kubeconfig_base64 }} + namespace: default + registry: ${{ secrets.AZURE_ACR_HOSTNAME }} + operator_tag: ${{ needs.operator-build.outputs.operator_default_tag }} + hf_token: ${{ secrets.HF_TOKEN }} + framework: sglang + test_name: sglang_efa + test_file: tests/deploy/test_deploy_efa_sglang.py + extra_pytest_args: -m framework_with_efa + # Straight from ECR, as for vLLM above. + image: ${{ vars.ECR_REGISTRY }}/${{ vars.ECR_REPOSITORY }}:${{ needs.resolve-source-sha.outputs.source_sha }}-${{ needs.sglang-efa-build.outputs.target_tag_plain }}-nightly + platform_arch: amd64 + efa-deploy-cleanup: - name: vLLM EFA Deploy Cleanup + name: EFA Deploy Cleanup if: always() - needs: [efa-deploy-operator, efa-deploy-test-vllm] + needs: [efa-deploy-operator, efa-deploy-test-vllm, efa-deploy-test-sglang] runs-on: prod-deploy-tester-v1 permissions: contents: read @@ -816,7 +907,13 @@ jobs: enable_coverage: true run_cpu_only_tests: true cpu_only_test_markers: vllm and gpu_0 - gpu_test_markers: vllm and gpu_1 + # not sidecar: sidecar-marked tests need the dynamo-vllm-sidecar binary + # installed (see pr.yaml's sidecar-vllm-test), which this plain + # vllm-runtime image doesn't have. + # TODO: this ad-hoc exclusion is duplicated across pr.yaml, post-merge-ci.yml, + # and nightly-ci.yml's gpu_test_markers; consolidate into shared-test.yml + # once there's a second marker that needs the same treatment. + gpu_test_markers: vllm and gpu_1 and not sidecar gpu_test_timeout_minutes: 240 # Profiled tests run in the parallel stage; unprofiled fall through to sequential. # 24 GiB admits all currently profiled vLLM tests (max is ~20.4 GiB) on a 48 GiB GPU. @@ -891,7 +988,7 @@ jobs: platform: '["amd64"]' enable_coverage: true run_sanity_check: false - gpu_test_markers: vllm and (gpu_1 or gpu_2) + gpu_test_markers: vllm and (gpu_1 or gpu_2) and not sidecar gpu_test_timeout_minutes: 240 # 80 GiB H100 admits more of the profiled pool per slot than the 24 GiB lanes. run_gpu_parallel_tests: true @@ -919,7 +1016,7 @@ jobs: run_cpu_only_tests: true cpu_only_test_markers: sglang and gpu_0 # Exclude h100 tests: they need >24 GiB and run on the sglang-h100-test lane below. - gpu_test_markers: sglang and gpu_1 and not h100 + gpu_test_markers: sglang and gpu_1 and not h100 and not sidecar gpu_test_timeout_minutes: 240 # Profiled tests run in the VRAM-aware GPU stage; unprofiled fall through to sequential. # Current single-GPU runners are 24 GiB, so this cap admits the profiled SGLang pool @@ -994,7 +1091,7 @@ jobs: platform: '["amd64"]' enable_coverage: true run_sanity_check: false - gpu_test_markers: sglang and (gpu_1 or gpu_2) + gpu_test_markers: sglang and (gpu_1 or gpu_2) and not sidecar gpu_test_timeout_minutes: 240 # 80 GiB H100 admits more of the profiled pool per slot than the 24 GiB lanes. run_gpu_parallel_tests: true @@ -1021,7 +1118,7 @@ jobs: enable_coverage: true run_cpu_only_tests: true cpu_only_test_markers: trtllm and gpu_0 - gpu_test_markers: trtllm and gpu_1 + gpu_test_markers: trtllm and gpu_1 and not sidecar gpu_test_timeout_minutes: 240 # Profiled tests run in the VRAM-aware GPU stage; unprofiled fall through to sequential. # Current single-GPU runners are 24 GiB, so this cap admits the profiled TRT-LLM pool @@ -1073,7 +1170,7 @@ jobs: platform: '["amd64"]' enable_coverage: true run_sanity_check: false - gpu_test_markers: trtllm and (gpu_1 or gpu_2) + gpu_test_markers: trtllm and (gpu_1 or gpu_2) and not sidecar gpu_test_timeout_minutes: 240 # 80 GiB H100 admits more of the profiled pool per slot than the 24 GiB lanes. run_gpu_parallel_tests: true diff --git a/.github/workflows/post-merge-ci.yml b/.github/workflows/post-merge-ci.yml index 18b7f01958bd..52e13f5fec30 100644 --- a/.github/workflows/post-merge-ci.yml +++ b/.github/workflows/post-merge-ci.yml @@ -215,6 +215,11 @@ jobs: extra_tags: | ${{ github.ref_name == 'main' && 'main-vllm-runtime' || '' }} ${{ github.ref_name == 'main' && format('main-vllm-runtime-{0}', github.sha) || '' }} + # Lets sidecar-vllm-test (pr.yaml) fall back to this floating tag via its + # source_ref input instead of rebuilding vllm-runtime on every + # lib/sidecar/**-only PR. + test_extra_tags: | + ${{ github.ref_name == 'main' && 'main-vllm-runtime-test' || '' }} archive_sources: false secrets: inherit @@ -280,6 +285,11 @@ jobs: extra_tags: | ${{ github.ref_name == 'main' && 'main-sglang-runtime' || '' }} ${{ github.ref_name == 'main' && format('main-sglang-runtime-{0}', github.sha) || '' }} + # Lets sidecar-sglang-test (pr.yaml) fall back to this floating tag via + # its source_ref input instead of rebuilding sglang-runtime on every + # lib/sidecar/**-only PR. + test_extra_tags: | + ${{ github.ref_name == 'main' && 'main-sglang-runtime-test' || '' }} archive_sources: false secrets: inherit @@ -345,6 +355,11 @@ jobs: extra_tags: | ${{ github.ref_name == 'main' && 'main-trtllm-runtime' || '' }} ${{ github.ref_name == 'main' && format('main-trtllm-runtime-{0}', github.sha) || '' }} + # Lets sidecar-trtllm-test (pr.yaml) fall back to this floating tag via + # its source_ref input instead of rebuilding trtllm-runtime on every + # lib/sidecar/**-only PR. + test_extra_tags: | + ${{ github.ref_name == 'main' && 'main-trtllm-runtime-test' || '' }} archive_sources: false secrets: inherit @@ -441,7 +456,7 @@ jobs: platform: 'linux/amd64,linux/arm64' builder_name: b-${{ github.run_id }}-${{ github.run_attempt }} inline_compliance: true - build_timeout_minutes: 45 + build_timeout_minutes: 90 # TEMPORARY CHANGE - TODO: INVESTIGATE AND SPEEDUP SLOW BUILDS # OSRB CSV diff baseline: post-merge context. main -> previous commit; # release/* -> prior release tag (see resolve_diff_base.py). diff_event_context: push @@ -471,7 +486,10 @@ jobs: platform: '["amd64", "arm64"]' # arm64 for CPU tests, single GPU tests are skipped run_cpu_only_tests: true cpu_only_test_markers: '(pre_merge or post_merge) and vllm and gpu_0' - gpu_test_markers: '(pre_merge or post_merge) and vllm and gpu_1' + # not sidecar: sidecar-marked tests need dynamo-vllm-sidecar layered onto + # the image (see pr.yaml's sidecar-vllm-test), which this plain + # vllm-runtime image doesn't have. + gpu_test_markers: '(pre_merge or post_merge) and vllm and gpu_1 and not sidecar' gpu_test_timeout_minutes: 120 # Profiled tests (those carrying @pytest.mark.profiled_vram_gib(N)) run concurrently in # the parallel stage; the rest fall through to the sequential stage. 24 GiB admits all @@ -549,7 +567,10 @@ jobs: platform: '["amd64", "arm64"]' # arm64 for CPU tests, single GPU tests are skipped run_cpu_only_tests: true cpu_only_test_markers: '(pre_merge or post_merge) and sglang and gpu_0' - gpu_test_markers: '(pre_merge or post_merge) and sglang and gpu_1' + # not sidecar: sidecar-marked tests need dynamo-sglang-sidecar layered + # onto the image (see pr.yaml's sidecar-sglang-test), which this plain + # sglang-runtime image doesn't have. + gpu_test_markers: '(pre_merge or post_merge) and sglang and gpu_1 and not sidecar' gpu_test_timeout_minutes: 120 # Profiled tests run in the VRAM-aware GPU stage; unprofiled fall through to sequential. # Current single-GPU runners are 24 GiB, so this cap admits the profiled SGLang pool @@ -591,7 +612,10 @@ jobs: platform: '["amd64", "arm64"]' # arm64 for CPU tests, single GPU tests are skipped run_cpu_only_tests: true cpu_only_test_markers: '(pre_merge or post_merge) and trtllm and gpu_0' - gpu_test_markers: '(pre_merge or post_merge) and trtllm and gpu_1' + # not sidecar: sidecar-marked tests need dynamo-trtllm-sidecar layered + # onto the image (see pr.yaml's sidecar-trtllm-test), which this plain + # trtllm-runtime image doesn't have. + gpu_test_markers: '(pre_merge or post_merge) and trtllm and gpu_1 and not sidecar' gpu_test_timeout_minutes: 120 # Profiled tests run in the VRAM-aware GPU stage; unprofiled fall through to sequential. # Current single-GPU runners are 24 GiB, so this cap admits the profiled TRT-LLM pool diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 09428070b404..61608ce05052 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -134,6 +134,9 @@ jobs: - snapshot-placeholder-trtllm - power-agent - sidecar-build + - sidecar-vllm-test + - sidecar-sglang-test + - sidecar-trtllm-test - vllm-build - vllm-local-dev-build - sglang-local-dev-build @@ -355,12 +358,99 @@ jobs: uses: $/.github/workflows/shared-build-sidecar.yml with: builder_name: ${{ needs.changed-files.outputs.builder_name }} + # Default (60) is tight for a cold cache: multi-arch cargo build of 3 + # sidecar crates + dynamo-llm/dynamo-runtime, plus the compliance and + # binary-extraction steps that follow it. Observed a legitimate cold + # build exceed 60m with room to spare in run 34547963529. + build_timeout_minutes: 90 artifact_retention_days: '7' diff_event_context: pr diff_current_branch: ${{ github.ref_name }} diff_base_branch: ${{ needs.changed-files.outputs.base_ref }} diff_merge_base_sha: ${{ needs.changed-files.outputs.merge_base_sha }} + # Reuse the backend test image and install the sidecar artifact before pytest. + + sidecar-vllm-test: + name: sidecar-runtime # This name overlaps with other sidecar jobs to group them in the UI + needs: [changed-files, vllm-build, sidecar-build] + # vllm-build must have succeeded or been intentionally skipped (a + # lib/sidecar-only PR) -- not failed or cancelled -- or the fallback below + # would silently test this PR's sidecar binary against main's vllm image + # and report green while this PR's own vllm-runtime build is red. + if: ${{ !cancelled() && needs.changed-files.outputs.sidecar == 'true' && needs.sidecar-build.result == 'success' && (needs.vllm-build.result == 'success' || needs.vllm-build.result == 'skipped') }} + permissions: + actions: read + contents: read + uses: $/.github/workflows/shared-test.yml + with: + test_suite_name: vllm + framework: vllm # matches the framework passed to vllm-build + test_type: Sidecar E2E Test + amd_runner: prod-tester-amd-gpu-v2 + target_tag_plain: ${{ needs.vllm-build.result == 'success' && needs.vllm-build.outputs.target_tag_plain || 'vllm-runtime' }} + source_ref: ${{ needs.vllm-build.result == 'success' && github.sha || 'main' }} + cuda_version: '["13.0"]' + platform: '["amd64"]' + run_sanity_check: false + gpu_test_markers: pre_merge and sidecar and vllm and gpu_1 + gpu_test_timeout_minutes: 45 + sidecar_backend: vllm + sidecar_binary_artifact: ${{ needs.sidecar-build.outputs.binary_artifact }} + secrets: inherit + + sidecar-sglang-test: + name: sidecar-runtime # This name overlaps with other sidecar jobs to group them in the UI + needs: [changed-files, sglang-build, sidecar-build] + # See sidecar-vllm-test above: sglang-build must have succeeded or been + # intentionally skipped, not failed or cancelled. + if: ${{ !cancelled() && needs.changed-files.outputs.sidecar == 'true' && needs.sidecar-build.result == 'success' && (needs.sglang-build.result == 'success' || needs.sglang-build.result == 'skipped') }} + permissions: + actions: read + contents: read + uses: $/.github/workflows/shared-test.yml + with: + test_suite_name: sglang + framework: sglang # matches the framework passed to sglang-build + test_type: Sidecar E2E Test + amd_runner: prod-tester-amd-gpu-v2 + target_tag_plain: ${{ needs.sglang-build.result == 'success' && needs.sglang-build.outputs.target_tag_plain || 'sglang-runtime' }} + source_ref: ${{ needs.sglang-build.result == 'success' && github.sha || 'main' }} + cuda_version: '["13.0"]' + platform: '["amd64"]' + run_sanity_check: false + gpu_test_markers: pre_merge and sidecar and sglang and gpu_1 + gpu_test_timeout_minutes: 45 + sidecar_backend: sglang + sidecar_binary_artifact: ${{ needs.sidecar-build.outputs.binary_artifact }} + secrets: inherit + + sidecar-trtllm-test: + name: sidecar-runtime # This name overlaps with other sidecar jobs to group them in the UI + needs: [changed-files, trtllm-build, sidecar-build] + # See sidecar-vllm-test above: trtllm-build must have succeeded or been + # intentionally skipped, not failed or cancelled. + if: ${{ !cancelled() && needs.changed-files.outputs.sidecar == 'true' && needs.sidecar-build.result == 'success' && (needs.trtllm-build.result == 'success' || needs.trtllm-build.result == 'skipped') }} + permissions: + actions: read + contents: read + uses: $/.github/workflows/shared-test.yml + with: + test_suite_name: trtllm + framework: trtllm # matches the framework passed to trtllm-build + test_type: Sidecar E2E Test + amd_runner: prod-tester-amd-gpu-v2 + target_tag_plain: ${{ needs.trtllm-build.result == 'success' && needs.trtllm-build.outputs.target_tag_plain || 'trtllm-runtime' }} + source_ref: ${{ needs.trtllm-build.result == 'success' && github.sha || 'main' }} + cuda_version: '["13.1"]' + platform: '["amd64"]' + run_sanity_check: false + gpu_test_markers: pre_merge and sidecar and trtllm and gpu_1 + gpu_test_timeout_minutes: 45 + sidecar_backend: trtllm + sidecar_binary_artifact: ${{ needs.sidecar-build.outputs.binary_artifact }} + secrets: inherit + # ============================================================================ # FRAMEWORK PIPELINES (Build โ†’ Test โ†’ Copy) # ============================================================================ @@ -813,7 +903,10 @@ jobs: # Keep the outer step budget above the test's 450s timeout so pytest can # finish teardown and upload diagnostics instead of being killed first. cpu_e2e_test_timeout_minutes: 15 - gpu_test_markers: pre_merge and vllm and gpu_1 + # not sidecar: sidecar-marked tests need dynamo-vllm-sidecar layered onto + # the image (see sidecar-vllm-test below), which this plain vllm-runtime + # image doesn't have. + gpu_test_markers: pre_merge and vllm and gpu_1 and not sidecar gpu_test_timeout_minutes: 45 # Profiled tests run in the parallel stage; unprofiled fall through to sequential. # 24 GiB admits all currently profiled vLLM tests (max is ~20.4 GiB) on a 48 GiB GPU. @@ -858,7 +951,10 @@ jobs: platform: '["amd64", "arm64"]' # arm64 for CPU tests, single GPU tests are skipped run_cpu_only_tests: true cpu_only_test_markers: pre_merge and sglang and gpu_0 - gpu_test_markers: pre_merge and sglang and gpu_1 + # not sidecar: sidecar-marked tests need dynamo-sglang-sidecar layered + # onto the image (see sidecar-sglang-test below), which this plain + # sglang-runtime image doesn't have. + gpu_test_markers: pre_merge and sglang and gpu_1 and not sidecar # Profiled tests run in the VRAM-aware GPU stage; unprofiled fall through to sequential. # Current single-GPU runners are 24 GiB, so this cap admits the profiled SGLang pool # while yielding one auto slot today. Larger runners will get more slots from the same markers. @@ -903,7 +999,10 @@ jobs: platform: '["amd64", "arm64"]' # arm64 for CPU tests, single GPU tests are skipped run_cpu_only_tests: true cpu_only_test_markers: pre_merge and trtllm and gpu_0 - gpu_test_markers: pre_merge and trtllm and gpu_1 + # not sidecar: sidecar-marked tests need dynamo-trtllm-sidecar layered + # onto the image (see sidecar-trtllm-test below), which this plain + # trtllm-runtime image doesn't have. + gpu_test_markers: pre_merge and trtllm and gpu_1 and not sidecar # Profiled tests run in the VRAM-aware GPU stage; unprofiled fall through to sequential. # Current single-GPU runners are 24 GiB, so this cap admits the profiled TRT-LLM pool # while yielding one auto slot today. Larger runners will get more slots from the same markers. diff --git a/.github/workflows/pre-merge.yml b/.github/workflows/pre-merge.yml index b04b69723b47..470c972ca81b 100644 --- a/.github/workflows/pre-merge.yml +++ b/.github/workflows/pre-merge.yml @@ -435,6 +435,13 @@ jobs: working-directory: ${{ matrix.dir }} # NOTE: --all-targets doesn't run doc tests run: cargo test --locked --all-targets + - name: Check isolated kv-router feature builds + if: matrix.dir == '.' + working-directory: ${{ matrix.dir }} + # Workspace feature unification hides these; check each on its own. + run: | + cargo check --locked -p dynamo-kv-router --no-default-features --features standalone-slot-tracker + cargo check --locked -p dynamo-kv-router --no-default-features --features standalone-indexer - name: Run standalone slot-tracker feature tests if: matrix.dir == '.' working-directory: ${{ matrix.dir }} @@ -534,7 +541,7 @@ jobs: working-directory: ${{ matrix.dir }} run: | # custom-policy is on by default; also lint it alongside the standalone selection - # service, and the build that links no worker-selection policies at all. + # service, and the build without the optional policy catalog. cargo clippy --locked --no-deps --features custom-policy,select-service --all-targets -- -D warnings cargo clippy --locked --no-deps --no-default-features --all-targets -- -D warnings - name: Lint feature-gated benchmark entrypoints diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 2c3958e48852..ef8993f2a033 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -554,6 +554,14 @@ jobs: if: env.NIGHTLY != 'true' env: NGC_HELM_REPO: https://helm.ngc.nvidia.com/${{ secrets.NGC_PUBLISH_ORG }}/ai-dynamo + # OCI mirror of the same staging namespace. ChartMuseum and OCI are + # SEPARATE registries in NGC โ€” a chart in one is invisible from the + # other โ€” so a GA release that only cm-pushes leaves the OCI copy + # missing, and the OCI publishing MR has nothing to promote. The + # nightly step below has dual-pushed since it was written; this step + # did not, which is why every GA chart through 1.4.0 is + # ChartMuseum-only. + NGC_HELM_OCI_REPO: oci://nvcr.io/${{ secrets.NGC_PUBLISH_ORG }}/ai-dynamo NGC_TOKEN: ${{ secrets.NGC_PUBLISH_TOKEN }} NGC_USERNAME: ${{ secrets.NGC_PUBLISH_USERNAME }} HELM_CHART_VERSION: ${{ needs.prepare-release.outputs.helm_chart_version }} @@ -592,6 +600,22 @@ jobs: helm cm-push "${CHART_FILE}" "${REPO_ALIAS}" echo "- \`${CHART_NAME}:${HELM_CHART_VERSION}\` pushed to NGC Helm registry" >> $GITHUB_STEP_SUMMARY + # Same tarball, second registry. helm does not read docker's login for + # OCI, so this needs its own registry login. Fail-soft: ChartMuseum is + # still the primary distribution channel and has already succeeded by + # this point, so an OCI hiccup must not fail the release. The GitLab + # publish:helm-oci:ondemand job probes for this copy and fails loudly + # if it is missing, so a soft failure here cannot ship silently. + echo "Pushing ${CHART_FILE} to NGC OCI registry..." + if echo "${NGC_TOKEN}" | helm registry login nvcr.io \ + --username "${NGC_USERNAME}" --password-stdin \ + && helm push "${CHART_FILE}" "${NGC_HELM_OCI_REPO}"; then + echo "- \`${CHART_NAME}:${HELM_CHART_VERSION}\` pushed to NGC OCI (\`${NGC_HELM_OCI_REPO}/${CHART_NAME}\`)" >> $GITHUB_STEP_SUMMARY + else + echo "::warning::OCI push of ${CHART_FILE} failed โ€” the ChartMuseum copy is unaffected. Re-run this job or push by hand before the OCI publishing MR." + echo "- :warning: \`${CHART_NAME}:${HELM_CHART_VERSION}\` OCI push FAILED (ChartMuseum copy is fine)" >> $GITHUB_STEP_SUMMARY + fi + helm repo remove "${REPO_ALIAS}" # Nightly: package the platform chart at the dated pre-release version diff --git a/.github/workflows/shared-build-image.yml b/.github/workflows/shared-build-image.yml index ace6d4d752fb..b3e44ed63342 100644 --- a/.github/workflows/shared-build-image.yml +++ b/.github/workflows/shared-build-image.yml @@ -37,6 +37,11 @@ on: required: false type: string default: '' + test_extra_tags: + description: 'Additional tags for the -test image only (newline-separated plain tag names, no ECR prefix). Unlike extra_tags (runtime image), the -test image otherwise gets no floating tag.' + required: false + type: string + default: '' no_cache: description: 'Disable Docker build cache' required: false @@ -603,6 +608,14 @@ jobs: - name: Build and Push Test Image if: ${{ !inputs.compliance_only && inputs.target != 'dev' && inputs.target != 'local-dev' && inputs.target != 'frontend' }} shell: bash + env: + ECR_REGISTRY: ${{ vars.ECR_REGISTRY }} + ECR_REPOSITORY: ${{ vars.ECR_REPOSITORY }} + TEST_EXTRA_TAGS: ${{ inputs.test_extra_tags }} + # Routed through env (not inlined via ${{ }}) so a value containing + # shell metacharacters can't be interpreted as script rather than data. + IMAGE_URI: ${{ steps.calculate-target-tag.outputs.image_uri }} + TEST_IMAGE_URI: ${{ steps.calculate-target-tag.outputs.test_image_uri }} run: | PUSH_ARGS="" if [ "${{ inputs.push_image }}" == "true" ]; then @@ -630,6 +643,15 @@ jobs: PYPI_SECRET_ARGS+=(--secret "id=pypi-netrc,src=${NETRC}") fi + TEST_TAG_ARGS=(-t "${TEST_IMAGE_URI}") + if [ -n "$TEST_EXTRA_TAGS" ]; then + while IFS= read -r tag; do + if [ -n "$tag" ]; then + TEST_TAG_ARGS+=(-t "${ECR_REGISTRY}/${ECR_REPOSITORY}:${tag}") + fi + done <<< "$TEST_EXTRA_TAGS" + fi + docker buildx build \ --progress=plain \ ${PUSH_ARGS} \ @@ -638,8 +660,8 @@ jobs: --platform ${{ inputs.platform }} \ --target ${TEST_TARGET} \ -f container/Dockerfile.test \ - --build-arg BASE_IMAGE=${{ steps.calculate-target-tag.outputs.image_uri }} \ - -t ${{ steps.calculate-target-tag.outputs.test_image_uri }} . + --build-arg BASE_IMAGE="${IMAGE_URI}" \ + "${TEST_TAG_ARGS[@]}" . - name: Show summary shell: bash if: ${{ inputs.push_image == 'true' && inputs.show_summary == 'true' }} diff --git a/.github/workflows/shared-build-sidecar.yml b/.github/workflows/shared-build-sidecar.yml index ecd6cb550287..6219ad142563 100644 --- a/.github/workflows/shared-build-sidecar.yml +++ b/.github/workflows/shared-build-sidecar.yml @@ -45,6 +45,13 @@ on: required: false type: string default: '' + outputs: + target_tag_plain: + description: 'Plain sidecar image tag prefix' + value: ${{ jobs.build.outputs.target_tag_plain }} + binary_artifact: + description: 'Name of the workflow artifact containing the amd64 dynamo-*-sidecar binaries' + value: ${{ jobs.build.outputs.binary_artifact }} jobs: build: @@ -54,6 +61,9 @@ jobs: permissions: actions: read contents: read + outputs: + target_tag_plain: ${{ steps.calculate-target-tag.outputs.target_tag_plain }} + binary_artifact: ${{ steps.binary-artifact.outputs.name }} steps: - name: Checkout repository uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -123,6 +133,42 @@ jobs: no_load: true use_sccache: false + # No derived per-backend test image: shared-test.yml pulls the backend's + # own -test image unmodified and installs the sidecar binary from this + # artifact via PATH (see shared-test.yml's sidecar_binary_artifact + # input). amd64-only: sidecar tests only run on amd64 GPU runners. + # + # buildx local output, not `docker pull` + `docker create` + `docker + # cp`: these builder runners route to a remote BuildKit worker (see + # init-dynamo-builder above) with no local docker daemon, so plain + # `docker` CLI commands against an image fail with "failed to connect + # to the docker API at unix:///var/run/docker.sock". `docker buildx + # build` talks to the remote builder directly and needs no daemon. + # Reuses the same builder and build cache as "Build and Push Image" + # above, so the sidecar-binaries target (which only depends on the + # {vllm,sglang,trtllm}-builder stages) is a near-instant cache hit. + - name: Extract sidecar binaries (amd64) + id: binary-artifact + shell: bash + run: | + set -euo pipefail + mkdir -p sidecar-binaries + docker buildx build \ + --target sidecar-binaries \ + --platform linux/amd64 \ + -f lib/sidecar/Dockerfile \ + --output type=local,dest=sidecar-binaries \ + . + name="dynamo-sidecar-binaries-${GITHUB_SHA}" + echo "name=${name}" >> "${GITHUB_OUTPUT}" + + - name: Upload sidecar binaries + uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0 + with: + name: ${{ steps.binary-artifact.outputs.name }} + path: sidecar-binaries/ + retention-days: 7 + - name: Resolve diff baseline id: resolve-diff-base if: ${{ inputs.diff_event_context != '' }} diff --git a/.github/workflows/shared-copy.yml b/.github/workflows/shared-copy.yml index 51f3fa8b9b81..a995b1728ea8 100644 --- a/.github/workflows/shared-copy.yml +++ b/.github/workflows/shared-copy.yml @@ -31,7 +31,7 @@ on: default: '' outputs: runtime_version: - description: 'Dynamo runtime compatibility version from Cargo.toml' + description: 'Dynamo runtime version from Cargo.toml' value: ${{ jobs.copy-to-acr.outputs.runtime_version }} image_tag: description: 'Image tag published to ACR' diff --git a/.github/workflows/shared-test.yml b/.github/workflows/shared-test.yml index c8f3b08037dd..30a8b58be273 100644 --- a/.github/workflows/shared-test.yml +++ b/.github/workflows/shared-test.yml @@ -11,7 +11,7 @@ on: required: true type: string framework: - description: 'Backend the test image was built with (vllm, sglang, trtllm, dynamo); labels test metadata and coverage/Allure artifacts' + description: 'Backend the test image was built with (dynamo, sglang, triton, trtllm, vllm); labels test metadata and coverage/Allure artifacts' required: true type: string test_type: @@ -117,7 +117,7 @@ on: type: boolean default: true source_ref: - description: 'Git ref/SHA to checkout and to resolve image tags against. Empty = caller SHA (github.sha).' + description: 'Git ref/SHA for image tags and checkout. Sidecar jobs always check out the caller SHA. Empty = caller SHA (github.sha).' required: false type: string default: '' @@ -136,6 +136,16 @@ on: required: false type: string default: 'true' + sidecar_backend: + description: 'Backend name (vllm/sglang/trtllm) to install a dynamo-*-sidecar binary for before running tests. Empty = skip (non-sidecar callers).' + required: false + type: string + default: '' + sidecar_binary_artifact: + description: 'Name of the workflow artifact (from shared-build-sidecar.yml) containing the amd64 dynamo-*-sidecar binaries. Required when sidecar_backend is set.' + required: false + type: string + default: '' secrets: AWS_DEFAULT_REGION: required: true @@ -193,9 +203,15 @@ jobs: DD_CIVISIBILITY_ENABLED: 'true' DD_CIVISIBILITY_AGENTLESS_ENABLED: 'true' DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER: github - PYTEST_ADDOPTS: --ddtrace + PYTEST_ADDOPTS: > + --ddtrace + ${{ inputs.framework == 'triton' && '--ignore=tests/fault_tolerance/test_canary_rank_pause.py' || '' }} + ${{ inputs.framework == 'triton' && '--ignore=tests/serve/test_vllm.py' || '' }} + ${{ inputs.framework == 'triton' && '--ignore=tests/serve/test_sglang.py' || '' }} + ${{ inputs.framework == 'triton' && '--ignore=tests/serve/test_trtllm.py' || '' }} permissions: contents: read + actions: read services: minio: image: ${{ vars.ECR_REGISTRY }}/quay/minio/minio:latest @@ -204,15 +220,44 @@ jobs: MINIO_ROOT_PASSWORD: minioadmin MINIO_VOLUMES: /tmp/minio-data MINIO_CONSOLE_ADDRESS: ':9001' - ports: - - 9000:9000 - - 9001:9001 steps: - name: Checkout repository id: checkout uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - ref: ${{ inputs.source_ref != '' && inputs.source_ref || github.sha }} + ref: ${{ inputs.sidecar_backend != '' && github.sha || inputs.source_ref || github.sha }} + - name: Download sidecar binary + if: ${{ inputs.sidecar_backend != '' }} + uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + with: + name: ${{ inputs.sidecar_binary_artifact }} + path: sidecar-artifact + - name: Install sidecar binary + if: ${{ inputs.sidecar_backend != '' }} + shell: bash + env: + SIDECAR_BACKEND: ${{ inputs.sidecar_backend }} + CONTAINER_WORKSPACE: ${{ inputs.container_workspace }} + run: | + set -euo pipefail + # Use a workspace-local directory writable by the container user. + bin_dir="${GITHUB_WORKSPACE}/.sidecar-bin" + mkdir -p "${bin_dir}" + install -m755 "sidecar-artifact/dynamo-${SIDECAR_BACKEND}-sidecar" \ + "${bin_dir}/dynamo-${SIDECAR_BACKEND}-sidecar" + echo "${bin_dir}" >> "${GITHUB_PATH}" + if [ -z "${CONTAINER_WORKSPACE}" ]; then + echo "::error::container_workspace is empty; refusing to replace test files" + exit 1 + fi + if [ ! "${GITHUB_WORKSPACE}" -ef "${CONTAINER_WORKSPACE}" ]; then + # Replace the test inputs so files deleted by the PR cannot linger. + for path in lib/sidecar tests examples/common conftest.py pyproject.toml Cargo.toml; do + rm -rf "${CONTAINER_WORKSPACE}/${path}" + mkdir -p "$(dirname "${CONTAINER_WORKSPACE}/${path}")" + cp -r "${GITHUB_WORKSPACE}/${path}" "${CONTAINER_WORKSPACE}/${path}" + done + fi - name: Export KVBM MLA backend if: ${{ inputs.kvbm_mla_backend != '' }} shell: bash @@ -238,6 +283,7 @@ jobs: fi - name: Run CPU-only tests (parallelized) + id: cpu-only-tests if: ${{ inputs.run_cpu_only_tests }} timeout-minutes: ${{ inputs.cpu_only_test_timeout_minutes }} uses: $/.github/actions/pytest-local @@ -253,6 +299,7 @@ jobs: enable_coverage: ${{ inputs.enable_coverage }} gpu_tests: 'false' - name: Run CPU end-to-end tests (isolated) + id: cpu-e2e-tests # Long-lived deployment tests should not compete with the xdist unit # pool or retry until the parent step kills the runner container. if: ${{ !cancelled() && steps.checkout.outcome == 'success' && inputs.run_cpu_e2e_tests }} @@ -271,6 +318,7 @@ jobs: gpu_tests: 'false' dd_flaky_retry_enabled: 'false' - name: Run GPU tests (parallel) + id: gpu-parallel-tests timeout-minutes: ${{ inputs.gpu_test_timeout_minutes }} if: ${{ matrix.platform == 'amd64' && inputs.run_gpu_tests && inputs.run_gpu_parallel_tests }} uses: $/.github/actions/pytest-local @@ -290,6 +338,7 @@ jobs: gpu_tests: 'true' verify_gpu: ${{ inputs.verify_gpu }} - name: Run GPU tests (sequential) + id: gpu-sequential-tests timeout-minutes: ${{ inputs.gpu_test_timeout_minutes }} # !cancelled() so a parallel-stage failure doesn't hide sequential results, # but require checkout to have succeeded so an infra/prereq failure (e.g. the @@ -324,3 +373,96 @@ jobs: path: test-results/pytest_test_report_*.xml retention-days: 7 if-no-files-found: warn + + # A stage cut off mid-run is still reported as a successful step, so real + # test failures go green. Take the verdict only from the status file + # pytest-local writes last, never from the step outcome. Placed after the + # upload so results survive a failure here, and once per job rather than + # per stage to avoid three more container-hook round trips. + - name: Verify test stages completed + if: ${{ !cancelled() && steps.checkout.outcome == 'success' }} + shell: bash + env: + STAGE_OUTCOMES: | + pre_merge_cpu ${{ steps.cpu-only-tests.outcome }} + pre_merge_cpu_e2e ${{ steps.cpu-e2e-tests.outcome }} + pre_merge_gpu_parallel ${{ steps.gpu-parallel-tests.outcome }} + pre_merge_gpu ${{ steps.gpu-sequential-tests.outcome }} + run: | + # Runs inside the job's test container. The planner image is + # distroless and ships only the binaries listed in + # container/Dockerfile.test, so this stays on bash builtins. + set -uo pipefail + status_dir="${GITHUB_WORKSPACE}/test-results/status" + failed=0 + checked="" + + while read -r stage outcome; do + [[ -z "${stage}" ]] && continue + exitcode_file="${status_dir}/${stage}.exitcode" + + case "${outcome}" in + ''|skipped) + # An unresolvable id renders empty, which would otherwise read + # as "skipped" and exempt a stage that really ran. + if [[ -f "${exitcode_file}" ]]; then + echo "โŒ ${stage}: wrote a status file but its outcome is" + echo " '${outcome:-}'. Its id in STAGE_OUTCOMES does not" + echo " resolve, so the stage ran unverified." + failed=1 + else + echo "โญ๏ธ ${stage}: did not run" + fi + ;; + success) + if [[ ! -f "${exitcode_file}" ]]; then + echo "โŒ ${stage}: reported success but wrote no status file." + echo " Its script did not run to completion, so the result is" + echo " unknown and cannot be treated as a pass." + failed=1 + else + # An empty file means the stage created it on entry and never + # reached the line that records an exit code. read is a + # builtin and returns non-zero on an empty file. + code="" + read -r code < "${exitcode_file}" || true + if [[ -z "${code}" ]]; then + echo "โŒ ${stage}: reported success but recorded no exit code." + echo " Its script was cut off mid-run, so the result is" + echo " unknown and cannot be treated as a pass." + failed=1 + elif [[ "${code}" != "0" ]]; then + echo "โŒ ${stage}: pytest exited ${code}" + failed=1 + else + echo "โœ… ${stage}: pytest exited 0" + fi + fi + ;; + *) + echo "โŒ ${stage}: step outcome '${outcome}'" + failed=1 + ;; + esac + checked+=" ${stage}" + done <<< "${STAGE_OUTCOMES}" + + # A stage added above but not listed here would escape this check. + shopt -s nullglob + for f in "${status_dir}"/*.exitcode; do + stage="${f##*/}" + stage="${stage%.exitcode}" + if [[ " ${checked} " != *" ${stage} "* ]]; then + echo "โŒ ${stage}: ran but is not listed in STAGE_OUTCOMES, so its" + echo " result is unverified. Add it to this step." + failed=1 + fi + done + + # Without this the checks tab shows only "Process completed with exit + # code 1" and the stage name stays buried in this step's log. + if [[ "${failed}" -ne 0 ]]; then + echo "::error title=Test stage verification failed::A pytest stage did not record a passing result. See this step's log for the stage." + fi + + exit "${failed}" diff --git a/.github/workflows/trigger-ci-approval-flow.yml b/.github/workflows/trigger-ci-approval-flow.yml index ab0d1589e6b9..70aeb775219d 100644 --- a/.github/workflows/trigger-ci-approval-flow.yml +++ b/.github/workflows/trigger-ci-approval-flow.yml @@ -1,53 +1,43 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# Intent -# ------ -# Safely approve trusted CI for the exact head commit of a pull request from a fork. Fork PRs do -# not receive this approval automatically because their code is external to the repository. This -# workflow evaluates who is responsible for the PR update, verifies every PR commit, and posts the -# `/ok to test ` command that starts the trusted CI flow only when all checks pass. +# Approve CI for a fork PR by posting `/ok to test `. +# Approval applies only to that commit. # -# When this workflow runs -# ----------------------- -# The workflow receives `pull_request_target` events when a fork PR is opened, synchronized by a -# new push, or reopened. The job is skipped for branches in this repository because those branches -# use the normal internal-contributor CI path. +# This workflow handles fork PRs when they are opened, pushed to, or reopened, +# and when glamr-agent posts a CI request. PRs from branches in this repository +# use the internal-contributor CI workflow instead. # -# A fork PR can reach the signature check through any of these approval paths: +# Automatic approval is allowed in these cases: # -# 1. The PR author has write, maintain, or admin access to this repository. This path applies when -# the PR is opened or reopened. A synchronized update must also be pushed by a maintainer. -# 2. The PR author is listed in the `APPROVED_EXTERNAL_CONTRIBUTORS` environment secret. This path -# applies when the PR is opened, synchronized, or reopened. -# 3. A user with write, maintain, or admin access to this repository synchronizes the PR branch. -# For example, a maintainer may merge `main` into a contributor's branch. This path approves -# only the head created by that synchronization; it does not make the PR or contributor trusted -# for later pushes. +# 1. A maintainer opens or reopens a PR. Here, a maintainer is someone with write, +# maintain, or admin access to this repository. +# 2. A PR author listed in APPROVED_EXTERNAL_CONTRIBUTORS opens, updates, or +# reopens a PR. +# 3. A maintainer pushes to a PR branch, for example by merging main into it. +# Approval covers that head commit. Later pushes need a new approval. +# 4. glamr-agent posts `@dynamo-ops please run full CI for ` on +# its own PR. The request may start with a label and colon. The commenter, +# PR author, and head repository owner must all have GitHub account ID +# 282374502, and the SHA must match the PR's current head. This case does not +# require signed commits. +# Set the repository Actions variable GLAMR_CI_REQUESTS_DISABLED=true to +# disable these requests. # -# All paths require every commit currently in the PR to have a signature that GitHub reports as -# verified. If any commit is unsigned, the workflow does not post `/ok to test`. For PRs authored -# by maintainers, it explains which commits failed verification. Approved external contributors -# receive guidance for either signing the commits or requesting manual approval. Other external -# contributors wait for manual approval without receiving an unsigned-commit comment. Repository- -# permission lookup failures fail closed for the maintainer paths, and malformed contributor-list -# configuration fails closed for the approved-author path. +# Cases 1-3 require GitHub to verify the signature of every commit in the recorded +# base-to-head comparison. If verification fails, the workflow tells maintainer +# authors which commits failed. Authors on the approved list receive instructions +# for signing or getting manual approval. Other contributors receive no comment. +# Failed permission lookups cannot grant maintainer approval. An invalid approved +# contributor list cannot grant approval through that list. # -# Why a contributor's PR may not trigger trusted CI -# ------------------------------------------------- -# - A new contributor who is not on the approved list needs a maintainer to approve the current -# head, either by synchronizing the branch or by posting the approval command manually. -# - A contributor push after a maintainer synchronization creates a new head and removes the -# one-head approval. The new head must be approved again. -# - Every PR commit must be signed and verified, including commits added before a maintainer update. -# - Contributors must enable "Allow edits from maintainers" before a maintainer can update a branch -# in a user-owned fork. +# A maintainer can also approve CI manually with `/ok to test `. +# To let a maintainer push to a user-owned fork, the contributor must enable +# "Allow edits from maintainers" on the PR. # -# Desired state -# ------------- -# The workflow posts an `/ok to test` command containing the current PR head SHA only after the -# current head has a trusted approval path and all commits are verified. Every later push creates a -# different head SHA and causes the complete approval process to run again. +# Before posting approval, the workflow checks that the PR is still open and its +# base and head have not changed. A push can arrive after this check, so the CI +# trigger must also check the command's SHA when approval is required. name: Trigger CI Approval Flow @@ -57,6 +47,8 @@ on: - opened - synchronize - reopened + issue_comment: + types: [created] permissions: contents: read @@ -65,15 +57,54 @@ permissions: jobs: ok-to-test: - if: github.event.pull_request.head.repo.fork + if: >- + (github.event_name == 'pull_request_target' && github.event.pull_request.head.repo.fork) || + (github.event_name == 'issue_comment' && github.event.issue.pull_request && + vars.GLAMR_CI_REQUESTS_DISABLED != 'true' && + github.event.comment.user.id == 282374502 && + contains(github.event.comment.body, '@dynamo-ops please run full CI for ')) runs-on: ubuntu-latest environment: external_collaborator + concurrency: + group: ci-approval-${{ github.event.pull_request.number || github.event.issue.number }} + cancel-in-progress: false + queue: max + env: + REPO_OWNER: ${{ github.repository_owner }} + REPO_NAME: ${{ github.event.repository.name }} + UNSIGNED_COMMENT_MARKER: "" + PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }} steps: + - name: Resolve current PR and requested commit + id: request + env: + GH_TOKEN: ${{ github.token }} + GLAMR_USER_ID: "282374502" + run: | + pr=$(gh api "repos/$GITHUB_REPOSITORY/pulls/$PR_NUMBER") + [[ $(jq -r '.state == "open" and .head.repo.fork == true' <<< "$pr") == true ]] || exit 0 + head=$(jq -r '.head.sha' <<< "$pr") + if [[ "$GITHUB_EVENT_NAME" == issue_comment ]]; then + [[ $(jq -r '.user.id' <<< "$pr") == "$GLAMR_USER_ID" ]] || exit 0 + [[ $(jq -r '.head.repo.owner.id' <<< "$pr") == "$GLAMR_USER_ID" ]] || exit 0 + body=$(jq -r '.comment.body' "$GITHUB_EVENT_PATH") + pattern='^([^[:space:]]+:[[:blank:]]+)?@dynamo-ops please run full CI for ([0-9a-f]{40})[[:space:]]*$' + [[ "$body" =~ $pattern ]] || exit 0 + [[ "${BASH_REMATCH[2]}" == "$head" ]] || exit 0 + echo "glamr_request_approved=true" >> "$GITHUB_OUTPUT" + else + [[ $(jq -r '.pull_request.head.sha' "$GITHUB_EVENT_PATH") == "$head" ]] || exit 0 + fi + echo "head=$head" >> "$GITHUB_OUTPUT" + echo "base=$(jq -r '.base.sha' <<< "$pr")" >> "$GITHUB_OUTPUT" + echo "author=$(jq -r '.user.login' <<< "$pr")" >> "$GITHUB_OUTPUT" + echo "current=true" >> "$GITHUB_OUTPUT" - name: Validate approved contributor id: author_check + if: steps.request.outputs.current == 'true' env: APPROVED_EXTERNAL_CONTRIBUTORS: ${{ secrets.APPROVED_EXTERNAL_CONTRIBUTORS }} - PR_AUTHOR: ${{ github.event.pull_request.user.login }} + PR_AUTHOR: ${{ steps.request.outputs.author }} run: | if ! echo "$APPROVED_EXTERNAL_CONTRIBUTORS" | jq -e 'if type == "array" then . else error("not an array") end' > /dev/null 2>&1; then echo "Warning: APPROVED_EXTERNAL_CONTRIBUTORS is not a valid JSON array." >&2 @@ -90,10 +121,11 @@ jobs: fi - name: Validate PR author permission id: author_permission_check + if: steps.request.outputs.current == 'true' env: GH_TOKEN: ${{ github.token }} REPOSITORY: ${{ github.repository }} - PR_AUTHOR: ${{ github.event.pull_request.user.login }} + PR_AUTHOR: ${{ steps.request.outputs.author }} run: | # This checks effective repository permission, not organization membership. if ! author_permission=$(gh api \ @@ -118,7 +150,7 @@ jobs: esac - name: Validate triggering actor id: actor_check - if: github.event.action == 'synchronize' && steps.author_check.outputs.is_approved != 'true' + if: steps.request.outputs.current == 'true' && github.event_name == 'pull_request_target' && github.event.action == 'synchronize' && steps.author_check.outputs.is_approved != 'true' env: GH_TOKEN: ${{ github.token }} REPOSITORY: ${{ github.repository }} @@ -149,13 +181,15 @@ jobs: esac - name: Determine CI approval id: approval_check + if: steps.request.outputs.current == 'true' env: + GLAMR_REQUEST_APPROVED: ${{ steps.request.outputs.glamr_request_approved }} AUTHOR_APPROVED: ${{ steps.author_check.outputs.is_approved }} AUTHOR_MAINTAINER: ${{ steps.author_permission_check.outputs.is_maintainer }} ACTOR_APPROVED: ${{ steps.actor_check.outputs.is_approved }} EVENT_ACTION: ${{ github.event.action }} run: | - if [[ "$AUTHOR_APPROVED" == "true" ]]; then + if [[ "$GLAMR_REQUEST_APPROVED" == "true" || "$AUTHOR_APPROVED" == "true" ]]; then echo "is_approved=true" >> "$GITHUB_OUTPUT" elif [[ "$EVENT_ACTION" == "synchronize" && "$ACTOR_APPROVED" == "true" ]]; then echo "is_approved=true" >> "$GITHUB_OUTPUT" @@ -166,16 +200,15 @@ jobs: fi - name: Verify commit signatures id: signature_check - if: steps.approval_check.outputs.is_approved == 'true' + if: steps.approval_check.outputs.is_approved == 'true' && steps.request.outputs.glamr_request_approved != 'true' env: GH_TOKEN: ${{ github.token }} - REPO_OWNER: ${{ github.repository_owner }} - REPO_NAME: ${{ github.event.repository.name }} - PR_NUMBER: ${{ github.event.pull_request.number }} + PR_BASE: ${{ steps.request.outputs.base }} + PR_SHA: ${{ steps.request.outputs.head }} run: | commit_shas=$(gh api --paginate \ - "repos/$REPO_OWNER/$REPO_NAME/pulls/$PR_NUMBER/commits?per_page=100" \ - --jq '.[].sha') + "repos/$REPO_OWNER/$REPO_NAME/compare/$PR_BASE...$PR_SHA?per_page=100" \ + --jq '.commits[].sha') if [[ -z "$commit_shas" ]]; then echo "Warning: No commits found in PR." >&2 @@ -217,14 +250,11 @@ jobs: if: steps.approval_check.outputs.is_approved == 'true' env: GH_TOKEN: ${{ secrets.AUTO_VERIFIER_CI_TOKEN }} - REPO_OWNER: ${{ github.repository_owner }} - REPO_NAME: ${{ github.event.repository.name }} - PR_NUMBER: ${{ github.event.pull_request.number }} run: | bot_login=$(gh api user --jq '.login') gh api "repos/$REPO_OWNER/$REPO_NAME/issues/$PR_NUMBER/comments" \ - | jq -r --arg login "$bot_login" \ - '.[] | select(.user.login == $login and (.body | contains(""))) | .id' \ + | jq -r --arg login "$bot_login" --arg marker "$UNSIGNED_COMMENT_MARKER" \ + '.[] | select(.user.login == $login and (.body | contains($marker))) | .id' \ | while read -r comment_id; do echo "Deleting stale unsigned-commits comment $comment_id" gh api --method DELETE "repos/$REPO_OWNER/$REPO_NAME/issues/comments/$comment_id" @@ -233,13 +263,10 @@ jobs: if: steps.approval_check.outputs.is_approved == 'true' && steps.author_permission_check.outputs.is_maintainer == 'true' && steps.signature_check.outputs.all_signed == 'false' env: GH_TOKEN: ${{ secrets.AUTO_VERIFIER_CI_TOKEN }} - REPO_OWNER: ${{ github.repository_owner }} - REPO_NAME: ${{ github.event.repository.name }} - PR_NUMBER: ${{ github.event.pull_request.number }} UNSIGNED_COMMITS: ${{ steps.signature_check.outputs.unsigned_commits }} run: | table=$(echo "$UNSIGNED_COMMITS" | jq -r '.[] | "| `\(.sha)` | \(.reason) |"') - body=" + body="$UNSIGNED_COMMENT_MARKER ## Unsigned Commits Detected The following commits are not GPG-signed and must be signed before CI can run: @@ -257,14 +284,11 @@ jobs: if: steps.author_check.outputs.is_approved == 'true' && steps.author_permission_check.outputs.is_maintainer != 'true' && steps.signature_check.outputs.all_signed == 'false' env: GH_TOKEN: ${{ secrets.AUTO_VERIFIER_CI_TOKEN }} - REPO_OWNER: ${{ github.repository_owner }} - REPO_NAME: ${{ github.event.repository.name }} - PR_NUMBER: ${{ github.event.pull_request.number }} - PR_SHA: ${{ github.event.pull_request.head.sha }} + PR_SHA: ${{ steps.request.outputs.head }} UNSIGNED_COMMITS: ${{ steps.signature_check.outputs.unsigned_commits }} run: | table=$(echo "$UNSIGNED_COMMITS" | jq -r '.[] | "| `\(.sha)` | \(.reason) |"') - body=" + body="$UNSIGNED_COMMENT_MARKER ## Automatic CI Approval Skipped Automatic approval requires every commit to have a signature that GitHub reports as verified. The following commits do not meet that requirement: @@ -279,19 +303,28 @@ jobs: --repo "$REPO_OWNER/$REPO_NAME" \ --body "$body" - name: Comment /ok to test - if: steps.approval_check.outputs.is_approved == 'true' && steps.signature_check.outputs.all_signed == 'true' + if: steps.approval_check.outputs.is_approved == 'true' && (steps.request.outputs.glamr_request_approved == 'true' || steps.signature_check.outputs.all_signed == 'true') env: GH_TOKEN: ${{ secrets.AUTO_VERIFIER_CI_TOKEN }} - REPO_OWNER: ${{ github.repository_owner }} - REPO_NAME: ${{ github.event.repository.name }} - PR_NUMBER: ${{ github.event.pull_request.number }} - PR_SHA: ${{ github.event.pull_request.head.sha }} + PR_BASE: ${{ steps.request.outputs.base }} + PR_SHA: ${{ steps.request.outputs.head }} run: | if [[ -z "$PR_SHA" ]]; then echo "Could not determine PR head SHA; aborting /ok to test comment." >&2 exit 1 fi + bot_login=$(gh api user --jq '.login') + comments=$(gh api --paginate --slurp "repos/$REPO_OWNER/$REPO_NAME/issues/$PR_NUMBER/comments?per_page=100") + if jq -e --arg login "$bot_login" --arg body "/ok to test $PR_SHA" \ + 'any(.[][]; .user.login == $login and .body == $body)' <<< "$comments" > /dev/null; then + echo "CI approval already posted for $PR_SHA." + exit 0 + fi + pr=$(gh api "repos/$REPO_OWNER/$REPO_NAME/pulls/$PR_NUMBER") + jq -e --arg sha "$PR_SHA" --arg base "$PR_BASE" \ + '.state == "open" and .head.sha == $sha and .base.sha == $base' <<< "$pr" > /dev/null || exit 0 + gh pr comment "$PR_NUMBER" \ --repo "$REPO_OWNER/$REPO_NAME" \ --body "/ok to test $PR_SHA" diff --git a/.gitignore b/.gitignore index c861a90dcc02..dc5f5124cca5 100644 --- a/.gitignore +++ b/.gitignore @@ -171,9 +171,9 @@ package-lock.json # HuggingFace hub cache downloaded by tests (never commit; committed fixtures use plain dir names) lib/llm/tests/data/sample-models/models--*/ -# Generated by .github/scripts/build_install_selector.py -docs/data/install-selector.json -fern/components/install-selector-data.ts +# Generated by docs/fern/scripts/gen_nightly_selector.py (rebuilt by the docs +# workflow on every publish) +docs/fern/components/nightly-selector-data.generated.ts # Local git worktrees .worktrees/ diff --git a/.lycheeignore b/.lycheeignore index d2437c912db9..cf4a75ff0d97 100644 --- a/.lycheeignore +++ b/.lycheeignore @@ -56,9 +56,13 @@ # (afterwards lychee validates the URL again, catching any later rename/move). ^https://github\.com/ai-dynamo/dynamo/blob/main/deploy/inference-gateway/ext-proc/examples/onramp/agg\.yaml -# The v1.4.0 release entry lands with the current-release promotion; the tag -# URL 404s from CI until v1.4.0 is cut. Remove after the tag exists. -^https://github\.com/ai-dynamo/dynamo/releases/tag/v1\.4\.0$ +# The v1.5.0 release entry lands with the current-release promotion; the tag +# URL 404s from CI until v1.5.0 is cut. Remove after the tag exists. +^https://github\.com/ai-dynamo/dynamo/releases/tag/v1\.5\.0$ + +# The classic dynamo-platform 1.5.0 chart publishes to helm.ngc.nvidia.com in a +# follow-up NGC publishing MR; the .tgz URL 404s until then. Remove once live. +^https://helm\.ngc\.nvidia\.com/nvidia/ai-dynamo/charts/dynamo-platform-1\.5\.0\.tgz$ # The Dynamo Enterprise NGC collection is not published on the public catalog # yet, so the canonical collection URL 404s from CI. The org and team segments @@ -80,3 +84,7 @@ # link, so the URL returns 404 until the docs deploy. Remove this entry after # the page is live. ^https://docs\.nvidia\.com/dynamo/dev/recipes/deepseek-v4-1-flash$ + +# Introduced by the same PR as its docs link: the K-EXAONE 2.0 recipe does not exist on main +# until this PR merges, so these blob/main URLs 404 during review. Remove after merge. +^https://github\.com/ai-dynamo/dynamo/(tree|blob)/main/recipes/k-exaone-2\.0 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 86ab70bbea98..cd932c6295be 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -155,6 +155,15 @@ repos: # regenerate leaves every committed generated span stale. files: ^docs/fern/(scripts/gen_llms_tables\.py|components/releases\.data\.ts|pages/reference/general/.*\.mdx)$ pass_filenames: false + - id: gen-llms-tables-determinism + name: Check the generated docs tables are reproducible + entry: docs/fern/scripts/check_llms_tables_determinism.sh + language: system + # Triggers on the script itself as well as the generator and its inputs: + # the reproducibility claim is only as good as the harness asserting it. + files: ^docs/fern/(scripts/(gen_llms_tables|check_llms_tables_determinism)\.(py|sh)|components/releases\.data\.ts|pages/reference/general/.*\.mdx)$ + pass_filenames: false + - id: check-component-imports name: Check Fern components trigger no rolldown bundling entry: python3 docs/fern/scripts/check_component_imports.py @@ -177,6 +186,7 @@ repos: - pytest-asyncio - tomli - pydantic + - redis>=6.2.0,<9.0.0 - filelock - pyyaml - prometheus_client>=0.23.1 @@ -198,6 +208,12 @@ repos: # be a second copy to keep in sync. Full scan is well under 0.1s. files: ^\.github/ pass_filenames: false + - id: compliance-extract-selftest + name: Test compliance export retries + entry: python3 .github/actions/compliance-extract/test_extract_with_retry.py + language: python + pass_filenames: false + files: ^(\.github/actions/compliance-extract/|\.pre-commit-config\.yaml$) - id: check-action-pins-selftest name: Self-test the action-pin matcher entry: python3 scripts/check_action_pins.py --test diff --git a/AGENTS.md b/AGENTS.md index 30f9a229993f..fe9864409553 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -232,10 +232,14 @@ cargo fmt --all && cargo clippy --workspace `style`, and `build`. - PR descriptions must include `Summary` and `Validation`. - Sign every commit with DCO: `git commit -s`. -- For fork PRs that qualify for automatic trusted-CI approval, every commit must have a - cryptographic signature that GitHub reports as `Verified`; a DCO sign-off alone does not - satisfy this requirement. Signing commits does not itself qualify a PR for automatic approval; - a maintainer can manually approve the current head with `/ok to test `. +- For fork PRs that qualify for automatic CI approval, GitHub must report every commit's + signature as `Verified`. A DCO sign-off alone is not enough. Signing commits does not + grant automatic approval; a maintainer can approve the current head manually with + `/ok to test `. + The only signature exception is a full-CI request from `glamr-agent` on its own PR. + The commenter, PR author, and head repository owner must all be `glamr-agent`, + checked by GitHub account ID. The request's full SHA must match the PR's current + head. Other automatic approvals still require verified signatures. - Do not hand-edit a generated artifact โ€” change its source and regenerate. A generated file says so in a `do not edit` marker, and its generator has a `--check` mode that fails when the committed output is stale. Resolve a @@ -256,8 +260,8 @@ cargo fmt --all && cargo clippy --workspace skill automates all of this. - Full CI on a PR runs only after a maintainer comments `/ok to test ` with the short SHA of the latest commit; copy-pr-bot then creates the `pull-request/N` branch that - triggers it. For an eligible fork PR, the automatic approval flow posts that command only - after every PR commit is GitHub-verified. Fix failures before requesting human review. + triggers it. For an eligible fork PR, automatic approval posts this command after the + checks above pass. Fix CI failures before requesting human review. - Architecture changes require a Dynamo Enhancement Proposal (DEP), filed as a GitHub issue on `ai-dynamo/dynamo` with `dep:*` labels (the `dep-create` skill automates this). diff --git a/CODEOWNERS b/CODEOWNERS index e91a11f2a59a..c92d47e0dda8 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -241,7 +241,6 @@ /components/src/dynamo/common/ @ai-dynamo/dynamo-runtime-codeowners /tests/test_models_dir_flag.py @ai-dynamo/dynamo-runtime-codeowners /tests/report_pytest_markers.py @ai-dynamo/dynamo-runtime-codeowners -/examples/backends/tritonserver/ @ai-dynamo/dynamo-runtime-codeowners /tests/serve/multimodal_profiles/ @ai-dynamo/dynamo-runtime-codeowners /tests/test_predownload_models.py @ai-dynamo/dynamo-runtime-codeowners /components/src/dynamo/common/lora/ @ai-dynamo/dynamo-runtime-codeowners @@ -385,6 +384,7 @@ # === rl (@ai-dynamo/dynamo-rl-codeowners) === /lib/rl/ @ai-dynamo/dynamo-rl-codeowners +/examples/rl/ @ai-dynamo/dynamo-rl-codeowners # === performance (@ai-dynamo/dynamo-performance-codeowners) === /recipes/ @ai-dynamo/dynamo-performance-codeowners @@ -645,6 +645,7 @@ requirements.txt /lib/llm/src/http/service/metrics.rs @ai-dynamo/dynamo-frontend-codeowners @ai-dynamo/dynamo-observability-codeowners /container/deps/requirements.vllm.txt @ai-dynamo/dynamo-ops-codeowners @ai-dynamo/dynamo-backend-vllm-codeowners /deploy/operator/internal/checkpoint/ @ai-dynamo/dynamo-gms-codeowners @ai-dynamo/dynamo-operator-codeowners +/lib/llm/src/preprocessor/mm_routing/ @ai-dynamo/dynamo-frontend-codeowners @ai-dynamo/dynamo-multimodal-codeowners @ai-dynamo/dynamo-router-codeowners /lib/llm/src/protocols/openai/audios/ @ai-dynamo/dynamo-frontend-codeowners @ai-dynamo/dynamo-multimodal-codeowners /lib/llm/src/protocols/openai/images/ @ai-dynamo/dynamo-frontend-codeowners @ai-dynamo/dynamo-diffusion-codeowners /lib/llm/src/protocols/openai/videos/ @ai-dynamo/dynamo-frontend-codeowners @ai-dynamo/dynamo-diffusion-codeowners @@ -680,7 +681,6 @@ requirements.txt /container/deps/requirements.frontend.txt @ai-dynamo/dynamo-ops-codeowners @ai-dynamo/dynamo-frontend-codeowners /lib/kv-router/src/services/shared_cache/ @ai-dynamo/dynamo-router-codeowners @ai-dynamo/dynamo-kv-memory-codeowners /lib/kv-router/src/zmq_wire/extra_keys.rs @ai-dynamo/dynamo-router-codeowners @ai-dynamo/dynamo-multimodal-codeowners -/lib/llm/src/preprocessor/lightseek_mm.rs @ai-dynamo/dynamo-frontend-codeowners @ai-dynamo/dynamo-multimodal-codeowners @ai-dynamo/dynamo-router-codeowners /lib/runtime/src/metrics/frontend_perf.rs @ai-dynamo/dynamo-observability-codeowners @ai-dynamo/dynamo-frontend-codeowners /tests/fault_tolerance/migration/utils.py @ai-dynamo/dynamo-fault-tolerance-codeowners @ai-dynamo/dynamo-observability-codeowners /tests/test_generate_kustomize_openapi.py @ai-dynamo/dynamo-runtime-codeowners @ai-dynamo/dynamo-ops-codeowners @@ -688,7 +688,6 @@ requirements.txt /components/src/dynamo/trtllm/publisher.py @ai-dynamo/dynamo-backend-trtllm-codeowners @ai-dynamo/dynamo-kv-memory-codeowners /container/deps/requirements.benchmark.txt @ai-dynamo/dynamo-ops-codeowners @ai-dynamo/dynamo-performance-codeowners /deploy/operator/internal/gpu/discovery.go @ai-dynamo/dynamo-operator-codeowners @ai-dynamo/dynamo-observability-codeowners -/examples/backends/tritonserver/Dockerfile @ai-dynamo/dynamo-runtime-codeowners @ai-dynamo/dynamo-ops-codeowners /lib/llm/src/block_manager/metrics_kvbm.rs @ai-dynamo/dynamo-kv-memory-codeowners @ai-dynamo/dynamo-observability-codeowners /components/src/dynamo/thunderagent_router/ @ai-dynamo/dynamo-router-codeowners @ai-dynamo/dynamo-agents-codeowners /components/src/dynamo/vllm/health_check.py @ai-dynamo/dynamo-backend-vllm-codeowners @ai-dynamo/dynamo-observability-codeowners diff --git a/Cargo.lock b/Cargo.lock index bd250b8af591..6ea9c4f9795f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,9 +34,9 @@ dependencies = [ [[package]] name = "aisimulate-core" -version = "0.12.0-dev.1" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4db6a085073ecb1c2569f796ba44742ed4681f8f4060daaed2456ef4ea9b652" +checksum = "7d266db788afdb866421aad39f8fbc7096dba9517105ff7d770ee83ad70f82bb" dependencies = [ "anyhow", "bincode 1.3.3", @@ -2392,7 +2392,7 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "dynamo-backend-common" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-stream", @@ -2411,6 +2411,7 @@ dependencies = [ "serde_json", "serial_test", "temp-env", + "tempfile", "tokio", "tokio-util", "tracing", @@ -2420,7 +2421,7 @@ dependencies = [ [[package]] name = "dynamo-bench" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "blake3", @@ -2462,16 +2463,24 @@ dependencies = [ [[package]] name = "dynamo-custom-policy-builtin" -version = "1.5.0" +version = "1.6.0" dependencies = [ + "async-trait", + "criterion", "dynamo-kv-router", + "fastrand", + "parking_lot", + "rustc-hash 2.1.2", "serde", "tempfile", + "thiserror 2.0.18", + "tokio", + "tracing", ] [[package]] name = "dynamo-custom-policy-example-catalog" -version = "1.5.0" +version = "1.6.0" dependencies = [ "dynamo-custom-policy-example-disagg-filter-score-pick", "dynamo-custom-policy-example-simple-filter-score-pick", @@ -2482,7 +2491,7 @@ dependencies = [ [[package]] name = "dynamo-custom-policy-example-disagg-filter-score-pick" -version = "1.5.0" +version = "1.6.0" dependencies = [ "dynamo-kv-router", "serde", @@ -2490,7 +2499,7 @@ dependencies = [ [[package]] name = "dynamo-custom-policy-example-epp" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "dynamo-custom-policy-example-catalog", @@ -2501,7 +2510,7 @@ dependencies = [ [[package]] name = "dynamo-custom-policy-example-simple-filter-score-pick" -version = "1.5.0" +version = "1.6.0" dependencies = [ "dynamo-kv-router", "serde", @@ -2509,7 +2518,7 @@ dependencies = [ [[package]] name = "dynamo-custom-policy-example-simple-stacked-score-pick" -version = "1.5.0" +version = "1.6.0" dependencies = [ "dynamo-kv-router", "serde", @@ -2517,7 +2526,7 @@ dependencies = [ [[package]] name = "dynamo-custom-policy-example-soft-pin-repin" -version = "1.5.0" +version = "1.6.0" dependencies = [ "dynamo-kv-router", "serde", @@ -2525,7 +2534,7 @@ dependencies = [ [[package]] name = "dynamo-data-gen" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "blake3", @@ -2539,7 +2548,7 @@ dependencies = [ [[package]] name = "dynamo-epp-sidecar" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-trait", @@ -2558,12 +2567,14 @@ dependencies = [ [[package]] name = "dynamo-ext-proc" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", + "async-trait", "axum 0.8.4", "bytes", "dashmap", + "dynamo-custom-policy-builtin", "dynamo-kv-router", "dynamo-llm", "dynamo-protocols", @@ -2598,7 +2609,7 @@ dependencies = [ [[package]] name = "dynamo-kv-hashing" -version = "1.5.0" +version = "1.6.0" dependencies = [ "derive_builder", "dynamo-tokens", @@ -2610,7 +2621,7 @@ dependencies = [ [[package]] name = "dynamo-kv-router" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "arc-swap", @@ -2628,6 +2639,7 @@ dependencies = [ "dynamo-truthy", "fastrand", "flume 0.12.0", + "futures-util", "indexmap 2.14.0", "once_cell", "ordered-float 4.6.0", @@ -2658,7 +2670,7 @@ dependencies = [ [[package]] name = "dynamo-llm" -version = "1.5.0" +version = "1.6.0" dependencies = [ "aligned-vec", "anyhow", @@ -2689,6 +2701,7 @@ dependencies = [ "derive_builder", "dialoguer", "dynamo-bench", + "dynamo-custom-policy-builtin", "dynamo-kv-router", "dynamo-memory", "dynamo-mocker", @@ -2777,7 +2790,7 @@ dependencies = [ [[package]] name = "dynamo-memory" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "cudarc", @@ -2796,7 +2809,7 @@ dependencies = [ [[package]] name = "dynamo-mocker" -version = "1.5.0" +version = "1.6.0" dependencies = [ "aisimulate-core", "anyhow", @@ -2804,6 +2817,7 @@ dependencies = [ "bytes", "dashmap", "derive_builder", + "dynamo-custom-policy-builtin", "dynamo-data-gen", "dynamo-kv-router", "dynamo-tokens", @@ -2830,7 +2844,7 @@ dependencies = [ [[package]] name = "dynamo-mocker-backend" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-stream", @@ -2851,9 +2865,9 @@ dependencies = [ [[package]] name = "dynamo-parsers" -version = "8.1.0" +version = "9.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e78ad3396a91e388a6ab6e4e9f66236238793076b5d0b80c37671134f54bcfa5" +checksum = "e316d4d0f5082257f1629ca61973eafc17cee3e6ee0c24f7a55a61beeff40d97" dependencies = [ "aho-corasick", "anyhow", @@ -2874,9 +2888,9 @@ dependencies = [ [[package]] name = "dynamo-parsers-v2" -version = "0.3.2" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e2277648c57af9f41943c1c6abf573ed42ddbad9d05f2e4053e9d66e9a41182" +checksum = "e3dba2a81c6382dc217fb4caa2b5b03a7e80d5bda6694a3c002ee9eeefa84336" dependencies = [ "anyhow", "num-traits", @@ -2891,9 +2905,9 @@ dependencies = [ [[package]] name = "dynamo-protocols" -version = "5.4.1" +version = "6.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dbaf3179de1accba21fcad16eb22b5e113722084265e0315e808e32150088ee7" +checksum = "9ef02e31ca224ca8db24b88a990b7db3eb2b99c77697b0d8ae44feaa32709e38" dependencies = [ "async-openai", "derive_builder", @@ -2908,9 +2922,9 @@ dependencies = [ [[package]] name = "dynamo-renderer" -version = "5.1.2" +version = "5.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ae2eaa139651c535aeaad8856c5546709608931ccd4d24d3a529f6c5233c4b6" +checksum = "06464b1deb9b6e185476e2e4280cbbbc5f16d0d6fd8732ff8b5732d43ec0567b" dependencies = [ "anyhow", "chrono", @@ -2926,7 +2940,7 @@ dependencies = [ [[package]] name = "dynamo-rl" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "axum 0.8.4", @@ -2934,6 +2948,8 @@ dependencies = [ "futures", "serde", "serde_json", + "serial_test", + "temp-env", "tokio", "tracing", "url", @@ -2941,7 +2957,7 @@ dependencies = [ [[package]] name = "dynamo-runtime" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "arc-swap", @@ -2978,6 +2994,7 @@ dependencies = [ "log", "lru 0.16.4", "mio 1.2.2", + "nix 0.30.1", "notify", "nuid", "once_cell", @@ -3026,7 +3043,7 @@ dependencies = [ [[package]] name = "dynamo-sglang-mocker" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-stream", @@ -3046,7 +3063,7 @@ dependencies = [ [[package]] name = "dynamo-sglang-sidecar" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-stream", @@ -3070,7 +3087,7 @@ dependencies = [ [[package]] name = "dynamo-sidecar-common" -version = "1.5.0" +version = "1.6.0" dependencies = [ "clap", "dynamo-backend-common", @@ -3084,9 +3101,9 @@ dependencies = [ [[package]] name = "dynamo-tokenizers" -version = "1.8.1" +version = "1.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed232a9d254149a4e2e14f0a7d77f59bd64161964c0ba55267beea3955650faa" +checksum = "cd9116db4a17d4ac186a13719b1eea8b41669cec2d95c8c93363869c9163c8aa" dependencies = [ "aho-corasick", "anyhow", @@ -3107,7 +3124,7 @@ dependencies = [ [[package]] name = "dynamo-tokens" -version = "1.5.0" +version = "1.6.0" dependencies = [ "bs58", "bytemuck", @@ -3125,7 +3142,7 @@ dependencies = [ [[package]] name = "dynamo-trtllm-sidecar" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-stream", @@ -3146,21 +3163,24 @@ dependencies = [ [[package]] name = "dynamo-truthy" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", ] [[package]] name = "dynamo-vllm-mocker" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-stream", "blake3", "clap", "dynamo-backend-common", + "dynamo-kv-router", + "dynamo-llm", "dynamo-mocker", + "dynamo-runtime", "dynamo-vllm-sidecar", "futures", "prost-types 0.14.3", @@ -3175,17 +3195,22 @@ dependencies = [ [[package]] name = "dynamo-vllm-sidecar" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-stream", "async-trait", + "base64 0.22.1", "clap", "dynamo-backend-common", + "dynamo-llm", + "dynamo-runtime", "dynamo-sidecar-common", "futures", "prost-types 0.14.3", + "serde", "serde_json", + "tempfile", "tokio", "tokio-stream", "tokio-util", @@ -3197,7 +3222,7 @@ dependencies = [ [[package]] name = "dynamo-worker-selection-policy-catalog" -version = "1.5.0" +version = "1.6.0" dependencies = [ "dynamo-kv-router", ] @@ -3476,9 +3501,9 @@ dependencies = [ [[package]] name = "fastokens" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35ff8cf975d8f772e7d4c3be439659b0b23b1ae1086ebee7204573807053ff7f" +checksum = "102017e5eb15d2d2d5308f5bcb537390e38b0547174000f408d1f9d6aa62a369" dependencies = [ "daachorse", "fancy-regex 0.17.0", @@ -5322,7 +5347,7 @@ dependencies = [ [[package]] name = "kvbm-common" -version = "1.5.0" +version = "1.6.0" dependencies = [ "dynamo-tokens", "serde", @@ -5330,7 +5355,7 @@ dependencies = [ [[package]] name = "kvbm-config" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "dynamo-memory", @@ -5350,7 +5375,7 @@ dependencies = [ [[package]] name = "kvbm-consolidator" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "dynamo-kv-hashing", @@ -5375,7 +5400,7 @@ dependencies = [ [[package]] name = "kvbm-engine" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-nats", @@ -5415,7 +5440,7 @@ dependencies = [ [[package]] name = "kvbm-kernels" -version = "1.5.0" +version = "1.6.0" dependencies = [ "clap", "cudarc", @@ -5426,7 +5451,7 @@ dependencies = [ [[package]] name = "kvbm-logical" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-stream", @@ -5455,7 +5480,7 @@ dependencies = [ [[package]] name = "kvbm-physical" -version = "1.5.0" +version = "1.6.0" dependencies = [ "aligned-vec", "anyhow", @@ -5505,7 +5530,7 @@ checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "libdynamo_llm" -version = "1.5.0" +version = "1.6.0" dependencies = [ "anyhow", "async-once-cell", @@ -7779,9 +7804,9 @@ dependencies = [ [[package]] name = "quinn" -version = "0.11.9" +version = "0.11.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" +checksum = "4051e23e9185c255a7e33ef59cdbca87a22d359052eecd22fc6b901fb37d9d11" dependencies = [ "bytes", "cfg_aliases", @@ -7799,9 +7824,9 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.17" +version = "0.11.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04759210543be93709136e28212294a659ef5001836ff4eab4d663e4529bba83" +checksum = "a9746dbde176634f4f2f1faf2404e30a31b2bc1e9cafb5329c95d8177a18c9fc" dependencies = [ "bytes", "getrandom 0.4.2", @@ -11072,9 +11097,9 @@ checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1" [[package]] name = "vllm-proto" -version = "0.1.0" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520830f75f6a97b9423e5f9155fd4a6aa5578b247dea7d31adb49a5ce5f6e817" +checksum = "77383505a016d5090ecf460297c815752339669aee187e29111baf444cd6e171" dependencies = [ "prost 0.14.3", "prost-types 0.14.3", diff --git a/Cargo.toml b/Cargo.toml index e76399a77681..721f8de1beb8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,7 +45,7 @@ members = [ resolver = "3" [workspace.package] -version = "1.5.0" +version = "1.6.0" edition = "2024" description = "Dynamo Inference Framework" authors = ["NVIDIA Inc. "] @@ -55,42 +55,42 @@ repository = "https://github.com/ai-dynamo/dynamo.git" keywords = ["llm", "genai", "inference", "nvidia", "distributed"] [workspace.dependencies] +dynamo-custom-policy-builtin = { path = "lib/router-plugins/builtin", version = "1.6.0" } # Published separately by https://github.com/ai-dynamo/aisimulate. -aisimulate-core = { version = "=0.12.0-dev.1" } +aisimulate-core = { version = "=0.12.0" } # Local crates -dynamo-runtime = { path = "lib/runtime", version = "1.5.0" } -dynamo-llm = { path = "lib/llm", version = "1.5.0" } -dynamo-rl = { path = "lib/rl", version = "1.5.0" } -dynamo-tokenizers = { version = "=1.8.1" } -dynamo-renderer = { version = "=5.1.2" } -dynamo-tokens = { path = "lib/tokens", version = "1.5.0" } -dynamo-truthy = { path = "lib/truthy", version = "1.5.0" } -dynamo-data-gen = { path = "lib/data-gen", version = "1.5.0" } -dynamo-kv-hashing = { path = "lib/kv-hashing", version = "1.5.0" } -dynamo-memory = { path = "lib/memory", version = "1.5.0" } -dynamo-mocker = { path = "lib/mocker", version = "1.5.0" } -dynamo-kv-router = { path = "lib/kv-router", version = "1.5.0" } -dynamo-protocols = { version = "=5.4.1" } -dynamo-parsers = { version = "8.1.0" } +dynamo-runtime = { path = "lib/runtime", version = "1.6.0" } +dynamo-llm = { path = "lib/llm", version = "1.6.0" } +dynamo-rl = { path = "lib/rl", version = "1.6.0" } +dynamo-tokenizers = { version = "=1.9.1" } +dynamo-renderer = { version = "=5.4.0" } +dynamo-tokens = { path = "lib/tokens", version = "1.6.0" } +dynamo-truthy = { path = "lib/truthy", version = "1.6.0" } +dynamo-data-gen = { path = "lib/data-gen", version = "1.6.0" } +dynamo-kv-hashing = { path = "lib/kv-hashing", version = "1.6.0" } +dynamo-memory = { path = "lib/memory", version = "1.6.0" } +dynamo-mocker = { path = "lib/mocker", version = "1.6.0" } +dynamo-kv-router = { path = "lib/kv-router", version = "1.6.0" } +dynamo-protocols = { version = "=6.1.0" } +dynamo-parsers = { version = "9.0.0" } # Published on crates.io. To see all available versions: # web: https://crates.io/crates/dynamo-parsers-v2/versions # sparse index: curl https://index.crates.io/dy/na/dynamo-parsers-v2 (one JSON line per release; read the "vers" field) -dynamo-parsers-v2 = { version = "0.3.2" } -dynamo-backend-common = { path = "lib/backend-common", version = "1.5.0" } -dynamo-sidecar-common = { path = "lib/sidecar/common", version = "1.5.0" } -fastokens = { version = "0.2.0" } +dynamo-parsers-v2 = { version = "0.5.3" } +dynamo-backend-common = { path = "lib/backend-common", version = "1.6.0" } +dynamo-sidecar-common = { path = "lib/sidecar/common", version = "1.6.0" } # kvbm -kvbm-common = { path = "lib/kvbm-common", version = "1.5.0" } -kvbm-config = { path = "lib/kvbm-config", version = "1.5.0" } -kvbm-consolidator = { path = "lib/kvbm-consolidator", version = "1.5.0" } -kvbm-engine = { path = "lib/kvbm-engine", version = "1.5.0" } -kvbm-kernels = { path = "lib/kvbm-kernels", version = "1.5.0" } -kvbm-logical = { path = "lib/kvbm-logical", version = "1.5.0" } -kvbm-physical = { path = "lib/kvbm-physical", version = "1.5.0" } +kvbm-common = { path = "lib/kvbm-common", version = "1.6.0" } +kvbm-config = { path = "lib/kvbm-config", version = "1.6.0" } +kvbm-consolidator = { path = "lib/kvbm-consolidator", version = "1.6.0" } +kvbm-engine = { path = "lib/kvbm-engine", version = "1.6.0" } +kvbm-kernels = { path = "lib/kvbm-kernels", version = "1.6.0" } +kvbm-logical = { path = "lib/kvbm-logical", version = "1.6.0" } +kvbm-physical = { path = "lib/kvbm-physical", version = "1.6.0" } # velo velo = { version = "0.12.0" } @@ -199,11 +199,11 @@ tonic-build = { version = "0.13.1" } prost-types-v14 = { package = "prost-types", version = "0.14.3" } tonic-v14 = { package = "tonic", version = "0.14.6" } tonic-health-v14 = { package = "tonic-health", version = "0.14.6" } -vllm-proto = { version = "=0.1.0" } +vllm-proto = { version = "=0.3.0" } kube = { version = "2.0.1", default-features = false, features = ["client", "rustls-tls", "aws-lc-rs", "runtime"] } k8s-openapi = { version = "0.26.0", features = ["v1_32"] } rcgen = { version = "0.13" } -quinn = { version = "0.11.9", default-features = false, features = ["runtime-tokio", "rustls-ring"] } +quinn = { version = "0.11.12", default-features = false, features = ["runtime-tokio", "rustls-ring"] } sha2 = { version = "0.10" } tokio-rustls = { version = "0.26" } rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } diff --git a/README.md b/README.md index 2b54b2fd1202..385c058c9f5a 100644 --- a/README.md +++ b/README.md @@ -42,8 +42,8 @@ Built in Rust for performance, Python for extensibility. | Date | Event | Location | |:-----|:------|:---------| +| Wed, Sep 23, 2026 | ~~[Dynamo community meeting](https://calendar.google.com/calendar/render?action=TEMPLATE&text=Dynamo+community+meeting&dates=20260923T173000Z%2F20260923T181500Z&location=https%3A%2F%2Fmeet.google.com%2Fheb-demu-qok)~~ | [Online](https://meet.google.com/heb-demu-qok) | | Thu, Sep 10, 2026 | ~~[Baseten x Dynamo x SGLang RL post training meetup](https://calendar.google.com/calendar/render?action=TEMPLATE&text=Baseten+x+Dynamo+x+SGLang+RL+post+training+meetup&dates=20260911T010000Z%2F20260911T040000Z&location=https%3A%2F%2Fluma.com%2FBaseDynSGL)~~ | [Luma](https://luma.com/BaseDynSGL) | -| Mon, Aug 24, 2026 | ~~[vLLM x Dynamo meetup](https://calendar.google.com/calendar/render?action=TEMPLATE&text=vLLM+x+Dynamo+meetup&dates=20260825T010000Z%2F20260825T040000Z&location=https%3A%2F%2Fluma.com%2Fr8o604o0)~~ | [Luma](https://luma.com/r8o604o0) | > Events are updated automatically. Subscribe to our [public calendar](https://calendar.google.com/calendar/embed?src=c_c2448d2efb09eac2ddee1f34524124135bd3f4554868769059105e18e1b97e8f%40group.calendar.google.com). @@ -138,15 +138,37 @@ the Gateway API setup, supported features, and configuration. ### Option A: Container (fastest) +Choose either vLLM or SGLang: + +#### vLLM + +```bash +# Pull a prebuilt container +docker run --gpus all --network host --rm -it nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.5.0 + +# Inside the container โ€” start frontend and worker +python3 -m dynamo.frontend --http-port 8000 --discovery-backend file > /dev/null 2>&1 & +python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --discovery-backend file & + +# Once the worker is ready, send a request +curl -s localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{ + "model": "Qwen/Qwen3-0.6B", + "messages": [{"role": "user", "content": "Hello!"}], + "max_tokens": 100 +}' | jq +``` + +#### SGLang + ```bash -# Pull a prebuilt container (SGLang example) -docker run --gpus all --network host --rm -it nvcr.io/nvidia/ai-dynamo/sglang-runtime:1.4.2 +# Pull a prebuilt container +docker run --gpus all --network host --rm -it nvcr.io/nvidia/ai-dynamo/sglang-runtime:1.5.0 # Inside the container โ€” start frontend and worker python3 -m dynamo.frontend --http-port 8000 --discovery-backend file > /dev/null 2>&1 & python3 -m dynamo.sglang --model-path Qwen/Qwen3-0.6B --discovery-backend file & -# Send a request +# Once the worker is ready, send a request curl -s localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{ "model": "Qwen/Qwen3-0.6B", "messages": [{"role": "user", "content": "Hello!"}], @@ -154,19 +176,27 @@ curl -s localhost:8000/v1/chat/completions -H "Content-Type: application/json" - }' | jq ``` -Also available: [`tensorrtllm-runtime:1.4.2`](https://docs.nvidia.com/dynamo/resources/release-artifacts) and [`vllm-runtime:1.4.2`](https://docs.nvidia.com/dynamo/resources/release-artifacts). +Also available: [`tensorrtllm-runtime:1.5.0`](https://docs.nvidia.com/dynamo/resources/release-artifacts). ### Option B: Install from PyPI -Install [uv](https://github.com/astral-sh/uv) (`curl -LsSf https://astral.sh/uv/install.sh | sh`), then: +Install [uv](https://github.com/astral-sh/uv) (`curl -LsSf https://astral.sh/uv/install.sh | sh`), then choose your backend: + +#### vLLM + +```bash +uv pip install --prerelease=allow "ai-dynamo[vllm]" +``` + +#### SGLang ```bash -uv pip install --prerelease=allow "ai-dynamo[sglang]" # or [vllm] +uv pip install --prerelease=allow "ai-dynamo[sglang]" ``` > **Note:** TensorRT-LLM requires `pip` with `--extra-index-url https://pypi.nvidia.com`. See the [install guide](docs/fern/pages/cli/installation/install-dynamo.mdx) for TRT-LLM-specific instructions. -Then start the frontend and a worker as shown above. See the [full installation guide](docs/fern/pages/cli/installation/install-dynamo.mdx) for system dependencies and backend-specific notes. +Then start the frontend and the matching worker as shown above. See the [full installation guide](docs/fern/pages/cli/installation/install-dynamo.mdx) for system dependencies and backend-specific notes. ### Option C: Kubernetes (recommended) diff --git a/README.zh-CN.md b/README.zh-CN.md index 038ce35006c6..619787cc9d45 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -113,7 +113,7 @@ limitations under the License. ```bash # ๆ‹‰ๅ–้ข„ๆž„ๅปบๅฎนๅ™จ๏ผˆSGLang ็คบไพ‹๏ผ‰ -docker run --gpus all --network host --rm -it nvcr.io/nvidia/ai-dynamo/sglang-runtime:1.4.0 +docker run --gpus all --network host --rm -it nvcr.io/nvidia/ai-dynamo/sglang-runtime:1.5.0 # ๅœจๅฎนๅ™จๅ†…ๅฏๅŠจ frontend ๅ’Œ worker python3 -m dynamo.frontend --http-port 8000 --discovery-backend file > /dev/null 2>&1 & @@ -127,7 +127,7 @@ curl -s localhost:8000/v1/chat/completions -H "Content-Type: application/json" - }' | jq ``` -ๅฆๆœ‰ [`tensorrtllm-runtime:1.4.0`](https://docs.nvidia.com/dynamo/resources/release-artifacts) ๅ’Œ [`vllm-runtime:1.4.0`](https://docs.nvidia.com/dynamo/resources/release-artifacts) ๅฏ็”จใ€‚ +ๅฆๆœ‰ [`tensorrtllm-runtime:1.5.0`](https://docs.nvidia.com/dynamo/resources/release-artifacts) ๅ’Œ [`vllm-runtime:1.5.0`](https://docs.nvidia.com/dynamo/resources/release-artifacts) ๅฏ็”จใ€‚ ### ้€‰้กน B๏ผšไปŽ PyPI ๅฎ‰่ฃ… diff --git a/benchmarks/pyproject.toml b/benchmarks/pyproject.toml index 221802fd115e..612b122a3011 100644 --- a/benchmarks/pyproject.toml +++ b/benchmarks/pyproject.toml @@ -39,7 +39,7 @@ classifiers = [ ] dependencies = [ - "aisimulate==0.12.0.dev1; python_version >= '3.11' and python_version < '3.14'", + "aisimulate==0.12.0; python_version >= '3.11' and python_version < '3.14'", "aiperf==0.12.0", "matplotlib", "networkx", diff --git a/benchmarks/router/common.py b/benchmarks/router/common.py index ae25d805fbcc..eae8872aaff1 100755 --- a/benchmarks/router/common.py +++ b/benchmarks/router/common.py @@ -10,8 +10,6 @@ import logging import os -from prefix_data_generator.synthesizer import Synthesizer - # Default values DEFAULT_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B" DEFAULT_URL = "http://localhost:8000" @@ -326,6 +324,10 @@ def prepare_trace_dataset(args, output_dir, logger): ) logger.info(f" Random seed: {args.seed}") + # Synthetic generation has optional dependencies absent from slim CPU test images. + # Keep the request helpers and trace-file path usable without those dependencies. + from prefix_data_generator.synthesizer import Synthesizer + synthesizer = Synthesizer( args.input_dataset, block_size=args.block_size, diff --git a/benchmarks/router/tests/test_aiperf_commands.py b/benchmarks/router/tests/test_aiperf_commands.py new file mode 100644 index 000000000000..7e397a0680fe --- /dev/null +++ b/benchmarks/router/tests/test_aiperf_commands.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression coverage for trace command arguments (NVBug 6414352).""" + +import sys + +import pytest + +from benchmarks.router import common + +pytestmark = [ + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.unit, + pytest.mark.parallel, +] + + +@pytest.fixture +def command_args(tmp_path): + return { + "model": "test-model", + "tokenizer": "test-tokenizer", + "input_dataset": str(tmp_path / "input trace.jsonl"), + "artifact_dir": str(tmp_path / "aiperf artifacts"), + "seed": 17, + "url": "http://router.test", + } + + +@pytest.fixture +def agent_command_builder(monkeypatch): + # The standalone script imports common by its unqualified module name. + # Import inside the fixture so the temporary alias exists first and is + # restored at teardown without changing imports for unrelated scripts. + monkeypatch.setitem(sys.modules, "common", common) + from benchmarks.router import agent_benchmark + + return agent_benchmark.get_aiperf_cmd + + +def assert_trace_command(cmd, args): + assert cmd[:2] == ["aiperf", "profile"] + flags = {token.split("=", 1)[0] for token in cmd if token.startswith("--")} + assert "--prompt-input-tokens-block-size" not in flags + assert "--isl-block-size" not in flags + assert not any(flag.startswith("--synthetic-input-tokens-") for flag in flags) + + expected_values = { + "--model": args["model"], + "--tokenizer": args["tokenizer"], + "--url": args["url"], + "--input-file": args["input_dataset"], + "--custom-dataset-type": "mooncake_trace", + "--random-seed": str(args["seed"]), + "--artifact-dir": args["artifact_dir"], + } + for flag, value in expected_values.items(): + assert cmd.count(flag) == 1 + index = cmd.index(flag) + assert cmd[index : index + 2] == [flag, value] + + # The block size is unused by the other controlled inputs. + assert str(args["block_size"]) not in cmd + assert "None" not in cmd + common_flags = common.get_common_aiperf_flags() + assert cmd[-len(common_flags) :] == common_flags + + +def test_fixed_schedule_trace_command_omits_block_size(command_args): + args = {**command_args, "block_size": 512} + cmd = common.get_aiperf_cmd_for_trace( + args["model"], + args["tokenizer"], + args["input_dataset"], + args["artifact_dir"], + args["seed"], + args["block_size"], + args["url"], + ) + + assert_trace_command(cmd, args) + assert cmd.count("--fixed-schedule") == 1 + assert cmd.count("--fixed-schedule-auto-offset") == 1 + + +def test_agent_trace_command_omits_block_size(command_args, agent_command_builder): + args = {**command_args, "block_size": 64, "concurrency": 3, "request_count": 23} + cmd = agent_command_builder( + args["model"], + args["tokenizer"], + args["input_dataset"], + args["artifact_dir"], + args["seed"], + args["concurrency"], + args["block_size"], + args["request_count"], + args["url"], + ) + + assert_trace_command(cmd, args) + for flag, value in ( + ("--concurrency", str(args["concurrency"])), + ("--request-count", str(args["request_count"])), + ): + assert cmd.count(flag) == 1 + index = cmd.index(flag) + assert cmd[index : index + 2] == [flag, value] diff --git a/benchmarks/router/tests/test_common.py b/benchmarks/router/tests/test_common.py index 4945b9599d29..903ff95627d1 100644 --- a/benchmarks/router/tests/test_common.py +++ b/benchmarks/router/tests/test_common.py @@ -3,7 +3,14 @@ import pytest -from benchmarks.router.common import add_expected_osl, tag_requests_with_priority +# Collection precedes marker filtering, including in the lean Planner image +# which does not install the router benchmark's graph-synthesis dependency. +pytest.importorskip("networkx", reason="router benchmark synthesis needs networkx") + +from benchmarks.router.common import ( # noqa: E402 -- optional dependency check + add_expected_osl, + tag_requests_with_priority, +) pytestmark = [ pytest.mark.pre_merge, diff --git a/components/src/dynamo/common/backend/agent_context.py b/components/src/dynamo/common/backend/agent_context.py new file mode 100644 index 000000000000..1c39b00bc17b --- /dev/null +++ b/components/src/dynamo/common/backend/agent_context.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Helpers for reading agent context from backend requests.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + + +def agent_context_from_request( + request: Mapping[str, Any], +) -> Mapping[str, Any] | None: + """Return a request's serialized agent context when it is object-shaped.""" + agent_context = request.get("agent_context") + return agent_context if isinstance(agent_context, Mapping) else None + + +def session_id_from_request(request: Mapping[str, Any]) -> str | None: + """Return the request's non-empty ``agent_context.session_id``.""" + agent_context = agent_context_from_request(request) + if agent_context is None: + return None + session_id = agent_context.get("session_id") + if not isinstance(session_id, str): + return None + session_id = session_id.strip() + return session_id or None diff --git a/components/src/dynamo/common/backend/engine.py b/components/src/dynamo/common/backend/engine.py index 57b69287f8e3..476f54ca1035 100644 --- a/components/src/dynamo/common/backend/engine.py +++ b/components/src/dynamo/common/backend/engine.py @@ -44,11 +44,13 @@ class GenerateRequest(TypedDict, total=False): Multimodal keys (``multi_modal_data``, ``mm_processor_kwargs``, ``mm_routing_info``) are populated by the frontend preprocessor when - the request carries media. ``encoder_result`` is set by the - frontend when forwarding a request from an Encode worker - to a downstream Prefill/Aggregated peer; engines read it via - :func:`dynamo.common.backend.multimodal.require_encoder_result`. All - four are object-shaped (``dict``) by contract. + the request carries media. ``encoder_result`` may be set by the frontend + when forwarding a request from an Encode worker or by an application + orchestrator before calling an Aggregated peer; engines read it via + :func:`dynamo.common.backend.multimodal.require_encoder_result`. + ``encoder_result`` is deliberately opaque at this common layer; its + producer and consuming engine own the concrete schema. All four are + object-shaped (``dict``) by contract. ``model`` carries the requested model name (set by the Rust preprocessor). Engines that support dynamic LoRA read it to route a @@ -132,6 +134,7 @@ class LlmRegistration: bootstrap_host: Optional[str] = None bootstrap_port: Optional[int] = None enable_eagle: bool = False + max_gpu_lora_count: Optional[int] = None @dataclass diff --git a/components/src/dynamo/common/backend/tests/test_backend_bindings.py b/components/src/dynamo/common/backend/tests/test_backend_bindings.py index 6edd9aff4aa4..714922192b92 100644 --- a/components/src/dynamo/common/backend/tests/test_backend_bindings.py +++ b/components/src/dynamo/common/backend/tests/test_backend_bindings.py @@ -88,6 +88,7 @@ def test_engine_config_full_kwargs_round_trip_through_getters(): max_num_seqs=64, max_num_batched_tokens=2048, enable_eagle=True, + max_gpu_lora_count=4, ), ) assert cfg.model == "m2" @@ -101,6 +102,7 @@ def test_engine_config_full_kwargs_round_trip_through_getters(): assert llm.max_num_seqs == 64 assert llm.max_num_batched_tokens == 2048 assert llm.enable_eagle is True + assert llm.max_gpu_lora_count == 4 @pytest.mark.unified @@ -109,6 +111,12 @@ def test_llm_registration_preserves_legacy_positional_arguments(): assert llm.bootstrap_host == "host" assert llm.bootstrap_port == 9000 assert llm.enable_eagle is False + assert llm.max_gpu_lora_count is None + extended = backend.LlmRegistration( + 2048, 16, 1000, 64, 2048, 2, 1, "host", 9000, True, 8 + ) + assert extended.enable_eagle is True + assert extended.max_gpu_lora_count == 8 @pytest.mark.parametrize("kwargs", [{}, {"enable_eagle": True}]) diff --git a/components/src/dynamo/common/configuration/groups/http_args.py b/components/src/dynamo/common/configuration/groups/http_args.py index 1ad1f74e31b8..c0de78d6481a 100644 --- a/components/src/dynamo/common/configuration/groups/http_args.py +++ b/components/src/dynamo/common/configuration/groups/http_args.py @@ -22,10 +22,6 @@ import os from typing import Optional -# typing.Self is 3.11+; pyproject declares requires-python>=3.10 so use the -# typing_extensions back-port (added as an explicit dep in pyproject.toml). -from typing_extensions import Self - from dynamo.common.configuration.arg_group import ArgGroup from dynamo.common.configuration.config_base import ConfigBase from dynamo.common.configuration.utils import add_argument, nullable_float @@ -47,6 +43,13 @@ "DYN_HTTP_CONCURRENCY": ("DYN_MM_HTTP_CONCURRENCY",), } +# Canonical names that are now accepted-but-ignored (they configured the removed +# httpx backend). Their legacy aliases warn that they're ignored rather than +# pointing operators at a live knob to migrate to. +_IGNORED_HTTP_KNOBS = frozenset( + {"DYN_HTTP_MAX_KEEPALIVE", "DYN_HTTP_POOL_TIMEOUT", "DYN_HTTP_CONCURRENCY"} +) + _legacy_warned: set[str] = set() @@ -68,7 +71,16 @@ def _apply_legacy_env_aliases() -> None: os.environ[canonical] = value if legacy not in _legacy_warned: _legacy_warned.add(legacy) - logger.warning("%s is deprecated; use %s instead.", legacy, canonical) + if canonical in _IGNORED_HTTP_KNOBS: + logger.warning( + "%s is deprecated and ignored; it configured the removed " + "httpx backend.", + legacy, + ) + else: + logger.warning( + "%s is deprecated; use %s instead.", legacy, canonical + ) break @@ -103,22 +115,15 @@ class HttpConfigBase(ConfigBase): # aiohttp: ``ClientTimeout.sock_connect`` connect_timeout: float - # --- httpx-only ------------------------------------------------------- - - # Cap on idle keepalive connections kept warm in the pool. Raising - # this to match ``max_connections`` prevents TLS re-handshake churn - # under fan-out. Maps to ``Limits.max_keepalive_connections``. + # --- Deprecated no-ops (retained for backward compatibility) ---------- + # These configured the removed httpx backend and have no aiohttp + # equivalent that we apply: aiohttp's connector queues natively (no + # semaphore/pool-timeout) and idle keepalive is governed by time + # (``keepalive_timeout``) rather than a connection count. The flags and + # env vars are still accepted so existing configs don't break, but the + # values are ignored. max_keepalive: int - - # Wait-for-free-slot timeout; ``Timeout.pool``. Decoupled from the - # read budget so a saturated pool surfaces quickly. pool_timeout: float - - # Process-wide cap on concurrent in-flight HTTP fetches via the - # httpx backend. The semaphore acts as backpressure in front of the - # pool so a burst of requests can't push ``PoolTimeout`` up the - # stack. aiohttp has no equivalent because its connector queues - # natively. concurrency: int # --- aiohttp-only ----------------------------------------------------- @@ -127,17 +132,6 @@ class HttpConfigBase(ConfigBase): # ``TCPConnector.keepalive_timeout``. keepalive_timeout: float - @classmethod - def from_cli_args(cls, args: argparse.Namespace) -> Self: - config = super().from_cli_args(args) - # ``--http-max-keepalive 0`` is the documented sentinel for - # "match max_connections" โ€” resolve here so the httpx backend - # always gets a non-zero keepalive cap regardless of which - # construction path the operator used. - if config.max_keepalive == 0: - config.max_keepalive = config.max_connections - return config - class HttpArgGroup(ArgGroup): """CLI / env-var registration for the HTTP fetch client.""" @@ -152,7 +146,7 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: default=100, arg_type=int, dest="max_connections", - help="Total pool size cap (httpx Limits.max_connections / aiohttp TCPConnector.limit).", + help="Pool size cap per connect-time policy (aiohttp TCPConnector.limit). The aiohttp client keeps one pool per connect-time policy outcome, at most two, so a deployment that sets DYN_MM_ALLOW_INTERNAL=1 and also issues stricter per-request policies can reach twice this value in total.", ) add_argument( g, @@ -183,12 +177,13 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: g, flag_name="--http-max-keepalive", env_var="DYN_HTTP_MAX_KEEPALIVE", - default=0, # 0 โ†’ match max_connections; resolved in from_env. + default=0, arg_type=int, dest="max_keepalive", help=( - "[httpx-only] Cap on idle keepalive connections in the pool. " - "0 โ†’ match --http-max-connections." + "[deprecated, no-op] Configured the removed httpx backend; " + "aiohttp governs keepalive by time (--http-keepalive-timeout). " + "Accepted but ignored." ), ) add_argument( @@ -198,7 +193,10 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: default=60.0, arg_type=float, dest="pool_timeout", - help="[httpx-only] Wait-for-free-slot timeout (seconds).", + help=( + "[deprecated, no-op] Configured the removed httpx backend; " + "aiohttp's connector queues natively. Accepted but ignored." + ), ) add_argument( g, @@ -208,8 +206,9 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: arg_type=int, dest="concurrency", help=( - "[httpx-only] Process-wide cap on concurrent in-flight fetches. " - "Acts as backpressure in front of the pool." + "[deprecated, no-op] Configured the removed httpx backend; " + "aiohttp's connector queues natively, so no front semaphore is " + "needed. Accepted but ignored." ), ) add_argument( diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 1a43c56a900a..0fc85150cfec 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -57,6 +57,7 @@ "router_queue_policy", "use_remote_indexer", "serve_indexer", + "enable_session_prefix_index", "shared_cache_multiplier", "shared_cache_type", "conditional_disagg_enabled", @@ -218,6 +219,7 @@ class KvRouterConfigBase(ConfigBase): router_queue_policy: str use_remote_indexer: bool = False serve_indexer: bool = False + enable_session_prefix_index: bool = False shared_cache_multiplier: float = 0.0 shared_cache_type: str = "none" conditional_disagg_enabled: bool = False @@ -293,11 +295,12 @@ def add_arguments(self, parser) -> None: default=False, dest="load_aware", help=( - "KV Router: Enable load-aware routing without cache-reuse signals. " + "KV Router: Enable the load-aware routing preset. " "On the frontend, this implies --router-mode kv. " "This preset sets overlap_score_credit=0, disables KV events and " "KV-reuse assumptions, enables active-block " - "and prefill-token load tracking, and disables remote/shared cache indexers." + "and prefill-token load tracking, and disables remote/shared cache indexers. " + "The builtin policy does not request cache inputs in this mode." ), ) add_argument( @@ -657,6 +660,18 @@ def add_arguments(self, parser) -> None: ), dest="use_remote_indexer", ) + add_negatable_bool_argument( + g, + flag_name="--enable-session-prefix-index", + env_var="DYN_ENABLE_SESSION_PREFIX_INDEX", + default=False, + help=( + "[EXPERIMENTAL] KV Router: Track per-session block lineage in a logical " + "prefix index that outlives engine cache eviction. Nothing consumes the " + "index for routing decisions yet." + ), + dest="enable_session_prefix_index", + ) add_argument( g, flag_name="--shared-cache-multiplier", diff --git a/components/src/dynamo/common/configuration/utils.py b/components/src/dynamo/common/configuration/utils.py index 79c07d33a28a..af2dae1d8a38 100644 --- a/components/src/dynamo/common/configuration/utils.py +++ b/components/src/dynamo/common/configuration/utils.py @@ -160,7 +160,7 @@ def add_negatable_bool_argument( *, flag_name: str, env_var: str, - default: bool, + default: Optional[bool], help: str, dest: Optional[str] = None, obsolete_flag: Optional[str] = None, @@ -173,7 +173,7 @@ def add_negatable_bool_argument( parser: ArgumentParser or argument group flag_name: Primary flag (must start with '--', e.g. "--enable-feature") env_var: Environment variable name (e.g., "DYN_ENABLE_FEATURE") - default: Default value + default: Default value. Use None when an omitted flag must remain unset. help: Help text dest: Optional destination name for the parsed value obsolete_flag: Optional obsolete/legacy flag (for help msg only, must start with '--') diff --git a/components/src/dynamo/common/http/README.md b/components/src/dynamo/common/http/README.md index de026f13decc..f8c66cfe7168 100644 --- a/components/src/dynamo/common/http/README.md +++ b/components/src/dynamo/common/http/README.md @@ -1,72 +1,57 @@ # `dynamo.common.http` -HTTP fetch client. Backend-neutral facade -(`fetch_bytes` / `close_http_client`) over an `HttpClient` ABC with -two concrete subclasses: `AiohttpClient` (default) and `HttpxClient`. -Backend selection: `DYN_HTTP_BACKEND={aiohttp,httpx}`. - -## Why aiohttp is the default - -Under high concurrency (e.g. one request fanning out to 100 image -URLs) the httpx backend hits `httpx.PoolTimeout`. aiohttp scales -markedly better on the same workload โ€” see the +HTTP fetch client: a facade (`fetch_bytes` / `close_http_client`) over +an `HttpClient` ABC with a single concrete subclass, `AiohttpClient` +(over `aiohttp.ClientSession`). `DYN_HTTP_BACKEND` is retained for +back-compat but only `aiohttp` is supported; any other value warns and +uses aiohttp. + +## Why aiohttp + +`AiohttpClient` is the single supported backend. Its connector queues +pending connections in `O(1)`, so latency stays close to the offered +rate when one request fans out to many URLs (e.g. 100 image fetches), +and it exposes a `TCPConnector(resolver=...)` DNS hook that pins the +validated DNS answers as a connect-time SSRF backstop against DNS +rebinding โ€” the default client wires a `BlocklistResolver` (see +`_ssrf_resolver.py`). The connector may return private addresses only when the +`DYN_MM_ALLOW_INTERNAL` deployment baseline **and** the request policy both +allow it, so neither side alone can weaken the other. The client keeps one +session per outcome and closes each resolver itself, because aiohttp closes +only a resolver it created. + +> [!IMPORTANT] +> The backstop governs **direct** connections only. When a proxy applies +> (`HTTP_PROXY` / `HTTPS_PROXY` โ€” the session runs `trust_env=True`), the +> connector dials the proxy and the *proxy* resolves the origin, out of this +> resolver's sight, so the check cannot govern the destination. +> +> A policy-protected fetch that a proxy would carry therefore **fails closed**. +> Set `DYN_MM_TRUST_EGRESS_PROXY=1` to assert that the proxy enforces +> destination policy itself, and the fetch proceeds. The gate asks aiohttp +> which proxy applies to that specific URL, so `NO_PROXY` is honored and a +> fetch that goes direct is never refused. It does not apply when +> `DYN_MM_ALLOW_INTERNAL=1`, which already permits private destinations. +> +> Note aiohttp never calls a resolver for an IP literal, so literal blocked +> addresses are `validate_url`'s job rather than the backstop's. + +See the [NeMo Gym aiohttp vs httpx note](https://docs.nvidia.com/nemo/gym/latest/infrastructure/engineering-notes/aiohttp-vs-httpx.html) -and [openai-python#1596](https://github.com/openai/openai-python/issues/1596). - -Root cause is in httpx's pool. The maintenance routine in -[`httpcore._async.connection_pool`](https://github.com/encode/httpcore/blob/master/httpcore/_async/connection_pool.py#L303-L309) -runs *"whenever a new request is added or removed from the pool"* and -is `O(queue_size ร— pool_size)` per call, so cost grows quadratically -with backlog. aiohttp's connector queues natively in `O(1)`, which is -why our httpx backend needed a process-wide semaphore -(`DYN_HTTP_CONCURRENCY`) in front of the pool to keep -`PoolTimeout` from leaking up the stack. That semaphore is redundant -under aiohttp. - -### Benchmark - -500 rps ร— 10k requests across server-processing-time buckets (lower -is better, all latencies in ms): - -``` -== request_rate=500 rps requests=10000 == - -[sweep] mean_ms=50 order=['aiohttp', 'httpx'] -backend wall(s) avg p50 p90 p99 -httpx 20.1 75.7 51.0 158.7 241.1 -aiohttp 20.1 50.7 50.6 50.8 51.1 - -[sweep] mean_ms=100 order=['httpx', 'aiohttp'] -httpx 23.2 1550.4 1484.3 2430.2 3164.9 -aiohttp 20.1 101.4 100.6 100.9 101.2 - -[sweep] mean_ms=200 order=['aiohttp', 'httpx'] -httpx 43.2 11833.2 11762.6 21167.1 23025.2 -aiohttp 20.8 320.5 231.2 733.3 830.3 - -[sweep] mean_ms=300 order=['httpx', 'aiohttp'] -httpx 62.2 21229.8 21318.0 37834.0 41728.4 -aiohttp 32.2 6355.0 6395.1 11119.6 12071.7 -``` - -Both backends are equivalent below saturation (`mean_ms=50`). Past -saturation httpx degrades super-linearly while aiohttp stays close to -the offered rate. - -Reproduce: +for the fan-out latency comparison. -```bash -python -m benchmarks.multimodal.http.sweep \ - --server-processing-time-means-ms 50,100,200,300 \ - --request-rate 500 \ - --requests 10000 \ - --timeout 60 -``` +> [!NOTE] +> **Deprecated:** `DYN_HTTP_BACKEND` now accepts only `aiohttp` (any other +> value warns and falls back). The httpx-only knobs โ€” `DYN_HTTP_MAX_KEEPALIVE`, +> `DYN_HTTP_POOL_TIMEOUT`, and `DYN_HTTP_CONCURRENCY` (and their `--http-*` +> flags) โ€” are still accepted for backward compatibility but **ignored**; +> aiohttp consumes none of them. They configured a second HTTP backend that +> has been removed. ## Operator-tunable knobs See [`http_args.py`](../configuration/groups/http_args.py) for the full `DYN_HTTP_*` env-var / `--http-*` CLI-flag reference (pool size, -per-call timeout override, httpx-only semaphore concurrency, aiohttp -keepalive, etc.). Legacy `DYN_MM_HTTP_*` env vars are still honored. +per-call timeout override, aiohttp keepalive, etc.). Legacy +`DYN_MM_HTTP_*` env vars are still honored. diff --git a/components/src/dynamo/common/http/__init__.py b/components/src/dynamo/common/http/__init__.py index e28d1ffee93a..a8eed9266746 100644 --- a/components/src/dynamo/common/http/__init__.py +++ b/components/src/dynamo/common/http/__init__.py @@ -1,25 +1,26 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Backend-neutral facade for Dynamo's HTTP fetch client. - -Env var ``DYN_HTTP_BACKEND`` selects the implementation behind the -process-wide singleton used by :func:`fetch_bytes`: - - - ``aiohttp`` (default) โ€” :class:`AiohttpClient` over an - ``aiohttp.ClientSession``. - - ``httpx`` โ€” :class:`HttpxClient` over an ``httpx.AsyncClient`` - (fixed config: decoupled pool timeout, ``max_keepalive_connections`` - matches ``max_connections``). - -Callers use :func:`fetch_bytes` and catch the unified exception classes -(``HttpTimeoutError``, ``HttpConnectionError``, ``HttpStatusError``). -For SSRF-safe fetches (e.g. from ``ImageLoader``) pass a -``UrlValidationPolicy`` โ€” :func:`fetch_bytes` then follows redirects -manually and revalidates each hop against the policy. - -Tests / advanced callers can also instantiate :class:`AiohttpClient` or -:class:`HttpxClient` directly and bypass the singleton. +"""Facade for Dynamo's HTTP fetch client (aiohttp). + +:func:`fetch_bytes` runs over a process-wide :class:`AiohttpClient` singleton. +Callers catch the unified exception classes (``HttpTimeoutError``, +``HttpConnectionError``, ``HttpStatusError``). For SSRF-safe fetches (e.g. from +``ImageLoader``) pass a ``UrlValidationPolicy`` โ€” :func:`fetch_bytes` then +follows redirects manually and revalidates each hop against the policy. + +``DYN_HTTP_BACKEND`` accepts only ``aiohttp``; any other value logs a warning +and uses aiohttp. aiohttp scales well under fan-out and exposes a +``TCPConnector(resolver=...)`` hook. The default client wires a +``BlocklistResolver`` (``_ssrf_resolver.py``) into it that pins validated DNS +answers at connect time โ€” the SSRF backstop against DNS rebinding. The +connector may return private addresses only when the ``DYN_MM_ALLOW_INTERNAL`` +deployment baseline and the request policy both allow it. It governs +**direct** connections: when a proxy applies, the proxy resolves the origin, +so a policy-protected fetch that would be proxied fails closed unless +``DYN_MM_TRUST_EGRESS_PROXY=1`` asserts that the proxy enforces destination +policy. ``NO_PROXY`` is honored, so a fetch that goes direct is never refused. +IP literals never reach a resolver in aiohttp and stay ``validate_url``'s job. """ from __future__ import annotations @@ -38,36 +39,35 @@ from .aiohttp_client import AiohttpClient from .base import ( HttpClient, + HttpConfigurationError, HttpConnectionError, HttpError, HttpStatusError, HttpTimeoutError, ) -from .httpx_client import HttpxClient logger = logging.getLogger(__name__) -_VALID_BACKENDS = ("aiohttp", "httpx") _default: Optional[HttpClient] = None def _create_client() -> HttpClient: - """Pick a concrete client class based on ``DYN_HTTP_BACKEND``. + """Instantiate the aiohttp client (the only supported backend). Mirrors any legacy ``DYN_MM_HTTP_*`` env vars to their canonical ``DYN_HTTP_*`` names first, so the client's lazy ``from_env()`` - inside ``HttpClient.__init__`` sees the migrated values. + inside ``HttpClient.__init__`` sees the migrated values. A stale + ``DYN_HTTP_BACKEND`` other than ``aiohttp`` is ignored with a warning. """ _apply_legacy_env_aliases() name = os.environ.get("DYN_HTTP_BACKEND", "aiohttp").lower() - if name == "aiohttp": - return AiohttpClient() - if name == "httpx": - return HttpxClient() - raise ValueError( - f"DYN_HTTP_BACKEND={name!r} is invalid; must be one of {_VALID_BACKENDS}" - ) + if name not in ("", "aiohttp"): + logger.warning( + "DYN_HTTP_BACKEND=%r is not supported; using aiohttp.", + name, + ) + return AiohttpClient() def get_default_client() -> HttpClient: @@ -79,9 +79,11 @@ def get_default_client() -> HttpClient: return _default -async def fetch_bytes(url, timeout, *, policy=None) -> bytes: +async def fetch_bytes(url, timeout, *, policy=None, max_bytes=None) -> bytes: """Singleton-backed convenience wrapper over :meth:`HttpClient.fetch_bytes`.""" - return await get_default_client().fetch_bytes(url, timeout, policy=policy) + return await get_default_client().fetch_bytes( + url, timeout, policy=policy, max_bytes=max_bytes + ) async def close_http_client() -> None: @@ -101,9 +103,9 @@ async def close_http_client() -> None: __all__ = [ "HttpClient", "AiohttpClient", - "HttpxClient", "HttpError", "HttpTimeoutError", + "HttpConfigurationError", "HttpConnectionError", "HttpStatusError", "HttpConfigBase", diff --git a/components/src/dynamo/common/http/_ssrf_resolver.py b/components/src/dynamo/common/http/_ssrf_resolver.py new file mode 100644 index 000000000000..8d550fbd6544 --- /dev/null +++ b/components/src/dynamo/common/http/_ssrf_resolver.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Connect-time SSRF backstop for the aiohttp backend. + +``validate_url`` checks a hostname's resolved IPs, but the client re-resolves at +connect, so a DNS-rebinding server can return a public IP on the check and an +internal one at connect. aiohttp's ``TCPConnector(resolver=...)`` hook lets us +resolve + filter once and hand the connector the validated addresses to dial, +while the hostname is still used for TLS SNI / certificate verification โ€” the +same mechanism the Rust frontend uses on reqwest's ``dns_resolver`` +(``lib/llm/src/preprocessor/media/loader.rs``). Reuses +:func:`url_validator.is_blocked_ip`. + +Scope: this governs **direct** connections. When an egress proxy is configured, +the proxy resolves the origin, so SSRF must be enforced at the proxy / network +layer instead (true of the Rust path as well). +""" + +from __future__ import annotations + +import errno +import socket + +from aiohttp.abc import AbstractResolver, ResolveResult +from aiohttp.helpers import proxies_from_env +from aiohttp.resolver import DefaultResolver + +from .url_validator import describe_media_source, is_blocked_ip + + +class SsrfBlockedAddress(OSError): + """Raised at connect time when every resolved IP is in a blocked range.""" + + +def _env_proxy_hosts() -> frozenset[str]: + """Hosts of any egress proxy configured in the environment. + + The session runs with ``trust_env=True``, so when a proxy is configured the + connector dials *the proxy*, and it is the proxy's own address that reaches + this resolver -- the origin is resolved by the proxy, out of our sight. A + corporate proxy on a private address would therefore be filtered here and + take out every fetch: measured, HTTP_PROXY at a private host turns each one + into ClientConnectorDNSError. Exempt the configured proxy so proxied + deployments keep working, and see the module docstring for what that means + for enforcement. + """ + return frozenset( + info.proxy.host + for info in proxies_from_env().values() + if info.proxy.host is not None + ) + + +class BlocklistResolver(AbstractResolver): + """aiohttp resolver that drops blocked IPs before the connector dials. + + Returns the full set of non-blocked addresses (not just the first) so + aiohttp keeps its normal multi-address / Happy-Eyeballs fallback. + + Note aiohttp short-circuits IP literals in ``TCPConnector._resolve_host`` + and never calls a resolver for them, so literal blocked addresses are + ``validate_url``'s job, not this one. This covers the hostname case, which + is the one that rebinds. + """ + + def __init__(self, *, allow_private_ips: bool) -> None: + self._inner = DefaultResolver() + self._allow_private_ips = allow_private_ips + self._proxy_hosts = _env_proxy_hosts() + + async def resolve( + self, + host: str, + port: int = 0, + family: socket.AddressFamily = socket.AF_INET, + ) -> list[ResolveResult]: + hosts = await self._inner.resolve(host, port, family) + if self._allow_private_ips or host in self._proxy_hosts: + return hosts + allowed = [h for h in hosts if not is_blocked_ip(h["host"])] + if not allowed: + # ``host`` is caller-supplied and unbounded; this text reaches an + # error response and a log line through the facade's connection + # error, so bound it the way every other message on this path is. + # + # Two-argument OSError on purpose: aiohttp wraps this in + # ClientConnectorDNSError and renders ``strerror``, so the + # one-argument form reaches the caller as a bare "[None]" and the + # reason is lost exactly where an operator would read it. + raise SsrfBlockedAddress( + errno.EHOSTUNREACH, + f"host {describe_media_source(host)} resolves only to blocked IPs", + ) + return allowed + + async def close(self) -> None: + await self._inner.close() diff --git a/components/src/dynamo/common/http/aiohttp_client.py b/components/src/dynamo/common/http/aiohttp_client.py index 1257a09e2a64..1294183c8904 100644 --- a/components/src/dynamo/common/http/aiohttp_client.py +++ b/components/src/dynamo/common/http/aiohttp_client.py @@ -12,32 +12,60 @@ import asyncio import logging +import os from typing import Optional import aiohttp +from aiohttp.helpers import get_env_proxy_for_url from yarl import URL -from .base import HttpClient, HttpConnectionError, HttpStatusError, HttpTimeoutError +from ._ssrf_resolver import BlocklistResolver +from .base import ( + HttpClient, + HttpConfigurationError, + HttpConnectionError, + HttpStatusError, + HttpTimeoutError, + collect_capped, +) +from .url_validator import ( + UrlValidationPolicy, + describe_error_detail, + describe_media_source, +) logger = logging.getLogger(__name__) _REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308}) +# Read granularity for the capped reader. Chunks are joined, so this only +# bounds how far past the limit a single read can carry. +_READ_CHUNK = 64 * 1024 + +# Set to "1" to assert that the configured egress proxy enforces destination +# policy itself. Spelled like DYN_MM_ALLOW_INTERNAL, which it sits beside. +DYN_MM_TRUST_EGRESS_PROXY = "DYN_MM_TRUST_EGRESS_PROXY" + class AiohttpClient(HttpClient): """aiohttp-backed concrete client.""" def __init__(self, config=None) -> None: super().__init__(config) - self._session: Optional[aiohttp.ClientSession] = None + # Keyed by whether the connector may return private addresses. At most + # two entries, so a request that asks for a stricter policy than the + # deployment baseline gets a connector that actually enforces it. + self._sessions: dict[bool, aiohttp.ClientSession] = {} + # aiohttp marks an injected resolver as externally owned, so closing a + # session never closes it. Hold each one and close it ourselves. + self._resolvers: dict[bool, BlocklistResolver] = {} def _effective_timeout(self, timeout: float) -> aiohttp.ClientTimeout: # The override caps ``total`` (whole request). ``sock_connect`` - # bounds the TCP+TLS handshake independently โ€” without it a - # stuck origin would burn the full ``total`` budget before any - # byte arrives, matching what httpx's ``Timeout.connect`` gives - # us on the other backend. + # bounds just the TCP+TLS handshake, so a stuck origin fast-fails at + # the connect budget instead of burning the full ``total`` before any + # byte arrives. total = ( self._config.per_call_timeout_override if self._config.per_call_timeout_override is not None @@ -47,7 +75,57 @@ def _effective_timeout(self, timeout: float) -> aiohttp.ClientTimeout: total=total, sock_connect=self._config.connect_timeout ) - def _build_session(self) -> aiohttp.ClientSession: + @staticmethod + def _connect_allows_private(policy: Optional[UrlValidationPolicy]) -> bool: + """Whether the connector may return private addresses for this fetch. + + The intersection of the deployment baseline and the request policy. + Taking the baseline alone lets ``DYN_MM_ALLOW_INTERNAL=1`` override a + caller that explicitly asked for ``allow_private_ips=False``, and + taking the request alone lets any caller opt out of the deployment + setting. Only when both allow it does the backstop stand down. + """ + env_allows = UrlValidationPolicy.from_env().allow_private_ips + if policy is None: + return env_allows + return env_allows and policy.allow_private_ips + + @staticmethod + async def _require_trusted_egress_proxy(url: str) -> None: + """Refuse a protected fetch that a proxy would put out of our reach. + + When a proxy applies, the connector dials the proxy and the proxy + resolves the origin, so the origin hostname never reaches the + connect-time check. Measured: the resolver sees only the proxy host, + while the proxy receives the absolute URI. Documenting that boundary + does not enforce it, so fail closed unless an operator asserts that + the proxy enforces destination policy itself. + + Asks aiohttp which proxy applies to *this* URL rather than whether any + proxy variable is set, so ``NO_PROXY`` is honored and a fetch that + would go direct is not refused. + """ + if os.getenv(DYN_MM_TRUST_EGRESS_PROXY, "").strip() == "1": + return + try: + # aiohttp runs this same helper through asyncio.to_thread because + # it does proxy-bypass discovery and .netrc file reads. Match that + # rather than repeating the blocking work on the event loop. + await asyncio.to_thread(get_env_proxy_for_url, URL(url)) + except LookupError: + # No proxy for this URL, so aiohttp dials the origin and the + # connect-time check governs it. + return + raise HttpConfigurationError( + f"{describe_media_source(url)} would be fetched through an egress " + "proxy, so the connect-time address check cannot govern the " + f"origin; set {DYN_MM_TRUST_EGRESS_PROXY}=1 to assert that the " + "proxy enforces destination policy, or unset the proxy" + ) + + def _build_session(self, allow_private_ips: bool) -> aiohttp.ClientSession: + resolver = BlocklistResolver(allow_private_ips=allow_private_ips) + self._resolvers[allow_private_ips] = resolver connector = aiohttp.TCPConnector( limit=self._config.max_connections, # Single-origin fan-out is the whole point; capping per-host @@ -55,13 +133,19 @@ def _build_session(self) -> aiohttp.ClientSession: limit_per_host=0, keepalive_timeout=self._config.keepalive_timeout, enable_cleanup_closed=True, + # Connect-time SSRF backstop against DNS rebinding (see + # _ssrf_resolver). See _connect_allows_private for how the flag is + # derived. DYN_MM_ALLOW_INTERNAL is the deployment knob. + resolver=resolver, ) return aiohttp.ClientSession(connector=connector, trust_env=True) - async def _get_session(self) -> aiohttp.ClientSession: + async def _get_session(self, allow_private_ips: bool) -> aiohttp.ClientSession: async with self._lock: - if self._session is None or self._session.closed: - self._session = self._build_session() + session = self._sessions.get(allow_private_ips) + if session is None or session.closed: + session = self._build_session(allow_private_ips) + self._sessions[allow_private_ips] = session logger.info( "aiohttp backend initialized: limit=%d, limit_per_host=0, " "keepalive_timeout=%.1fs%s", @@ -71,34 +155,64 @@ async def _get_session(self) -> aiohttp.ClientSession: if self._config.per_call_timeout_override is not None else "; total timeout set per-request", ) - return self._session + return session - async def _fetch_simple(self, url: str, timeout: float) -> bytes: - session = await self._get_session() + async def _fetch_simple( + self, + url: str, + timeout: float, + *, + max_bytes: Optional[int] = None, + policy: Optional[UrlValidationPolicy] = None, + ) -> bytes: + allow_private = self._connect_allows_private(policy) + # Only when the check is meant to bite. If private destinations are + # already permitted for this fetch, the proxy gate protects nothing. + if not allow_private: + await self._require_trusted_egress_proxy(url) + session = await self._get_session(allow_private) client_timeout = self._effective_timeout(timeout) try: async with session.get( url, timeout=client_timeout, allow_redirects=True ) as response: response.raise_for_status() - return await response.read() + return await collect_capped( + response.content.iter_chunked(_READ_CHUNK), url, max_bytes + ) except aiohttp.ClientResponseError as e: raise HttpStatusError(e.status, e.message or "", url) from e except (asyncio.TimeoutError, aiohttp.ServerTimeoutError) as e: - raise HttpTimeoutError(f"Timeout loading {url}") from e + raise HttpTimeoutError( + f"Timeout loading {describe_media_source(url)}" + ) from e except ( aiohttp.ClientConnectionError, aiohttp.ClientConnectorError, aiohttp.ServerDisconnectedError, ) as e: - raise HttpConnectionError(f"Connection error loading {url}: {e}") from e + raise HttpConnectionError( + f"Connection error loading {describe_media_source(url)}: {describe_error_detail(str(e))}" + ) from e except aiohttp.ClientError as e: - raise HttpConnectionError(f"HTTP error loading {url}: {e}") from e + raise HttpConnectionError( + f"HTTP error loading {describe_media_source(url)}: {describe_error_detail(str(e))}" + ) from e async def _fetch_body_or_redirect( - self, url: str, timeout: float + self, + url: str, + timeout: float, + *, + max_bytes: Optional[int] = None, + policy: Optional[UrlValidationPolicy] = None, ) -> tuple[bytes | None, str | None]: - session = await self._get_session() + allow_private = self._connect_allows_private(policy) + # Only when the check is meant to bite. If private destinations are + # already permitted for this fetch, the proxy gate protects nothing. + if not allow_private: + await self._require_trusted_egress_proxy(url) + session = await self._get_session(allow_private) client_timeout = self._effective_timeout(timeout) try: async with session.get( @@ -109,26 +223,49 @@ async def _fetch_body_or_redirect( if location: next_url = str(response.url.join(URL(location))) return None, next_url - return await response.read(), None + return ( + await collect_capped( + response.content.iter_chunked(_READ_CHUNK), url, max_bytes + ), + None, + ) try: response.raise_for_status() except aiohttp.ClientResponseError as e: raise HttpStatusError(e.status, e.message or "", url) from e - return await response.read(), None + return ( + await collect_capped( + response.content.iter_chunked(_READ_CHUNK), url, max_bytes + ), + None, + ) except (asyncio.TimeoutError, aiohttp.ServerTimeoutError) as e: - raise HttpTimeoutError(f"Timeout loading {url}") from e + raise HttpTimeoutError( + f"Timeout loading {describe_media_source(url)}" + ) from e except ( aiohttp.ClientConnectionError, aiohttp.ClientConnectorError, aiohttp.ServerDisconnectedError, ) as e: - raise HttpConnectionError(f"Connection error loading {url}: {e}") from e + raise HttpConnectionError( + f"Connection error loading {describe_media_source(url)}: {describe_error_detail(str(e))}" + ) from e except aiohttp.ClientError as e: - raise HttpConnectionError(f"HTTP error loading {url}: {e}") from e + raise HttpConnectionError( + f"HTTP error loading {describe_media_source(url)}: {describe_error_detail(str(e))}" + ) from e async def close(self) -> None: async with self._lock: - if self._session is not None and not self._session.closed: - await self._session.close() - self._session = None + for session in self._sessions.values(): + if not session.closed: + await session.close() + self._sessions = {} + # TCPConnector sets _resolver_owner=False for an injected resolver + # and only closes the one it made itself, so closing the session + # above leaves ours open. + for resolver in self._resolvers.values(): + await resolver.close() + self._resolvers = {} diff --git a/components/src/dynamo/common/http/base.py b/components/src/dynamo/common/http/base.py index 22c5724def33..c928de3de6b6 100644 --- a/components/src/dynamo/common/http/base.py +++ b/components/src/dynamo/common/http/base.py @@ -3,8 +3,8 @@ """Abstract base class + unified exception classes for the HTTP facade. -Each concrete subclass (:class:`AiohttpClient`, :class:`HttpxClient`) -owns a backend-specific session/client singleton on the instance. +The concrete subclass (:class:`AiohttpClient`) +owns a backend-specific session singleton on the instance. Backend-neutral logic (the SSRF redirect loop) lives here so it isn't duplicated across subclasses. """ @@ -13,7 +13,7 @@ import abc import asyncio -from typing import Optional +from typing import AsyncIterator, Optional from dynamo.common.configuration.groups.http_args import HttpConfigBase, from_env @@ -21,9 +21,19 @@ _MAX_REDIRECTS, UrlValidationError, UrlValidationPolicy, + describe_error_detail, + describe_media_source, validate_url, ) +# Longest a client-supplied reason may render as inside an HttpStatusError. +# Matches ``dynamo.llm.exceptions.HttpError._MAX_MESSAGE_LENGTH``, which the +# binding applies to the other class it forwards on a 4xx. Generous on purpose: +# callers build real guidance here -- the trtllm decoder hint is ~480 characters +# and its test allows 2000 -- so a tight bound silently deletes the actionable +# part and leaves only a truncated prefix. +_MAX_MESSAGE_LENGTH = 8192 + class HttpError(Exception): """Base class for all HTTP fetch failures.""" @@ -37,16 +47,60 @@ class HttpConnectionError(HttpError): """Network-layer failure: DNS, refused, reset, half-close.""" +class HttpConfigurationError(HttpError): + """The deployment is configured so the fetch cannot be made safely. + + An operator fault, not a verdict on the caller's URL, so it must not reach + the client as a 4xx. Loaders that convert unknown exceptions into + ``ValueError`` have to let this one through: ``py_err_to_dynamo`` maps + ``ValueError`` to ``InvalidArgument``, which is a client error. + """ + + class HttpStatusError(HttpError): """Server responded with a non-2xx status.""" def __init__(self, status: int, message: str, url: str) -> None: - super().__init__(f"HTTP {status} for {url}: {message}") + # Both halves are client-supplied: ``url`` directly, and ``message`` + # because aiohttp's own ClientResponseError text repeats the URL. The video + # diffusion handler puts str(exc) in its response body. + # + # ``message`` is bounded in the *attribute*, not just in the rendered + # string: errors.rs::extract_http_like_error reads ``.status`` and + # ``.message`` off this class by name and, per the SECURITY note there, + # forwards ``.message`` verbatim to the client on a 4xx. Nothing in the + # rendered text reaches that path. ``url`` is not part of that protocol, + # so it keeps its full value for debugging. + message = describe_error_detail(message, _MAX_MESSAGE_LENGTH) + super().__init__(f"HTTP {status} for {describe_media_source(url)}: {message}") self.status = status self.message = message self.url = url +async def collect_capped( + chunks: AsyncIterator[bytes], url: str, max_bytes: Optional[int] +) -> bytes: + """Join ``chunks`` into one body, refusing to buffer past ``max_bytes``. + + The check is per chunk rather than on the finished body: a declared + Content-Length is attacker-controlled and absent on a chunked response, and + a single capped read is not enough either โ€” aiohttp's ``read(n)`` returns + *at most* n bytes, so a short read would look like a body under the limit. + """ + out: list[bytes] = [] + total = 0 + async for chunk in chunks: + total += len(chunk) + if max_bytes is not None and total > max_bytes: + raise UrlValidationError( + f"Media exceeds the {max_bytes} byte download limit: " + f"{describe_media_source(url)}" + ) + out.append(chunk) + return b"".join(out) + + class HttpClient(abc.ABC): """Backend-neutral HTTP client. @@ -66,28 +120,38 @@ async def fetch_bytes( timeout: float, *, policy: Optional[UrlValidationPolicy] = None, + max_bytes: Optional[int] = None, ) -> bytes: """Fetch ``url`` and return the response body. Single-shot: no retries. Raises one of the unified exception - classes above; callers never see native httpx/aiohttp classes. + classes above; callers never see native aiohttp classes. ``policy=None``: use the backend's built-in redirect handling. + ``max_bytes`` set: refuse a body larger than that while it is being + read, so an attacker-chosen URL cannot buffer an unbounded response. + Raises :class:`UrlValidationError` โ€” it is a verdict on a + client-supplied source, like the redirect cap below. + ``policy`` set: follow redirects manually and revalidate each hop against the policy via :func:`url_validator.validate_url`. This is the SSRF-safe path; raises :class:`UrlValidationError` if any hop fails or the chain exceeds ``_MAX_REDIRECTS``. """ if policy is None: - return await self._fetch_simple(url, timeout) - return await self._fetch_with_revalidation(url, timeout, policy) + return await self._fetch_simple(url, timeout, max_bytes=max_bytes) + return await self._fetch_with_revalidation( + url, timeout, policy, max_bytes=max_bytes + ) async def _fetch_with_revalidation( self, url: str, timeout: float, policy: UrlValidationPolicy, + *, + max_bytes: Optional[int] = None, ) -> bytes: """Manual redirect loop with per-hop SSRF validation (backend-neutral).""" current = url @@ -97,30 +161,51 @@ async def _fetch_with_revalidation( await validate_url(current, policy) visited.append(current) - body, redirect_to = await self._fetch_body_or_redirect(current, timeout) + body, redirect_to = await self._fetch_body_or_redirect( + current, timeout, max_bytes=max_bytes, policy=policy + ) if redirect_to is None: if body is None: raise HttpError( - f"Backend returned (None, None) for {current}; " + f"Backend returned (None, None) for {describe_media_source(current)}; " "expected bytes on a terminal (2xx or 3xx-without-Location) response" ) return body if hops_remaining <= 0: + # ``visited`` holds attacker-chosen URLs; describe each one. + chain = [describe_media_source(hop) for hop in visited] raise UrlValidationError( - f"Too many redirects (max={_MAX_REDIRECTS}); chain={visited}" + f"Too many redirects (max={_MAX_REDIRECTS}); chain={chain}" ) hops_remaining -= 1 current = redirect_to @abc.abstractmethod - async def _fetch_simple(self, url: str, timeout: float) -> bytes: - """Backend's native redirect-following GET (no SSRF policy applied).""" + async def _fetch_simple( + self, + url: str, + timeout: float, + *, + max_bytes: Optional[int] = None, + policy: Optional[UrlValidationPolicy] = None, + ) -> bytes: + """Backend's native redirect-following GET (no SSRF policy applied). + + ``policy`` is not applied to the URL here. It is passed so a backend + can honor a request that is stricter than the deployment baseline at + connect time; see ``AiohttpClient._connect_allows_private``. + """ @abc.abstractmethod async def _fetch_body_or_redirect( - self, url: str, timeout: float + self, + url: str, + timeout: float, + *, + max_bytes: Optional[int] = None, + policy: Optional[UrlValidationPolicy] = None, ) -> tuple[bytes | None, str | None]: """Single hop with redirects disabled. diff --git a/components/src/dynamo/common/http/httpx_client.py b/components/src/dynamo/common/http/httpx_client.py deleted file mode 100644 index 8a5c13cf1aa1..000000000000 --- a/components/src/dynamo/common/http/httpx_client.py +++ /dev/null @@ -1,158 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""httpx implementation of :class:`HttpClient`. - -Behavior notes (operator-tunable knobs live in :mod:`.args`): - - - ``httpx.Timeout`` is built per-request so each fetch sees its - caller-supplied read timeout instead of the first caller's value - getting baked into the singleton. ``write`` is unbounded โ€” fan-out - GET workloads have no meaningful write phase. - - ``max_keepalive_connections`` is sized to match ``max_connections`` - so idle connections get reused under fan-out instead of churning - TLS handshakes. - - A process-wide semaphore (``DYN_HTTP_CONCURRENCY``) wraps every - fetch so a burst can't push ``PoolTimeout`` up the stack. aiohttp - has no equivalent because its connector queues natively. - -An operator who flips ``DYN_HTTP_BACKEND=httpx`` gets a working backend, -not the original PoolTimeout bug. -""" - -from __future__ import annotations - -import asyncio -import logging -from typing import Optional - -import httpx - -from .base import HttpClient, HttpConnectionError, HttpStatusError, HttpTimeoutError - -logger = logging.getLogger(__name__) - - -class HttpxClient(HttpClient): - """httpx-backed concrete client.""" - - def __init__(self, config=None) -> None: - super().__init__(config) - self._client: Optional[httpx.AsyncClient] = None - self._semaphore: Optional[asyncio.Semaphore] = None - - def _get_semaphore(self) -> asyncio.Semaphore: - """Process-wide concurrency cap on in-flight httpx fetches. - - Lazy-init so the bound (``DYN_HTTP_CONCURRENCY``) is read once - and the ``asyncio.Semaphore`` binds to whichever loop first - awaits it. - """ - if self._semaphore is None: - self._semaphore = asyncio.Semaphore(self._config.concurrency) - return self._semaphore - - def _per_call_timeout(self, read_timeout: float) -> httpx.Timeout: - # Override targets ``read`` only โ€” ``connect`` / ``pool`` stay - # independent so a stuck handshake or saturated pool still - # fast-fails on its own budget. - effective_read = ( - self._config.per_call_timeout_override - if self._config.per_call_timeout_override is not None - else read_timeout - ) - return httpx.Timeout( - connect=self._config.connect_timeout, - read=effective_read, - write=None, - pool=self._config.pool_timeout, - ) - - def _build_client(self) -> httpx.AsyncClient: - return httpx.AsyncClient( - timeout=self._per_call_timeout(60.0), - follow_redirects=True, - limits=httpx.Limits( - max_connections=self._config.max_connections, - max_keepalive_connections=self._config.max_keepalive, - ), - ) - - async def _get_client(self) -> httpx.AsyncClient: - async with self._lock: - if self._client is None or self._client.is_closed: - self._client = self._build_client() - logger.info( - "httpx backend initialized: max_connections=%d, max_keepalive=%d, " - "timeout(connect=%.1fs, write=None, pool=%.1fs); read timeout %s", - self._config.max_connections, - self._config.max_keepalive, - self._config.connect_timeout, - self._config.pool_timeout, - f"forced to {self._config.per_call_timeout_override:.1f}s via env" - if self._config.per_call_timeout_override is not None - else "set per-request", - ) - return self._client - - async def _fetch_simple(self, url: str, timeout: float) -> bytes: - client = await self._get_client() - async with self._get_semaphore(): - try: - response = await client.get( - url, timeout=self._per_call_timeout(timeout) - ) - response.raise_for_status() - return response.content - except httpx.HTTPStatusError as e: - raise HttpStatusError(e.response.status_code, str(e), url) from e - except httpx.TimeoutException as e: - raise HttpTimeoutError(f"Timeout loading {url}") from e - except (httpx.ConnectError, httpx.NetworkError) as e: - raise HttpConnectionError(f"Connection error loading {url}: {e}") from e - except httpx.HTTPError as e: - raise HttpConnectionError(f"HTTP error loading {url}: {e}") from e - - async def _fetch_body_or_redirect( - self, url: str, timeout: float - ) -> tuple[bytes | None, str | None]: - client = await self._get_client() - async with self._get_semaphore(): - try: - request = client.build_request( - "GET", url, timeout=self._per_call_timeout(timeout) - ) - response = await client.send( - request, - follow_redirects=False, - ) - except httpx.TimeoutException as e: - raise HttpTimeoutError(f"Timeout loading {url}") from e - except (httpx.ConnectError, httpx.NetworkError) as e: - raise HttpConnectionError(f"Connection error loading {url}: {e}") from e - except httpx.HTTPError as e: - raise HttpConnectionError(f"HTTP error loading {url}: {e}") from e - - try: - if response.is_redirect: - location = response.headers.get("location") - if location: - next_url = str(response.url.join(location)) - return None, next_url - # 3xx without Location: treat as terminal, surface the body. - return response.content, None - - try: - response.raise_for_status() - except httpx.HTTPStatusError as e: - raise HttpStatusError(e.response.status_code, str(e), url) from e - return response.content, None - finally: - await response.aclose() - - async def close(self) -> None: - async with self._lock: - if self._client is not None and not self._client.is_closed: - await self._client.aclose() - self._client = None - self._semaphore = None diff --git a/components/src/dynamo/common/http/media_reference.py b/components/src/dynamo/common/http/media_reference.py new file mode 100644 index 000000000000..83c91196505c --- /dev/null +++ b/components/src/dynamo/common/http/media_reference.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Materialize a client media reference to a trusted local path. + +`validate_media_reference` blocks the *initial* URL, but a validated URL must +still not be handed to a downstream generator that fetches it and follows +redirects on its own โ€” an allowed origin can `302` to an internal address +(redirect SSRF). This fetches URL references through the SSRF-safe client (which +revalidates every redirect hop against the policy) into a temp file and yields +that local path; local references pass through unchanged. Callers hand the +generator only a trusted local path, never a URL. +""" + +from __future__ import annotations + +import logging +import os +import tempfile +from contextlib import asynccontextmanager +from pathlib import Path +from typing import AsyncIterator +from urllib.parse import urlparse + +from .url_validator import ( + UrlValidationError, + UrlValidationPolicy, + describe_media_source, + validate_media_reference, +) + +logger = logging.getLogger(__name__) + +# A temp-file suffix is derived from the client's URL path, so it has to be +# bounded: mkstemp() raises ENAMETOOLONG past the filesystem's limit, and that +# OSError carries the server's temp directory back to the caller. Long enough +# for any real media extension. +_MAX_SUFFIX_LEN = 16 + +# Before this module existed, a URL input_reference was downloaded by SGLang's +# own get_image_bytes -> download_remote_media, which streams under +# ``media_url_max_file_size_mb`` (default 64 MiB, sglang 0.5.18 +# server_args.py:2804). Handing the generator a local path moves it to the +# open() branch, where that cap no longer applies -- so carry the same default +# here rather than silently dropping the bound. +# +# That knob was operator-tunable as a SGLang server arg, so carry the knob too: +# ``DYN_MM_MAX_FILE_SIZE_MB`` overrides the default, in megabytes, matching both +# the SGLang arg it replaces and trtllm's ``DYN_TRTLLM_MAX_FILE_SIZE_MB``. +DYN_MM_MAX_FILE_SIZE_MB = "DYN_MM_MAX_FILE_SIZE_MB" +DEFAULT_MAX_MEDIA_MB = 64 +MAX_MEDIA_BYTES = DEFAULT_MAX_MEDIA_MB * 1024 * 1024 + +# A parameter default binds at definition time, so it cannot call the resolver +# below -- the env has to be read per call for a worker that is configured after +# import and for tests that monkeypatch it. This sentinel means "resolve it". +_FROM_ENV = -1 + + +def max_media_bytes() -> int: + """Media download cap in bytes, from ``DYN_MM_MAX_FILE_SIZE_MB``. + + Mirrors ``media_decoder._video_num_frames``: read at call time, and an + empty, unparseable or non-positive value falls back to the default with a + warning rather than raising. A malformed operator value must not take the + worker down, and must not silently remove the bound either. + """ + raw = os.getenv(DYN_MM_MAX_FILE_SIZE_MB, "").strip() + if not raw: + return MAX_MEDIA_BYTES + try: + value = int(raw) + except ValueError: + logger.warning( + "Ignoring invalid %s=%r; using %s MB", + DYN_MM_MAX_FILE_SIZE_MB, + raw, + DEFAULT_MAX_MEDIA_MB, + ) + return MAX_MEDIA_BYTES + if value <= 0: + logger.warning( + "Ignoring non-positive %s=%r; using %s MB", + DYN_MM_MAX_FILE_SIZE_MB, + raw, + DEFAULT_MAX_MEDIA_MB, + ) + return MAX_MEDIA_BYTES + return value * 1024 * 1024 + + +def _temp_suffix(url: str) -> str: + """Extension to give the temp file, bounded and stripped of path separators. + + Only a plausible extension survives: the generator picks a decoder from it, + and everything else in a client-supplied path is noise at best. + """ + suffix = Path(urlparse(url).path).suffix + if len(suffix) > _MAX_SUFFIX_LEN or not suffix[1:].isalnum(): + return "" + return suffix + + +@asynccontextmanager +async def local_media_reference( + reference: str, + policy: UrlValidationPolicy, + *, + timeout: float = 30.0, + max_bytes: int | None = _FROM_ENV, +) -> AsyncIterator[str]: + """Yield a trusted local filesystem path for ``reference``. + + URL references are fetched through the policy-aware client (per-hop redirect + revalidation) into a temp file that is removed on exit. Local references are + validated by ``validate_media_reference`` and yielded as-is. + + ``max_bytes`` bounds the download while it is read; exceeding it raises + ``UrlValidationError``, as an oversized client-chosen source is a bad + request rather than a server fault. Left unset it resolves from + ``DYN_MM_MAX_FILE_SIZE_MB`` per call; ``None`` disables the bound. + + Raises ``UrlValidationError`` for anything that is not a local path or an + http(s) URL โ€” notably ``data:``, which ``validate_url`` allows but which is + a URI, not a path, and would reach the generator as one. + """ + if max_bytes == _FROM_ENV: + max_bytes = max_media_bytes() + + resolved = await validate_media_reference(reference, policy) + scheme = urlparse(resolved).scheme + if scheme in ("http", "https"): + # Imported here, not at module scope, so tests can monkeypatch + # ``dynamo.common.http.fetch_bytes`` and see the patched function. + from . import fetch_bytes + + data = await fetch_bytes(resolved, timeout, policy=policy, max_bytes=max_bytes) + fd, tmp = tempfile.mkstemp(suffix=_temp_suffix(resolved)) + try: + with os.fdopen(fd, "wb") as fh: + fh.write(data) + yield tmp + finally: + try: + os.unlink(tmp) + except OSError: + pass + elif scheme: + raise UrlValidationError( + f"Media reference scheme '{scheme}' cannot be resolved to a local " + f"file: {describe_media_source(reference)}" + ) + else: + yield resolved diff --git a/components/src/dynamo/common/http/url_validator.py b/components/src/dynamo/common/http/url_validator.py index 32242c48198e..b336901ed8ca 100644 --- a/components/src/dynamo/common/http/url_validator.py +++ b/components/src/dynamo/common/http/url_validator.py @@ -20,7 +20,8 @@ import socket from dataclasses import dataclass from pathlib import Path -from urllib.parse import urlparse +from typing import Final +from urllib.parse import unquote, urlparse class UrlValidationError(ValueError): @@ -72,6 +73,62 @@ class UrlValidationError(ValueError): ) +# Longest a media source may render as inside an error message or log line. +# Generous enough to keep an ordinary URL intact and identifiable. +SOURCE_LABEL_LIMIT: Final = 120 + + +def describe_media_source(source: str, limit: int = SOURCE_LABEL_LIMIT) -> str: + """Render ``source`` as a bounded label safe to put in an error or log. + + A ``data:`` URI carries the whole media payload inline, so echoing one into + an error message serializes megabytes of base64 -- to the client, and to + every log sink that records the failure. Describe those by media type and + size instead, never by content. Other sources are truncated, since a URL + identifies the request without being unbounded. + + Lives here rather than in ``multimodal.media_source`` so the validators + below can bound their own messages: importing that package pulls in torch. + """ + if not isinstance(source, str): + return "" + if source.startswith("data:"): + meta = source[len("data:") :].partition(",")[0] + media_type = meta.split(";")[0] or "application/octet-stream" + # The media-type field is client-supplied and unbounded: a reference of + # ``"data:" + "A" * 200_000 + ",AAAA"`` puts all of it here, so eliding + # the payload alone still renders a 200 KB label. Bound it like any + # other source. Omitting the comma takes the same path. + if len(media_type) > limit: + media_type = f"{media_type[:limit]}... ({len(media_type)} chars)" + return f"data:{media_type} ({len(source)} chars, payload elided)" + if len(source) > limit: + return f"{source[:limit]}... ({len(source)} chars)" + return source + + +def describe_error_detail(detail: str, limit: int = SOURCE_LABEL_LIMIT) -> str: + """Bound a backend exception's text, keeping both ends. + + Unlike a media source, the useful part of one of these is usually at the + end: aiohttp renders the client-supplied host *before* the errno, so a + head-only truncation would keep the attacker's string and drop the + diagnosis. + """ + if not isinstance(detail, str): + return "" + if len(detail) <= limit: + return detail + # The marker comes out of the budget, not on top of it: a caller that sizes + # a buffer by ``limit`` should not be handed something longer. + marker = f"... ({len(detail)} chars) ..." + budget = limit - len(marker) + if budget <= 0: + return marker + head = budget // 2 + return f"{detail[:head]}{marker}{detail[-(budget - head):]}" + + def is_blocked_ip(ip_text: str) -> bool: """Return True if ``ip_text`` parses as an IP inside one of the blocked ranges.""" try: @@ -109,7 +166,7 @@ async def validate_url(url: str, policy: UrlValidationPolicy) -> str: For URLs with a hostname, we resolve it here (off the event loop via ``loop.getaddrinfo``) and check the resulting IPs against the blocked ranges. This catches obvious DNS rebinding but not an attacker who - changes their answer between this lookup and httpx's actual connect. + changes their answer between this lookup and the client's actual connect. Raises ``UrlValidationError`` on any policy violation. """ @@ -119,11 +176,21 @@ async def validate_url(url: str, policy: UrlValidationPolicy) -> str: parsed = urlparse(url) scheme = parsed.scheme.lower() + # Before the label: describe_media_source copies the source, and a data: + # URI carries the whole payload inline, so building one for the branch that + # returns without using it dominates the call (98% of it at 32 MiB). if scheme == "data": return url + # Every message below is surfaced to the caller (the diffusion handlers + # turn it into a 400 body) and logged, so nothing client-supplied goes in + # at full length. + label = describe_media_source(url) + if scheme not in ("http", "https"): - raise UrlValidationError(f"URL scheme '{scheme}' not allowed") + raise UrlValidationError( + f"URL scheme '{describe_media_source(scheme)}' not allowed" + ) if scheme == "http" and not policy.allow_http: raise UrlValidationError( @@ -132,11 +199,12 @@ async def validate_url(url: str, policy: UrlValidationPolicy) -> str: host = (parsed.hostname or "").lower() if not host: - raise UrlValidationError(f"URL has no host component: {url!r}") + raise UrlValidationError(f"URL has no host component: {label!r}") if not policy.allow_private_ips and host in _BLOCKED_HOSTS: raise UrlValidationError( - f"Host '{host}' is blocked (resolves to internal service)" + f"Host '{describe_media_source(host)}' is blocked " + "(resolves to internal service)" ) try: @@ -155,11 +223,15 @@ async def validate_url(url: str, policy: UrlValidationPolicy) -> str: try: infos = await loop.getaddrinfo(host, None) except socket.gaierror as exc: - raise UrlValidationError(f"Could not resolve host '{host}': {exc}") from exc + raise UrlValidationError( + f"Could not resolve host '{describe_media_source(host)}': {exc}" + ) from exc for info in infos: addr = info[4][0] if is_blocked_ip(addr): - raise UrlValidationError(f"Host '{host}' resolves to blocked IP '{addr}'") + raise UrlValidationError( + f"Host '{describe_media_source(host)}' resolves to blocked IP '{addr}'" + ) return url @@ -178,25 +250,46 @@ def validate_local_path(path: str, policy: UrlValidationPolicy) -> Path: "Local media paths are not permitted; set " "DYN_MM_LOCAL_PATH to enable" ) + # ``path`` is client-supplied and unbounded: describe_media_source keeps it + # out of an error response and a log line at full length. Short, ordinary + # paths render unchanged. + label = describe_media_source(path) + try: resolved = Path(path).expanduser().resolve(strict=True) except FileNotFoundError as exc: - raise UrlValidationError(f"File not found: {path}") from exc + raise UrlValidationError(f"File not found: {label}") from exc except OSError as exc: - raise UrlValidationError(f"Could not resolve path '{path}': {exc}") from exc + # ``exc`` renders the offending filename in full -- ENAMETOOLONG on a + # 200,000-character name gives a 200,069-character string -- which walks + # straight past ``label``'s bound into the error response and the log. + # ``strerror`` is the errno text alone ("File name too long"), with no + # path in it; the bounded ``str(exc)`` is the fallback when it is None. + detail = exc.strerror or describe_error_detail(str(exc)) + raise UrlValidationError(f"Could not resolve path '{label}': {detail}") from exc + except ValueError as exc: + # An embedded NUL makes lstat() raise ValueError, not OSError. Reachable + # since file:// paths are percent-decoded, so %00 becomes a real NUL. + # Callers map UrlValidationError to a client error; a bare ValueError + # would reach them as a server error instead. + raise UrlValidationError(f"Invalid path '{label}': {exc}") from exc try: allowed = Path(policy.allowed_local_path).expanduser().resolve(strict=True) except FileNotFoundError as exc: + # Same reason as the message below: callers put this in a client + # response, and the configured directory is deployment detail. raise UrlValidationError( - f"Configured allowed_local_path does not exist: {policy.allowed_local_path}" + "Configured allowed_local_path does not exist" ) from exc try: resolved.relative_to(allowed) except ValueError as exc: + # The configured directory is deployment detail; naming it here puts + # it in the client's error body, since callers surface this message. raise UrlValidationError( - f"Path '{path}' is outside the allowed directory '{policy.allowed_local_path}'" + f"Path '{label}' is outside the allowed directory" ) from exc return resolved @@ -220,11 +313,28 @@ async def validate_media_url(url: str, policy: UrlValidationPolicy) -> str: scheme = parsed.scheme.lower() if scheme in ("", "file"): - raw_path = parsed.path if scheme == "file" else url + # file:// paths are percent-encoded; a bare path is literal. + raw_path = unquote(parsed.path) if scheme == "file" else url resolved = validate_local_path(raw_path, policy) return resolved.as_uri() return await validate_url(url, policy) +async def validate_media_reference(reference: str, policy: UrlValidationPolicy) -> str: + """Like :func:`validate_media_url` but return a plain filesystem path for + local references instead of a ``file://`` URI, for callers that pass the + result to a loader expecting a bare path. + """ + if not reference: + raise UrlValidationError("Media reference is empty") + + parsed = urlparse(reference) + if parsed.scheme.lower() in ("", "file"): + # file:// paths are percent-encoded; a bare path is literal. + raw_path = unquote(parsed.path) if parsed.scheme else reference + return str(validate_local_path(raw_path, policy)) + return await validate_url(reference, policy) + + _MAX_REDIRECTS = 3 diff --git a/components/src/dynamo/common/multimodal/audio_loader.py b/components/src/dynamo/common/multimodal/audio_loader.py index b9a5941389bd..79d4c3df631b 100644 --- a/components/src/dynamo/common/multimodal/audio_loader.py +++ b/components/src/dynamo/common/multimodal/audio_loader.py @@ -8,7 +8,7 @@ import numpy as np -from dynamo.common.http import HttpStatusError, fetch_bytes +from dynamo.common.http import HttpConfigurationError, HttpStatusError, fetch_bytes from dynamo.common.http.url_validator import ( UrlValidationError, UrlValidationPolicy, @@ -21,6 +21,8 @@ from dynamo.common.utils import nvtx_utils as _nvtx from dynamo.common.utils.runtime import run_async +from dynamo.common.http.media_reference import max_media_bytes # isort: skip + logger = logging.getLogger(__name__) # Constants for multimodal data variants @@ -128,7 +130,10 @@ async def _load_audio_with_vllm(self, audio_url: str) -> tuple[np.ndarray, float # data: and file:// never touch the network, so vLLM can handle them. if urlparse(normalized_url).scheme in ("http", "https"): content = await fetch_bytes( - normalized_url, self._http_timeout, policy=self._url_policy + normalized_url, + self._http_timeout, + policy=self._url_policy, + max_bytes=max_media_bytes(), ) return await asyncio.to_thread(media_io.load_bytes, content) @@ -153,7 +158,7 @@ async def load_audio(self, audio_url: str) -> tuple[np.ndarray, float]: return waveform, sr except FileNotFoundError: raise - except (UrlValidationError, HttpStatusError): + except (UrlValidationError, HttpStatusError, HttpConfigurationError): # Preserve deliberate client-error verdicts. UrlValidationError is # a ValueError, so the generic handler below would otherwise erase # its type and prevent the frontend from returning a 4xx. diff --git a/components/src/dynamo/common/multimodal/codec_errors.py b/components/src/dynamo/common/multimodal/codec_errors.py index a5473288a93f..8d5ff5ed308d 100644 --- a/components/src/dynamo/common/multimodal/codec_errors.py +++ b/components/src/dynamo/common/multimodal/codec_errors.py @@ -20,8 +20,11 @@ from __future__ import annotations +import importlib +import importlib.metadata + from dynamo.common.multimodal.nvdec_decoder import HW_ROUTED_CODECS, nvdec_available -from dynamo.common.utils.install_media_decoders import VALIDATED_SPECS +from dynamo.common.utils.install_media_decoders import VALIDATED_SPECS, installer_covers INSTALLER_CMD = "python -m dynamo.common.utils.install_media_decoders" @@ -35,14 +38,71 @@ class MissingMediaDecoderError(RuntimeError): """ -def _install_hint(backend: str, package: str) -> str: - spec = VALIDATED_SPECS[package] +def _carrier_present(module: str) -> bool: + """Return whether the carrier imports; keep import errors out of diagnostics.""" + try: + importlib.import_module(module) + except Exception: # noqa: BLE001 - any import failure means unusable here + return False + return True + + +def _needs_binary_wheel(backend: str, package: str) -> bool: + """Whether this package must come from a wheel: vLLM's OpenCV sdist has no codecs.""" + return backend == "vllm" and package == "opencv-python-headless" + + +def _is_vllm_source_built_cv2(backend: str, package: str, module: str) -> bool: + """Whether this is the codec-free OpenCV build shipped by vLLM.""" return ( - f"install the validated decoder with `pip install --no-deps '{spec}'` " - f"(or `{INSTALLER_CMD} {backend}`)" + _needs_binary_wheel(backend, package) + and module == "cv2" + and _carrier_present(module) ) +def _installed_version(package: str) -> str | None: + """The installed distribution version, or None when it has no metadata.""" + try: + return importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + return None + + +def _install_hint(backend: str, package: str, module: str) -> str: + """Use the installed OpenCV version for vLLM replacements, validated bounds otherwise.""" + if _is_vllm_source_built_cv2(backend, package, module): + installed = _installed_version(package) + if installed: + hint = ( + "replace the shipped build with the binary wheel of the same " + f"version: `pip install --no-deps --force-reinstall " + f"--only-binary {package} '{package}=={installed}'`" + ) + else: + # cv2 omits the distribution packaging revision. + hint = ( + "first determine the shipped cv2 version with " + '`python -c "import cv2; print(cv2.__version__)"`, then replace ' + "it with the binary wheel built from that same version -- the " + "distribution adds a packaging revision, so match on the prefix: " + f"`pip install --no-deps --force-reinstall --only-binary {package} " + f"'{package}==.*'`" + ) + else: + only_binary = ( + f"--only-binary {package} " if _needs_binary_wheel(backend, package) else "" + ) + hint = ( + "install the validated decoder with " + f"`pip install --no-deps --force-reinstall " + f"{only_binary}'{VALIDATED_SPECS[package]}'`" + ) + if installer_covers(backend, package): + hint += f" (or `{INSTALLER_CMD} {backend}`)" + return hint + + def _with_cause(message: str, cause: str | None) -> str: """Append the underlying decoder text so diagnostics survive the wrap. @@ -79,6 +139,13 @@ def video_decoder_missing( "but NVDEC is unavailable in this container. Grant the 'video' " "driver capability (NVIDIA_DRIVER_CAPABILITIES) to enable it, or " ) + elif _is_vllm_source_built_cv2(backend, package, module): + lead = ( + f"this video ({codec_desc}) has no decoder in this image: shipped " + "images decode only H.264/H.265 (in hardware, via NVDEC), and the " + f"'{module}' they ship is built without a video backend. " + "Re-encode the input to H.264/H.265, or " + ) else: lead = ( f"this video ({codec_desc}) has no decoder in this image: shipped " @@ -88,7 +155,10 @@ def video_decoder_missing( ) return MissingMediaDecoderError( _with_cause( - "Cannot decode video: " + lead + _install_hint(backend, package) + ".", + "Cannot decode video: " + + lead + + _install_hint(backend, package, module) + + ".", cause, ) ) @@ -106,7 +176,7 @@ def audio_decoder_missing( _with_cause( "Cannot decode audio: this input needs the PyAV decoder ('av'), which " "this image deliberately does not ship, and NVDEC does not decode " - "audio. To enable audio input, " + _install_hint(backend, "av") + ".", + "audio. To enable audio input, " + _install_hint(backend, "av", "av") + ".", cause, ) ) diff --git a/components/src/dynamo/common/multimodal/image_loader.py b/components/src/dynamo/common/multimodal/image_loader.py index 1cc22929d18c..034b9e0fd9d5 100644 --- a/components/src/dynamo/common/multimodal/image_loader.py +++ b/components/src/dynamo/common/multimodal/image_loader.py @@ -4,9 +4,11 @@ import asyncio import base64 import binascii +import hashlib import logging import os from collections import OrderedDict +from collections.abc import Mapping from io import BytesIO from typing import Any, Coroutine, Dict, Final, List, Literal, overload from urllib.parse import urlparse @@ -23,11 +25,13 @@ HttpTimeoutError, fetch_bytes, ) +from ..http.media_reference import max_media_bytes from ..http.url_validator import ( UrlValidationError, UrlValidationPolicy, validate_media_url, ) +from .shared_image_cache import SharedImageCache, SharedImageCacheStats logger = logging.getLogger(__name__) @@ -35,6 +39,47 @@ URL_VARIANT_KEY: Final = "Url" DECODED_VARIANT_KEY: Final = "Decoded" UUID_ONLY_VARIANT_KEY: Final = "UuidOnly" +IMAGE_CACHE_SCOPE_KEY: Final = "image_cache_scope" +IMAGE_CACHE_SESSION_SCOPED_ENV: Final = "DYN_MM_IMAGE_CACHE_SESSION_SCOPED" + + +def image_cache_scope_from_request(request: Mapping[str, Any]) -> str | None: + """Return a non-empty frontend-derived image-cache scope, if present.""" + scope = request.get(IMAGE_CACHE_SCOPE_KEY) + if not isinstance(scope, str): + return None + scope = scope.strip() + return scope or None + + +def image_cache_session_scoped_from_env() -> bool: + """Return whether multimodal caches must be partitioned by request scope.""" + value = os.environ.get(IMAGE_CACHE_SESSION_SCOPED_ENV, "0") + if value not in ("0", "1"): + raise ValueError( + f"{IMAGE_CACHE_SESSION_SCOPED_ENV} must be '0' or '1', got {value!r}" + ) + return value == "1" + + +def scope_image_cache_key( + cache_key: str, + cache_scope: str | None, + *, + session_scoped_cache: bool, +) -> str | None: + """Apply the common session-scope policy to a multimodal cache key. + + When session scoping is disabled, the original key remains stable. When it + is enabled, a valid scope partitions the key and a missing scope returns + ``None`` so the caller bypasses its cache. + """ + if not session_scoped_cache: + return cache_key + if cache_scope is None or not cache_scope.strip(): + return None + scope_digest = hashlib.sha256(cache_scope.strip().encode("utf-8")).hexdigest() + return f"{scope_digest}:{cache_key}" def _create_nixl_connector() -> Any: @@ -71,6 +116,8 @@ def __init__( http_timeout: float = 30.0, enable_frontend_decoding: bool = False, url_policy: UrlValidationPolicy | None = None, + max_bytes: int | None = None, + session_scoped_cache: bool | None = None, ): """ Initialize the ImageLoader with caching, HTTP settings, and optional NIXL config for @@ -85,11 +132,22 @@ def __init__( decoded images directly from frontend memory, bypassing standard network transport. Defaults to False. url_policy: Policy for validating URLs. Defaults to UrlValidationPolicy.from_env(). + max_bytes: Maximum remote image size in bytes. When omitted, resolve + DYN_MM_MAX_FILE_SIZE_MB for each request. + session_scoped_cache: Include a caller-provided cache scope in local, + in-flight, and shared-cache keys. ``None`` (default) reads + DYN_MM_IMAGE_CACHE_SESSION_SCOPED (default off). When enabled, + HTTP loads without a scope bypass all image caches. """ self._http_timeout = http_timeout self._cache_size = cache_size + self._configured_max_bytes = max_bytes + if session_scoped_cache is None: + session_scoped_cache = image_cache_session_scoped_from_env() + self._session_scoped_cache = session_scoped_cache self._image_cache: OrderedDict[str, Image.Image] = OrderedDict() self._inflight: dict[str, asyncio.Task[Image.Image]] = {} + self._shared_image_cache = SharedImageCache.from_env() self._enable_frontend_decoding = enable_frontend_decoding self._url_policy = url_policy or UrlValidationPolicy.from_env() # Lazy-init NIXL connector only when frontend decoding is enabled @@ -100,6 +158,46 @@ def __init__( self._nixl_connector.initialize ) # Synchronously wait for async init + def _max_bytes(self) -> int: + if self._configured_max_bytes is not None: + return self._configured_max_bytes + return max_media_bytes() + + @property + def cache_entries(self) -> int: + """Current number of cached images (read by metrics export).""" + return len(self._image_cache) + + @property + def session_scoped_cache(self) -> bool: + """Whether all multimodal cache keys require a request scope.""" + return self._session_scoped_cache + + @property + def shared_image_cache_stats(self) -> SharedImageCacheStats | None: + """Latency samples for the optional shared encoded-image cache.""" + if self._shared_image_cache is None: + return None + return self._shared_image_cache.stats + + def _cache_key( + self, normalized_url: str, cache_scope: str | None = None + ) -> str | None: + """Return the LRU/inflight/shared-cache key for a validated http(s) URL. + + ``None`` means the caller must bypass all image caches: session + scoping is enabled but no valid scope was provided for this load. + """ + # Preserve path and query case: unlike the scheme and host, both may be + # case-sensitive and therefore identify different origin objects. + url_key = normalized_url + + return scope_image_cache_key( + url_key, + cache_scope, + session_scoped_cache=self._session_scoped_cache, + ) + @staticmethod def _open_image_sync(image_data: BytesIO) -> Image.Image: """Open, validate, and decode an image from raw bytes. Runs in a thread.""" @@ -126,22 +224,48 @@ def _cache_put(self, key: str, image: Image.Image) -> None: self._image_cache.popitem(last=False) self._image_cache[key] = image - async def _fetch_and_process(self, image_url: str) -> Image.Image: + async def _fetch_and_process(self, key: str | None, image_url: str) -> Image.Image: """Fetch image via HTTP(S), decode with PIL, return RGB Image. - All exception normalization happens here so shared callers - see identical error types. + Checks the optional shared encoded-image cache before hitting the + origin and refills it after a successful origin fetch. A ``None`` key + bypasses the shared cache (session scoping without a valid scope). + All exception normalization happens here so shared callers see + identical error types. """ try: + if self._shared_image_cache is not None and key is not None: + cached_content = await self._shared_image_cache.get(key) + if cached_content is not None: + try: + return await self._open_image(BytesIO(cached_content)) + except (OSError, ValueError) as exc: + logger.warning( + "Discarding invalid shared image cache entry for '%s': %s", + image_url[:80], + exc, + ) + await self._shared_image_cache.delete(key) + with _nvtx.annotate("mm:img:http_fetch", color="lime"): content = await fetch_bytes( - image_url, self._http_timeout, policy=self._url_policy + image_url, + self._http_timeout, + policy=self._url_policy, + max_bytes=self._max_bytes(), ) if not content: raise ValueError("Empty response content from image URL") image_data = BytesIO(content) - return await self._open_image(image_data) + image = await self._open_image(image_data) + if self._shared_image_cache is not None and key is not None: + # Cache fills are awaited deliberately so Redis acknowledges a + # successful write before this request returns. A cache outage + # remains fail-open, but can add up to the configured Redis I/O + # timeout to a miss; a bounded async queue could avoid that later. + await self._shared_image_cache.put(key, content) + return image except HttpStatusError as e: logger.error(f"HTTP {e.status} loading image: '{image_url}'") @@ -185,12 +309,12 @@ async def _fetch_and_process(self, image_url: str) -> Image.Image: raise ValueError(f"Failed to load image: '{image_url}': {e}") from e except Exception as e: logger.error(f"{type(e).__name__} loading image: '{image_url}': {e}") - raise ValueError(f"Failed to load image: '{image_url}': {e}") from e + raise async def _fetch_and_cache(self, key: str, image_url: str) -> Image.Image: """Shared task: fetch, cache, then remove from _inflight.""" try: - image = await self._fetch_and_process(image_url) + image = await self._fetch_and_process(key, image_url) self._cache_put(key, image) return image finally: @@ -208,7 +332,9 @@ async def _read_and_convert_nixl_image( return image if image.mode == "RGB" else image.convert("RGB") @_nvtx.annotate("mm:img:load_image", color="lime") - async def load_image(self, image_url: str) -> Image.Image: + async def load_image( + self, image_url: str, *, cache_scope: str | None = None + ) -> Image.Image: parsed_url = urlparse(image_url) if parsed_url.scheme in ("", "file"): raise ValueError( @@ -218,7 +344,10 @@ async def load_image(self, image_url: str) -> Image.Image: parsed_url = urlparse(normalized_url) if parsed_url.scheme in ("http", "https"): - key = normalized_url.lower() + key = self._cache_key(normalized_url, cache_scope) + + if key is None: + return await self._fetch_and_process(None, normalized_url) if key in self._image_cache: logger.debug(f"Image found in cache for URL: {image_url}") @@ -254,7 +383,7 @@ async def load_image(self, image_url: str) -> Image.Image: except Image.UnidentifiedImageError as e: logger.error(f"Unsupported image format decoding: '{image_url}'") raise HttpStatusError(415, "Unsupported Media Type", image_url) from e - except Exception as e: + except ValueError as e: if "Unsupported image format" in str(e): logger.error(f"Unsupported image format decoding: '{image_url}'") raise HttpStatusError( @@ -262,6 +391,14 @@ async def load_image(self, image_url: str) -> Image.Image: ) from e logger.error(f"{type(e).__name__} decoding image: '{image_url}': {e}") raise ValueError(f"Failed to decoding image: '{image_url}': {e}") from e + except OSError as e: + logger.error(f"Invalid or truncated image data: '{image_url}'") + raise HttpStatusError( + 400, "Invalid or truncated image data", image_url + ) from e + except Exception: + logger.error(f"Unexpected error decoding image: '{image_url}'") + raise # It's not file:, http:, https:, or data: raise ValueError(f"Invalid image source scheme: {parsed_url.scheme}") @@ -272,6 +409,7 @@ async def load_image_batch( image_mm_items: List[Dict[str, Any]], *, preserve_uuid_slots: Literal[False] = False, + cache_scope: str | None = None, ) -> list[Image.Image]: ... @@ -281,6 +419,7 @@ async def load_image_batch( image_mm_items: List[Dict[str, Any]], *, preserve_uuid_slots: Literal[True], + cache_scope: str | None = None, ) -> list[Image.Image | None]: ... @@ -289,6 +428,7 @@ async def load_image_batch( image_mm_items: List[Dict[str, Any]], *, preserve_uuid_slots: bool = False, + cache_scope: str | None = None, ) -> list[Any]: """ Load a batch of images from multimodal data items. @@ -303,6 +443,8 @@ async def load_image_batch( image_mm_items: List of multimodal data items for images preserve_uuid_slots: Allow UUID-only items and preserve their positions as None. This is enabled only by backends that resolve such slots. + cache_scope: Stable session or request identifier used when + DYN_MM_IMAGE_CACHE_SESSION_SCOPED is enabled. Returns: Loaded images, with None for UUID-only cache slots @@ -315,6 +457,9 @@ async def load_image_batch( client-error code instead of 500. UrlValidationError: If a media URL is rejected by the SSRF policy; preserved as a ValueError so the frontend returns a 4xx, not 500. + ValueError: If any image fails client-side input validation (e.g. + a malformed data: URI โ€” invalid base64, missing ;base64 marker, + or non-image MIME). Exception: If any image fails to load for any other reason ValueError: If enable_frontend_decoding=True but nixl_connector is None ValueError: If a UUID-only slot is received without opting in @@ -327,7 +472,7 @@ async def load_image_batch( # URL path: download and decode in Python backend url = item[URL_VARIANT_KEY] slot_to_future_idx.append(len(image_futures)) - image_futures.append(self.load_image(url)) + image_futures.append(self.load_image(url, cache_scope=cache_scope)) logger.debug(f"Preparing to load image from URL: {url[:80]}...") elif isinstance(item, dict) and DECODED_VARIANT_KEY in item: if self._enable_frontend_decoding: @@ -360,6 +505,7 @@ async def load_image_batch( collective_exceptions = "" status_error: HttpStatusError | None = None url_error: UrlValidationError | None = None + value_error: ValueError | None = None for media_item, future_idx in zip( image_mm_items, slot_to_future_idx, strict=True ): @@ -389,6 +535,10 @@ async def load_image_batch( # preserve it so the frontend still gets a 4xx, not a 500. elif url_error is None and isinstance(result, UrlValidationError): url_error = result + # The bindings map ValueError to Backend(InvalidArgument) โ†’ 400, + # so preserve the type so the frontend still gets a 4xx. + elif value_error is None and isinstance(result, ValueError): + value_error = result continue loaded_images.append(result) @@ -398,6 +548,9 @@ async def load_image_batch( if url_error is not None: raise url_error + if value_error is not None: + raise value_error + if collective_exceptions: raise Exception(collective_exceptions) diff --git a/components/src/dynamo/common/multimodal/media_connector.py b/components/src/dynamo/common/multimodal/media_connector.py index 31eb097ce51f..bc9b918107fc 100644 --- a/components/src/dynamo/common/multimodal/media_connector.py +++ b/components/src/dynamo/common/multimodal/media_connector.py @@ -19,7 +19,7 @@ from PIL import Image -from dynamo.common.http import HttpStatusError +from dynamo.common.http import HttpConfigurationError, HttpStatusError from dynamo.common.http.url_validator import UrlValidationError from dynamo.common.multimodal.image_loader import ImageLoader @@ -70,7 +70,7 @@ async def fetch_image_async( """Async image fetch via dynamo's ImageLoader with LRU cache.""" try: img = await self._image_loader.load_image(image_url) - except (UrlValidationError, HttpStatusError): + except (UrlValidationError, HttpStatusError, HttpConfigurationError): # These are deliberate client-error verdicts, not signals that # another loader should retry the same rejected URL. raise diff --git a/components/src/dynamo/common/multimodal/media_source.py b/components/src/dynamo/common/multimodal/media_source.py index 37e82740c5d6..1aa37e99936f 100644 --- a/components/src/dynamo/common/multimodal/media_source.py +++ b/components/src/dynamo/common/multimodal/media_source.py @@ -23,18 +23,31 @@ import base64 import binascii import logging -from typing import Final from urllib.parse import unquote, urlparse from urllib.request import url2pathname from dynamo.common.http.url_validator import ( + SOURCE_LABEL_LIMIT, UrlValidationError, UrlValidationPolicy, + describe_media_source, validate_local_path, ) logger = logging.getLogger(__name__) +# describe_media_source now lives beside the validators, which have to bound a +# client-supplied source in their own messages and cannot import this package +# (it pulls in torch). Re-exported here because callers import it from here. +__all__ = [ + "LOCAL_MEDIA_SCHEMES", + "SOURCE_LABEL_LIMIT", + "decode_data_uri", + "describe_media_source", + "is_local_media_url", + "read_local_media_bytes", +] + # Schemes this module can turn into bytes. http(s) is deliberately absent: it # belongs to fetch_bytes, which applies SSRF revalidation on every redirect hop. LOCAL_MEDIA_SCHEMES = frozenset({"file", "data"}) @@ -45,31 +58,6 @@ def is_local_media_url(url: str) -> bool: return urlparse(url).scheme in LOCAL_MEDIA_SCHEMES -# Longest a media source may render as inside an error message or log line. -# Generous enough to keep an ordinary URL intact and identifiable. -SOURCE_LABEL_LIMIT: Final = 120 - - -def describe_media_source(source: str, limit: int = SOURCE_LABEL_LIMIT) -> str: - """Render ``source`` as a bounded label safe to put in an error or log. - - A ``data:`` URI carries the whole media payload inline, so echoing one into - an error message serializes megabytes of base64 -- to the client, and to - every log sink that records the failure. Describe those by media type and - size instead, never by content. Other sources are truncated, since a URL - identifies the request without being unbounded. - """ - if not isinstance(source, str): - return "" - if source.startswith("data:"): - meta = source[len("data:") :].partition(",")[0] - media_type = meta.split(";")[0] or "application/octet-stream" - return f"data:{media_type} ({len(source)} chars, payload elided)" - if len(source) > limit: - return f"{source[:limit]}... ({len(source)} chars)" - return source - - def decode_data_uri(url: str) -> bytes: """Decode a ``data:`` URI body to bytes. diff --git a/components/src/dynamo/common/multimodal/shared_image_cache.py b/components/src/dynamo/common/multimodal/shared_image_cache.py new file mode 100644 index 000000000000..297df02b3120 --- /dev/null +++ b/components/src/dynamo/common/multimodal/shared_image_cache.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Optional Redis Cluster-backed cache for encoded multimodal image bytes.""" + +import hashlib +import logging +import os +import threading +import time +from collections import defaultdict +from typing import cast + +from redis.asyncio.cluster import RedisCluster +from redis.asyncio.retry import Retry +from redis.backoff import NoBackoff +from redis.exceptions import RedisClusterException, RedisError + +logger = logging.getLogger(__name__) + +_ENABLED_ENV = "DYN_MM_SHARED_IMAGE_CACHE_ENABLED" +_URL_ENV = "DYN_MM_SHARED_IMAGE_CACHE_URL" +_TTL_ENV = "DYN_MM_SHARED_IMAGE_CACHE_TTL_SECS" +_CONNECT_TIMEOUT_ENV = "DYN_MM_SHARED_IMAGE_CACHE_CONNECT_TIMEOUT_SECS" +_IO_TIMEOUT_ENV = "DYN_MM_SHARED_IMAGE_CACHE_IO_TIMEOUT_SECS" +_MAX_PENDING_DURATIONS = 10_000 +_ERROR_WARNING_INTERVAL_SECONDS = 60.0 +# redis-py defines RedisClusterException outside the RedisError hierarchy. +_REDIS_OPERATION_ERRORS = (RedisError, RedisClusterException) + + +def _size_bucket(size_bytes: int | None) -> str: + if size_bytes is None: + return "unknown" + for upper_bound, label in ( + (64 * 1024, "le_64kib"), + (256 * 1024, "le_256kib"), + (1024 * 1024, "le_1mib"), + (4 * 1024 * 1024, "le_4mib"), + (16 * 1024 * 1024, "le_16mib"), + (32 * 1024 * 1024, "le_32mib"), + ): + if size_bytes <= upper_bound: + return label + return "gt_32mib" + + +class SharedImageCacheStats: + """Thread-safe buffer of Redis operation latency samples.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._durations: defaultdict[tuple[str, str, str], list[float]] = defaultdict( + list + ) + self._pending_count = 0 + + def record( + self, operation: str, outcome: str, size_bucket: str, duration_s: float + ) -> None: + with self._lock: + if self._pending_count >= _MAX_PENDING_DURATIONS: + return + self._durations[(operation, outcome, size_bucket)].append(duration_s) + self._pending_count += 1 + + def snapshot_and_drain(self) -> dict[tuple[str, str, str], list[float]]: + """Atomically take all pending operation latency samples.""" + with self._lock: + durations = dict(self._durations) + self._durations.clear() + self._pending_count = 0 + return durations + + +class _RateLimitedErrorWarnings: + """Track cache failures and allow one aggregate warning per interval.""" + + def __init__(self, interval_seconds: float) -> None: + self._interval_seconds = interval_seconds + self._lock = threading.Lock() + self._last_warning_at: float | None = None + self._failures_since_warning = 0 + + def record(self, now: float) -> int | None: + """Return the failure count when a warning is due, otherwise ``None``.""" + with self._lock: + self._failures_since_warning += 1 + if ( + self._last_warning_at is not None + and now - self._last_warning_at < self._interval_seconds + ): + return None + failure_count = self._failures_since_warning + self._failures_since_warning = 0 + self._last_warning_at = now + return failure_count + + +class SharedImageCache: + """Fail-open Redis Cluster cache for compressed image payloads.""" + + def __init__(self, client: RedisCluster, ttl_seconds: int) -> None: + self._client = client + self._ttl_seconds = ttl_seconds + self.stats = SharedImageCacheStats() + self._error_warnings = _RateLimitedErrorWarnings( + _ERROR_WARNING_INTERVAL_SECONDS + ) + + @classmethod + def from_env(cls) -> "SharedImageCache | None": + """Build the cache when explicitly enabled, otherwise return ``None``.""" + enabled = os.environ.get(_ENABLED_ENV, "0") + if enabled not in ("0", "1"): + raise ValueError(f"{_ENABLED_ENV} must be '0' or '1', got {enabled!r}") + if enabled == "0": + return None + + url = os.environ.get(_URL_ENV) + if not url: + raise ValueError(f"{_URL_ENV} must be set when {_ENABLED_ENV}=1") + + ttl_seconds = int(os.environ.get(_TTL_ENV, "3600")) + if ttl_seconds <= 0: + raise ValueError(f"{_TTL_ENV} must be greater than zero") + + connect_timeout_seconds = float(os.environ.get(_CONNECT_TIMEOUT_ENV, "0.1")) + if connect_timeout_seconds <= 0: + raise ValueError(f"{_CONNECT_TIMEOUT_ENV} must be greater than zero") + + io_timeout_seconds = float(os.environ.get(_IO_TIMEOUT_ENV, "2.0")) + if io_timeout_seconds <= 0: + raise ValueError(f"{_IO_TIMEOUT_ENV} must be greater than zero") + + client = RedisCluster.from_url( + url, + decode_responses=False, + socket_connect_timeout=connect_timeout_seconds, + socket_timeout=io_timeout_seconds, + # Rediscover through the configured endpoint, typically a Kubernetes + # Service, instead of replacing it with the node addresses the + # cluster announces; those go stale once every node's IP changes. + # The async client accepts this from redis-py 6.2, our floor. + dynamic_startup_nodes=False, + # Fail fast: an origin fetch is cheaper than retrying a cache call. + retry=Retry(NoBackoff(), 0), + ) + # Do not log the URL: Redis URLs may contain credentials. + logger.info("Shared image cache enabled") + return cls(client, ttl_seconds) + + @staticmethod + def _key(cache_identity: str) -> str: + digest = hashlib.sha256(cache_identity.encode("utf-8")).hexdigest() + return f"dynamo:mm:image:{{{digest}}}:bytes" + + def _log_operation_error(self, operation: str, exc: Exception) -> None: + """Log each cache failure at debug and rate-limit aggregate warnings.""" + logger.debug("Shared image cache %s failed: %s", operation, exc) + failure_count = self._error_warnings.record(time.monotonic()) + if failure_count is not None: + logger.warning( + "Shared image cache unavailable; observed %d operation failure(s) " + "since the last warning or startup. Requests continue without " + "the shared cache; repeated failures are logged at debug level " + "for %.0f seconds.", + failure_count, + _ERROR_WARNING_INTERVAL_SECONDS, + ) + + async def get(self, cache_identity: str) -> bytes | None: + """Return cached encoded bytes, or ``None`` on a miss or cache error.""" + key = self._key(cache_identity) + started = time.perf_counter() + outcome = "error" + size_bucket = "unknown" + try: + value = await self._client.get(key) + except _REDIS_OPERATION_ERRORS as exc: + self._log_operation_error("read", exc) + return None + else: + outcome = "hit" if value is not None else "miss" + size_bucket = _size_bucket(len(value) if value is not None else None) + # RedisCluster.get() is typed for both decoded and binary clients, + # but this client is always configured with decode_responses=False. + return cast(bytes | None, value) + finally: + self.stats.record( + "get", outcome, size_bucket, time.perf_counter() - started + ) + + async def put(self, cache_identity: str, content: bytes) -> None: + """Store validated encoded bytes, ignoring cache availability errors.""" + key = self._key(cache_identity) + started = time.perf_counter() + outcome = "error" + size_bucket = _size_bucket(len(content)) + try: + await self._client.set(key, content, ex=self._ttl_seconds) + except _REDIS_OPERATION_ERRORS as exc: + self._log_operation_error("write", exc) + else: + outcome = "success" + finally: + self.stats.record( + "set", outcome, size_bucket, time.perf_counter() - started + ) + + async def delete(self, cache_identity: str) -> None: + """Remove a corrupt entry, ignoring cache availability errors.""" + key = self._key(cache_identity) + try: + await self._client.delete(key) + except _REDIS_OPERATION_ERRORS as exc: + self._log_operation_error("delete", exc) diff --git a/components/src/dynamo/common/multimodal/video_loader.py b/components/src/dynamo/common/multimodal/video_loader.py index 6c39c0d51fc9..4c40cb75cb38 100644 --- a/components/src/dynamo/common/multimodal/video_loader.py +++ b/components/src/dynamo/common/multimodal/video_loader.py @@ -14,14 +14,16 @@ # limitations under the License. import asyncio +import functools import logging import os +import re from typing import Any, Awaitable, Dict, Final, List from urllib.parse import urlparse import numpy as np -from dynamo.common.http import HttpStatusError, fetch_bytes +from dynamo.common.http import HttpConfigurationError, HttpStatusError, fetch_bytes from dynamo.common.http.url_validator import ( UrlValidationError, UrlValidationPolicy, @@ -42,9 +44,39 @@ ) from dynamo.common.utils.runtime import run_async +from dynamo.common.http.media_reference import max_media_bytes # isort: skip + logger = logging.getLogger(__name__) +def _attributable_to_cv2(exc: BaseException, media_io: Any | None = None) -> bool: + """Identify OpenCV errors without relabeling failures from other video backends.""" + if getattr(exc, "name", None) == "cv2" or "cv2" in str(exc): + return True + return ( + isinstance(exc, ValueError) + and str(exc) == "Could not open video stream" + and (getattr(media_io, "kwargs", None) or {}).get("backend", "opencv") + == "opencv" + ) + + +@functools.lru_cache(maxsize=1) +def _cv2_lacks_video_backend() -> bool: + """Return whether OpenCV is broken or lacks FFmpeg and GStreamer.""" + try: + import cv2 + except ImportError: + return False + except Exception: # noqa: BLE001 - a broken import has no usable video backend + return True + build_info = cv2.getBuildInformation() + return not any( + re.search(rf"^\s*{backend}:\s*YES", build_info, re.MULTILINE | re.IGNORECASE) + for backend in ("FFMPEG", "GSTREAMER") + ) + + URL_VARIANT_KEY: Final = "Url" DECODED_VARIANT_KEY: Final = "Decoded" @@ -141,7 +173,10 @@ async def _load_video_with_vllm( # data: and file:// never touch the network, so vLLM can handle them. if urlparse(normalized_url).scheme in ("http", "https"): content = await fetch_bytes( - normalized_url, self._http_timeout, policy=self._url_policy + normalized_url, + self._http_timeout, + policy=self._url_policy, + max_bytes=max_media_bytes(), ) return await self._decode_video_bytes(content, media_io) @@ -165,13 +200,8 @@ async def _decode_video_bytes( ) -> tuple[np.ndarray, Dict[str, Any]]: """Decode video bytes: H.264/H.265 on NVDEC, all else software. - The runtime images purge the software decode wheels (opencv/av/decord/ - torchcodec) for codec compliance, so the software fallback only - resolves where a decoder was installed separately. When it is absent, - vLLM's lazy import surfaces a bare ``No module named 'cv2'`` with no - codec and no remedy -- convert that into the actionable - unsupported-codec error, which can name the codec because the probe - already ran here. + vLLM ships OpenCV for still images without video backends. Translate + its decode failures into install guidance; preserve other backends. The NVDEC path honors the two ``media_io_kwargs`` that decide *which* frames come back -- ``num_frames`` and ``fps`` -- because it samples @@ -198,9 +228,22 @@ async def _decode_video_bytes( try: return await asyncio.to_thread(media_io.load_bytes, content) except ImportError as exc: + if not _attributable_to_cv2(exc): + raise raise video_decoder_missing( "vllm", "opencv-python-headless", "cv2", codec, cause=str(exc) ) from exc + except (OSError, SystemError, ValueError) as exc: + # Confirm both the error source and missing video support. + if not (_attributable_to_cv2(exc, media_io) and _cv2_lacks_video_backend()): + raise + raise video_decoder_missing( + "vllm", + "opencv-python-headless", + "cv2", + codec, + cause=str(exc), + ) from exc def _extract_nvdec_args(self, media_io: Any) -> dict[str, Any]: num_frames = getattr(media_io, "num_frames", self._num_frames) @@ -233,7 +276,7 @@ async def load_video( return np.ascontiguousarray(frames), metadata except FileNotFoundError: raise - except (UrlValidationError, HttpStatusError): + except (UrlValidationError, HttpStatusError, HttpConfigurationError): # Preserve deliberate client-error verdicts. UrlValidationError is # a ValueError, so the generic handler below would otherwise erase # its type and prevent the frontend from returning a 4xx. diff --git a/components/src/dynamo/common/protocols/audio_protocol.py b/components/src/dynamo/common/protocols/audio_protocol.py index 17ea95da55d0..9365c9c6dee2 100644 --- a/components/src/dynamo/common/protocols/audio_protocol.py +++ b/components/src/dynamo/common/protocols/audio_protocol.py @@ -27,6 +27,13 @@ class AudioNvExt(BaseModel): leaves the N-2 compatibility window in v1.7. """ + cfg_scale: Optional[float] = None + """Classifier-free guidance strength (Audex only, which is why it is an + extension rather than a top-level OpenAI field; vLLM-Omni likewise takes it + under ``extra_params``). 1.0 disables guidance; 1.5 is the recommended TTS + quality setting, and 3.0 the official TTA setting (applied by default for + text-to-audio).""" + class NvCreateAudioSpeechRequest(BaseModel): """Request for audio speech generation (/v1/audio/speech endpoint). diff --git a/components/src/dynamo/common/protocols/video_protocol.py b/components/src/dynamo/common/protocols/video_protocol.py index 644cea28e161..13d061344d64 100644 --- a/components/src/dynamo/common/protocols/video_protocol.py +++ b/components/src/dynamo/common/protocols/video_protocol.py @@ -108,6 +108,12 @@ class VideoData(BaseModel): b64_json: Optional[str] = None """Base64-encoded video (if response_format is 'b64_json').""" + fps: Optional[int] = None + """Actual video frame rate when reported by the model.""" + + audio_sample_rate: Optional[int] = None + """Muxed audio sample rate when the generated video contains audio.""" + class NvVideosResponse(BaseModel): """Response structure for video generation. diff --git a/components/src/dynamo/common/tests/configuration/test_kv_router_args.py b/components/src/dynamo/common/tests/configuration/test_kv_router_args.py index 6e0a7cc9e634..0080001a55f2 100644 --- a/components/src/dynamo/common/tests/configuration/test_kv_router_args.py +++ b/components/src/dynamo/common/tests/configuration/test_kv_router_args.py @@ -191,6 +191,42 @@ def test_decode_active_request_weight_flows_to_binding_kwargs() -> None: assert kwargs["decode_active_request_weight"] == 64.0 +def test_session_prefix_index_is_opt_in() -> None: + parser = argparse.ArgumentParser() + KvRouterArgGroup().add_arguments(parser) + + default_kwargs = KvRouterConfigBase.from_cli_args( + parser.parse_args([]) + ).kv_router_kwargs() + assert default_kwargs["enable_session_prefix_index"] is False + + enabled_kwargs = KvRouterConfigBase.from_cli_args( + parser.parse_args(["--enable-session-prefix-index"]) + ).kv_router_kwargs() + assert enabled_kwargs["enable_session_prefix_index"] is True + + disabled_kwargs = KvRouterConfigBase.from_cli_args( + parser.parse_args(["--no-enable-session-prefix-index"]) + ).kv_router_kwargs() + assert disabled_kwargs["enable_session_prefix_index"] is False + + +def test_session_prefix_index_environment_flows_to_binding_kwargs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DYN_ENABLE_SESSION_PREFIX_INDEX", "true") + parser = argparse.ArgumentParser() + KvRouterArgGroup().add_arguments(parser) + + kwargs = KvRouterConfigBase.from_cli_args(parser.parse_args([])).kv_router_kwargs() + assert kwargs["enable_session_prefix_index"] is True + + overridden = KvRouterConfigBase.from_cli_args( + parser.parse_args(["--no-enable-session-prefix-index"]) + ).kv_router_kwargs() + assert overridden["enable_session_prefix_index"] is False + + def test_load_aware_cli_applies_no_cache_load_balancing_preset() -> None: parser = argparse.ArgumentParser() KvRouterArgGroup().add_arguments(parser) diff --git a/components/src/dynamo/common/tests/http/conftest.py b/components/src/dynamo/common/tests/http/conftest.py index 0bd56223c909..c80c344c694b 100644 --- a/components/src/dynamo/common/tests/http/conftest.py +++ b/components/src/dynamo/common/tests/http/conftest.py @@ -3,17 +3,41 @@ """Shared fixtures for HTTP facade tests. -Autouse fixture ensures the shared HTTP client singleton is closed after -each test so ``DYN_HTTP_BACKEND`` changes take effect between runs and -no "Unclosed client session" warning bleeds across tests. +One autouse fixture closes the shared HTTP client singleton after each test, +so ``DYN_HTTP_BACKEND`` changes take effect between runs and no "Unclosed +client session" warning bleeds across tests. + +A second clears the proxy environment. The client refuses a policy-protected +fetch that an egress proxy would carry (see ``DYN_MM_TRUST_EGRESS_PROXY``), so +a developer machine or CI runner with an ambient ``http_proxy`` would other- +wise fail tests that have nothing to do with proxies. A test that wants a +proxy sets one itself with ``monkeypatch``. """ from __future__ import annotations +import pytest import pytest_asyncio from dynamo.common.http import close_http_client +_PROXY_VARS = ( + "http_proxy", + "https_proxy", + "all_proxy", + "no_proxy", + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "NO_PROXY", +) + + +@pytest.fixture(autouse=True) +def _no_ambient_proxy(monkeypatch): + for name in _PROXY_VARS: + monkeypatch.delenv(name, raising=False) + @pytest_asyncio.fixture(autouse=True) async def _close_shared_http_client(): diff --git a/components/src/dynamo/common/tests/http/test_aiohttp_client.py b/components/src/dynamo/common/tests/http/test_aiohttp_client.py index fa467e2c2aed..5a0a1a37e73d 100644 --- a/components/src/dynamo/common/tests/http/test_aiohttp_client.py +++ b/components/src/dynamo/common/tests/http/test_aiohttp_client.py @@ -24,6 +24,8 @@ from dynamo.common import http as mm_http from dynamo.common.http import AiohttpClient +from dynamo.common.http._ssrf_resolver import BlocklistResolver +from dynamo.common.http.base import HttpConfigurationError from dynamo.common.http.url_validator import UrlValidationPolicy pytestmark = [ @@ -34,6 +36,17 @@ ] +class _FakeContent: + """``response.content`` stand-in: the reader the size cap streams from.""" + + def __init__(self, body: bytes) -> None: + self._body = body + + async def iter_chunked(self, n: int): + for i in range(0, len(self._body), n) or [0]: + yield self._body[i : i + n] + + class _FakeResponse: """Minimal aiohttp response stand-in for ``async with session.get(...) as r``.""" @@ -42,6 +55,7 @@ def __init__(self, *, status=200, headers=None, url=None, body=b"") -> None: self.headers = headers or {} self.url = url self._body = body + self.content = _FakeContent(body) def raise_for_status(self) -> None: return None @@ -106,11 +120,21 @@ def _get(url, **kwargs): def _make_client_with_session(session) -> AiohttpClient: client = AiohttpClient() - client._session = session + # Both strictness keys map to the same double, so a test does not have to + # care which session the policy under test selects. + client._sessions = {True: session, False: session} return client _PERMISSIVE = UrlValidationPolicy(allow_http=True, allow_private_ips=True) +_STRICT = UrlValidationPolicy(allow_http=True, allow_private_ips=False) + +# Building the real connector trips aiohttp's notice that ``enable_cleanup_closed`` +# is a no-op on Python >= 3.12.7. That flag is pre-existing on main, so this is +# scoped to the tests that build a real connector rather than ignored repo-wide. +_allows_cleanup_closed_notice = pytest.mark.filterwarnings( + "ignore:enable_cleanup_closed ignored because:DeprecationWarning" +) async def test_fetch_bytes_returns_body_on_200() -> None: @@ -175,3 +199,247 @@ async def test_redirect_resolved_through_policy_path() -> None: body = await client.fetch_bytes("https://h/x.png", 30.0, policy=_PERMISSIVE) assert body == b"final" + + +@_allows_cleanup_closed_notice +async def test_build_session_installs_the_connect_time_resolver() -> None: + """The shared connector carries the connect-time resolver. + + Every fetch goes through a connector built here, so losing the + ``resolver=`` wiring silently drops the connect-time check for the whole + process while the resolver's own unit tests stay green. + """ + client = AiohttpClient() + # Through _get_session, so client.close() owns it. _build_session does not + # register the session, and closing the client would leave it open. + session = await client._get_session(False) + try: + assert isinstance(session.connector._resolver, BlocklistResolver) + finally: + await client.close() + assert session.closed + + +@_allows_cleanup_closed_notice +async def test_connect_policy_is_the_intersection(monkeypatch) -> None: + """Neither side alone decides whether the connector may return private IPs. + + A permissive request must not loosen a strict deployment, and a permissive + deployment must not override a request that explicitly asked for strict. + """ + strict = UrlValidationPolicy(allow_http=True, allow_private_ips=False) + + monkeypatch.delenv("DYN_MM_ALLOW_INTERNAL", raising=False) + client = AiohttpClient() + assert client._connect_allows_private(None) is False + assert client._connect_allows_private(_PERMISSIVE) is False + assert client._connect_allows_private(strict) is False + + monkeypatch.setenv("DYN_MM_ALLOW_INTERNAL", "1") + assert client._connect_allows_private(None) is True + assert client._connect_allows_private(_PERMISSIVE) is True + # The case that regressed: the deployment allows private, the caller does + # not, and the caller wins. + assert client._connect_allows_private(strict) is False + + +@_allows_cleanup_closed_notice +async def test_a_strict_request_gets_a_filtering_connector(monkeypatch) -> None: + """The intersection reaches the connector, not just the helper.""" + monkeypatch.setenv("DYN_MM_ALLOW_INTERNAL", "1") + strict = UrlValidationPolicy(allow_http=True, allow_private_ips=False) + client = AiohttpClient() + try: + permissive_session = await client._get_session( + client._connect_allows_private(_PERMISSIVE) + ) + strict_session = await client._get_session( + client._connect_allows_private(strict) + ) + assert permissive_session is not strict_session + assert permissive_session.connector._resolver._allow_private_ips is True + assert strict_session.connector._resolver._allow_private_ips is False + finally: + await client.close() + + +@_allows_cleanup_closed_notice +async def test_close_closes_every_resolver(monkeypatch) -> None: + """aiohttp will not close an injected resolver, so the client must. + + ``TCPConnector`` sets ``_resolver_owner=False`` for a resolver it did not + create, and only closes its own. Without this the optional aiodns resolver + stays registered for the loop after the session is gone. + """ + monkeypatch.delenv("DYN_MM_ALLOW_INTERNAL", raising=False) + client = AiohttpClient() + session = await client._get_session(False) + resolver = session.connector._resolver + assert session.connector._resolver_owner is False + + closed = [] + original = resolver._inner.close + + async def _spy() -> None: + closed.append(1) + await original() + + resolver._inner.close = _spy + await client.close() + + assert closed, "the injected resolver was never closed" + assert client._sessions == {} + assert client._resolvers == {} + + +async def test_a_proxied_fetch_fails_closed_without_the_opt_in(monkeypatch) -> None: + """A proxied fetch puts the origin out of the resolver's reach. + + Documenting the boundary does not enforce it, so without an operator + assertion the fetch must not quietly run unchecked. Driven through + ``fetch_bytes`` so it also pins that the gate is wired into the fetch + path, not merely present as a helper. + """ + monkeypatch.delenv("DYN_MM_ALLOW_INTERNAL", raising=False) + monkeypatch.delenv("DYN_MM_TRUST_EGRESS_PROXY", raising=False) + for name in ("http_proxy", "HTTP_PROXY", "https_proxy", "HTTPS_PROXY"): + monkeypatch.setenv(name, "http://proxy.internal:3128") + + client = AiohttpClient() + try: + with pytest.raises(HttpConfigurationError) as excinfo: + # No policy, so no URL validation runs first and the gate is the + # only thing that can reject this. + await client.fetch_bytes("https://example.com/x.png", 5.0) + message = str(excinfo.value) + assert "DYN_MM_TRUST_EGRESS_PROXY" in message + # An operator configuration fault, not a verdict on the caller's URL, + # so it must not arrive as the ValueError subclass that maps to a 400. + assert not isinstance(excinfo.value, ValueError) + finally: + await client.close() + + +async def test_the_gate_guards_the_revalidating_path_too(monkeypatch) -> None: + """Both fetch seams must be gated, not just the simple one.""" + monkeypatch.delenv("DYN_MM_ALLOW_INTERNAL", raising=False) + monkeypatch.delenv("DYN_MM_TRUST_EGRESS_PROXY", raising=False) + for name in ("http_proxy", "HTTP_PROXY", "https_proxy", "HTTPS_PROXY"): + monkeypatch.setenv(name, "http://proxy.internal:3128") + + client = AiohttpClient() + try: + with pytest.raises(HttpConfigurationError) as excinfo: + await client._fetch_body_or_redirect( + "https://example.com/x.png", 5.0, policy=_STRICT + ) + assert "DYN_MM_TRUST_EGRESS_PROXY" in str(excinfo.value) + finally: + await client.close() + + +async def test_no_proxy_exempts_a_host_from_the_gate(monkeypatch) -> None: + """NO_PROXY means the fetch goes direct, so the gate must not fire. + + Asking only "is any proxy variable set" would refuse this fetch even + though the connect-time check governs it perfectly well. Asserted on the + helper directly, because driving it through a fetch reaches the transport + and passes whether or not the gate is precise. + """ + monkeypatch.delenv("DYN_MM_ALLOW_INTERNAL", raising=False) + monkeypatch.delenv("DYN_MM_TRUST_EGRESS_PROXY", raising=False) + for name in ("http_proxy", "HTTP_PROXY", "https_proxy", "HTTPS_PROXY"): + monkeypatch.setenv(name, "http://proxy.internal:3128") + for name in ("no_proxy", "NO_PROXY"): + monkeypatch.setenv(name, "example.com") + + client = AiohttpClient() + # Must not raise: NO_PROXY sends this host direct. + await client._require_trusted_egress_proxy("https://example.com/x.png") + # A host NOT covered by NO_PROXY is still gated, which is the control. + with pytest.raises(HttpConfigurationError): + await client._require_trusted_egress_proxy("https://other.invalid/x.png") + await client.close() + + +async def test_the_proxy_gate_does_not_fire_when_private_is_allowed( + monkeypatch, +) -> None: + """DYN_MM_ALLOW_INTERNAL=1 already permits private destinations. + + Nothing is left for the proxy gate to protect, so it must not demand a + second variable. + """ + monkeypatch.setenv("DYN_MM_ALLOW_INTERNAL", "1") + monkeypatch.delenv("DYN_MM_TRUST_EGRESS_PROXY", raising=False) + for name in ("http_proxy", "HTTP_PROXY"): + monkeypatch.setenv(name, "http://proxy.internal:3128") + + client = AiohttpClient() + assert client._connect_allows_private(None) is True + # The gate is keyed to the effective policy, so it never runs here. + client._require_trusted_egress_proxy # present, but not reached + await client.close() + + +async def test_the_opt_in_allows_a_proxied_fetch(monkeypatch) -> None: + """With the assertion set, the fetch proceeds and the operator owns it.""" + monkeypatch.delenv("DYN_MM_ALLOW_INTERNAL", raising=False) + monkeypatch.setenv("DYN_MM_TRUST_EGRESS_PROXY", "1") + for name in ("http_proxy", "HTTP_PROXY"): + monkeypatch.setenv(name, "http://proxy.internal:3128") + + client = AiohttpClient() + # Returns without raising, which is the whole assertion. + await client._require_trusted_egress_proxy("https://example.com/x.png") + await client.close() + + +@_allows_cleanup_closed_notice +async def test_each_policy_pool_gets_the_full_connection_limit(monkeypatch) -> None: + """The cap is per connect-time policy, and the help text says so. + + Two pools can coexist, each with the configured limit, so a deployment + that enables internal access and also issues stricter per-request policies + can reach twice the value. Pinned here so the number and the documented + meaning cannot drift apart silently. + """ + monkeypatch.setenv("DYN_MM_ALLOW_INTERNAL", "1") + client = AiohttpClient() + try: + permissive = await client._get_session(True) + strict = await client._get_session(False) + limit = client._config.max_connections + assert permissive is not strict + assert permissive.connector.limit == limit + assert strict.connector.limit == limit + finally: + await client.close() + + +async def test_the_default_configuration_only_ever_builds_one_pool(monkeypatch) -> None: + """Without DYN_MM_ALLOW_INTERNAL the connect policy is always strict. + + So the second pool, and the doubled cap with it, cannot appear in the + default deployment whatever a caller passes. + """ + monkeypatch.delenv("DYN_MM_ALLOW_INTERNAL", raising=False) + client = AiohttpClient() + strict = UrlValidationPolicy(allow_http=True, allow_private_ips=False) + assert client._connect_allows_private(None) is False + assert client._connect_allows_private(_PERMISSIVE) is False + assert client._connect_allows_private(strict) is False + await client.close() + + +async def test_the_proxy_gate_is_awaitable() -> None: + """The gate must not run aiohttp's proxy discovery on the event loop. + + ``get_env_proxy_for_url`` does proxy-bypass discovery and ``.netrc`` file + reads, which is why aiohttp itself calls it through ``asyncio.to_thread``. + A synchronous gate repeats that blocking work inline on every protected + hop. + """ + import inspect + + assert inspect.iscoroutinefunction(AiohttpClient._require_trusted_egress_proxy) diff --git a/components/src/dynamo/common/tests/http/test_http_facade.py b/components/src/dynamo/common/tests/http/test_http_facade.py index 49cc941f33fc..0d3f64086b8c 100644 --- a/components/src/dynamo/common/tests/http/test_http_facade.py +++ b/components/src/dynamo/common/tests/http/test_http_facade.py @@ -4,7 +4,7 @@ """Facade-level tests: backend resolution + SSRF revalidation loop. Per-backend exception mapping lives in -``test_aiohttp_client.py`` / ``test_httpx_client.py``. +``test_aiohttp_client.py``. """ from __future__ import annotations @@ -14,7 +14,7 @@ import pytest from dynamo.common import http as mm_http -from dynamo.common.http import AiohttpClient, HttpxClient, from_env +from dynamo.common.http import AiohttpClient, base, from_env from dynamo.common.http.url_validator import UrlValidationError, UrlValidationPolicy pytestmark = [ @@ -47,15 +47,15 @@ async def test_default_backend_is_aiohttp(monkeypatch) -> None: assert isinstance(mm_http.get_default_client(), AiohttpClient) -async def test_httpx_backend_selected(monkeypatch) -> None: +async def test_stale_backend_falls_back_to_aiohttp(monkeypatch) -> None: + # The httpx backend was removed; any non-aiohttp value warns and uses aiohttp. monkeypatch.setenv("DYN_HTTP_BACKEND", "httpx") - assert isinstance(mm_http.get_default_client(), HttpxClient) + assert isinstance(mm_http.get_default_client(), AiohttpClient) -async def test_invalid_backend_raises(monkeypatch) -> None: +async def test_unknown_backend_falls_back_to_aiohttp(monkeypatch) -> None: monkeypatch.setenv("DYN_HTTP_BACKEND", "requests") - with pytest.raises(ValueError, match="DYN_HTTP_BACKEND"): - mm_http.get_default_client() + assert isinstance(mm_http.get_default_client(), AiohttpClient) async def test_legacy_mm_http_env_var_still_honored(monkeypatch) -> None: @@ -81,10 +81,10 @@ async def test_legacy_mm_http_env_var_still_honored(monkeypatch) -> None: def _client_for(name: str): - return {"aiohttp": AiohttpClient, "httpx": HttpxClient}[name] + return {"aiohttp": AiohttpClient}[name] -@pytest.mark.parametrize("backend_name", ["aiohttp", "httpx"]) +@pytest.mark.parametrize("backend_name", ["aiohttp"]) async def test_fetch_with_policy_returns_first_response( monkeypatch, backend_name ) -> None: @@ -94,7 +94,7 @@ async def test_fetch_with_policy_returns_first_response( call_count = {"n": 0} - async def _fake(url, timeout): + async def _fake(url, timeout, *, max_bytes=None, policy=None): call_count["n"] += 1 return b"body-bytes", None @@ -106,7 +106,7 @@ async def _fake(url, timeout): assert call_count["n"] == 1 -@pytest.mark.parametrize("backend_name", ["aiohttp", "httpx"]) +@pytest.mark.parametrize("backend_name", ["aiohttp"]) async def test_fetch_with_policy_follows_safe_redirect( monkeypatch, backend_name ) -> None: @@ -115,7 +115,7 @@ async def test_fetch_with_policy_follows_safe_redirect( hops: list[str] = [] - async def _fake(url, timeout): + async def _fake(url, timeout, *, max_bytes=None, policy=None): hops.append(url) if url == "https://example.com/x.png": return None, "https://example.com/final.png" @@ -129,7 +129,7 @@ async def _fake(url, timeout): assert hops == ["https://example.com/x.png", "https://example.com/final.png"] -@pytest.mark.parametrize("backend_name", ["aiohttp", "httpx"]) +@pytest.mark.parametrize("backend_name", ["aiohttp"]) async def test_fetch_with_policy_blocks_redirect_to_private_ip( monkeypatch, backend_name ) -> None: @@ -138,7 +138,7 @@ async def test_fetch_with_policy_blocks_redirect_to_private_ip( strict = UrlValidationPolicy(allow_private_ips=False) - async def _fake(url, timeout): + async def _fake(url, timeout, *, max_bytes=None, policy=None): return None, "http://169.254.169.254/latest/meta-data/" with patch.object(client, "_fetch_body_or_redirect", _fake): @@ -146,7 +146,7 @@ async def _fake(url, timeout): await mm_http.fetch_bytes("https://8.8.8.8/x.png", 30.0, policy=strict) -@pytest.mark.parametrize("backend_name", ["aiohttp", "httpx"]) +@pytest.mark.parametrize("backend_name", ["aiohttp"]) async def test_fetch_with_policy_enforces_redirect_limit( monkeypatch, backend_name ) -> None: @@ -161,9 +161,110 @@ async def test_fetch_with_policy_enforces_redirect_limit( "https://example.com/d": "https://example.com/e", } - async def _fake(url, timeout): + async def _fake(url, timeout, *, max_bytes=None, policy=None): return None, chain[url] with patch.object(client, "_fetch_body_or_redirect", _fake): with pytest.raises(UrlValidationError, match="Too many redirects"): await mm_http.fetch_bytes("https://example.com/a", 30.0, policy=_PERMISSIVE) + + +# --- Download cap (collect_capped) --- +# +# The cap exists because moving the diffusion download out of SGLang left its +# media_url_max_file_size_mb behind. It has to hold while the body streams: a +# declared Content-Length is attacker-controlled and absent when chunked. + + +async def _chunks(*sizes: int): + for n in sizes: + yield b"x" * n + + +async def test_collect_capped_joins_a_body_under_the_limit() -> None: + assert await base.collect_capped(_chunks(4, 4, 2), "u", 100) == b"x" * 10 + + +async def test_collect_capped_without_a_limit_reads_everything() -> None: + assert len(await base.collect_capped(_chunks(50, 50), "u", None)) == 100 + + +async def test_collect_capped_refuses_a_body_over_the_limit() -> None: + with pytest.raises(UrlValidationError, match="download limit"): + await base.collect_capped(_chunks(60, 60), "u", 100) + + +async def test_collect_capped_counts_across_short_chunks() -> None: + """Each chunk is well under the limit; only the running total exceeds it. + + aiohttp's ``read(n)`` returns at most n bytes and in practice returns far + fewer, so a single capped read would let this body through. + """ + with pytest.raises(UrlValidationError, match="download limit"): + await base.collect_capped(_chunks(*([10] * 20)), "u", 100) + + +async def test_collect_capped_does_not_echo_an_unbounded_url() -> None: + url = "https://example.com/" + "a" * 200_000 + with pytest.raises(UrlValidationError) as excinfo: + await base.collect_capped(_chunks(200), url, 100) + + assert len(str(excinfo.value)) < 500 + + +async def test_http_status_error_bounds_both_halves_of_its_message() -> None: + """The video diffusion handler puts str(exc) straight in its response body, + and aiohttp's own status-error text repeats the URL, so neither the url nor + the backend message can go in at full length.""" + url = "https://example.com/" + "u" * 200_000 + err = base.HttpStatusError(404, "Not Found " + "m" * 200_000, url) + + # The binding reads .message off this class by name and forwards it + # verbatim to the client on a 4xx (errors.rs extract_http_like_error), so + # bounding only the rendered string leaves the client-facing path open. + assert len(err.message) <= base._MAX_MESSAGE_LENGTH + assert len(str(err)) < 2 * base._MAX_MESSAGE_LENGTH + # .url is not part of that protocol and keeps its full value. + assert err.url == url + + +async def test_http_status_error_keeps_both_ends_of_a_long_detail() -> None: + """aiohttp renders the client-supplied host before the errno, so a + head-only bound would keep the attacker's string and drop the reason.""" + detail = "Cannot connect to host " + "h" * 40_000 + ":80 [nodename not known]" + err = base.HttpStatusError(400, detail, "https://example.com/x") + + assert "Cannot connect" in err.message + assert "nodename not known" in err.message + + +async def test_http_status_error_keeps_a_real_guidance_message_intact() -> None: + """Callers build actionable guidance in ``message``; bounding must not eat it. + + The trtllm video-decoder hint is ~480 characters of spec, installer command + and vendor cause, and its own test allows 2000. A tight bound here deletes + all of that and leaves a truncated prefix, which is how this class was + broken once already. + """ + guidance = ( + "Cannot decode video: this video (an undetected codec) has no decoder " + "in this image: shipped images decode only H.264/H.265 (in hardware, " + "via NVDEC), and the software decoder 'cv2' is deliberately not " + "installed. Re-encode the input to H.264/H.265, or install the " + "validated decoder with `pip install --no-deps " + "'opencv-python-headless==4.13.0.90'` (or `python -m " + "dynamo.common.utils.install_media_decoders trtllm`) " + "(decoder reported: OpenCV (cv2) is required for video decoding)" + ) + err = base.HttpStatusError( + 500, guidance, "data:video/mp4 (90022 chars, payload elided)" + ) + + assert err.message == guidance, "legitimate guidance must survive untouched" + for needle in ( + "install_media_decoders trtllm", + "opencv-python-headless", + "OpenCV (cv2) is required for video decoding", + ): + assert needle in str(err), f"{needle!r} was truncated away" + assert len(str(err)) < 2_000 diff --git a/components/src/dynamo/common/tests/http/test_httpx_client.py b/components/src/dynamo/common/tests/http/test_httpx_client.py deleted file mode 100644 index b152e63a67d1..000000000000 --- a/components/src/dynamo/common/tests/http/test_httpx_client.py +++ /dev/null @@ -1,141 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Unit tests for ``HttpxClient`` exception mapping + redirect parsing. - -The client singleton is held on the instance (``client._client``) and -replaced via ``patch.object``; the autouse ``_close_shared_http_client`` -fixture in ``conftest.py`` resets the process-wide singleton between -tests. - -Redirect parsing is exercised through the public ``fetch_bytes(..., -policy=...)`` path so we don't reach into the -``_fetch_body_or_redirect`` abstract-method seam from tests. -""" - -from __future__ import annotations - -from unittest.mock import AsyncMock, MagicMock - -import httpx -import pytest - -from dynamo.common import http as mm_http -from dynamo.common.http import HttpxClient -from dynamo.common.http.url_validator import UrlValidationPolicy - -pytestmark = [ - pytest.mark.asyncio, - pytest.mark.unit, - pytest.mark.gpu_0, - pytest.mark.pre_merge, -] - - -def _async_returning(value): - """Coroutine factory used as ``side_effect`` for an awaited async call.""" - - async def _coro(*args, **kwargs): - return value - - return _coro - - -def _async_raising(exc): - async def _coro(*args, **kwargs): - raise exc - - return _coro - - -def _make_client_with_inner(inner) -> HttpxClient: - client = HttpxClient() - client._client = inner - return client - - -_PERMISSIVE = UrlValidationPolicy(allow_http=True, allow_private_ips=True) - - -async def test_fetch_bytes_returns_body_on_200() -> None: - response = MagicMock(spec=httpx.Response) - response.content = b"hello" - response.raise_for_status = MagicMock(return_value=None) - inner = MagicMock(spec=httpx.AsyncClient) - inner.is_closed = False - inner.get = MagicMock(side_effect=_async_returning(response)) - client = _make_client_with_inner(inner) - result = await client.fetch_bytes("https://h/x", 30.0) - assert result == b"hello" - - -async def test_fetch_bytes_maps_timeout() -> None: - inner = MagicMock(spec=httpx.AsyncClient) - inner.is_closed = False - inner.get = MagicMock(side_effect=_async_raising(httpx.ConnectTimeout("timeout"))) - client = _make_client_with_inner(inner) - with pytest.raises(mm_http.HttpTimeoutError) as exc: - await client.fetch_bytes("https://h/x", 30.0) - assert isinstance(exc.value.__cause__, httpx.ConnectTimeout) - - -async def test_fetch_bytes_maps_status() -> None: - response = MagicMock(spec=httpx.Response) - response.status_code = 404 - response.raise_for_status = MagicMock( - side_effect=httpx.HTTPStatusError( - "404 Not Found", request=MagicMock(), response=response - ) - ) - inner = MagicMock(spec=httpx.AsyncClient) - inner.is_closed = False - inner.get = MagicMock(side_effect=_async_returning(response)) - client = _make_client_with_inner(inner) - with pytest.raises(mm_http.HttpStatusError) as exc: - await client.fetch_bytes("https://h/x", 30.0) - assert exc.value.status == 404 - - -async def test_fetch_bytes_maps_connection_error() -> None: - inner = MagicMock(spec=httpx.AsyncClient) - inner.is_closed = False - inner.get = MagicMock(side_effect=_async_raising(httpx.ConnectError("refused"))) - client = _make_client_with_inner(inner) - with pytest.raises(mm_http.HttpConnectionError): - await client.fetch_bytes("https://h/x", 30.0) - - -async def test_redirect_resolved_through_policy_path() -> None: - """302 โ†’ absolute next URL is parsed correctly when the SSRF policy - drives the redirect loop. Verifies relative-Location resolution - against the response URL through the public API.""" - redirect_response = MagicMock(spec=httpx.Response) - redirect_response.is_redirect = True - redirect_response.headers = {"location": "/next.png"} - redirect_response.url = httpx.URL("https://h/x.png") - redirect_response.aclose = AsyncMock(return_value=None) - - final_response = MagicMock(spec=httpx.Response) - final_response.is_redirect = False - final_response.content = b"final" - final_response.raise_for_status = MagicMock(return_value=None) - final_response.aclose = AsyncMock(return_value=None) - - responses_by_url = { - "https://h/x.png": redirect_response, - "https://h/next.png": final_response, - } - - async def _send(request, *args, **kwargs): - return responses_by_url[str(request.url)] - - inner = MagicMock(spec=httpx.AsyncClient) - inner.is_closed = False - inner.build_request = MagicMock( - side_effect=lambda method, url, **kw: MagicMock(url=httpx.URL(url)) - ) - inner.send = MagicMock(side_effect=_send) - - client = _make_client_with_inner(inner) - body = await client.fetch_bytes("https://h/x.png", 30.0, policy=_PERMISSIVE) - assert body == b"final" diff --git a/components/src/dynamo/common/tests/http/test_media_reference.py b/components/src/dynamo/common/tests/http/test_media_reference.py new file mode 100644 index 000000000000..3c7bca67086e --- /dev/null +++ b/components/src/dynamo/common/tests/http/test_media_reference.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for ``local_media_reference``. + +The redirect-SSRF fix: a URL reference must be fetched through the policy-aware +client (which revalidates each redirect hop) into a temp file, so the caller +hands a downstream generator only a trusted local path โ€” never a URL it would +fetch and redirect itself. Local references pass through unchanged. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from dynamo.common.http.media_reference import ( + DEFAULT_MAX_MEDIA_MB, + DYN_MM_MAX_FILE_SIZE_MB, + MAX_MEDIA_BYTES, + local_media_reference, + max_media_bytes, +) +from dynamo.common.http.url_validator import UrlValidationError, UrlValidationPolicy + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.gpu_0, +] + + +async def test_url_reference_is_fetched_to_a_temp_path(monkeypatch) -> None: + seen = {} + + async def fake_fetch(url, timeout, *, policy=None, max_bytes=None): + seen["url"] = url + seen["policy"] = policy + seen["max_bytes"] = max_bytes + return b"PNGDATA" + + # patched on the package: local_media_reference imports it lazily as `fetch_bytes` + monkeypatch.setattr("dynamo.common.http.fetch_bytes", fake_fetch) + + # allow_private_ips=True keeps validate_url off the network in the test. + policy = UrlValidationPolicy(allow_http=True, allow_private_ips=True) + url = "https://example.com/img.png" + + async with local_media_reference(url, policy) as path: + assert path != url # a local path, not the URL + assert os.path.exists(path) + assert Path(path).read_bytes() == b"PNGDATA" + assert path.endswith(".png") # suffix preserved for the generator + tmp = path + + assert seen["url"] == url + assert not os.path.exists(tmp) # temp cleaned on exit + # Without the policy, fetch_bytes takes the backend's own redirect path and + # no hop is revalidated โ€” the entire redirect-SSRF guard would be gone with + # every other assertion here still passing. + assert seen["policy"] is policy + assert seen["max_bytes"] == MAX_MEDIA_BYTES + + +async def test_local_reference_passes_through(monkeypatch, tmp_path) -> None: + monkeypatch.setenv("DYN_MM_LOCAL_PATH", str(tmp_path)) + ref = tmp_path / "x.png" + ref.write_bytes(b"x") + + async with local_media_reference(str(ref), UrlValidationPolicy.from_env()) as path: + assert path == str(ref.resolve()) # yielded as-is, no temp file + + assert ref.exists() # a local reference is never deleted + + +async def test_a_data_uri_is_rejected_not_yielded_as_a_path() -> None: + """`data:` passes validate_url, but it is a URI, not a filesystem path. + + Yielding it unchanged hands the generator a string it will try to open as a + file, and puts the whole inline payload โ€” which can be megabytes โ€” wherever + that failure is reported. + """ + policy = UrlValidationPolicy() + + with pytest.raises(UrlValidationError, match="cannot be resolved to a local file"): + async with local_media_reference("data:image/png;base64,iVBORw0KGgo=", policy): + pass + + +async def test_data_uri_rejection_does_not_echo_the_payload() -> None: + payload = "A" * 200_000 + + with pytest.raises(UrlValidationError) as excinfo: + async with local_media_reference( + f"data:image/png;base64,{payload}", UrlValidationPolicy() + ): + pass + + assert payload not in str(excinfo.value) + assert len(str(excinfo.value)) < 500 + + +async def test_an_overlong_url_extension_does_not_become_a_temp_filename( + monkeypatch, +) -> None: + """The temp-file suffix comes from the client's URL path. + + Passed through, a 300-character extension makes mkstemp raise ENAMETOOLONG + โ€” an OSError the caller does not expect, carrying the server's temp + directory back to the client. + """ + + async def fake_fetch(url, timeout, *, policy=None, max_bytes=None): + return b"PNGDATA" + + monkeypatch.setattr("dynamo.common.http.fetch_bytes", fake_fetch) + policy = UrlValidationPolicy(allow_http=True, allow_private_ips=True) + + async with local_media_reference( + "https://example.com/img." + "a" * 300, policy + ) as path: + assert Path(path).read_bytes() == b"PNGDATA" + assert len(Path(path).name) < 64 + + +async def test_a_non_alphanumeric_url_extension_is_dropped(monkeypatch) -> None: + async def fake_fetch(url, timeout, *, policy=None, max_bytes=None): + return b"PNGDATA" + + monkeypatch.setattr("dynamo.common.http.fetch_bytes", fake_fetch) + policy = UrlValidationPolicy(allow_http=True, allow_private_ips=True) + + async with local_media_reference( + "https://example.com/img.%2F..%2F..%2Fetc%2Fpasswd", policy + ) as path: + assert Path(path).suffix == "" + + +async def test_the_download_cap_reaches_the_client(monkeypatch) -> None: + """The caller's limit must be handed to fetch_bytes, which enforces it. + + Enforcement itself is collect_capped's job (test_http_facade.py); what can + silently regress here is the limit never being passed at all. + """ + seen = {} + + async def fake_fetch(url, timeout, *, policy=None, max_bytes=None): + seen["max_bytes"] = max_bytes + return b"x" + + monkeypatch.setattr("dynamo.common.http.fetch_bytes", fake_fetch) + policy = UrlValidationPolicy(allow_http=True, allow_private_ips=True) + + async with local_media_reference( + "https://example.com/big.png", policy, max_bytes=512 + ): + pass + + assert seen["max_bytes"] == 512 + + +def test_max_media_bytes_defaults_when_unset(monkeypatch) -> None: + """Unset keeps the SGLang default this path carried forward.""" + monkeypatch.delenv(DYN_MM_MAX_FILE_SIZE_MB, raising=False) + + assert max_media_bytes() == MAX_MEDIA_BYTES + assert MAX_MEDIA_BYTES == DEFAULT_MAX_MEDIA_MB * 1024 * 1024 + + +def test_max_media_bytes_honours_the_override(monkeypatch) -> None: + """The point of the knob: an operator with media over 64 MiB can raise it.""" + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, "128") + + assert max_media_bytes() == 128 * 1024 * 1024 + + +def test_max_media_bytes_tolerates_surrounding_whitespace(monkeypatch) -> None: + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, " 256 ") + + assert max_media_bytes() == 256 * 1024 * 1024 + + +@pytest.mark.parametrize("raw", ["abc", "", " ", "12.5"]) +def test_max_media_bytes_falls_back_on_an_unparseable_value(monkeypatch, raw) -> None: + """A malformed operator value must not take the worker down.""" + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, raw) + + assert max_media_bytes() == MAX_MEDIA_BYTES + + +@pytest.mark.parametrize("raw", ["0", "-5"]) +def test_max_media_bytes_refuses_to_drop_the_bound(monkeypatch, raw) -> None: + """Non-positive must not read as "unlimited" -- that removes the bound.""" + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, raw) + + assert max_media_bytes() == MAX_MEDIA_BYTES + + +async def test_local_media_reference_uses_the_configured_cap(monkeypatch) -> None: + """The resolved value has to reach fetch_bytes, not just exist.""" + seen = {} + + async def fake_fetch(url, timeout, *, policy=None, max_bytes=None): + seen["max_bytes"] = max_bytes + return b"PNGDATA" + + monkeypatch.setattr("dynamo.common.http.fetch_bytes", fake_fetch) + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, "128") + policy = UrlValidationPolicy(allow_http=True, allow_private_ips=True) + + async with local_media_reference("https://example.com/img.png", policy): + pass + + assert seen["max_bytes"] == 128 * 1024 * 1024 + + +async def test_an_explicit_max_bytes_still_wins_over_the_env(monkeypatch) -> None: + """Control: the knob must not override a caller that asked for a value.""" + seen = {} + + async def fake_fetch(url, timeout, *, policy=None, max_bytes=None): + seen["max_bytes"] = max_bytes + return b"PNGDATA" + + monkeypatch.setattr("dynamo.common.http.fetch_bytes", fake_fetch) + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, "128") + policy = UrlValidationPolicy(allow_http=True, allow_private_ips=True) + + async with local_media_reference( + "https://example.com/img.png", policy, max_bytes=512 + ): + pass + + assert seen["max_bytes"] == 512 diff --git a/components/src/dynamo/common/tests/http/test_ssrf_resolver.py b/components/src/dynamo/common/tests/http/test_ssrf_resolver.py new file mode 100644 index 000000000000..3f2a0729d091 --- /dev/null +++ b/components/src/dynamo/common/tests/http/test_ssrf_resolver.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the aiohttp connect-time SSRF resolver (``_ssrf_resolver``). + +Exercise the DNS-rebinding case deterministically: the resolver is fed a mix of +public and blocked answers (as if a rebinding server flipped between check and +connect) and must drop the blocked ones, keep the rest, and fail closed when +none remain. No network. +""" + +from __future__ import annotations + +import errno + +import pytest + +from dynamo.common.http._ssrf_resolver import BlocklistResolver, SsrfBlockedAddress + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.gpu_0, +] + + +class _FakeInner: + """Stand-in for aiohttp's DefaultResolver returning canned entries.""" + + def __init__(self, ips: list[str]) -> None: + self._ips = ips + + async def resolve(self, host, port=0, family=0): + return [{"hostname": host, "host": ip, "port": port} for ip in self._ips] + + async def close(self): + pass + + +def _resolver_with(ips: list[str]) -> BlocklistResolver: + r = BlocklistResolver(allow_private_ips=False) + r._inner = _FakeInner(ips) + return r + + +async def test_resolver_drops_blocked_and_keeps_the_rest() -> None: + # Rebinding answer (public + metadata IP): the blocked one is dropped, and + # every non-blocked address survives for multi-address fallback. + resolver = _resolver_with(["93.184.216.34", "169.254.169.254", "8.8.8.8"]) + out = await resolver.resolve("evil.example.com") + assert [h["host"] for h in out] == ["93.184.216.34", "8.8.8.8"] + + +async def test_resolver_fails_closed_when_only_blocked() -> None: + resolver = _resolver_with(["169.254.169.254"]) + with pytest.raises(SsrfBlockedAddress): + await resolver.resolve("evil.example.com") + + +async def test_allow_private_ips_bypasses_filtering() -> None: + # The env baseline (DYN_MM_ALLOW_INTERNAL=1) opts the deployment into + # internal targets, so the resolver returns every answer unfiltered. + resolver = BlocklistResolver(allow_private_ips=True) + resolver._inner = _FakeInner(["169.254.169.254", "10.0.0.5"]) + out = await resolver.resolve("internal.svc") + assert [h["host"] for h in out] == ["169.254.169.254", "10.0.0.5"] + + +async def test_configured_egress_proxy_is_not_filtered(monkeypatch) -> None: + """A corporate proxy on a private address must not be filtered out. + + The session runs with ``trust_env=True``, so with a proxy configured the + connector dials the proxy and it is the proxy's own address that reaches + this resolver. Filtering it turns every fetch into a connection error -- + measured before the exemption, HTTP_PROXY at a private host failed them all. + """ + # aiohttp's proxies_from_env() prefers the lowercase name, and CI images can + # carry one already -- clear both so this test controls what is configured. + monkeypatch.delenv("http_proxy", raising=False) + monkeypatch.setenv("HTTP_PROXY", "http://proxy.internal:3128") + + resolver = BlocklistResolver(allow_private_ips=False) + resolver._inner = _FakeInner(["10.1.2.3"]) + + hosts = await resolver.resolve("proxy.internal") + + assert [h["host"] for h in hosts] == ["10.1.2.3"] + + +async def test_a_non_proxy_host_is_still_filtered(monkeypatch) -> None: + """Control: the exemption is for the configured proxy only.""" + monkeypatch.delenv("http_proxy", raising=False) + monkeypatch.setenv("HTTP_PROXY", "http://proxy.internal:3128") + + resolver = BlocklistResolver(allow_private_ips=False) + resolver._inner = _FakeInner(["10.1.2.3"]) + + with pytest.raises(SsrfBlockedAddress): + await resolver.resolve("origin.example.com") + + +async def test_blocked_message_bounds_the_hostname() -> None: + """The host is caller-supplied; this text reaches a response and a log.""" + resolver = _resolver_with(["169.254.169.254"]) + + with pytest.raises(SsrfBlockedAddress) as excinfo: + await resolver.resolve("h" * 200_000 + ".example.com") + + assert len(str(excinfo.value)) < 500 + assert "h" * 200 not in str(excinfo.value) + + +async def test_blocked_reason_survives_into_strerror() -> None: + """aiohttp renders ``strerror``, not ``str(exc)``, for a connector error. + + A one-argument OSError leaves ``strerror`` as None, and the caller then + receives ``Cannot connect to host ... [None]`` with the reason dropped at + the one place an operator reads it. + """ + resolver = _resolver_with(["169.254.169.254"]) + + with pytest.raises(SsrfBlockedAddress) as excinfo: + await resolver.resolve("evil.example.com") + + assert excinfo.value.strerror == ( + "host evil.example.com resolves only to blocked IPs" + ) + assert excinfo.value.errno == errno.EHOSTUNREACH diff --git a/components/src/dynamo/common/tests/http/test_url_validator.py b/components/src/dynamo/common/tests/http/test_url_validator.py index d3cff55ebcbb..5484150f76e8 100644 --- a/components/src/dynamo/common/tests/http/test_url_validator.py +++ b/components/src/dynamo/common/tests/http/test_url_validator.py @@ -22,8 +22,11 @@ from dynamo.common.http.url_validator import ( UrlValidationError, UrlValidationPolicy, + describe_error_detail, + describe_media_source, is_blocked_ip, validate_local_path, + validate_media_reference, validate_media_url, validate_url, ) @@ -359,8 +362,8 @@ def test_policy_from_env_allow_internal(monkeypatch) -> None: # Fetch-with-revalidation tests now live in test_http_backends.py where -# they exercise the backend-neutral facade path against both httpx and -# aiohttp. See ``test_fetch_with_policy_*``. +# they exercise the backend-neutral facade path against the aiohttp +# client. See ``test_fetch_with_policy_*``. # --------------------------------------------------------------------------- @@ -430,3 +433,180 @@ async def test_validate_media_url_rejects_file_uri_outside_prefix(tmp_path) -> N with pytest.raises(UrlValidationError, match="outside the allowed directory"): await validate_media_url(other.resolve().as_uri(), policy) + + +async def test_validate_media_url_rejects_a_percent_encoded_nul(tmp_path) -> None: + """%00 must stay a UrlValidationError, not leak out as a bare ValueError. + + file:// paths are percent-decoded, so %00 reaches Path.resolve() as a real + NUL and lstat() raises ValueError rather than OSError. Callers key their + 4xx-vs-5xx decision on the UrlValidationError type (see image_loader and + video_loader), so letting a plain ValueError escape turns a rejected input + into a server error. + """ + policy = UrlValidationPolicy(allowed_local_path=str(tmp_path)) + + with pytest.raises(UrlValidationError): + await validate_media_url(f"file://{tmp_path}/x%00.png", policy) + + +def test_validate_local_path_bounds_the_path_in_its_message(tmp_path) -> None: + """The rejected path is client-supplied and unbounded; the message is not. + + It lands in an error response and in a log line, so one request would + otherwise amplify into that much text at every sink. + """ + policy = UrlValidationPolicy(allowed_local_path=str(tmp_path)) + path = "/nope/" + "A" * 200_000 + ".png" + + with pytest.raises(UrlValidationError) as excinfo: + validate_local_path(path, policy) + + assert len(str(excinfo.value)) < 500 + assert "200010 chars" in str(excinfo.value) # true size stays visible + + +def test_validate_local_path_bounds_the_oserror_diagnostic(tmp_path) -> None: + """The errno text repeats the filename, which the label bound does not cover. + + The ``/nope/`` case above stops at the missing parent and never reaches the + filesystem. Under an *existing* parent the name is handed to lstat(), + ENAMETOOLONG comes back, and ``str(exc)`` carries the whole + 200,000-character filename into the same error response and log line the + label bound was there to protect. + """ + policy = UrlValidationPolicy(allowed_local_path=str(tmp_path)) + path = str(tmp_path / ("n" * 200_000)) + + with pytest.raises(UrlValidationError) as excinfo: + validate_local_path(path, policy) + + assert len(str(excinfo.value)) < 500 + assert "n" * 200 not in str(excinfo.value) + + +def test_describe_media_source_bounds_data_uri_metadata() -> None: + """Eliding the payload is not enough: the media-type field is client-supplied. + + Everything before the comma is the metadata, so a reference of + ``"data:" + "A" * 200_000 + ",AAAA"`` rendered a 200,036-character label with + the payload already elided. Omitting the comma takes the same branch. + """ + for source in ("data:" + "A" * 200_000 + ",AAAA", "data:" + "A" * 200_000): + label = describe_media_source(source) + + assert len(label) < 500 + assert "A" * 200 not in label + assert f"({len(source)} chars" in label # true size stays visible + + +def test_describe_media_source_keeps_an_ordinary_data_uri_intact() -> None: + """Control: a real media type is short and must survive the new bound.""" + label = describe_media_source("data:image/png;base64," + "A" * 50_000) + + assert label.startswith("data:image/png (") + assert "payload elided" in label + + +def test_validate_local_path_keeps_an_ordinary_path_intact(tmp_path) -> None: + """Control: bounding must not change the message for a normal path.""" + policy = UrlValidationPolicy(allowed_local_path=str(tmp_path)) + + with pytest.raises(UrlValidationError, match=r"Path '/etc/passwd' is outside"): + validate_local_path("/etc/passwd", policy) + + +async def test_validate_media_reference_rejects_empty(tmp_path) -> None: + """Parity with validate_media_url, which guards the empty string.""" + policy = UrlValidationPolicy(allowed_local_path=str(tmp_path)) + + with pytest.raises(UrlValidationError, match="empty"): + await validate_media_reference("", policy) + + +@pytest.mark.parametrize( + "url", + [ + "https:///" + "A" * 200_000, # no host component + "A" * 200_000 + "://x", # scheme is client-supplied too + ], +) +async def test_validate_url_bounds_the_url_in_its_message(url) -> None: + """These messages became client-visible once the diffusion handlers + started mapping UrlValidationError to a 400 that carries the reason.""" + with pytest.raises(UrlValidationError) as excinfo: + await validate_url(url, STRICT_HTTPS) + + assert len(str(excinfo.value)) < 500 + + +async def test_redirect_chain_in_the_limit_message_is_bounded() -> None: + """The chain is entirely attacker-chosen URLs.""" + from dynamo.common.http.base import HttpClient + + long_hop = "https://example.com/" + "A" * 200_000 + + class _Client(HttpClient): + async def _fetch_simple(self, url, timeout, *, max_bytes=None, policy=None): + raise AssertionError("unused") + + async def _fetch_body_or_redirect( + self, url, timeout, *, max_bytes=None, policy=None + ): + return None, long_hop + + async def close(self): + return None + + policy = UrlValidationPolicy(allow_http=True, allow_private_ips=True) + with pytest.raises(UrlValidationError, match="Too many redirects") as excinfo: + await _Client().fetch_bytes(long_hop, 1.0, policy=policy) + + assert len(str(excinfo.value)) < 1000 + + +async def test_a_data_uri_does_not_pay_for_a_label_it_cannot_use(monkeypatch) -> None: + """describe_media_source copies the source, and a data: URI is the payload. + + The data branch returns before any message is built, so building the label + first made validate_url O(payload): 1.35 ms for a 32 MiB URI against + 0.03 ms after. + """ + calls = [] + monkeypatch.setattr( + url_validator, + "describe_media_source", + lambda src, *a, **kw: calls.append(src) or src, + ) + + url = "data:image/png;base64,AAAA" + assert await url_validator.validate_url(url, STRICT_HTTPS) == url + assert calls == [] + + +def test_missing_allowed_dir_does_not_name_it(tmp_path) -> None: + """Same disclosure as the 'outside the allowed directory' message below it: + this one reaches the client too, on an ordinary misconfiguration.""" + policy = UrlValidationPolicy(allowed_local_path=str(tmp_path / "gone")) + # The input itself must resolve, or "File not found" fires first. + media = tmp_path / "x.png" + media.write_bytes(b"x") + + with pytest.raises(UrlValidationError, match="does not exist") as excinfo: + validate_local_path(str(media), policy) + + assert "gone" not in str(excinfo.value) + + +def test_describe_error_detail_keeps_a_short_detail_intact() -> None: + assert describe_error_detail("[Errno 8] not known") == "[Errno 8] not known" + + +def test_describe_error_detail_bounds_from_both_ends() -> None: + detail = "head-of-the-message " + "x" * 40_000 + " tail-of-the-message" + label = describe_error_detail(detail) + + assert len(label) < 300 + assert label.startswith("head-of-the-message") + assert label.endswith("tail-of-the-message") + assert str(len(detail)) in label # true size stays visible diff --git a/components/src/dynamo/common/tests/multimodal/conftest.py b/components/src/dynamo/common/tests/multimodal/conftest.py index 5fa2acd9b033..c9b7e8524975 100644 --- a/components/src/dynamo/common/tests/multimodal/conftest.py +++ b/components/src/dynamo/common/tests/multimodal/conftest.py @@ -15,10 +15,15 @@ import importlib +import pytest import pytest_asyncio from dynamo.common.http import close_http_client +# The decode carriers ``codec_errors`` probes. Only these names are answered by +# the ``carrier_imports`` fixture; every other import runs for real. +_MEDIA_CARRIERS = ("cv2", "av", "decord") + # Cached results of probing optional deps used by multimodal unit tests. # `None` = not attempted, `True` = importable, `False` = raised. _torch_importable: bool | None = None @@ -70,6 +75,32 @@ def pytest_ignore_collect(collection_path, config) -> bool | None: return None +@pytest.fixture +def carrier_imports(monkeypatch): + """Mock media-carrier imports while leaving unrelated imports intact.""" + # Import after collection's optional-dependency guard: the multimodal + # package imports torch, which CPU-only runtime images may not ship. + from dynamo.common.multimodal import codec_errors + + def _set( + present: tuple[str, ...] = (), error: str | BaseException | None = None + ) -> None: + real = importlib.import_module + + def _fake(name: str, *args, **kwargs): + if name in _MEDIA_CARRIERS: + if name in present: + return object() + if isinstance(error, BaseException): + raise error + raise ImportError(error or f"No module named '{name}'") + return real(name, *args, **kwargs) + + monkeypatch.setattr(codec_errors.importlib, "import_module", _fake) + + return _set + + @pytest_asyncio.fixture(autouse=True) async def _close_shared_http_client(): yield diff --git a/components/src/dynamo/common/tests/multimodal/test_audio_loader.py b/components/src/dynamo/common/tests/multimodal/test_audio_loader.py index 81497dd38f54..d8a8cf9ac946 100644 --- a/components/src/dynamo/common/tests/multimodal/test_audio_loader.py +++ b/components/src/dynamo/common/tests/multimodal/test_audio_loader.py @@ -6,13 +6,15 @@ import numpy as np import pytest -import dynamo.common.multimodal.audio_loader as audio_loader_module -from dynamo.common.http import HttpStatusError +from dynamo.common.http import HttpConfigurationError, HttpStatusError from dynamo.common.http.url_validator import UrlValidationError, UrlValidationPolicy +from dynamo.common.multimodal import audio_loader as audio_loader_module from dynamo.common.multimodal.audio_loader import AudioLoader from dynamo.common.multimodal.codec_errors import MissingMediaDecoderError from dynamo.common.utils.install_media_decoders import VALIDATED_SPECS +from dynamo.common.http.media_reference import DYN_MM_MAX_FILE_SIZE_MB # isort: skip + pytestmark = [ pytest.mark.unit, pytest.mark.gpu_0, @@ -94,6 +96,25 @@ async def test_load_audio_rejects_empty_waveform(): await loader.load_audio("https://example.com/empty.wav") +@pytest.mark.asyncio +async def test_http_fetch_honors_configured_media_limit(monkeypatch): + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, "1") + loader = AudioLoader(url_policy=_permissive_http_policy()) + waveform = np.zeros(8, dtype=np.float32) + + class _MediaIO: + def load_bytes(self, content): + return waveform, 16000.0 + + mock_fetch = AsyncMock(return_value=b"audio") + monkeypatch.setattr(audio_loader_module, "fetch_bytes", mock_fetch) + monkeypatch.setattr(loader, "_create_vllm_audio_io", lambda: _MediaIO()) + + await loader._load_audio_with_vllm("https://example.com/limited.wav") + + assert mock_fetch.await_args.kwargs["max_bytes"] == 1024 * 1024 + + @pytest.mark.asyncio async def test_load_audio_batch_uses_url_loader(): loader = AudioLoader() @@ -187,11 +208,14 @@ async def test_load_audio_batch_reads_decoded_variant(monkeypatch): @pytest.mark.asyncio -async def test_load_audio_missing_decoder_is_actionable(): +async def test_load_audio_missing_decoder_is_actionable(carrier_imports): """vLLM's own hint here is `pip install vllm[audio]`, which drags in an unpinned stack; the wrap must point at the validated bounded install and say there is no hardware alternative for audio.""" loader = AudioLoader() + # 'av' genuinely absent, independent of what the test machine has + # installed: _install_hint asks the running interpreter. + carrier_imports() loader._load_audio_with_vllm = AsyncMock( # type: ignore[method-assign] side_effect=ImportError("Please install vllm[audio] for audio support") ) @@ -217,3 +241,26 @@ async def test_load_audio_batch_preserves_missing_decoder_error(): await loader.load_audio_batch([{"Url": "https://example.com/x.mp3"}]) assert exc_info.value is err + + +@pytest.mark.asyncio +async def test_load_audio_preserves_a_configuration_error(monkeypatch): + """An operator fault must not reach the client as a 4xx. + + ``load_audio`` converts unknown exceptions into ``ValueError``, and + ``py_err_to_dynamo`` maps ``ValueError`` to ``InvalidArgument``. Measured + before this was preserved: an ``HttpError`` from the egress-proxy gate + arrived as ``builtins.ValueError``, so audio clients saw a 4xx for a + deployment misconfiguration. + """ + loader = AudioLoader.__new__(AudioLoader) + + async def _boom(url): + raise HttpConfigurationError("egress proxy is not trusted") + + monkeypatch.setattr(loader, "_load_audio_with_vllm", _boom, raising=False) + + with pytest.raises(HttpConfigurationError) as excinfo: + await loader.load_audio("https://example.com/a.wav") + + assert not isinstance(excinfo.value, ValueError) diff --git a/components/src/dynamo/common/tests/multimodal/test_codec_errors.py b/components/src/dynamo/common/tests/multimodal/test_codec_errors.py index 5e2d45755a35..1a4c60341a3e 100644 --- a/components/src/dynamo/common/tests/multimodal/test_codec_errors.py +++ b/components/src/dynamo/common/tests/multimodal/test_codec_errors.py @@ -5,7 +5,7 @@ import pytest -import dynamo.common.multimodal.codec_errors as codec_errors +from dynamo.common.multimodal import codec_errors from dynamo.common.multimodal.codec_errors import ( MissingMediaDecoderError, audio_decoder_missing, @@ -16,7 +16,13 @@ pytestmark = [pytest.mark.unit, pytest.mark.pre_merge, pytest.mark.gpu_0] -def test_video_message_names_codec_spec_and_installer(monkeypatch): +@pytest.fixture(autouse=True) +def _carrier_absent(carrier_imports): + """Make missing-carrier messages independent of installed packages.""" + carrier_imports() + + +def test_video_message_names_codec_and_spec(monkeypatch): monkeypatch.setattr(codec_errors, "nvdec_available", lambda: True) err = video_decoder_missing("vllm", "opencv-python-headless", "cv2", "vp9") @@ -32,6 +38,61 @@ def test_video_message_names_codec_spec_and_installer(monkeypatch): assert "H.264/H.265" in msg +def test_present_but_unusable_carrier_pins_the_installed_version( + monkeypatch, carrier_imports +): + """Replacing codec-free OpenCV must preserve the installed version.""" + carrier_imports(present=("cv2",)) + monkeypatch.setattr( + codec_errors.importlib.metadata, "version", lambda p: "5.0.0.93" + ) + msg = str(video_decoder_missing("vllm", "opencv-python-headless", "cv2", "vp9")) + assert "--force-reinstall" in msg + assert "--only-binary opencv-python-headless" in msg + assert "opencv-python-headless==5.0.0.93" in msg + assert VALIDATED_SPECS["opencv-python-headless"] not in msg + + +def test_present_carrier_without_metadata_requests_its_version( + monkeypatch, carrier_imports +): + """Without metadata, match the library version plus its packaging revision.""" + carrier_imports(present=("cv2",)) + + def _missing(_package): + raise codec_errors.importlib.metadata.PackageNotFoundError + + monkeypatch.setattr(codec_errors.importlib.metadata, "version", _missing) + msg = str(video_decoder_missing("vllm", "opencv-python-headless", "cv2", "vp9")) + assert "print(cv2.__version__)" in msg + assert "opencv-python-headless==.*" in msg + assert VALIDATED_SPECS["opencv-python-headless"] not in msg + assert "--force-reinstall" in msg + + +def test_broken_carrier_is_not_treated_as_the_source_build(carrier_imports): + """Broken imports need a reinstall, not a missing-video-backend diagnosis.""" + carrier_imports(error=OSError("libGL.so.1: cannot open shared object file")) + + msg = str( + video_decoder_missing( + "vllm", + "opencv-python-headless", + "cv2", + "vp9", + cause="libGL.so.1: cannot open shared object file", + ) + ) + assert "built without a video backend" not in msg + assert VALIDATED_SPECS["opencv-python-headless"] in msg + assert "--force-reinstall" in msg + # Without this, pip may rebuild the codec-free sdist the operator is + # trying to replace. + assert "--only-binary opencv-python-headless" in msg + # The carrier's own words are what identify the real fault. + assert "libGL.so.1" in msg + + def test_hw_codec_without_nvdec_points_at_driver_capability(monkeypatch): """H.264 with NVDEC unavailable is a capability problem first. diff --git a/components/src/dynamo/common/tests/multimodal/test_image_loader.py b/components/src/dynamo/common/tests/multimodal/test_image_loader.py index 647e43b0a401..6b296e0aa111 100644 --- a/components/src/dynamo/common/tests/multimodal/test_image_loader.py +++ b/components/src/dynamo/common/tests/multimodal/test_image_loader.py @@ -17,16 +17,28 @@ import asyncio import base64 +import logging from io import BytesIO from unittest.mock import AsyncMock, patch import numpy as np import pytest from PIL import Image +from redis.exceptions import RedisClusterException, RedisError from dynamo.common.http import HttpConnectionError, HttpStatusError, HttpTimeoutError +from dynamo.common.http.media_reference import DYN_MM_MAX_FILE_SIZE_MB from dynamo.common.http.url_validator import UrlValidationError, UrlValidationPolicy -from dynamo.common.multimodal.image_loader import URL_VARIANT_KEY, ImageLoader +from dynamo.common.multimodal.image_loader import ( + URL_VARIANT_KEY, + ImageLoader, + image_cache_scope_from_request, +) +from dynamo.common.multimodal.shared_image_cache import ( + SharedImageCache, + _RateLimitedErrorWarnings, + _size_bucket, +) pytestmark = [ pytest.mark.asyncio, @@ -36,6 +48,9 @@ ] _FETCH_BYTES_PATH = "dynamo.common.multimodal.image_loader.fetch_bytes" +_REDIS_CLUSTER_FACTORY_PATH = ( + "dynamo.common.multimodal.shared_image_cache.RedisCluster.from_url" +) def _make_png_bytes() -> bytes: @@ -73,7 +88,7 @@ def _mock_fetch_bytes( side_effect: If set, the mock raises this exception instead of returning. """ - async def _fetch(url, timeout, *, policy=None): + async def _fetch(url, timeout, *, policy=None, max_bytes=None): if delay > 0: await asyncio.sleep(delay) if side_effect is not None: @@ -207,6 +222,34 @@ async def test_http_timeout_raises_408(loader: ImageLoader) -> None: assert "https://example.com/img.png" in exc_info.value.url +async def test_http_fetch_honors_configured_media_limit( + loader: ImageLoader, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, "1") + mock_fetch = _mock_fetch_bytes() + + with patch(_FETCH_BYTES_PATH, mock_fetch): + await loader.load_image("https://example.com/limited.png") + + assert mock_fetch.await_args.kwargs["max_bytes"] == 1024 * 1024 + + +async def test_explicit_media_limit_overrides_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, "1") + loader = ImageLoader( + max_bytes=2 * 1024 * 1024, + url_policy=_permissive_policy(), + ) + mock_fetch = _mock_fetch_bytes() + + with patch(_FETCH_BYTES_PATH, mock_fetch): + await loader.load_image("https://example.com/explicit-limit.png") + + assert mock_fetch.await_args.kwargs["max_bytes"] == 2 * 1024 * 1024 + + async def test_http_connection_error_raises_400(loader: ImageLoader) -> None: """An unreachable user-supplied URL is a client error, so the frontend must surface 400 instead of an opaque 500.""" @@ -405,6 +448,89 @@ async def test_image_batch_prioritizes_typed_client_error( assert exc_info.value is client_error +@pytest.mark.parametrize( + "bad_url", + [ + "data:image/png;base64,not!!!base64!!!", + "data:image/png,abc", + ], + ids=["invalid-base64", "missing-base64-marker"], +) +async def test_malformed_data_url_batch_raises_value_error( + loader: ImageLoader, bad_url: str +) -> None: + """A malformed data: URI is a client error. The batch path must preserve + ValueError โ€” the bindings map it to Backend(InvalidArgument) โ†’ 4xx.""" + with pytest.raises(ValueError) as exc_info: + await loader.load_image_batch([{URL_VARIANT_KEY: bad_url}]) + + assert not isinstance(exc_info.value, UrlValidationError) + assert "Failed to decoding image" in str(exc_info.value) + + +async def test_unexpected_decoder_error_not_wrapped_as_value_error( + loader: ImageLoader, +) -> None: + """An unexpected decoder failure must not be + classified as a client validation error: it propagates unchanged from + load_image and the batch path folds it into a generic Exception, so the + frontend answers 500, not 400.""" + loader._open_image = AsyncMock( # type: ignore[method-assign] + side_effect=RuntimeError("decode engine fault") + ) + + with pytest.raises(RuntimeError, match="decode engine fault"): + await loader.load_image("data:image/png;base64,aGVsbG8=") + + with pytest.raises(Exception) as exc_info: + await loader.load_image_batch( + [{URL_VARIANT_KEY: "data:image/png;base64,aGVsbG8="}] + ) + assert not isinstance(exc_info.value, ValueError) + assert "decode engine fault" in str(exc_info.value) + + +async def test_truncated_data_url_batch_raises_400(loader: ImageLoader) -> None: + """Truncated image bytes are malformed client input, not a server fault: + the batch path should return 400.""" + img = Image.new("RGB", (64, 64), color="red") + buf = BytesIO() + img.save(buf, format="JPEG") + truncated_url = ( + f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()[:-20]).decode()}" + ) + + with pytest.raises(HttpStatusError) as exc_info: + await loader.load_image_batch([{URL_VARIANT_KEY: truncated_url}]) + assert exc_info.value.status == 400 + assert "Invalid or truncated image data" in exc_info.value.message + + +async def test_unexpected_http_decoder_error_not_wrapped_as_value_error( + loader: ImageLoader, +) -> None: + """An unexpected decoder failure on a fetched image must not be + classified as a client validation error: it propagates unchanged + from load_image and the batch path folds it into a generic Exception, + so the frontend answers 500, not 400. + """ + loader._open_image = AsyncMock( # type: ignore[method-assign] + side_effect=OSError("image file is truncated") + ) + mock_fetch = _mock_fetch_bytes() + + with patch(_FETCH_BYTES_PATH, mock_fetch): + with pytest.raises(OSError, match="truncated"): + await loader.load_image("https://example.com/img.png") + + with pytest.raises(Exception) as exc_info: + await loader.load_image_batch( + [{URL_VARIANT_KEY: "https://example.com/img.png"}] + ) + assert not isinstance(exc_info.value, ValueError) + assert "truncated" in str(exc_info.value) + + async def test_cache_is_lru_not_fifo(loader: ImageLoader) -> None: """Accessing a cached entry should protect it from eviction (LRU, not FIFO).""" loader._cache_size = 3 @@ -426,3 +552,448 @@ async def test_cache_is_lru_not_fifo(loader: ImageLoader) -> None: assert "https://example.com/b.png" not in loader._image_cache assert "https://example.com/c.png" in loader._image_cache assert "https://example.com/d.png" in loader._image_cache + + +# --- Shared encoded-image cache --- + + +def _enable_shared_image_cache(monkeypatch, ttl_seconds: int = 3600) -> None: + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_ENABLED", "1") + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_URL", "redis://dragonfly.invalid") + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_TTL_SECS", str(ttl_seconds)) + + +@pytest.mark.parametrize( + ("size_bytes", "expected_bucket"), + [ + (10 * 1024, "le_64kib"), + (30 * 1024 * 1024, "le_32mib"), + (33 * 1024 * 1024, "gt_32mib"), + (None, "unknown"), + ], +) +async def test_shared_cache_size_buckets( + size_bytes: int | None, expected_bucket: str +) -> None: + assert _size_bucket(size_bytes) == expected_bucket + + +async def test_shared_cache_disabled_does_not_create_client(monkeypatch) -> None: + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_ENABLED", "0") + with patch(_REDIS_CLUSTER_FACTORY_PATH) as client_factory: + ImageLoader(cache_size=4, url_policy=_permissive_policy()) + + client_factory.assert_not_called() + + +async def test_shared_cache_requires_url_when_enabled(monkeypatch) -> None: + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_ENABLED", "1") + monkeypatch.delenv("DYN_MM_SHARED_IMAGE_CACHE_URL", raising=False) + + with pytest.raises(ValueError, match="DYN_MM_SHARED_IMAGE_CACHE_URL"): + ImageLoader(cache_size=4, url_policy=_permissive_policy()) + + +async def test_shared_cache_uses_split_timeouts(monkeypatch) -> None: + _enable_shared_image_cache(monkeypatch) + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_CONNECT_TIMEOUT_SECS", "0.25") + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_IO_TIMEOUT_SECS", "3.5") + + with patch(_REDIS_CLUSTER_FACTORY_PATH) as client_factory: + ImageLoader(cache_size=4, url_policy=_permissive_policy()) + + client_factory.assert_called_once() + args, kwargs = client_factory.call_args + retry = kwargs.pop("retry") + assert args == ("redis://dragonfly.invalid",) + assert kwargs == { + "decode_responses": False, + "socket_connect_timeout": 0.25, + "socket_timeout": 3.5, + "dynamic_startup_nodes": False, + } + # Retry has no __eq__ at the redis-py 6.2 floor; check the behavior instead. + assert retry.get_retries() == 0 + + +async def test_shared_cache_hit_skips_origin_fetch(monkeypatch) -> None: + _enable_shared_image_cache(monkeypatch) + client = AsyncMock() + client.get.return_value = PNG_BYTES + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + image = await shared_loader.load_image("https://example.com/img.png?sig=one") + + assert image.size == (2, 2) + origin_fetch.assert_not_awaited() + client.get.assert_awaited_once() + client.set.assert_not_awaited() + + +async def test_shared_cache_miss_writes_validated_origin_bytes(monkeypatch) -> None: + _enable_shared_image_cache(monkeypatch, ttl_seconds=123) + client = AsyncMock() + client.get.return_value = None + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + image = await shared_loader.load_image("https://example.com/img.png?sig=one") + + assert image.size == (2, 2) + origin_fetch.assert_awaited_once() + client.set.assert_awaited_once() + _, stored_content = client.set.await_args.args + assert stored_content == PNG_BYTES + assert client.set.await_args.kwargs == {"ex": 123} + + +async def test_shared_cache_does_not_write_invalid_origin_bytes(monkeypatch) -> None: + """Origin bytes must decode successfully before they enter the shared cache.""" + _enable_shared_image_cache(monkeypatch) + client = AsyncMock() + client.get.return_value = None + origin_fetch = _mock_fetch_bytes(content=b"not an image") + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + with pytest.raises(HttpStatusError) as exc_info: + await shared_loader.load_image("https://example.com/img.png") + + assert exc_info.value.status == 415 + origin_fetch.assert_awaited_once() + client.set.assert_not_awaited() + + +@pytest.mark.parametrize( + ("session_scoped_cache", "cache_scope"), + [(False, None), (True, "session-a"), (True, None)], + ids=["unscoped", "scoped", "missing-scope-cache-bypass"], +) +async def test_shared_cache_enabled_origin_fetch_uses_media_size_limit( + monkeypatch, + session_scoped_cache: bool, + cache_scope: str | None, +) -> None: + """Shared-cache configuration must cap even cache-bypassed downloads.""" + _enable_shared_image_cache(monkeypatch) + monkeypatch.setenv("DYN_MM_MAX_FILE_SIZE_MB", "7") + client = AsyncMock() + client.get.return_value = None + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader( + cache_size=4, + url_policy=_permissive_policy(), + session_scoped_cache=session_scoped_cache, + ) + await shared_loader.load_image( + "https://example.com/img.png", cache_scope=cache_scope + ) + + assert origin_fetch.await_args.kwargs["max_bytes"] == 7 * 1024 * 1024 + + +@pytest.mark.parametrize( + "cache_error", + [ + RedisError("cache unavailable"), + RedisClusterException("Timeout connecting to server"), + ], + ids=["redis-error", "redis-cluster-error"], +) +async def test_shared_cache_errors_fall_back_to_origin( + monkeypatch, cache_error: Exception +) -> None: + _enable_shared_image_cache(monkeypatch) + client = AsyncMock() + client.get.side_effect = cache_error + client.set.side_effect = cache_error + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + image = await shared_loader.load_image("https://example.com/img.png") + + assert image.size == (2, 2) + origin_fetch.assert_awaited_once() + client.set.assert_awaited_once() + + +async def test_shared_cache_outage_warnings_are_rate_limited(caplog) -> None: + client = AsyncMock() + client.get.side_effect = RedisError("read unavailable") + client.set.side_effect = RedisError("write unavailable") + client.delete.side_effect = RedisError("delete unavailable") + cache = SharedImageCache(client, ttl_seconds=3600) + + with caplog.at_level( + logging.DEBUG, + logger="dynamo.common.multimodal.shared_image_cache", + ): + assert await cache.get("image") is None + await cache.put("image", PNG_BYTES) + await cache.delete("image") + + warning_records = [ + record for record in caplog.records if record.levelno == logging.WARNING + ] + debug_messages = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.DEBUG + ] + assert len(warning_records) == 1 + assert ( + "Requests continue without the shared cache" in warning_records[0].getMessage() + ) + assert any("read failed" in message for message in debug_messages) + assert any("write failed" in message for message in debug_messages) + assert any("delete failed" in message for message in debug_messages) + + +async def test_shared_cache_outage_warning_repeats_periodically() -> None: + warnings = _RateLimitedErrorWarnings(interval_seconds=60.0) + + assert warnings.record(100.0) == 1 + assert warnings.record(101.0) is None + assert warnings.record(159.9) is None + assert warnings.record(160.0) == 3 + + +async def test_shared_cache_delete_cluster_error_falls_back_to_origin( + monkeypatch, +) -> None: + _enable_shared_image_cache(monkeypatch) + client = AsyncMock() + client.get.return_value = b"not an image" + client.delete.side_effect = RedisClusterException("Timeout connecting to server") + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + image = await shared_loader.load_image("https://example.com/img.png") + + assert image.size == (2, 2) + client.delete.assert_awaited_once() + origin_fetch.assert_awaited_once() + client.set.assert_awaited_once() + + +async def test_invalid_shared_cache_entry_is_deleted_and_refetched(monkeypatch) -> None: + _enable_shared_image_cache(monkeypatch) + client = AsyncMock() + client.get.return_value = b"not an image" + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + image = await shared_loader.load_image("https://example.com/img.png") + + assert image.size == (2, 2) + client.delete.assert_awaited_once() + origin_fetch.assert_awaited_once() + client.set.assert_awaited_once() + + +async def test_truncated_shared_cache_entry_is_deleted_and_refetched( + monkeypatch, +) -> None: + """Pillow reports some recognized-but-truncated images as plain OSError.""" + _enable_shared_image_cache(monkeypatch) + client = AsyncMock() + client.get.return_value = PNG_BYTES[:-30] + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + image = await shared_loader.load_image("https://example.com/img.png") + + assert image.size == (2, 2) + client.delete.assert_awaited_once() + origin_fetch.assert_awaited_once() + client.set.assert_awaited_once() + + +async def test_cache_key_preserves_url_path_and_query_case(monkeypatch) -> None: + _enable_shared_image_cache(monkeypatch) + client = AsyncMock() + client.get.return_value = None + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + await shared_loader.load_image("https://example.com/Image.png?sig=AbC") + await shared_loader.load_image("https://example.com/image.png?sig=abc") + + assert origin_fetch.await_count == 2 + assert client.get.await_count == 2 + assert ( + client.get.await_args_list[0].args[0] != client.get.await_args_list[1].args[0] + ) + + +# --- Session-scoped image cache --- + + +async def test_session_scoped_cache_partitions_local_and_shared_keys( + monkeypatch, +) -> None: + _enable_shared_image_cache(monkeypatch) + client = AsyncMock() + client.get.return_value = None + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader( + cache_size=4, + url_policy=_permissive_policy(), + session_scoped_cache=True, + ) + await shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-a" + ) + await shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-b" + ) + await shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-a" + ) + + assert origin_fetch.await_count == 2 + assert client.get.await_count == 2 + assert client.set.await_count == 2 + assert ( + client.get.await_args_list[0].args[0] != client.get.await_args_list[1].args[0] + ) + + +async def test_session_scoped_cache_partitions_inflight_dedup(monkeypatch) -> None: + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_ENABLED", "0") + origin_fetch = _mock_fetch_bytes(delay=0.05) + shared_loader = ImageLoader( + cache_size=4, + url_policy=_permissive_policy(), + session_scoped_cache=True, + ) + + with patch(_FETCH_BYTES_PATH, origin_fetch): + await asyncio.gather( + shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-a" + ), + shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-a" + ), + shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-b" + ), + ) + + assert origin_fetch.await_count == 2 + + +async def test_session_scope_is_ignored_when_flag_is_disabled(monkeypatch) -> None: + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_ENABLED", "0") + origin_fetch = _mock_fetch_bytes() + shared_loader = ImageLoader( + cache_size=4, + url_policy=_permissive_policy(), + session_scoped_cache=False, + ) + + with patch(_FETCH_BYTES_PATH, origin_fetch): + await shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-a" + ) + await shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-b" + ) + + origin_fetch.assert_awaited_once() + + +async def test_session_scoped_cache_can_be_enabled_by_env(monkeypatch) -> None: + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_ENABLED", "0") + monkeypatch.setenv("DYN_MM_IMAGE_CACHE_SESSION_SCOPED", "1") + origin_fetch = _mock_fetch_bytes() + shared_loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + + with patch(_FETCH_BYTES_PATH, origin_fetch): + await shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-a" + ) + await shared_loader.load_image( + "https://example.com/img.png", cache_scope="session-b" + ) + + assert origin_fetch.await_count == 2 + + +async def test_session_scoped_cache_bypasses_when_scope_is_missing(monkeypatch) -> None: + _enable_shared_image_cache(monkeypatch) + client = AsyncMock() + origin_fetch = _mock_fetch_bytes() + + with ( + patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client), + patch(_FETCH_BYTES_PATH, origin_fetch), + ): + shared_loader = ImageLoader( + cache_size=4, + url_policy=_permissive_policy(), + session_scoped_cache=True, + ) + await shared_loader.load_image("https://example.com/img.png") + await shared_loader.load_image("https://example.com/img.png") + + assert origin_fetch.await_count == 2 + assert shared_loader.cache_entries == 0 + client.get.assert_not_awaited() + client.set.assert_not_awaited() + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"image_cache_scope": " session-a "}, "session-a"), + ({"image_cache_scope": ""}, None), + ({"image_cache_scope": 123}, None), + ({}, None), + ], +) +async def test_image_cache_scope_from_request(payload, expected) -> None: + assert image_cache_scope_from_request(payload) == expected diff --git a/components/src/dynamo/common/tests/multimodal/test_image_loader_metrics.py b/components/src/dynamo/common/tests/multimodal/test_image_loader_metrics.py new file mode 100644 index 000000000000..ce42c14a5982 --- /dev/null +++ b/components/src/dynamo/common/tests/multimodal/test_image_loader_metrics.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the shared-cache portion of register_image_loader_metrics.""" + +import asyncio +from io import BytesIO +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from PIL import Image +from redis.exceptions import RedisError + +from dynamo.common.http.url_validator import UrlValidationPolicy +from dynamo.common.multimodal.image_loader import ImageLoader +from dynamo.common.utils.prometheus import register_image_loader_metrics + +pytestmark = [ + pytest.mark.asyncio, + pytest.mark.unit, + pytest.mark.gpu_0, + pytest.mark.pre_merge, +] + +_FETCH_BYTES_PATH = "dynamo.common.multimodal.image_loader.fetch_bytes" +_REDIS_CLUSTER_FACTORY_PATH = ( + "dynamo.common.multimodal.shared_image_cache.RedisCluster.from_url" +) + + +def _make_png_bytes() -> bytes: + img = Image.new("RGB", (2, 2), color="red") + buf = BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +PNG_BYTES = _make_png_bytes() + + +def _permissive_policy() -> UrlValidationPolicy: + return UrlValidationPolicy(allow_http=True, allow_private_ips=True) + + +def _mock_fetch_bytes( + content: bytes = PNG_BYTES, + delay: float = 0.0, + side_effect: Exception | None = None, +) -> AsyncMock: + async def _fetch(url, timeout, *, policy=None, max_bytes=None): + if delay > 0: + await asyncio.sleep(delay) + if side_effect is not None: + raise side_effect + return content + + return AsyncMock(side_effect=_fetch) + + +def _parse_metric( + text: str, name: str, label_substr: str | None = None +) -> float | None: + """Parse a metric value from Prometheus expfmt text, optionally filtered + by a label substring (e.g. 'outcome="miss"').""" + for line in text.split("\n"): + if line.startswith(name + "{") or line.startswith(name + " "): + if label_substr is not None and label_substr not in line: + continue + parts = line.rsplit(" ", 1) + if len(parts) == 2: + return float(parts[1]) + return None + + +async def test_shared_cache_get_set_duration_histograms(monkeypatch) -> None: + """Shared-cache timings are exported with bounded outcome labels.""" + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_ENABLED", "1") + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_URL", "redis://dragonfly.invalid") + client = AsyncMock() + client.get.return_value = None + + endpoint = MagicMock() + with patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client): + loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + register_image_loader_metrics(endpoint, loader, "m", "c") + endpoint.metrics.register_prometheus_typed_callback.assert_called_once() + callback = endpoint.metrics.register_prometheus_expfmt_callback.call_args[0][0] + typed_callback = endpoint.metrics.register_prometheus_typed_callback.call_args[0][0] + + with patch(_FETCH_BYTES_PATH, _mock_fetch_bytes()): + await loader.load_image("https://example.com/shared.png") + + typed_metrics = typed_callback() + assert any( + family[0] == "dynamo_component_image_shared_cache_get_duration_seconds" + for family in typed_metrics + ) + text = callback() + assert ( + _parse_metric( + text, + "dynamo_component_image_shared_cache_get_duration_seconds_count", + 'outcome="miss",size_bucket="unknown"', + ) + == 1.0 + ) + assert ( + _parse_metric( + text, + "dynamo_component_image_shared_cache_set_duration_seconds_count", + 'outcome="success",size_bucket="le_64kib"', + ) + == 1.0 + ) + + +async def test_shared_cache_error_duration_histograms(monkeypatch) -> None: + """Redis failures are timed while the loader still fails open.""" + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_ENABLED", "1") + monkeypatch.setenv("DYN_MM_SHARED_IMAGE_CACHE_URL", "redis://dragonfly.invalid") + client = AsyncMock() + client.get.side_effect = RedisError("read unavailable") + client.set.side_effect = RedisError("write unavailable") + + endpoint = MagicMock() + with patch(_REDIS_CLUSTER_FACTORY_PATH, return_value=client): + loader = ImageLoader(cache_size=4, url_policy=_permissive_policy()) + register_image_loader_metrics(endpoint, loader, "m", "c") + callback = endpoint.metrics.register_prometheus_expfmt_callback.call_args[0][0] + + with patch(_FETCH_BYTES_PATH, _mock_fetch_bytes()): + await loader.load_image("https://example.com/shared-error.png") + + text = callback() + assert ( + _parse_metric( + text, + "dynamo_component_image_shared_cache_get_duration_seconds_count", + 'outcome="error",size_bucket="unknown"', + ) + == 1.0 + ) + assert ( + _parse_metric( + text, + "dynamo_component_image_shared_cache_set_duration_seconds_count", + 'outcome="error",size_bucket="le_64kib"', + ) + == 1.0 + ) diff --git a/components/src/dynamo/common/tests/multimodal/test_video_loader.py b/components/src/dynamo/common/tests/multimodal/test_video_loader.py index f45c47a5ab8f..d7a4bd42d175 100644 --- a/components/src/dynamo/common/tests/multimodal/test_video_loader.py +++ b/components/src/dynamo/common/tests/multimodal/test_video_loader.py @@ -1,18 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import sys from unittest.mock import AsyncMock import numpy as np import pytest -import dynamo.common.multimodal.video_loader as video_loader_module -from dynamo.common.http import HttpStatusError +from dynamo.common.http import HttpConfigurationError, HttpStatusError from dynamo.common.http.url_validator import UrlValidationError, UrlValidationPolicy +from dynamo.common.multimodal import codec_errors +from dynamo.common.multimodal import video_loader as video_loader_module from dynamo.common.multimodal.codec_errors import MissingMediaDecoderError from dynamo.common.multimodal.video_loader import VideoLoader from dynamo.common.utils.install_media_decoders import VALIDATED_SPECS +from dynamo.common.http.media_reference import DYN_MM_MAX_FILE_SIZE_MB # isort: skip + pytestmark = [ pytest.mark.unit, pytest.mark.pre_merge, @@ -73,6 +77,25 @@ async def test_load_video_uses_vllm_media_connector(): assert loaded_metadata == metadata +@pytest.mark.asyncio +async def test_http_fetch_honors_configured_media_limit(monkeypatch): + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, "1") + loader = VideoLoader( + url_policy=UrlValidationPolicy(allow_http=True, allow_private_ips=True) + ) + frames = np.zeros((1, 2, 2, 3), dtype=np.uint8) + metadata = {"fps": 1.0, "frames_indices": [0], "total_num_frames": 1} + mock_fetch = AsyncMock(return_value=b"video") + mock_decode = AsyncMock(return_value=(frames, metadata)) + monkeypatch.setattr(video_loader_module, "fetch_bytes", mock_fetch) + monkeypatch.setattr(loader, "_create_vllm_video_io", lambda kwargs=None: object()) + monkeypatch.setattr(loader, "_decode_video_bytes", mock_decode) + + await loader._load_video_with_vllm("https://example.com/limited.mp4") + + assert mock_fetch.await_args.kwargs["max_bytes"] == 1024 * 1024 + + @pytest.mark.asyncio async def test_load_video_batch_uses_url_loader(): loader = VideoLoader() @@ -287,7 +310,9 @@ def load_bytes(self, content: bytes): @pytest.mark.asyncio -async def test_decode_video_bytes_missing_decoder_is_actionable(monkeypatch): +async def test_decode_video_bytes_missing_decoder_is_actionable( + monkeypatch, carrier_imports +): """A bare `No module named 'cv2'` must become the actionable codec error. Reproduced on a real runtime image before this existed: the user-visible @@ -297,6 +322,8 @@ async def test_decode_video_bytes_missing_decoder_is_actionable(monkeypatch): loader = VideoLoader() monkeypatch.setattr(video_loader_module, "probe_video_codec", lambda b: "vp9") monkeypatch.setattr(video_loader_module, "should_use_nvdec", lambda c: False) + # cv2 genuinely absent, independent of what the test machine has installed. + carrier_imports() with pytest.raises(MissingMediaDecoderError) as exc_info: await loader._decode_video_bytes(b"vp9-bytes", _ImportErrorMediaIO()) @@ -304,10 +331,197 @@ async def test_decode_video_bytes_missing_decoder_is_actionable(monkeypatch): msg = str(exc_info.value) assert "'vp9'" in msg # names the codec assert VALIDATED_SPECS["opencv-python-headless"] in msg # bounded spec - assert "install_media_decoders vllm" in msg # installer command + assert "pip install" in msg assert "cv2" in msg +class _SystemErrorMediaIO: + """What a cv2 with no video backend does: fail inside VideoCapture.""" + + def load_bytes(self, content: bytes): + raise SystemError( + " returned a result with an exception set" + ) + + +@pytest.mark.asyncio +async def test_decode_video_bytes_backendless_cv2_is_actionable( + monkeypatch, carrier_imports +): + """Translate codec-free OpenCV failures into actionable decoder guidance.""" + loader = VideoLoader() + monkeypatch.setattr(video_loader_module, "probe_video_codec", lambda b: "vp9") + monkeypatch.setattr(video_loader_module, "should_use_nvdec", lambda c: False) + monkeypatch.setattr(video_loader_module, "_cv2_lacks_video_backend", lambda: True) + + carrier_imports(present=("cv2",)) + monkeypatch.setattr( + codec_errors.importlib.metadata, "version", lambda p: "5.0.0.93" + ) + + with pytest.raises(MissingMediaDecoderError) as exc_info: + await loader._decode_video_bytes(b"vp9-bytes", _SystemErrorMediaIO()) + + msg = str(exc_info.value) + assert "'vp9'" in msg + assert "opencv-python-headless==5.0.0.93" in msg + assert "--force-reinstall" in msg + assert "video backend" in msg + + +class _FailedOpenMediaIO: + """What OpenCV's own video backend raises when VideoCapture opens nothing.""" + + def load_bytes(self, content: bytes): + raise ValueError("Could not open video stream") + + +@pytest.mark.asyncio +async def test_decode_video_bytes_failed_open_value_error_is_actionable( + monkeypatch, carrier_imports +): + """Translate OpenCV failed-open errors when its video backend is absent.""" + loader = VideoLoader() + monkeypatch.setattr(video_loader_module, "probe_video_codec", lambda b: "vp9") + monkeypatch.setattr(video_loader_module, "should_use_nvdec", lambda c: False) + monkeypatch.setattr(video_loader_module, "_cv2_lacks_video_backend", lambda: True) + + carrier_imports(present=("cv2",)) + monkeypatch.setattr( + codec_errors.importlib.metadata, "version", lambda p: "5.0.0.93" + ) + + with pytest.raises(MissingMediaDecoderError) as exc_info: + await loader._decode_video_bytes(b"vp9-bytes", _FailedOpenMediaIO()) + + msg = str(exc_info.value) + assert "'vp9'" in msg + assert "opencv-python-headless==5.0.0.93" in msg + assert "--force-reinstall" in msg + assert "video backend" in msg + + +class _InitErrorMediaIO: + """What a lazily initialised cv2 raises when its own libraries fail to load.""" + + def load_bytes(self, content: bytes): + raise OSError("cv2: libavcodec.so.60: cannot open shared object file") + + +@pytest.mark.asyncio +async def test_decode_video_bytes_cv2_os_error_is_actionable( + monkeypatch, carrier_imports +): + """Keep the remedy visible when cv2 initialisation fails inside the decode.""" + loader = VideoLoader() + monkeypatch.setattr(video_loader_module, "probe_video_codec", lambda b: "vp9") + monkeypatch.setattr(video_loader_module, "should_use_nvdec", lambda c: False) + monkeypatch.setattr(video_loader_module, "_cv2_lacks_video_backend", lambda: True) + + carrier_imports(present=("cv2",)) + monkeypatch.setattr( + codec_errors.importlib.metadata, "version", lambda p: "5.0.0.93" + ) + + with pytest.raises(MissingMediaDecoderError) as exc_info: + await loader._decode_video_bytes(b"vp9-bytes", _InitErrorMediaIO()) + + msg = str(exc_info.value) + assert "'vp9'" in msg + assert "opencv-python-headless==5.0.0.93" in msg + assert "--force-reinstall" in msg + + +@pytest.mark.asyncio +async def test_unrelated_os_error_is_not_blamed_on_opencv(monkeypatch): + """An I/O failure that never mentions the carrier keeps its own error.""" + loader = VideoLoader() + monkeypatch.setattr(video_loader_module, "probe_video_codec", lambda b: "vp9") + monkeypatch.setattr(video_loader_module, "should_use_nvdec", lambda c: False) + monkeypatch.setattr(video_loader_module, "_cv2_lacks_video_backend", lambda: True) + + class _DiskFullMediaIO: + def load_bytes(self, content: bytes): + raise OSError("[Errno 28] No space left on device: '/tmp/clip.mp4'") + + with pytest.raises(OSError, match="No space left"): + await loader._decode_video_bytes(b"vp9-bytes", _DiskFullMediaIO()) + + +@pytest.mark.asyncio +async def test_backendless_cv2_does_not_pre_empt_a_working_decode(monkeypatch): + """Allow configured video backends to decode before checking OpenCV.""" + loader = VideoLoader() + monkeypatch.setattr(video_loader_module, "probe_video_codec", lambda b: "vp9") + monkeypatch.setattr(video_loader_module, "should_use_nvdec", lambda c: False) + monkeypatch.setattr(video_loader_module, "_cv2_lacks_video_backend", lambda: True) + + frames = np.zeros((2, 4, 4, 3), dtype=np.uint8) + + class _WorkingMediaIO: + def load_bytes(self, content: bytes): + return frames, {"frames_indices": [0, 1]} + + decoded, metadata = await loader._decode_video_bytes( + b"vp9-bytes", _WorkingMediaIO() + ) + assert decoded is frames + assert metadata == {"frames_indices": [0, 1]} + + +@pytest.mark.asyncio +async def test_other_backend_import_error_is_not_blamed_on_opencv(monkeypatch): + """Preserve import errors from other video backends.""" + loader = VideoLoader() + monkeypatch.setattr(video_loader_module, "probe_video_codec", lambda b: "vp9") + monkeypatch.setattr(video_loader_module, "should_use_nvdec", lambda c: False) + + class _PyAvMediaIO: + def load_bytes(self, content: bytes): + raise ModuleNotFoundError("No module named 'av'", name="av") + + with pytest.raises(ModuleNotFoundError, match="'av'"): + await loader._decode_video_bytes(b"vp9-bytes", _PyAvMediaIO()) + + +@pytest.mark.asyncio +async def test_other_backend_system_error_is_not_blamed_on_opencv(monkeypatch): + """Preserve other backend failures even when OpenCV lacks video support.""" + loader = VideoLoader() + monkeypatch.setattr(video_loader_module, "probe_video_codec", lambda b: "vp9") + monkeypatch.setattr(video_loader_module, "should_use_nvdec", lambda c: False) + monkeypatch.setattr(video_loader_module, "_cv2_lacks_video_backend", lambda: True) + + class _OtherBackendMediaIO: + def load_bytes(self, content: bytes): + raise SystemError("torchcodec decoder failed to initialise") + + with pytest.raises(SystemError, match="torchcodec"): + await loader._decode_video_bytes(b"vp9-bytes", _OtherBackendMediaIO()) + + +@pytest.mark.asyncio +async def test_unrelated_system_error_keeps_its_own_message(monkeypatch): + """A SystemError from anywhere else must not be blamed on the codec.""" + loader = VideoLoader() + monkeypatch.setattr(video_loader_module, "probe_video_codec", lambda b: "vp9") + monkeypatch.setattr(video_loader_module, "should_use_nvdec", lambda c: False) + monkeypatch.setattr(video_loader_module, "_cv2_lacks_video_backend", lambda: False) + + with pytest.raises(SystemError, match="VideoCapture"): + await loader._decode_video_bytes(b"vp9-bytes", _SystemErrorMediaIO()) + + +def test_cv2_lacks_video_backend_is_false_when_cv2_absent(monkeypatch): + """An absent cv2 is not this case: the ImportError path names it better.""" + video_loader_module._cv2_lacks_video_backend.cache_clear() + monkeypatch.setitem(sys.modules, "cv2", None) + try: + assert video_loader_module._cv2_lacks_video_backend() is False + finally: + video_loader_module._cv2_lacks_video_backend.cache_clear() + + @pytest.mark.asyncio async def test_load_video_batch_preserves_missing_decoder_error(): """The batch aggregate wraps failures in a generic Exception; the @@ -343,3 +557,25 @@ async def test_load_video_preserves_missing_decoder_error(monkeypatch): await loader.load_video("https://example.com/x.mp4") assert exc_info.value is err + + +@pytest.mark.asyncio +async def test_load_video_preserves_a_configuration_error(monkeypatch): + """An operator fault must not reach the client as a 4xx. + + Same reasoning as the audio loader: the generic ``except Exception`` there + rewrites unknown errors as ``ValueError``, which maps to ``InvalidArgument``. + """ + from dynamo.common.multimodal.video_loader import VideoLoader + + loader = VideoLoader.__new__(VideoLoader) + + async def _boom(*args, **kwargs): + raise HttpConfigurationError("egress proxy is not trusted") + + monkeypatch.setattr(loader, "_load_video_with_vllm", _boom, raising=False) + + with pytest.raises(HttpConfigurationError) as excinfo: + await loader.load_video("https://example.com/v.mp4") + + assert not isinstance(excinfo.value, ValueError) diff --git a/components/src/dynamo/common/tests/test_install_media_decoders.py b/components/src/dynamo/common/tests/test_install_media_decoders.py index 03576888e075..b9e608a693a3 100644 --- a/components/src/dynamo/common/tests/test_install_media_decoders.py +++ b/components/src/dynamo/common/tests/test_install_media_decoders.py @@ -35,6 +35,11 @@ def sandboxed(monkeypatch): monkeypatch.setattr( install_media_decoders, "_modules_missing_fresh", lambda mods: list(mods) ) + monkeypatch.setattr( + install_media_decoders, + "_vllm_opencv_spec", + lambda: install_media_decoders.VALIDATED_SPECS["opencv-python-headless"], + ) return monkeypatch @@ -133,21 +138,26 @@ def test_already_present_installs_nothing(sandboxed): assert calls == [] -def test_vllm_installs_bounded_video_and_audio_specs(sandboxed): +def test_vllm_installs_video_and_audio(sandboxed): present: set[str] = set() _set_available(sandboxed, present) + sandboxed.setattr( + install_media_decoders, + "_vllm_opencv_spec", + lambda: "opencv-python-headless==5.0.0.93", + ) calls = _record_pip_and_mark(sandboxed, present, "cv2", "av") installed = install_media_decoders.install_media_decoders("vllm") assert installed == [ - install_media_decoders.VALIDATED_SPECS["opencv-python-headless"], + "opencv-python-headless==5.0.0.93", install_media_decoders.VALIDATED_SPECS["av"], ] (cmd,) = calls assert "--break-system-packages" in cmd + assert "--force-reinstall" in cmd + assert cmd[cmd.index("--only-binary") + 1] == "opencv-python-headless" for spec in installed: assert spec in cmd - # Never installed: pynvvideocodec because the image already ships it as the - # NVDEC path, and the rest because no vLLM decode path imports them. for banned in ("torchcodec", "pynvvideocodec", "decord2", "libx264"): assert not any(banned in part for part in cmd) @@ -174,16 +184,6 @@ def test_trtllm_installs_opencv_only(sandboxed): assert install_media_decoders.VALIDATED_SPECS["av"] not in cmd -def test_installs_only_missing_modules(sandboxed): - present = {"cv2"} # video carrier present, audio missing - _set_available(sandboxed, present) - calls = _record_pip_and_mark(sandboxed, present, "av") - installed = install_media_decoders.install_media_decoders("vllm") - assert installed == [install_media_decoders.VALIDATED_SPECS["av"]] - (cmd,) = calls - assert install_media_decoders.VALIDATED_SPECS["opencv-python-headless"] not in cmd - - def test_every_install_uses_no_deps(sandboxed): """--no-deps is unconditional now that custom specs are out of scope -- the installer must never be able to shift the image's pinned stack.""" @@ -256,29 +256,19 @@ def test_dry_run_reports_without_installing(sandboxed): assert calls == [] -def test_pending_subset_installs_only_still_missing(sandboxed): - """A racing process may install part of the set while we wait on the lock. - - Probe round 1 (pre-check) sees both vLLM carriers missing; round 2 - (post-lock re-check) sees cv2 already installed by the racing process, so - only the audio carrier installs; round 3 (post-verify) sees everything. - """ +def test_racing_install_between_probe_and_lock_runs_no_pip(sandboxed): + """Skip pip when a concurrent installer makes both decoders usable.""" rounds = {"n": 0} def probe(mods): rounds["n"] += 1 - if rounds["n"] == 1: - return list(mods) - if rounds["n"] == 2: - return [m for m in mods if m != "cv2"] - return [] + return list(mods) if rounds["n"] == 1 else [] sandboxed.setattr(install_media_decoders, "_modules_missing_fresh", probe) calls = _record_pip(sandboxed) - installed = install_media_decoders.install_media_decoders("vllm") - assert installed == [install_media_decoders.VALIDATED_SPECS["av"]] - (cmd,) = calls - assert install_media_decoders.VALIDATED_SPECS["opencv-python-headless"] not in cmd + + assert install_media_decoders.install_media_decoders("vllm") == [] + assert calls == [] def test_modules_missing_fresh_real_probe(): @@ -310,6 +300,32 @@ def test_probe_treats_present_but_broken_package_as_missing(tmp_path, monkeypatc assert install_media_decoders._modules_missing_fresh(["brokenmod"]) == ["brokenmod"] +def test_opencv_spec_falls_back_when_the_version_probe_fails(monkeypatch): + """An unreadable OpenCV version must not abort the install. + + A cv2 that raises on import (rather than being absent) kills the probe + child, and check=True turns that into CalledProcessError in the caller. + The bounded spec is the right answer there, exactly as for an empty probe. + """ + + def _boom(cmd, **kwargs): + raise subprocess.CalledProcessError(1, cmd) + + monkeypatch.setattr(install_media_decoders.subprocess, "run", _boom) + assert ( + install_media_decoders._vllm_opencv_spec() + == install_media_decoders.VALIDATED_SPECS["opencv-python-headless"] + ) + + +def test_opencv_version_probe_survives_a_broken_cv2(tmp_path, monkeypatch): + """The probe child must exit cleanly when importing cv2 raises.""" + (tmp_path / "cv2.py").write_text("raise RuntimeError('native libs gone')\n") + monkeypatch.setenv("PYTHONPATH", str(tmp_path)) + spec = install_media_decoders._vllm_opencv_spec() + assert spec.startswith("opencv-python-headless") + + def test_lock_refuses_symlink_and_preserves_target(tmp_path, monkeypatch): """A pre-planted symlink at the fixed lock path must not be followed. diff --git a/components/src/dynamo/common/tests/test_video_protocol.py b/components/src/dynamo/common/tests/test_video_protocol.py index cc537ec657dd..524e5c63eaf8 100644 --- a/components/src/dynamo/common/tests/test_video_protocol.py +++ b/components/src/dynamo/common/tests/test_video_protocol.py @@ -56,3 +56,19 @@ def test_video_response_wire_shape(): "created": 0, "data": [{"output_format": "mp4", "url": "http://example.com/v.mp4"}], } + + +def test_video_response_reports_media_metadata(): + response = VideoData( + output_format="mp4", + url="http://example.com/v.mp4", + fps=24, + audio_sample_rate=32000, + ) + + assert response.model_dump(exclude_none=True) == { + "output_format": "mp4", + "url": "http://example.com/v.mp4", + "fps": 24, + "audio_sample_rate": 32000, + } diff --git a/components/src/dynamo/common/tests/test_video_utils.py b/components/src/dynamo/common/tests/test_video_utils.py index 7371965c6f8c..e60783677e80 100644 --- a/components/src/dynamo/common/tests/test_video_utils.py +++ b/components/src/dynamo/common/tests/test_video_utils.py @@ -20,6 +20,20 @@ def make_frames(n=3, h=8, w=8) -> np.ndarray: return np.zeros((n, h, w, 3), dtype=np.uint8) +# --------------------------------------------------------------------------- +# compute_num_frames +# --------------------------------------------------------------------------- + + +def test_compute_num_frames_rounds_fractional_frame_rate(): + from dynamo.common.utils.video_utils import compute_num_frames + + result = compute_num_frames(seconds=10, fps=23.976) + + assert result == 240 + assert isinstance(result, int) + + # --------------------------------------------------------------------------- # encode_to_video_bytes # --------------------------------------------------------------------------- diff --git a/components/src/dynamo/common/utils/input_params.py b/components/src/dynamo/common/utils/input_params.py index 1e8c38cde9e4..9c88cb7046e2 100644 --- a/components/src/dynamo/common/utils/input_params.py +++ b/components/src/dynamo/common/utils/input_params.py @@ -4,6 +4,18 @@ from typing import Any, Optional +def resolve_thinking_token_budget(request: dict) -> Optional[Any]: + """Resolve the thinking-token budget from an OpenAI-compatible request. + + Supports both the OpenAI-compatible root-level field + ``thinking_token_budget`` and the legacy Dynamo extension + ``nvext.max_thinking_tokens``. The root-level field takes precedence. + """ + root = request.get("thinking_token_budget") + legacy = (request.get("nvext") or {}).get("max_thinking_tokens") + return root if root is not None else legacy + + def _inject_reasoning_content(messages: list) -> None: """Inject reasoning_content as blocks into content. diff --git a/components/src/dynamo/common/utils/install_media_decoders.py b/components/src/dynamo/common/utils/install_media_decoders.py index 4a4140c7d2b6..2261ffabd205 100644 --- a/components/src/dynamo/common/utils/install_media_decoders.py +++ b/components/src/dynamo/common/utils/install_media_decoders.py @@ -38,8 +38,9 @@ broadens the codec surface of the running container; review that against your organization's distribution and security policies before rolling it out. -Each package installs at a version validated against Dynamo's multimodal test -suite (see the specs below): the lower bound is the exact validated version and +For vLLM, codec-free OpenCV is replaced with a binary wheel matching the installed +version. Other packages install at a version validated against Dynamo's multimodal +test suite (see the specs below): the lower bound is the exact validated version and the upper bound excludes the next major, so a fresh install cannot silently pick up an unvalidated release. This installer deliberately covers ONLY those tested combinations; to pin different versions or install a custom subset, run @@ -49,7 +50,7 @@ The default install runs with ``--no-deps`` so it cannot change the image's pinned dependency stack (e.g. numpy under torch/vLLM); the carriers need only numpy, which the backend already provides. The install is idempotent (skipped -when the module already imports), serialized across processes with a file +when imports and OpenCV video support are available), serialized with a file lock, and bounded by a timeout so a stalled index cannot hang forever. The encode path (video output) is unaffected and stays on the in-tree FFmpeg. The optional Rust frontend decoder links FFmpeg's compiled-in decoders and is @@ -140,33 +141,25 @@ class _Decoder: } +def installer_covers(backend: str, package: str) -> bool: + """Return whether the backend installer covers this package.""" + return any(d.package == package for d in _BACKEND_DECODERS.get(backend, ())) + + def _modules_missing_fresh(modules: Sequence[str]) -> list[str]: - """Return the subset of `modules` a FRESH interpreter cannot IMPORT. - - Fresh interpreter, not this process: running as a non-root user, pip - defaults to a user-site install, and a user-site directory created after - the parent interpreter started is never added to its ``sys.path`` - (``site.py`` only does that at startup, and - ``importlib.invalidate_caches()`` cannot add path entries). Verified on - all three runtime images: the same-process check reported the install - missing while a fresh process imported it fine. What matters - operationally is the worker process launched after this command -- - which is exactly a fresh interpreter. - - Real import, not ``find_spec``: a package whose files are present but - whose native libraries cannot load (a broken or partially removed wheel) - has a spec and would be treated as installed, silently skipping the - install and deferring the failure to request time. Importing in the - probe subprocess keeps this process's module state untouched. - """ + """Check imports and OpenCV video support in a fresh worker-like interpreter.""" if not modules: return [] probe = ( - "import importlib, sys\n" + "import importlib, re, sys\n" "missing = []\n" "for m in sys.argv[1:]:\n" " try:\n" - " importlib.import_module(m)\n" + " module = importlib.import_module(m)\n" + " if m == 'cv2' and not re.search(\n" + " r'^\\s*(FFMPEG|GSTREAMER):\\s*YES',\n" + " module.getBuildInformation(), re.M | re.I):\n" + " missing.append(m)\n" " except Exception:\n" " missing.append(m)\n" "print(' '.join(missing))\n" @@ -185,6 +178,42 @@ def _modules_missing_fresh(modules: Sequence[str]) -> list[str]: return out.stdout.split() +def _vllm_opencv_spec() -> str: + """Preserve the image's OpenCV version when replacing its source build.""" + probe = ( + "import importlib.metadata\n" + "version = ''\n" + "for package in ('opencv-python-headless', 'opencv-python'):\n" + " try:\n" + " version = importlib.metadata.version(package)\n" + " break\n" + " except importlib.metadata.PackageNotFoundError:\n" + " pass\n" + "if not version:\n" + " try:\n" + " import cv2\n" + # A cv2 that raises anything is as unusable as a missing one. + " except Exception:\n" + " pass\n" + " else:\n" + " version = cv2.__version__ + '.*'\n" + "print(version)\n" + ) + try: + out = subprocess.run( + [sys.executable, "-c", probe], + capture_output=True, + text=True, + timeout=120, + check=True, + ) + except Exception as exc: # noqa: BLE001 - an unreadable version is not fatal + logger.warning("OpenCV version probe failed to run: %s", exc) + return _OPENCV.spec + version = out.stdout.strip() + return f"{_OPENCV.package}=={version}" if version else _OPENCV.spec + + def _redact(text: str) -> str: """Mask URL userinfo (user:token@host) so pip args never leak secrets.""" return _CRED_RE.sub(r"\1***@", text) @@ -253,9 +282,9 @@ def install_media_decoders( ) -> list[str]: """Install `backend`'s media-decoder package(s); return the specs installed. - Installs the backend's validated, version-bounded specs with ``--no-deps``, - skipping any package whose module already imports. Deliberately covers only - the tested combinations: for custom versions or subsets, run pip directly. + Installs with ``--no-deps`` and skips usable decoders. vLLM OpenCV is + replaced with a same-version binary wheel; other packages use validated + bounds. For custom versions or subsets, run pip directly. Returns the list of pip specs actually installed (empty when everything was already present). ``dry_run`` returns what would install without running @@ -267,9 +296,7 @@ def install_media_decoders( raise ValueError( f"unknown backend {backend!r}; expected one of {sorted(_BACKEND_DECODERS)}" ) - # Install only what does not already import cleanly. The probe does a - # real import in a fresh interpreter, so a present-but-broken package - # counts as missing rather than being silently skipped. + # Include import failures and OpenCV builds without a video backend. missing_names = set(_modules_missing_fresh([d.module for d in decoders])) missing = [d for d in decoders if d.module in missing_names] if not missing: @@ -278,10 +305,18 @@ def install_media_decoders( backend, ) return [] - specs = [d.spec for d in missing] + specs = [ + _vllm_opencv_spec() if backend == "vllm" and d.module == "cv2" else d.spec + for d in missing + ] verify_modules = [d.module for d in missing] if dry_run: + if backend == "vllm" and "cv2" in verify_modules: + logger.info( + "dry run: OpenCV replacement requires " + "--force-reinstall --only-binary opencv-python-headless" + ) logger.info( "dry run: would install for %s: %s", backend, _redact(" ".join(specs)) ) @@ -306,13 +341,17 @@ def install_media_decoders( return [] specs = [spec for spec, _ in pending] verify_modules = [mod for _, mod in pending] - _pip_install(specs, pip_args, timeout_s=timeout_s) + extra_args = list(pip_args) + if backend == "vllm" and "cv2" in verify_modules: + specs[verify_modules.index("cv2")] = _vllm_opencv_spec() + extra_args += ["--force-reinstall", "--only-binary", _OPENCV.package] + _pip_install(specs, extra_args, timeout_s=timeout_s) importlib.invalidate_caches() still_missing = _modules_missing_fresh(verify_modules) if still_missing: raise RuntimeError( - "media decoder module(s) still not importable after install: " + "media decoder module(s) still not importable or video-capable after install: " + " ".join(still_missing) + ". If a package is present but broken, pip may have treated the " "requirement as already satisfied -- pip uninstall it and rerun, " @@ -329,8 +368,8 @@ def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( prog="python -m dynamo.common.utils.install_media_decoders", description=( - "Install a Dynamo backend's media-decoder package(s) at validated, " - "version-bounded releases. Explicit by design: nothing installs " + "Install a Dynamo backend's media decoders, preserving vLLM's OpenCV " + "version and using validated bounds otherwise. Nothing installs " "unless an operator runs this." ), ) @@ -348,7 +387,6 @@ def main(argv: list[str] | None = None) -> int: 'string. Use the = form -- --pip-args="--no-index --find-links ' '/wheels" -- so a value starting with a dash is not mistaken ' "for an option (for air-gapped hosts)" - "air-gapped hosts)" ), ) parser.add_argument( diff --git a/components/src/dynamo/common/utils/nixl_telemetry.py b/components/src/dynamo/common/utils/nixl_telemetry.py index e4855061e5f9..a98366600f8a 100644 --- a/components/src/dynamo/common/utils/nixl_telemetry.py +++ b/components/src/dynamo/common/utils/nixl_telemetry.py @@ -15,11 +15,16 @@ This module deliberately imports no inference engine: the derivation is pure arithmetic over a base port and a rank index, and the callers that know how to find a rank index live in the per-backend packages. + +Only environment configuration is inspected here. NIXL can also read these +settings from ``NIXL_CONFIG_FILE``, ``~/.nixl.cfg``, or ``/etc/nixl.cfg``; +Dynamo cannot derive or declare per-rank ports from those files. """ from __future__ import annotations import os +import re from collections.abc import Mapping MAX_PORT = 65535 @@ -28,7 +33,14 @@ NIXL_TELEMETRY_EXPORTER_ENV = "NIXL_TELEMETRY_EXPORTER" NIXL_TELEMETRY_PROMETHEUS_PORT_ENV = "NIXL_TELEMETRY_PROMETHEUS_PORT" -DEFAULT_NIXL_PROMETHEUS_PORT = 19090 +# NIXL's Prometheus exporter uses this port when no override is configured. +DEFAULT_NIXL_PROMETHEUS_PORT = 9090 + +# NIXL compares these tokens case-insensitively, without trimming whitespace. +_NIXL_TRUE_VALUES = frozenset({"y", "1", "yes", "on", "true", "enable"}) +_NIXL_FALSE_VALUES = frozenset({"n", "0", "no", "off", "false", "disable"}) + +_NIXL_UINT16 = re.compile(r"(?:0[xX][0-9a-fA-F]+|[0-9]+)") # Keep in sync with DynamoMaxNixlPorts in deploy/operator/internal/consts/consts.go: # a rank deriving a port past the reserved range would bind a port nothing scrapes. @@ -45,11 +57,7 @@ def configured_fixed_port( default: int | None = None, env: Mapping[str, str] | None = None, ) -> int | None: - """Return a configured fixed TCP port, ignoring disabled/invalid values. - - Mirrors ``dynamo.vllm.backend_args._configured_fixed_port`` so that a value - this package accepts is exactly a value that package accepts. - """ + """Return a configured fixed TCP port, ignoring disabled/invalid values.""" environ = os.environ if env is None else env raw = environ.get(env_name) if raw is None: @@ -61,18 +69,65 @@ def configured_fixed_port( return port if 0 < port <= MAX_PORT else None +def _parse_nixl_uint16(env_name: str, raw: str) -> int: + """Parse one present value with NIXL's unsigned 16-bit syntax.""" + # Unlike int(), NIXL rejects signs, whitespace, separators, and partial parses. + if _NIXL_UINT16.fullmatch(raw) is None: + raise ValueError(f"{env_name}={raw!r} is not a valid unsigned 16-bit integer") + + is_hex = raw.startswith(("0x", "0X")) + digits = (raw[2:] if is_hex else raw).lstrip("0") or "0" + max_digits = format(MAX_PORT, "x") if is_hex else str(MAX_PORT) + if len(digits) > len(max_digits) or ( + len(digits) == len(max_digits) and digits.lower() > max_digits + ): + raise ValueError(f"{env_name}={raw!r} is outside the range 0-{MAX_PORT}") + # Parse only the significant digits so Python's configurable integer-string + # limit cannot misclassify a valid value with many leading zeroes. + return int(digits, 16 if is_hex else 10) + + +def _parse_nixl_bool(env_name: str, raw: str) -> bool: + """Parse one present value with NIXL's case-insensitive boolean syntax.""" + normalized = raw.lower() + if normalized in _NIXL_TRUE_VALUES: + return True + if normalized in _NIXL_FALSE_VALUES: + return False + raise ValueError(f"{env_name}={raw!r} is not a boolean value recognized by NIXL") + + def nixl_prometheus_base_port(env: Mapping[str, str] | None = None) -> int | None: - """Return the base NIXL Prometheus port, or None when it is not in use.""" + """Return the fixed Prometheus port selected by the NIXL environment. + + ``None`` means this environment does not select the Prometheus exporter. + Invalid NIXL values raise, as they do during NIXL agent construction. + """ environ = os.environ if env is None else env - enabled = environ.get(NIXL_TELEMETRY_ENABLE_ENV, "").strip().lower() - exporter = environ.get(NIXL_TELEMETRY_EXPORTER_ENV, "prometheus") - if enabled != "y" or exporter.strip().lower() != "prometheus": + enabled = environ.get(NIXL_TELEMETRY_ENABLE_ENV) + if enabled is None or not _parse_nixl_bool(NIXL_TELEMETRY_ENABLE_ENV, enabled): + return None + + exporter = environ.get(NIXL_TELEMETRY_EXPORTER_ENV, "") + # TODO: NIXL main also has a prometheus_mp exporter that shares this port + # variable but aggregates all ranks behind one endpoint. Add separate + # reservation semantics when Dynamo moves to a NIXL release containing it. + if exporter != "prometheus": return None - return configured_fixed_port( - NIXL_TELEMETRY_PROMETHEUS_PORT_ENV, - default=DEFAULT_NIXL_PROMETHEUS_PORT, - env=environ, + + raw_port = environ.get(NIXL_TELEMETRY_PROMETHEUS_PORT_ENV) + port = ( + DEFAULT_NIXL_PROMETHEUS_PORT + if raw_port is None + else _parse_nixl_uint16(NIXL_TELEMETRY_PROMETHEUS_PORT_ENV, raw_port) ) + if port == 0: + raise ValueError( + f"{NIXL_TELEMETRY_PROMETHEUS_PORT_ENV}=0 asks NIXL for an ephemeral " + "port, which Dynamo cannot declare or configure as a scrape target. " + f"Set a fixed port between 1 and {MAX_PORT}." + ) + return port def reserved_port_ranges( diff --git a/components/src/dynamo/common/utils/prometheus.py b/components/src/dynamo/common/utils/prometheus.py index 63c64612a2fa..7f226347caa3 100644 --- a/components/src/dynamo/common/utils/prometheus.py +++ b/components/src/dynamo/common/utils/prometheus.py @@ -29,6 +29,7 @@ from prometheus_client import CollectorRegistry from dynamo.common.memory import MultimodalEmbeddingCacheManager + from dynamo.common.multimodal.image_loader import ImageLoader # Auto-label injection: always injects dynamo_namespace, dynamo_component, dynamo_endpoint labels # into engine metrics based on the endpoint hierarchy. @@ -52,6 +53,38 @@ class EmbeddingCacheMetrics(str, enum.Enum): ENTRIES = f"{EMBEDDING_CACHE_METRIC_PREFIX}_entries" +# Single source of truth for image loader metric names. +IMAGE_LOADER_METRIC_PREFIX = f"{name_prefix.COMPONENT}_image" + + +class ImageLoaderMetrics(str, enum.Enum): + """Prometheus metric names for the multimodal image loader.""" + + CACHE_ENTRIES = f"{IMAGE_LOADER_METRIC_PREFIX}_cache_entries" + SHARED_CACHE_GET_DURATION_SECONDS = ( + f"{IMAGE_LOADER_METRIC_PREFIX}_shared_cache_get_duration_seconds" + ) + SHARED_CACHE_SET_DURATION_SECONDS = ( + f"{IMAGE_LOADER_METRIC_PREFIX}_shared_cache_set_duration_seconds" + ) + + +_SHARED_IMAGE_CACHE_DURATION_BUCKETS = ( + 0.001, + 0.0025, + 0.005, + 0.01, + 0.025, + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.0, + 4.0, +) + + def register_engine_metrics_callback( endpoint: Endpoint, registry: "CollectorRegistry", @@ -674,3 +707,103 @@ def _collect_embedding_cache_typed() -> list: model_name, component_name, ) + + +def register_image_loader_metrics( + endpoint: "Endpoint", + loader: "ImageLoader", + model_name: str = "", + component_name: str = "", +) -> None: + """Register Prometheus metrics for an ImageLoader instance. + + Same mechanics as register_embedding_cache_metrics: a dedicated + CollectorRegistry (avoids prometheus_client import-ordering issues with + SGLang's multiprocess mode) and a threading.Lock against concurrent scrape + races. Shared-cache operation latencies are drained from the loader's + pending sample buffer and observed into the histograms once per scrape. + + Must be called AFTER engine initialization to ensure prometheus_client is + safe to import. + + Args: + endpoint: Dynamo Endpoint with metrics.register_prometheus_expfmt_callback(). + loader: The ImageLoader instance to observe. + model_name: Model name for the 'model' label. + component_name: Component name for the 'dynamo_component' label. + """ + from prometheus_client import CollectorRegistry, Gauge, Histogram, generate_latest + + registry = CollectorRegistry() + label_names = [labels.MODEL, labels.COMPONENT] + label_values = {labels.MODEL: model_name, labels.COMPONENT: component_name} + + ILM = ImageLoaderMetrics + + shared_cache_get_duration_histogram = Histogram( + ILM.SHARED_CACHE_GET_DURATION_SECONDS, + "Shared encoded-image cache GET latency in seconds by terminal outcome " + "and encoded payload size bucket.", + labelnames=label_names + ["outcome", "size_bucket"], + registry=registry, + buckets=_SHARED_IMAGE_CACHE_DURATION_BUCKETS, + ) + shared_cache_set_duration_histogram = Histogram( + ILM.SHARED_CACHE_SET_DURATION_SECONDS, + "Shared encoded-image cache SET latency in seconds by terminal outcome " + "and encoded payload size bucket.", + labelnames=label_names + ["outcome", "size_bucket"], + registry=registry, + buckets=_SHARED_IMAGE_CACHE_DURATION_BUCKETS, + ) + entries_gauge = Gauge( + ILM.CACHE_ENTRIES, + "Number of entries in the image cache.", + labelnames=label_names, + registry=registry, + ) + + lock = threading.Lock() + + def _refresh_locked() -> None: + shared_cache_stats = loader.shared_image_cache_stats + if shared_cache_stats is not None: + shared_cache_durations = shared_cache_stats.snapshot_and_drain() + for ( + operation, + outcome, + size_bucket, + ), samples in shared_cache_durations.items(): + histogram = ( + shared_cache_get_duration_histogram + if operation == "get" + else shared_cache_set_duration_histogram + ) + for duration in samples: + histogram.labels( + **label_values, + outcome=outcome, + size_bucket=size_bucket, + ).observe(duration) + + entries_gauge.labels(**label_values).set(loader.cache_entries) + + def _collect_image_loader_metrics() -> str: + """Callback invoked on each /metrics scrape.""" + with lock: + _refresh_locked() + return generate_latest(registry).decode("utf-8") + + def _collect_image_loader_typed() -> list: + """Return the same image-loader metrics for OTLP export.""" + with lock: + _refresh_locked() + return get_prometheus_typed(registry) + + endpoint.metrics.register_prometheus_expfmt_callback(_collect_image_loader_metrics) + endpoint.metrics.register_prometheus_typed_callback(_collect_image_loader_typed) + logging.info( + "Registered image loader metrics (model=%s, component=%s)", + model_name, + component_name, + ) diff --git a/components/src/dynamo/common/utils/tests/test_nixl_telemetry.py b/components/src/dynamo/common/utils/tests/test_nixl_telemetry.py index c8e0dcfe6cdd..dfee579a9a00 100644 --- a/components/src/dynamo/common/utils/tests/test_nixl_telemetry.py +++ b/components/src/dynamo/common/utils/tests/test_nixl_telemetry.py @@ -6,6 +6,7 @@ import pytest from dynamo.common.utils.nixl_telemetry import ( + DEFAULT_NIXL_PROMETHEUS_PORT, MAX_COLOCATED_NIXL_EXPORTERS, MAX_PORT, derive_nixl_prometheus_port, @@ -95,16 +96,72 @@ def test_a_launch_wider_than_the_pod_reserves_is_rejected(self): class TestNixlPrometheusBasePort: - def test_operator_defaults_are_recognized(self): - assert nixl_prometheus_base_port(OPERATOR_ENV) == 19090 + @pytest.mark.parametrize( + "enabled_value", ["y", "1", "yes", "on", "true", "enable", "TRUE"] + ) + def test_nixl_truthy_token_is_recognized(self, enabled_value): + env = {**OPERATOR_ENV, "NIXL_TELEMETRY_ENABLE": enabled_value} + assert nixl_prometheus_base_port(env) == 19090 + + @pytest.mark.parametrize( + "enabled_value", ["n", "0", "no", "off", "false", "disable", "FALSE"] + ) + def test_nixl_false_token_disables_telemetry(self, enabled_value): + env = {**OPERATOR_ENV, "NIXL_TELEMETRY_ENABLE": enabled_value} + assert nixl_prometheus_base_port(env) is None + + @pytest.mark.parametrize("enabled_value", ["", "maybe", " y", "y "]) + def test_invalid_enable_value_is_rejected(self, enabled_value): + env = {**OPERATOR_ENV, "NIXL_TELEMETRY_ENABLE": enabled_value} + with pytest.raises(ValueError, match="NIXL_TELEMETRY_ENABLE"): + nixl_prometheus_base_port(env) + + @pytest.mark.parametrize( + ("removed_name", "expected"), + [ + ("NIXL_TELEMETRY_EXPORTER", None), + ("NIXL_TELEMETRY_PROMETHEUS_PORT", DEFAULT_NIXL_PROMETHEUS_PORT), + ], + ) + def test_unset_value(self, removed_name, expected): + env = dict(OPERATOR_ENV) + del env[removed_name] + assert nixl_prometheus_base_port(env) == expected + + def test_hexadecimal_port_is_recognized(self): + env = {**OPERATOR_ENV, "NIXL_TELEMETRY_PROMETHEUS_PORT": "0x4A92"} + assert nixl_prometheus_base_port(env) == 19090 + + @pytest.mark.parametrize("port_value", ["abc", "99999", " 9090", "+9090"]) + def test_invalid_port_is_rejected(self, port_value): + env = {**OPERATOR_ENV, "NIXL_TELEMETRY_PROMETHEUS_PORT": port_value} + with pytest.raises(ValueError, match="NIXL_TELEMETRY_PROMETHEUS_PORT"): + nixl_prometheus_base_port(env) + + def test_oversized_port_is_reported_as_out_of_range(self): + env = {**OPERATOR_ENV, "NIXL_TELEMETRY_PROMETHEUS_PORT": "9" * 5000} + with pytest.raises(ValueError, match="outside the range"): + nixl_prometheus_base_port(env) + + def test_many_leading_zeroes_do_not_trigger_python_integer_limit(self): + env = { + **OPERATOR_ENV, + "NIXL_TELEMETRY_PROMETHEUS_PORT": "0" * 5000 + "19090", + } + assert nixl_prometheus_base_port(env) == 19090 + + def test_ephemeral_port_is_rejected_by_dynamo(self): + env = {**OPERATOR_ENV, "NIXL_TELEMETRY_PROMETHEUS_PORT": "0"} + with pytest.raises(ValueError, match="ephemeral"): + nixl_prometheus_base_port(env) @pytest.mark.parametrize( "override", [ - {"NIXL_TELEMETRY_ENABLE": "n"}, - {"NIXL_TELEMETRY_ENABLE": ""}, - {"NIXL_TELEMETRY_EXPORTER": "file"}, + {"NIXL_TELEMETRY_EXPORTER": "doca"}, + {"NIXL_TELEMETRY_EXPORTER": "PROMETHEUS"}, + {"NIXL_TELEMETRY_EXPORTER": "prometheus "}, ], ) - def test_disabled_telemetry_has_no_base_port(self, override): + def test_inactive_configuration_has_no_base_port(self, override): assert nixl_prometheus_base_port({**OPERATOR_ENV, **override}) is None diff --git a/components/src/dynamo/common/utils/tests/test_topology.py b/components/src/dynamo/common/utils/tests/test_topology.py index 8ae100b2375b..7d3fb27fc69f 100644 --- a/components/src/dynamo/common/utils/tests/test_topology.py +++ b/components/src/dynamo/common/utils/tests/test_topology.py @@ -114,6 +114,35 @@ def test_returns_empty_when_enabled_empty_string(self, monkeypatch): config = read_topology_config() assert not config.enabled + def test_alt_truthy_values_enable_topology(self, monkeypatch, tmp_path): + """parse_bool accepts '1', 'on', 'yes' as truthy; they enable topology like 'true'.""" + topology_dir = tmp_path / "topology" + topology_dir.mkdir() + (topology_dir / "zone").write_text("us-east-1a") + for value in ("1", "on", "yes"): + _enable_topology(monkeypatch, topology_dir) + monkeypatch.setenv("DYN_TOPOLOGY_ENABLED", value) + config = read_topology_config() + assert ( + config.enabled + ), f"expected enabled for DYN_TOPOLOGY_ENABLED={value!r}" + + def test_invalid_values_disable_with_warning(self, monkeypatch, caplog): + """Unrecognized values disable topology and log a warning instead of raising.""" + import logging + + for value in ("maybe", "2", "enabled", "truthy"): + monkeypatch.setenv("DYN_TOPOLOGY_ENABLED", value) + with caplog.at_level(logging.WARNING): + config = read_topology_config() + assert ( + not config.enabled + ), f"expected disabled for DYN_TOPOLOGY_ENABLED={value!r}" + assert any( + "DYN_TOPOLOGY_ENABLED" in r.message for r in caplog.records + ), f"expected warning for DYN_TOPOLOGY_ENABLED={value!r}" + caplog.clear() + def test_reads_all_non_hidden_topology_files(self, monkeypatch, tmp_path): """Reads every visible, non-empty file under the topology mount.""" topology_dir = tmp_path / "topology" diff --git a/components/src/dynamo/common/utils/topology.py b/components/src/dynamo/common/utils/topology.py index fb8590a9b378..112dbf60de12 100644 --- a/components/src/dynamo/common/utils/topology.py +++ b/components/src/dynamo/common/utils/topology.py @@ -10,7 +10,7 @@ derives canonical topology taints from the published topology domains. Environment variables: - DYN_TOPOLOGY_ENABLED: Set to "true" to enable topology reading. + DYN_TOPOLOGY_ENABLED: Set to a truthy value (true/1/on/yes) to enable; any other non-empty value is logged and ignored. DYN_TOPOLOGY_MOUNT_PATH: Directory containing topology domain files (default: /etc/dynamo/topology). DYN_KV_TRANSFER_DOMAIN: Which topology domain the router should enforce @@ -22,6 +22,7 @@ enforcement is "preferred". """ +import argparse import logging import os import sys @@ -29,6 +30,8 @@ from dataclasses import dataclass, field from pathlib import Path +from dynamo.common.configuration.utils import parse_bool + _TOPOLOGY_ENABLED_VAR = "DYN_TOPOLOGY_ENABLED" _TOPOLOGY_MOUNT_PATH_VAR = "DYN_TOPOLOGY_MOUNT_PATH" _KV_TRANSFER_DOMAIN_VAR = "DYN_KV_TRANSFER_DOMAIN" @@ -100,7 +103,7 @@ def _read_kv_transfer_policy() -> tuple[str, str | None, float | None]: if not kv_transfer_domain: logger.error( - "DYN_TOPOLOGY_ENABLED=true but %s is not set. The deployment " + "DYN_TOPOLOGY_ENABLED is set but %s is not set. The deployment " "environment must set the KV transfer domain when topology is " "enabled. Exiting.", _KV_TRANSFER_DOMAIN_VAR, @@ -119,7 +122,7 @@ def read_topology_config( The deployment environment injects env vars for topology location and transfer policy: - - DYN_TOPOLOGY_ENABLED=true + - DYN_TOPOLOGY_ENABLED=true # any truthy value (true/1/on/yes) enables - DYN_TOPOLOGY_MOUNT_PATH=/etc/dynamo/topology - DYN_KV_TRANSFER_DOMAIN=zone - DYN_KV_TRANSFER_ENFORCEMENT=required @@ -138,12 +141,21 @@ def read_topology_config( Empty config if topology is not enabled. Raises: - SystemExit: If DYN_TOPOLOGY_ENABLED=true but DYN_KV_TRANSFER_DOMAIN is + SystemExit: If DYN_TOPOLOGY_ENABLED is set but DYN_KV_TRANSFER_DOMAIN is not set, the transfer-domain topology file is still missing or empty after the timeout. """ - enabled = os.environ.get(_TOPOLOGY_ENABLED_VAR, "").strip().lower() - if enabled != "true": + raw = os.environ.get(_TOPOLOGY_ENABLED_VAR, "").strip() + try: + enabled = parse_bool(raw) + except argparse.ArgumentTypeError: + if raw: + logger.warning( + "Unrecognized DYN_TOPOLOGY_ENABLED=%r, treating as disabled; use 'true'/'false' (or 1/0, on/off, yes/no)", + raw, + ) + return TopologyConfig() + if not enabled: return TopologyConfig() ( @@ -173,7 +185,7 @@ def read_topology_config( if kv_transfer_domain not in topology_domains: logger.error( - "DYN_TOPOLOGY_ENABLED=true but topology file %s was not populated " + "DYN_TOPOLOGY_ENABLED is set but topology file %s was not populated " "within %.0fs. This indicates the configured topology source did " "not publish the selected transfer domain. Exiting.", transfer_domain_file, diff --git a/components/src/dynamo/common/utils/video_utils.py b/components/src/dynamo/common/utils/video_utils.py index c7cc59f1be0c..0f0cb1d5a141 100644 --- a/components/src/dynamo/common/utils/video_utils.py +++ b/components/src/dynamo/common/utils/video_utils.py @@ -45,7 +45,7 @@ def parse_size( def compute_num_frames( num_frames: int | None = None, seconds: int | None = None, - fps: int | None = None, + fps: float | None = None, default_fps: int = DEFAULT_VIDEO_FPS, default_num_frames: int = DEFAULT_VIDEO_NUM_FRAMES, ) -> int: @@ -58,7 +58,7 @@ def compute_num_frames( if seconds is not None or fps is not None: _seconds = seconds if seconds is not None else 4 _fps = fps if fps is not None else default_fps - return _seconds * _fps + return round(_seconds * _fps) return default_num_frames diff --git a/components/src/dynamo/frontend/cpu_affinity.py b/components/src/dynamo/frontend/cpu_affinity.py new file mode 100644 index 000000000000..3db1396c5b5e --- /dev/null +++ b/components/src/dynamo/frontend/cpu_affinity.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Warn when the frontend can run across multiple NUMA nodes.""" + +import logging +import os +from pathlib import Path + +_CPU_SYSFS_ROOT = Path("/sys/devices/system/cpu") + + +def _format_cpu_ranges(cpus: list[int]) -> str: + ranges: list[tuple[int, int]] = [] + for cpu in cpus: + if ranges and cpu == ranges[-1][1] + 1: + ranges[-1] = (ranges[-1][0], cpu) + else: + ranges.append((cpu, cpu)) + return ",".join( + str(start) if start == end else f"{start}-{end}" for start, end in ranges + ) + + +def warn_if_frontend_cpu_affinity_spans_numa_nodes( + logger: logging.Logger, +) -> None: + """Log a warning if the current CPU affinity covers multiple NUMA nodes.""" + + try: + cpus = sorted(os.sched_getaffinity(0)) + nodes = sorted( + { + node_path.name + for cpu in cpus + for node_path in (_CPU_SYSFS_ROOT / f"cpu{cpu}").glob("node[0-9]*") + } + ) + except (AttributeError, OSError) as error: + logger.debug("NUMA affinity check skipped: %s", error, exc_info=True) + return + + if not nodes: + logger.debug( + "NUMA affinity check skipped: no NUMA node links found for CPUs %s", + _format_cpu_ranges(cpus), + ) + return + + if len(nodes) == 1: + return + + border = "=" * 80 + logger.warning( + "%s\n" + "CRITICAL PERFORMANCE WARNING: Frontend CPU affinity spans " + "multiple NUMA nodes!\n" + "Running the frontend across multiple NUMA domains is known to cause " + "catastrophic performance issues.\n" + "Available CPUs: %s\n" + "NUMA nodes: %s\n" + "Pin the frontend to CPUs from a single NUMA node.\n" + "%s", + border, + _format_cpu_ranges(cpus), + ", ".join(nodes), + border, + ) diff --git a/components/src/dynamo/frontend/main.py b/components/src/dynamo/frontend/main.py index c53752323544..05dc905deb72 100644 --- a/components/src/dynamo/frontend/main.py +++ b/components/src/dynamo/frontend/main.py @@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Any, Optional import uvloop +from packaging.version import Version from dynamo.common.config_dump import dump_config from dynamo.common.configuration.groups.router_args import build_router_config @@ -41,6 +42,7 @@ from dynamo.runtime import DistributedRuntime from dynamo.runtime.logging import configure_dynamo_logging +from .cpu_affinity import warn_if_frontend_cpu_affinity_spans_numa_nodes from .frontend_args import FrontendArgGroup, FrontendConfig if TYPE_CHECKING: @@ -291,7 +293,11 @@ def parse_args() -> tuple[FrontendConfig, Optional[Namespace], Optional[Namespac try: from vllm.engine.arg_utils import AsyncEngineArgs - from vllm.entrypoints.openai.cli_args import FrontendArgs + + if Version(importlib.metadata.version("vllm")).release >= (0, 29): + from vllm.entrypoints.launchers.cli_args import FrontendArgs + else: + from vllm.entrypoints.openai.cli_args import FrontendArgs except ModuleNotFoundError: logger.exception("Flag '--chat-processor vllm' requires vllm be installed.") sys.exit(1) @@ -391,6 +397,7 @@ async def async_main(): # it connects to NATS eagerly, so NATS (m)TLS env vars must already be set or # the CLI flags are silently ignored (unlike the lazily-dialed TCP planes). _export_transport_tls_env(config) + warn_if_frontend_cpu_affinity_spans_numa_nodes(logger) runtime = DistributedRuntime( loop, config.discovery_backend, diff --git a/components/src/dynamo/frontend/prepost.py b/components/src/dynamo/frontend/prepost.py index e6a6f0b5f1b2..f7d4010e47a4 100644 --- a/components/src/dynamo/frontend/prepost.py +++ b/components/src/dynamo/frontend/prepost.py @@ -6,11 +6,12 @@ import json import os import re +import weakref from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import ThreadPoolExecutor from copy import deepcopy from dataclasses import dataclass -from functools import lru_cache +from functools import lru_cache, partial from typing import TYPE_CHECKING, Any, Protocol, TypeGuard, cast from vllm.entrypoints.chat_utils import make_tool_call_id @@ -19,11 +20,6 @@ ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, -) from vllm.reasoning import ReasoningParser from vllm.renderers import ChatParams, merge_kwargs from vllm.sampling_params import SamplingParams @@ -33,6 +29,7 @@ from vllm.utils.async_utils import make_async from dynamo.common.utils.guided_json import admits_only_empty_object +from dynamo.frontend.vllm_protocol import DeltaFunctionCall, DeltaMessage, DeltaToolCall from dynamo.llm.exceptions import InvalidArgument from .thinking import apply_default_thinking_mode_to_template_kwargs @@ -69,10 +66,35 @@ class PreprocessResult: uses_dynamo_json_tool_call_fallback: bool = False -_ASYNC_TOKENIZER_POOL: dict[int, Callable[..., Awaitable[Any]]] = {} +# One executor per live tokenizer. The id-keyed registry stores only a weak +# reference to the tokenizer, and every lookup verifies object identity. This +# avoids both id reuse and WeakKeyDictionary's referent-based equality. +_ASYNC_TOKENIZER_EXECUTORS: dict[ + int, tuple[weakref.ReferenceType[TokenizerLike], ThreadPoolExecutor] +] = {} +# Fallback for tokenizers that do not support weak references; retains +# entries for the process lifetime (the previous behavior for all tokenizers). +# The tokenizer is stored alongside its executor so its id() cannot be +# recycled by a different tokenizer while the entry lives. +_STRONG_ASYNC_TOKENIZER_EXECUTORS: dict[ + int, tuple[TokenizerLike, ThreadPoolExecutor] +] = {} SKIP_REQUEST_VALIDATION = os.getenv("DYN_VLLM_SKIP_REQUEST_VALIDATION", "1") == "1" +def _evict_async_tokenizer_executor( + key: int, + executor: ThreadPoolExecutor, + tokenizer_ref: weakref.ReferenceType[TokenizerLike], +) -> None: + entry = _ASYNC_TOKENIZER_EXECUTORS.get(key) + if entry is not None and entry[0] is tokenizer_ref: + del _ASYNC_TOKENIZER_EXECUTORS[key] + # The tokenizer's last reference may be released by this executor's own + # worker, so shutdown must not try to join the current thread. + executor.shutdown(wait=False) + + def _reject_non_finite_json(value: str) -> Any: raise ValueError(f"non-finite JSON number: {value}") @@ -332,14 +354,30 @@ def _build_assistant_guided_decoding( def _get_async_tokenizer(tokenizer: TokenizerLike) -> Callable[..., Awaitable[Any]]: - key = id(tokenizer) - async_tokenizer = _ASYNC_TOKENIZER_POOL.get(key) - if async_tokenizer is None: - async_tokenizer = make_async( - tokenizer, executor=ThreadPoolExecutor(max_workers=1) - ) - _ASYNC_TOKENIZER_POOL[key] = async_tokenizer - return async_tokenizer + try: + tokenizer_ref = weakref.ref(tokenizer) + except TypeError: + # Tokenizer does not support weak references. + key = id(tokenizer) + entry = _STRONG_ASYNC_TOKENIZER_EXECUTORS.get(key) + if entry is None: + executor = ThreadPoolExecutor(max_workers=1) + _STRONG_ASYNC_TOKENIZER_EXECUTORS[key] = (tokenizer, executor) + else: + executor = entry[1] + else: + key = id(tokenizer) + entry = _ASYNC_TOKENIZER_EXECUTORS.get(key) + if entry is None or entry[0]() is not tokenizer: + executor = ThreadPoolExecutor(max_workers=1) + tokenizer_ref = weakref.ref( + tokenizer, + partial(_evict_async_tokenizer_executor, key, executor), + ) + _ASYNC_TOKENIZER_EXECUTORS[key] = (tokenizer_ref, executor) + else: + executor = entry[1] + return make_async(tokenizer, executor=executor) def _materialize_assistant_tool_calls( diff --git a/components/src/dynamo/frontend/sglang_prepost.py b/components/src/dynamo/frontend/sglang_prepost.py index dde1eb2ef4e6..b04beadfe833 100644 --- a/components/src/dynamo/frontend/sglang_prepost.py +++ b/components/src/dynamo/frontend/sglang_prepost.py @@ -347,14 +347,31 @@ def named_closed_zero_arg_tool(request: dict[str, Any]) -> str | None: return None -def _guided_tool_choice_requires_reasoning( - request: dict[str, Any], force_reasoning: bool +def _guided_output_requires_reasoning( + request: dict[str, Any], + force_reasoning: bool, + reasoning_parser_name: str | None = None, + guided_decoding: dict[str, Any] | None = None, ) -> bool: - """Return whether SGLang should reason before guided tool-call JSON.""" + """Return whether SGLang should reason before guided output.""" + if not force_reasoning: + return False + tool_choice = request.get("tool_choice", "auto") - return force_reasoning and ( - tool_choice == "required" or _is_named_tool_choice(tool_choice) - ) + if tool_choice == "required" or _is_named_tool_choice(tool_choice): + return True + + # Explicit legacy constraints take precedence over response_format. + if legacy_guided_decoding(request): + return False + + response_format = request.get("response_format") + if isinstance(response_format, dict) and response_format.get("type") != "text": + return reasoning_parser_name != "gpt-oss" + + # An auto tool-call grammar forbids the end-of-thinking marker, so it must + # also wait for thinking to finish. + return guided_decoding is not None and "structural_tag" in guided_decoding def _normalize_deepseek_v4_hint(value: Any) -> str: diff --git a/components/src/dynamo/frontend/sglang_processor.py b/components/src/dynamo/frontend/sglang_processor.py index b25de327a1a6..726c66272e15 100644 --- a/components/src/dynamo/frontend/sglang_processor.py +++ b/components/src/dynamo/frontend/sglang_processor.py @@ -32,7 +32,7 @@ ToolCallParserType, _client_wants_separate_reasoning, _get_history_tool_calls_count, - _guided_tool_choice_requires_reasoning, + _guided_output_requires_reasoning, convert_tools, create_parsers, detect_force_reasoning_from_template, @@ -42,6 +42,7 @@ from .thinking import runtime_default_thinking_mode from .utils import ( PreprocessError, + as_error_envelope, extract_mm_urls, handle_engine_error, make_internal_error, @@ -357,8 +358,8 @@ def _preprocess_worker( pre.guided_decoding, pre.tool_call_parser, pre.reasoning_parser, - require_reasoning=_guided_tool_choice_requires_reasoning( - request, pre.force_reasoning + require_reasoning=_guided_output_requires_reasoning( + request, pre.force_reasoning, _w_reasoning_parser_name, pre.guided_decoding ), ) @@ -615,8 +616,11 @@ async def _generator_inner( pre.guided_decoding, pre.tool_call_parser, pre.reasoning_parser, - require_reasoning=_guided_tool_choice_requires_reasoning( - request, pre.force_reasoning + require_reasoning=_guided_output_requires_reasoning( + request, + pre.force_reasoning, + self.reasoning_parser_name, + pre.guided_decoding, ), ) except PreprocessError as exc: @@ -848,8 +852,12 @@ def flush_pending( cached_tokens = _cached_tokens_from_usage(usage_for_metrics) if cached_tokens is not None: metrics["cached_tokens"] = cached_tokens - envelope["event"] = "llm_metrics" - envelope["comment"] = [json.dumps(metrics)] + # Attach metrics to data when available; otherwise use an annotation. + if data := envelope.get("data"): + data["llm_metrics"] = metrics + else: + envelope["event"] = "llm_metrics" + envelope["comment"] = [json.dumps(metrics)] pending_token_ids = [] pending_log_probs = None @@ -867,7 +875,7 @@ def flush_pending( request_id, message, ) - yield make_internal_error(request_id, message) + yield as_error_envelope(make_internal_error(request_id, message)) break engine_response = dynamo_response.data() @@ -880,7 +888,9 @@ def flush_pending( not isinstance(engine_response, dict) or "token_ids" not in engine_response ): - yield handle_engine_error(engine_response, request_id, logger) + yield as_error_envelope( + handle_engine_error(engine_response, request_id, logger) + ) break new_ids = engine_response["token_ids"] @@ -943,7 +953,7 @@ def flush_pending( yield envelope if post.locally_finished: break - except Unknown: + except (InvalidArgument, Unknown): raise except Exception as e: logger.exception("Error generating response for request %s", request_id) diff --git a/components/src/dynamo/frontend/tests/test_cpu_affinity.py b/components/src/dynamo/frontend/tests/test_cpu_affinity.py new file mode 100644 index 000000000000..33f64cea1b79 --- /dev/null +++ b/components/src/dynamo/frontend/tests/test_cpu_affinity.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import os +from pathlib import Path + +import pytest + +from dynamo.frontend import cpu_affinity +from dynamo.frontend.cpu_affinity import warn_if_frontend_cpu_affinity_spans_numa_nodes + +pytestmark = [pytest.mark.unit, pytest.mark.gpu_0, pytest.mark.pre_merge] + + +def _add_cpu(root: Path, cpu: int, node: int) -> None: + (root / f"cpu{cpu}" / f"node{node}").mkdir(parents=True) + + +def test_warns_when_affinity_spans_multiple_numa_nodes( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +): + for cpu in (0, 1, 2, 3): + _add_cpu(tmp_path, cpu, 0) + for cpu in (8, 9): + _add_cpu(tmp_path, cpu, 1) + monkeypatch.setattr(cpu_affinity, "_CPU_SYSFS_ROOT", tmp_path) + monkeypatch.setattr(os, "sched_getaffinity", lambda _pid: {0, 1, 2, 3, 8, 9}) + logger = logging.getLogger("test.frontend.cpu_affinity.multiple") + + with caplog.at_level(logging.WARNING, logger=logger.name): + warn_if_frontend_cpu_affinity_spans_numa_nodes(logger) + + assert len(caplog.records) == 1 + message = caplog.records[0].getMessage() + assert "CRITICAL PERFORMANCE WARNING" in message + assert "Frontend CPU affinity spans multiple NUMA nodes" in message + assert ( + "Running the frontend across multiple NUMA domains is known to cause " + "catastrophic performance issues." + ) in message + assert "Available CPUs: 0-3,8-9" in message + assert "node0" in message + assert "node1" in message + + +@pytest.mark.parametrize( + ("cpus", "cpu_nodes", "debug_message"), + [ + pytest.param({0, 1}, [(0, 0), (1, 0)], None, id="same_node"), + pytest.param({8}, [], "no NUMA node links found", id="cpu_missing_from_sysfs"), + ], +) +def test_does_not_warn_without_multiple_detected_nodes( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + cpus: set[int], + cpu_nodes: list[tuple[int, int]], + debug_message: str | None, +): + for cpu, node in cpu_nodes: + _add_cpu(tmp_path, cpu, node) + monkeypatch.setattr(cpu_affinity, "_CPU_SYSFS_ROOT", tmp_path) + monkeypatch.setattr(os, "sched_getaffinity", lambda _pid: cpus) + logger = logging.getLogger("test.frontend.cpu_affinity.silent") + + with caplog.at_level(logging.DEBUG, logger=logger.name): + warn_if_frontend_cpu_affinity_spans_numa_nodes(logger) + + assert not [ + record for record in caplog.records if record.levelno >= logging.WARNING + ] + if debug_message is None: + assert caplog.records == [] + else: + assert len(caplog.records) == 1 + assert debug_message in caplog.records[0].getMessage() + + +def test_affinity_failure_logs_debug_traceback_without_warning( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + def fail_affinity(_pid: int) -> set[int]: + raise OSError("not permitted") + + monkeypatch.setattr(os, "sched_getaffinity", fail_affinity) + logger = logging.getLogger("test.frontend.cpu_affinity.failure") + + with caplog.at_level(logging.DEBUG, logger=logger.name): + warn_if_frontend_cpu_affinity_spans_numa_nodes(logger) + + assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.DEBUG + assert "not permitted" in caplog.records[0].getMessage() + assert caplog.records[0].exc_info is not None diff --git a/components/src/dynamo/frontend/tests/test_multimodal_utils.py b/components/src/dynamo/frontend/tests/test_multimodal_utils.py index 49a0245f0ce0..4aaf68fd4aa6 100644 --- a/components/src/dynamo/frontend/tests/test_multimodal_utils.py +++ b/components/src/dynamo/frontend/tests/test_multimodal_utils.py @@ -140,6 +140,22 @@ def test_extracts_multiple_items_per_modality(): ) +def test_extracts_tool_message_media(): + messages = [ + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/t.png"}} + ], + } + ] + assert extract_mm_urls(messages) == ( + {"image_url": [{"Url": "https://example.com/t.png"}]}, + None, + ) + + def test_ignores_non_user_messages(): messages = [ {"role": "system", "content": "You are helpful."}, diff --git a/components/src/dynamo/frontend/tests/test_sglang_processor_metrics_unit.py b/components/src/dynamo/frontend/tests/test_sglang_processor_metrics_unit.py index 2ea8e6dcd4a5..a8bddea30c31 100644 --- a/components/src/dynamo/frontend/tests/test_sglang_processor_metrics_unit.py +++ b/components/src/dynamo/frontend/tests/test_sglang_processor_metrics_unit.py @@ -5,7 +5,6 @@ import asyncio import importlib -import json import sys import types @@ -114,6 +113,11 @@ class _ReasoningParser: get_json_schema_constraint=lambda *args, **kwargs: None, ) _install_module(install_module, "sglang.srt.parser") + _install_module( + install_module, + "sglang.srt.parser.conversation", + chat_template_exists=lambda *_args, **_kwargs: False, + ) _install_module( install_module, "sglang.srt.parser.jinja_template_utils", @@ -150,6 +154,27 @@ def _load_processor_module(module_stubs): install_module, remove_module = module_stubs _install_sglang_stubs(install_module) _install_module(install_module, "dynamo._internal", ModelDeploymentCard=object) + _install_module(install_module, "dynamo.common.multimodal") + _install_module( + install_module, + "dynamo.common.multimodal.cache_uuid", + reject_unsupported_multimodal_uuids=lambda *_args, **_kwargs: None, + ) + _install_module( + install_module, + "dynamo.frontend.sglang_prepost", + ReasoningParser=object, + SglangStreamingPostProcessor=object, + ToolCallParserType=object, + _client_wants_separate_reasoning=lambda *_args, **_kwargs: False, + _get_history_tool_calls_count=lambda *_args, **_kwargs: 0, + _guided_output_requires_reasoning=lambda *_args, **_kwargs: False, + convert_tools=lambda *_args, **_kwargs: None, + create_parsers=lambda *_args, **_kwargs: (None, None), + detect_force_reasoning_from_template=lambda *_args, **_kwargs: False, + preprocess_chat_request=lambda *_args, **_kwargs: None, + resolve_skip_special_tokens=lambda *_args, **_kwargs: True, + ) _install_module( install_module, "dynamo.frontend.frontend_args", @@ -165,6 +190,7 @@ def _load_processor_module(module_stubs): _install_module( install_module, "dynamo.llm.exceptions", + HttpError=type("HttpError", (Exception,), {}), InvalidArgument=type("InvalidArgument", (Exception,), {}), Unknown=type("Unknown", (Exception,), {}), ) @@ -172,7 +198,7 @@ def _load_processor_module(module_stubs): return importlib.import_module("dynamo.frontend.sglang_processor") -def test_stream_emits_llm_metrics_annotation(module_stubs): +def test_stream_embeds_typed_llm_metrics(module_stubs): module = _load_processor_module(module_stubs) completion_usage = { "prompt_tokens": 10, @@ -209,15 +235,14 @@ async def collect(): ] items = asyncio.run(collect()) - metric_items = [item for item in items if item.get("event") == "llm_metrics"] - - assert len(metric_items) == 1 - envelope = metric_items[0] + assert len(items) == 1 + envelope = items[0] assert envelope["_dynamo_annotated"] is True assert envelope["data"]["usage"] == completion_usage - metrics = json.loads(envelope["comment"][0]) + assert "event" not in envelope + assert "comment" not in envelope # Zero counts are omitted (text-only request), mirroring the Rust skip-zero behavior. - assert metrics == { + assert envelope["data"]["llm_metrics"] == { "input_tokens": 10, "output_tokens": 3, "chunk_tokens": 3, @@ -278,9 +303,8 @@ async def collect(): ] items = asyncio.run(collect()) - metric_items = [item for item in items if item.get("event") == "llm_metrics"] - assert len(metric_items) == 1 - metrics = json.loads(metric_items[0]["comment"][0]) + assert len(items) == 1 + metrics = items[0]["data"]["llm_metrics"] assert metrics["image_count"] == 2 assert metrics["video_count"] == 1 # audio has zero parts, so the key is omitted from the emitted metrics. diff --git a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py index cbd97aee0205..f9bc8a1ea128 100644 --- a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py @@ -14,7 +14,7 @@ import json import sys import types -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor from contextlib import nullcontext import pytest @@ -37,7 +37,7 @@ SglangPreprocessResult, SglangStreamingPostProcessor, _flatten_message_content, - _guided_tool_choice_requires_reasoning, + _guided_output_requires_reasoning, _normalize_assistant_tool_call_arguments, _normalize_prompt_token_ids, _normalize_sglang_parser_name, @@ -1061,7 +1061,7 @@ def test_qwen_separate_reasoning_false_keeps_generation_gate(tokenizer, tool_cho assert result.force_reasoning is True assert result.reasoning_parser is None - assert _guided_tool_choice_requires_reasoning(request, result.force_reasoning) + assert _guided_output_requires_reasoning(request, result.force_reasoning) @pytest.mark.parametrize( @@ -1138,7 +1138,187 @@ def test_guided_tool_choice_requires_effective_reasoning( ): """Only reasoning-enabled required or named tools activate the gate.""" request = {"tool_choice": tool_choice} - assert _guided_tool_choice_requires_reasoning(request, force_reasoning) is expected + assert _guided_output_requires_reasoning(request, force_reasoning) is expected + + +@pytest.mark.core +def test_structured_response_requires_effective_reasoning(): + request = {"response_format": {"type": "json_schema"}} + + assert _guided_output_requires_reasoning(request, True, "deepseek-v4") is True + assert _guided_output_requires_reasoning(request, False, "deepseek-v4") is False + assert _guided_output_requires_reasoning(request, True, "gpt-oss") is False + assert not _guided_output_requires_reasoning( + {"response_format": {"type": "text"}}, True, "deepseek-v4" + ) + + +@pytest.mark.parametrize( + ("request_fields", "force_reasoning", "expected"), + [ + ({}, True, True), + ({"response_format": {"type": "text"}}, True, True), + # A structural_tag response_format keeps the gpt-oss exception. + ({"response_format": {"type": "structural_tag"}}, True, False), + ], +) +def test_auto_tool_structural_tag_requires_effective_reasoning( + request_fields, force_reasoning, expected +): + request = {"tool_choice": "auto", **request_fields} + guided_decoding = {"structural_tag": {}} + assert ( + _guided_output_requires_reasoning( + request, force_reasoning, "gpt-oss", guided_decoding + ) + is expected + ) + + +@pytest.mark.parametrize("thinking", [True, False]) +def test_kimi_k3_auto_tool_sets_reasoning_gate_pool(thinking, monkeypatch): + """Kimi K3's auto tool grammar forbids <|close|>think<|sep|>, so it must + only apply after thinking ends.""" + + class StubTokenizer: + chat_template = "template" + + def apply_chat_template(self, messages, **kwargs): + return [1, 2, 3] + + monkeypatch.setattr(sglang_processor_module, "_w_tokenizer", StubTokenizer()) + monkeypatch.setattr(sglang_processor_module, "_w_tool_call_parser_name", "kimi_k3") + monkeypatch.setattr(sglang_processor_module, "_w_reasoning_parser_name", "kimi_k3") + monkeypatch.setattr( + sglang_processor_module, "_w_exclude_tools_when_tool_choice_none", True + ) + monkeypatch.setattr(sglang_processor_module, "_w_template_force_reasoning", False) + + request = { + "model": "moonshotai/Kimi-K3", + "messages": [{"role": "user", "content": "Weather in San Francisco?"}], + "chat_template_kwargs": {"thinking": thinking}, + "tool_choice": "auto", + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + } + result = _preprocess_worker(request, MODEL, eos_token_ids=None) + + guided_decoding = result.dynamo_preproc["sampling_options"]["guided_decoding"] + assert "structural_tag" in guided_decoding + assert result.dynamo_preproc["require_reasoning"] is thinking + + +@pytest.mark.core +@pytest.mark.parametrize( + ("legacy_constraint", "expected"), + [ + ({"guided_regex": "trueish"}, False), + ({"guided_choice": []}, True), + ], +) +def test_structured_response_respects_legacy_constraint_precedence( + legacy_constraint, expected +): + request = {"response_format": {"type": "json_schema"}, **legacy_constraint} + + assert _guided_output_requires_reasoning(request, True, "qwen3") is expected + + +@pytest.mark.core +@pytest.mark.timeout(60) +@pytest.mark.parametrize( + ("use_pool", "thinking", "separate_reasoning", "legacy_regex"), + [ + pytest.param(False, True, True, False, id="inline-structured"), + pytest.param(False, False, True, False, id="inline-thinking-disabled"), + pytest.param(False, True, False, False, id="inline-separation-disabled"), + pytest.param(False, True, True, True, id="inline-legacy-regex"), + pytest.param(True, True, True, False, id="pool-structured"), + ], +) +def test_structured_response_generator_forwards_reasoning_gate( + tokenizer, monkeypatch, use_pool, thinking, separate_reasoning, legacy_regex +): + response_format = { + "type": "json_schema", + "json_schema": { + "name": "answer", + "schema": { + "type": "object", + "properties": {"answer": {"type": "integer"}}, + "required": ["answer"], + }, + }, + } + request = { + "model": MODEL, + "messages": [{"role": "user", "content": "Return answer 42 as JSON."}], + "response_format": response_format, + "chat_template_kwargs": {"enable_thinking": thinking}, + "separate_reasoning": separate_reasoning, + "stream": True, + } + if legacy_regex: + request["guided_regex"] = "trueish" + routed_engine = FakeRoutedEngine(items=[{"token_ids": [], "finish_reason": "stop"}]) + + class InlinePreprocessPool: + # Keep the real worker and pool generator path without spawning processes. + def submit(self, fn, *args): + future = Future() + future.set_result(fn(*args)) + return future + + if use_pool: + monkeypatch.setattr(sglang_processor_module, "_w_tokenizer", tokenizer) + monkeypatch.setattr(sglang_processor_module, "_w_tool_call_parser_name", None) + monkeypatch.setattr( + sglang_processor_module, "_w_reasoning_parser_name", "qwen3" + ) + monkeypatch.setattr( + sglang_processor_module, "_w_exclude_tools_when_tool_choice_none", True + ) + monkeypatch.setattr( + sglang_processor_module, "_w_template_force_reasoning", False + ) + monkeypatch.setattr(sglang_processor_module, "_w_default_thinking_mode", None) + processor = SglangProcessor( + tokenizer=tokenizer, + routed_engine=routed_engine, + tool_call_parser_name=None, + reasoning_parser_name="qwen3", + eos_token_ids=None, + preprocess_pool=InlinePreprocessPool() if use_pool else None, + ) + + async def collect(): + return [item async for item in processor.generator(request)] + + asyncio.run(collect()) + assert len(routed_engine.requests) == 1 + assert routed_engine.requests[0]["require_reasoning"] is ( + thinking and not legacy_regex + ) + if legacy_regex: + assert routed_engine.requests[0]["sampling_options"]["guided_decoding"] == { + "regex": "trueish" + } + else: + assert routed_engine.requests[0]["sampling_options"]["guided_decoding"] == { + "json": response_format["json_schema"]["schema"] + } class _CapturingReasoningParser: @@ -4049,16 +4229,70 @@ async def collect(): return asyncio.run(collect()) - def test_routed_engine_is_error_yields_internal_error(self, tokenizer): - """is_error() True yields a single internal_error chunk with the comment text.""" + @pytest.mark.parametrize("after_output", [False, True]) + def test_routed_engine_is_error_yields_error_envelope(self, after_output): + """An annotated engine error terminates with a binding-compatible error frame.""" + prefix = [{"token_ids": [ord("A")]}] if after_output else [] items = self._run_stream( - tokenizer, - [FakeRoutedItem(None, is_error=True, comments=["backend disconnected"])], + self.ByteTokenizer(), + [ + *prefix, + FakeRoutedItem(None, is_error=True, comments=["backend disconnected"]), + {"token_ids": [], "finish_reason": "stop"}, + ], ) - assert len(items) == 1 - err = items[0]["error"] - assert err["type"] == "internal_error" - assert "backend disconnected" in err["message"] + assert len(items) == len(prefix) + 1 + if after_output: + assert items[0]["choices"][0]["delta"]["content"] == "A" + assert items[0]["choices"][0]["finish_reason"] is None + assert items[-1] == { + "_dynamo_annotated": True, + "event": "error", + "comment": ["backend disconnected"], + } + + @pytest.mark.parametrize("after_output", [False, True]) + def test_routed_validation_exception_preserves_type_and_message(self, after_output): + """A typed iterator failure survives before or after an emitted delta.""" + message = "Failed to compile json grammar: unsupported schema type" + + class FailingEngine: + async def generate(self, preprocessed, **kwargs): + """Model the routed iterator raising a native validation exception.""" + + async def stream(): + if after_output: + yield FakeRoutedItem({"token_ids": [ord("A")]}) + raise InvalidArgument(message) + + return stream() + + async def check(): + tokenizer = self.ByteTokenizer() + processor = SglangProcessor( + tokenizer=tokenizer, + routed_engine=FailingEngine(), + tool_call_parser_name=None, + reasoning_parser_name=None, + eos_token_ids=None, + ) + post = SglangStreamingPostProcessor( + tokenizer=tokenizer, tool_call_parser=None, reasoning_parser=None + ) + stream = processor._generate_and_stream( + "req-invalid", {"model": "test-model"}, {}, [], post + ) + if after_output: + first = await anext(stream) + assert first["data"]["choices"][0]["delta"]["content"] == "A" + assert first["data"]["choices"][0]["finish_reason"] is None + with pytest.raises(InvalidArgument) as error: + await anext(stream) + assert str(error.value) == message + with pytest.raises(StopAsyncIteration): + await anext(stream) + + asyncio.run(check()) def test_routed_engine_none_data_is_skipped(self, tokenizer): """data() is None (e.g. comment-only event) is skipped, not yielded as error.""" @@ -4080,7 +4314,11 @@ def test_malformed_engine_response_yields_engine_error(self, tokenizer): [{"status": "error", "message": "kv cache exhausted"}], ) assert len(items) == 1 - assert "error" in items[0] + assert items[0] == { + "_dynamo_annotated": True, + "event": "error", + "comment": ["kv cache exhausted"], + } def test_completed_batches_replace_decode_context(self): """Completed batches replace context instead of accumulating history.""" @@ -4451,6 +4689,84 @@ def test_required_tool_distinguishes_bare_json_from_reasoning( "city": "New York" } + @pytest.mark.core + @pytest.mark.timeout(60) + @pytest.mark.parametrize( + ("chunk_size", "finish_with_tokens"), [(1, False), (10000, True)] + ) + @pytest.mark.parametrize( + ("text", "expected_reasoning", "expected_content", "finish_reason"), + [ + ('{"answer":42}', '{"answer":42}', "", "stop"), + ('{"answer":42}', '{"answer":42}', "", "length"), + ( + '{"answer":41}{"answer":42}', + '{"answer":41}', + '{"answer":42}', + "stop", + ), + ], + ) + def test_structured_response_preserves_reasoning_boundaries( + self, + tokenizer, + chunk_size, + finish_with_tokens, + text, + expected_reasoning, + expected_content, + finish_reason, + ): + _, reasoning_parser = create_parsers( + {}, + tool_call_parser_name=None, + reasoning_parser_name="qwen3", + force_reasoning=True, + ) + assert reasoning_parser is not None + case_tokenizer = copy.deepcopy(tokenizer) + detector = reasoning_parser.detector + case_tokenizer.add_special_tokens( + { + "additional_special_tokens": [ + detector.think_start_token, + detector.think_end_token, + ] + } + ) + post = SglangStreamingPostProcessor( + tokenizer=case_tokenizer, + tool_call_parser=None, + reasoning_parser=reasoning_parser, + ) + token_ids = case_tokenizer.encode(text, add_special_tokens=False) + responses = [] + for offset in range(0, len(token_ids), chunk_size): + is_last = offset + chunk_size >= len(token_ids) + choice = post.process_output( + { + "token_ids": token_ids[offset : offset + chunk_size], + "finish_reason": ( + finish_reason if finish_with_tokens and is_last else None + ), + } + ) + if choice: + responses.append(choice) + if not finish_with_tokens: + terminal = post.process_output( + {"token_ids": [], "finish_reason": finish_reason} + ) + assert terminal is not None + responses.append(terminal) + content = "".join(item["delta"].get("content", "") for item in responses) + reasoning = "".join( + item["delta"].get("reasoning_content", "") for item in responses + ) + assert content == expected_content + assert reasoning == expected_reasoning + assert responses[-1]["finish_reason"] == finish_reason + # --------------------------------------------------------------------------- # Utility functions diff --git a/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py b/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py index 60d8aeee2398..7dcbf4433c77 100644 --- a/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py @@ -1008,6 +1008,106 @@ def test_kv_router_copies_reasoning_metadata_to_extra_args(self): } +class TestReasoningParserOutputCapability: + def test_harmony_note_only_for_gptoss(self): + from dynamo.frontend.vllm_processor import ( + _ensure_reasoning_parser_output_capable, + ) + + class BoundaryOnlyParser: + def __init__(self, tokenizer, *args, **kwargs): + pass + + def extract_reasoning_streaming(self, *args): + raise NotImplementedError("boundary detection only") + + with pytest.raises(RuntimeError, match="boundary detection") as exc_info: + _ensure_reasoning_parser_output_capable( + "future_parser", BoundaryOnlyParser, object(), {}, None + ) + assert "HarmonyParser" not in str(exc_info.value) + + with pytest.raises(RuntimeError, match="HarmonyParser"): + _ensure_reasoning_parser_output_capable( + "openai_gptoss", BoundaryOnlyParser, object(), {}, None + ) + + def test_output_capable_parser_accepted(self): + from dynamo.frontend.vllm_processor import ( + _ensure_reasoning_parser_output_capable, + ) + + class WorkingParser: + def __init__(self, tokenizer, *args, **kwargs): + pass + + def extract_reasoning_streaming(self, *args): + return None + + _ensure_reasoning_parser_output_capable( + "fake", WorkingParser, object(), {}, None + ) + + def test_probe_tolerates_other_empty_input_failures(self): + from dynamo.frontend.vllm_processor import ( + _ensure_reasoning_parser_output_capable, + ) + + class PickyParser: + def __init__(self, tokenizer, *args, **kwargs): + pass + + def extract_reasoning_streaming(self, *args): + raise IndexError("empty input") + + _ensure_reasoning_parser_output_capable("fake", PickyParser, object(), {}, None) + + def test_incompatible_signature_rejected(self): + from dynamo.frontend.vllm_processor import ( + _ensure_reasoning_parser_output_capable, + ) + + class WrongSignatureParser: + def __init__(self, tokenizer, *args, **kwargs): + pass + + def extract_reasoning_streaming(self, delta_text): + return None + + with pytest.raises(RuntimeError, match="signature"): + _ensure_reasoning_parser_output_capable( + "fake", WrongSignatureParser, object(), {}, None + ) + + def test_real_gptoss_parser_rejected(self): + pytest.importorskip("vllm.reasoning.gptoss_reasoning_parser") + from vllm.reasoning import ReasoningParserManager + + from dynamo.frontend.vllm_processor import ( + _ensure_reasoning_parser_output_capable, + ) + + class GptOssTokenizer: + vocab = {"<|end|>": 1} + encoded = { + "<|channel|>final": [2], + "<|message|>": [3], + "<|start|>assistant<|channel|>final<|message|>": [4], + } + + def encode(self, text, *args, **kwargs): + return self.encoded[text] + + def get_vocab(self): + return self.vocab + + parser_class = ReasoningParserManager.get_reasoning_parser("openai_gptoss") + with pytest.raises(RuntimeError, match="openai_gptoss"): + _ensure_reasoning_parser_output_capable( + "openai_gptoss", parser_class, GptOssTokenizer(), {}, None + ) + + @pytest.mark.asyncio @pytest.mark.multimodal async def test_build_engine_inputs_preserves_multimodal_uuids( @@ -1586,8 +1686,6 @@ async def test_routed_stream_produces_openai_chunks(self, vllm_processor_module) chunks = await _run_generate(processor, _base_preproc()) - # One annotated envelope per iteration carries both data and the - # llm_metrics annotation; observer strips the annotation before SSE. assert len(chunks) == 1 envelope = chunks[0] @@ -1604,16 +1702,15 @@ async def test_routed_stream_produces_openai_chunks(self, vllm_processor_module) "created": envelope["data"]["created"], "model": MODEL, "object": "chat.completion.chunk", + "llm_metrics": { + "input_tokens": 3, + "output_tokens": 1, + "chunk_tokens": 1, + }, } - assert envelope["event"] == "llm_metrics" - assert len(envelope["comment"]) == 1 - # Zero counts are omitted (text-only request), mirroring the Rust skip-zero behavior. - assert json.loads(envelope["comment"][0]) == { - "input_tokens": 3, - "output_tokens": 1, - "chunk_tokens": 1, - } + assert "event" not in envelope + assert "comment" not in envelope @pytest.mark.asyncio async def test_routed_stream_emits_multimodal_counts(self, vllm_processor_module): @@ -1657,7 +1754,7 @@ async def test_routed_stream_emits_multimodal_counts(self, vllm_processor_module ) ] - metrics = json.loads(chunks[0]["comment"][0]) + metrics = chunks[0]["data"]["llm_metrics"] assert metrics["image_count"] == 2 assert metrics["video_count"] == 1 # audio has zero parts, so the key is omitted from the emitted metrics. diff --git a/components/src/dynamo/frontend/utils.py b/components/src/dynamo/frontend/utils.py index d9126a599bf1..9f667e7092af 100644 --- a/components/src/dynamo/frontend/utils.py +++ b/components/src/dynamo/frontend/utils.py @@ -165,7 +165,7 @@ def extract_mm_urls( has_user_uuid = False for msg in messages: - if not isinstance(msg, dict) or msg.get("role") != "user": + if not isinstance(msg, dict) or msg.get("role") not in ("user", "tool"): continue content = msg.get("content") if not isinstance(content, list): diff --git a/components/src/dynamo/frontend/vllm_processor.py b/components/src/dynamo/frontend/vllm_processor.py index ac17a343fe8f..7dc94fa03e15 100644 --- a/components/src/dynamo/frontend/vllm_processor.py +++ b/components/src/dynamo/frontend/vllm_processor.py @@ -6,6 +6,7 @@ # import asyncio +import inspect import json import logging import os @@ -37,6 +38,7 @@ ) from dynamo.common.multimodal.routing_utils import build_mm_routing_info_from_features from dynamo.common.utils import nvtx_utils as _nvtx +from dynamo.common.utils.input_params import resolve_thinking_token_budget from dynamo.frontend.frontend_args import FrontendConfig from dynamo.llm import ModelCardInstanceId, PythonAsyncEngine, RoutedEngine from dynamo.llm.exceptions import HttpError @@ -238,6 +240,57 @@ class _ReasoningParserMetadata: parser_kwargs: dict[str, Any] | None +def _ensure_reasoning_parser_output_capable( + parser_name: str, + parser_class: type[ReasoningParser], + tokenizer: TokenizerLike, + chat_template_kwargs: dict[str, Any], + model_config: Any, +) -> None: + # vLLM ships boundary-only parsers (e.g. GptOssReasoningParser) that raise + # NotImplementedError from every output-parsing method, while this + # processor calls extract_reasoning_streaming per request. Probe once here + # so the combination is rejected at engine setup instead of failing every + # request with a 500 (issue #14936). The probe constructor mirrors the + # production construction sites (chat_template_kwargs, model_config). + probe = parser_class( + tokenizer, + chat_template_kwargs=chat_template_kwargs, + model_config=model_config, + ) + try: + inspect.signature(probe.extract_reasoning_streaming).bind( + "", "", "", [], [], [] + ) + except TypeError as e: + raise RuntimeError( + f"reasoning_parser {parser_name!r} ({parser_class.__name__}) has an " + f"extract_reasoning_streaming signature this processor cannot call: {e}" + ) from e + try: + probe.extract_reasoning_streaming("", "", "", [], [], []) + except NotImplementedError as e: + msg = ( + f"reasoning_parser {parser_name!r} ({parser_class.__name__}) only " + "provides boundary detection; this processor needs a parser that " + "implements extract_reasoning_streaming (issue #14936)" + ) + if parser_name == "openai_gptoss": + msg += ( + "; gpt-oss output parsing requires HarmonyParser, which this " + "processor does not support yet" + ) + raise RuntimeError(msg) from e + except Exception as e: + # Parsers are not contracted to accept empty input; the probe only + # proves the method is implemented, so log and accept. + logger.debug( + "reasoning_parser %r probe raised %r on empty input; accepting", + parser_name, + e, + ) + + def _build_reasoning_parser_metadata( reasoning_parser_class: type[ReasoningParser] | None, tokenizer: TokenizerLike, @@ -658,7 +711,7 @@ async def _generator_inner( sampling_fields = ( set(getattr(SamplingParams, "__annotations__", ())) & set(type(request_for_sampling).model_fields) - ) - {"max_tokens", "logprobs", "output_kind"} + ) - {"max_tokens", "logprobs", "output_kind", "thinking_token_budget"} for k in sorted(sampling_fields): v = getattr(request_for_sampling, k, None) if v is not None: @@ -667,11 +720,10 @@ async def _generator_inner( # frontend's InputProcessor is built without reasoning_config (it only # tokenizes), so setting sampling_params.thinking_token_budget would # cause process_inputs._validate_params to reject the request. Pluck - # the value out of nvext and pass it directly into dynamo_preproc - # below. - nvext_max_thinking_tokens = (request.get("nvext") or {}).get( - "max_thinking_tokens" - ) + # the value out of the request and pass it directly into dynamo_preproc + # below. Prefer the OpenAI-compatible root-level field, fall back to the + # legacy nvext passthrough. + thinking_token_budget = resolve_thinking_token_budget(request) with _nvtx.annotate("mm_frontend:process_inputs", color="orange"): # render_messages_async returns a raw prompt. Convert it to a typed @@ -718,7 +770,7 @@ async def _generator_inner( "stop_token_ids": sp.stop_token_ids, "min_tokens": sp.min_tokens, "ignore_eos": sp.ignore_eos, - "max_thinking_tokens": nvext_max_thinking_tokens, + "max_thinking_tokens": thinking_token_budget, }, "sampling_options": { "n": sp.n, @@ -1052,8 +1104,12 @@ async def _generate_and_stream( metrics["video_count"] = video_count if audio_count: metrics["audio_count"] = audio_count - envelope["event"] = "llm_metrics" - envelope["comment"] = [json.dumps(metrics)] + # Attach metrics to data when available; otherwise use an annotation. + if data := envelope.get("data"): + data["llm_metrics"] = metrics + else: + envelope["event"] = "llm_metrics" + envelope["comment"] = [json.dumps(metrics)] yield envelope _nvtx.end_range(rng_stream) @@ -1225,6 +1281,13 @@ async def chat_engine_factory( reasoning_parser_class = ReasoningParserManager.get_reasoning_parser( reasoning_parser_name ) + _ensure_reasoning_parser_output_capable( + reasoning_parser_name, + reasoning_parser_class, + tokenizer, + getattr(self.flags, "default_chat_template_kwargs", None) or {}, + model_config, + ) else: reasoning_parser_class = None default_thinking_mode = runtime_default_thinking_mode(mdc.runtime_config()) diff --git a/components/src/dynamo/frontend/vllm_protocol.py b/components/src/dynamo/frontend/vllm_protocol.py new file mode 100644 index 000000000000..fb9ceaee7ab2 --- /dev/null +++ b/components/src/dynamo/frontend/vllm_protocol.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compatibility imports for vLLM's shared OpenAI protocol types.""" + +from importlib.util import find_spec + +try: + _HAS_GENERATE_PROTOCOL = ( + find_spec("vllm.entrypoints.generate.base.protocol") is not None + ) +except ModuleNotFoundError: + _HAS_GENERATE_PROTOCOL = False + +if _HAS_GENERATE_PROTOCOL: + # vLLM 0.29 moved the shared types out of the OpenAI engine package. + from vllm.entrypoints.generate.base.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + FunctionDefinition, + ) +else: + # vLLM 0.27 and 0.28 expose the same types at the legacy path. + from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + FunctionDefinition, + ) + +__all__ = [ + "DeltaFunctionCall", + "DeltaMessage", + "DeltaToolCall", + "FunctionDefinition", +] diff --git a/components/src/dynamo/global_planner/tests/unit/test_scale_request_handler.py b/components/src/dynamo/global_planner/tests/unit/test_scale_request_handler.py index 33a0e33e8b91..097b31d4adb6 100644 --- a/components/src/dynamo/global_planner/tests/unit/test_scale_request_handler.py +++ b/components/src/dynamo/global_planner/tests/unit/test_scale_request_handler.py @@ -418,11 +418,12 @@ def _install_connector(handler, dgd_key, dgd_spec_dict, parent_dgd_name="my-dgd" def _install_real_connector(handler, dgd_key, dgd_spec_dict, parent_dgd_name="my-dgd"): """Attach a real KubernetesConnector backed by a mocked Kubernetes API.""" - connector = KubernetesConnector.__new__(KubernetesConnector) - connector.parent_dgd_name = parent_dgd_name - connector.graph_deployment_name = parent_dgd_name - connector.raise_not_ready = True - connector.kube_api = MagicMock() + with patch("dynamo.planner.connectors.kubernetes.KubernetesAPI"): + connector = KubernetesConnector( + dynamo_namespace="test", + parent_dgd_name=parent_dgd_name, + raise_not_ready=True, + ) connector.kube_api.get_graph_deployment.return_value = dgd_spec_dict connector.kube_api.is_deployment_ready.return_value = True connector.kube_api.wait_for_graph_deployment_ready = AsyncMock() diff --git a/components/src/dynamo/kv_dc_relay/README.md b/components/src/dynamo/kv_dc_relay/README.md index 209d3c813474..b084f84c44b3 100644 --- a/components/src/dynamo/kv_dc_relay/README.md +++ b/components/src/dynamo/kv_dc_relay/README.md @@ -41,3 +41,28 @@ optional TLS/mTLS sidecar across trust boundaries, and do not expose the plainte - [Rust implementation](https://github.com/ai-dynamo/dynamo/blob/main/lib/llm/src/kv_dc_relay/docs/architecture.md) - [gRPC contract](https://github.com/ai-dynamo/dynamo/blob/main/lib/llm/src/kv_dc_relay/docs/grpc-contract.md) - [Protocol helpers and CBI1](https://github.com/ai-dynamo/dynamo/blob/main/lib/llm/src/kv_dc_relay/wan/grpc/protocol/README.md) + +## File mode and state + +Use `--sources-file /etc/relay/sources.json` instead of discovery filters to +select an explicit, live-reloaded set of logical namespaces: + +```json +{"version":1,"sources":[{"namespace":"production-llama"}]} +``` + +Without this option, discovery remains the default. An empty source list watches +nothing. Invalid initial files fail startup; invalid updates retain the last +working source set. Transport and credentials still come from runtime settings. +An optional `connectionRevision` must match `DYN_RELAY_CONNECTION_REVISION`; omit both for +standalone operation without a connection guard. Kubernetes operators supply both. +Both modes feed a common namespace reconciler. File mode polls the file once per +second; discovery mode follows model-registration events and periodically refreshes +the namespace set through DRT. Explicit `--namespaces` remain a fixed selection. +Endpoint discovery within the selected namespaces is shared by both modes. +Source updates preserve unchanged watches and do not restart the process. + +When system HTTP is enabled, `GET /engine/state` reports the mode, listener readiness, +selected namespace count (including namespaces without matching endpoints), and desired/applied revisions with the last reload error. File +revisions are null in discovery mode. The standard runtime health endpoints and +gRPC catalog retain their existing contracts. diff --git a/components/src/dynamo/kv_dc_relay/__main__.py b/components/src/dynamo/kv_dc_relay/__main__.py index 3231c919de8a..9d9bfd29567a 100644 --- a/components/src/dynamo/kv_dc_relay/__main__.py +++ b/components/src/dynamo/kv_dc_relay/__main__.py @@ -79,6 +79,8 @@ async def worker(runtime: DistributedRuntime) -> None: expected_unique_blocks=args.expected_unique_blocks, bind=args.bind, tuning=dict(args.tuning) or None, + sources_file=args.sources_file, + connection_revision=args.connection_revision, ) await relay.start() diagnostics = KvDcRelayDiagnostics(relay) diff --git a/components/src/dynamo/kv_dc_relay/cli.py b/components/src/dynamo/kv_dc_relay/cli.py index 5f12964af028..c084391a47ec 100644 --- a/components/src/dynamo/kv_dc_relay/cli.py +++ b/components/src/dynamo/kv_dc_relay/cli.py @@ -50,6 +50,8 @@ class KvDcRelayCliConfig: expected_unique_blocks: int bind: str | None tuning: tuple[tuple[str, int], ...] + sources_file: str | None = None + connection_revision: str | None = None class RelayShutdownWaiter(Protocol): @@ -130,6 +132,9 @@ def parse_args( environment = os.environ if environment is None else environment parser = argparse.ArgumentParser(description="Dynamo DC-scoped KV Relay") parser.add_argument("--dc-id") + parser.add_argument( + "--sources-file", help="Live source namespace file (selects file mode)" + ) parser.add_argument( "--namespaces", help="Comma-separated Dynamo namespaces containing inference DGDs", @@ -165,6 +170,27 @@ def parse_args( if dc_id != dc_id.strip(): parser.error("DC ID must not contain surrounding whitespace") + sources_file = parsed.sources_file + connection_revision = None + if sources_file is not None: + if not sources_file.strip(): + parser.error("--sources-file must not be empty") + connection_revision = environment.get("DYN_RELAY_CONNECTION_REVISION") + if connection_revision is not None and not connection_revision.strip(): + parser.error("DYN_RELAY_CONNECTION_REVISION must not be empty") + if any( + ( + parsed.namespaces is not None, + parsed.namespace_filter is not None, + parsed.watch_all, + parsed.endpoint_prefixes is not None, + "DYN_RELAY_NAMESPACES" in environment, + "DYN_RELAY_WATCH_ALL" in environment, + "DYN_RELAY_ENDPOINT_PREFIXES" in environment, + ) + ): + parser.error("--sources-file cannot be combined with discovery filters") + configured_cli_scopes = sum( ( parsed.namespaces is not None, @@ -176,7 +202,10 @@ def parse_args( parser.error( "--namespace-filter, --namespaces, and --watch-all are mutually exclusive" ) - if parsed.namespaces is not None: + if sources_file is not None: + namespaces: tuple[str, ...] = () + watch_all = False + elif parsed.namespaces is not None: namespaces = _csv_values(parsed.namespaces, "--namespaces", parser) watch_all = False elif parsed.namespace_filter is not None: @@ -272,6 +301,8 @@ def parse_args( ), bind=bind, tuning=tuple(tuning), + sources_file=sources_file, + connection_revision=connection_revision, ) diff --git a/components/src/dynamo/kv_dc_relay/tests/test_cli.py b/components/src/dynamo/kv_dc_relay/tests/test_cli.py index 286bfeba55cf..7e07567f33bf 100644 --- a/components/src/dynamo/kv_dc_relay/tests/test_cli.py +++ b/components/src/dynamo/kv_dc_relay/tests/test_cli.py @@ -358,3 +358,48 @@ def test_expected_unique_blocks_rejects_invalid_numbers( name = "DYN_RELAY_EXPECTED_UNIQUE_BLOCKS" environment = {name: value} assert_cli_error(["--dc-id", "dc-a", *argv], environment, capsys, name, diagnostic) + + +def test_sources_file_is_explicitly_empty_until_runtime_loads_sources() -> None: + config = parse_args( + ["--dc-id", "dc", "--sources-file", "/etc/relay/sources.json"], + {"DYN_RELAY_CONNECTION_REVISION": "connection"}, + ) + assert config.sources_file == "/etc/relay/sources.json" + assert config.connection_revision == "connection" + assert config.namespaces == () + assert config.watch_all is False + + +@pytest.mark.parametrize( + "extra,environment", + [ + ([], {"DYN_RELAY_CONNECTION_REVISION": ""}), + (["--watch-all"], {"DYN_RELAY_CONNECTION_REVISION": "connection"}), + (["--namespaces", "a"], {"DYN_RELAY_CONNECTION_REVISION": "connection"}), + ( + [], + { + "DYN_RELAY_CONNECTION_REVISION": "connection", + "DYN_RELAY_WATCH_ALL": "true", + }, + ), + ], +) +def test_sources_file_rejects_ambiguous_or_empty_guard_configuration( + extra: list[str], environment: dict[str, str] +) -> None: + with pytest.raises(SystemExit): + parse_args( + ["--dc-id", "dc", "--sources-file", "/etc/relay/sources.json", *extra], + environment, + ) + + +def test_standalone_sources_file_needs_no_connection_guard() -> None: + config = parse_args( + ["--dc-id", "dc", "--sources-file", "/etc/relay/sources.json"], {} + ) + assert config.connection_revision is None + assert config.watch_all is False + assert config.namespaces == () diff --git a/components/src/dynamo/kv_dc_relay/tests/test_native_sources_reload.py b/components/src/dynamo/kv_dc_relay/tests/test_native_sources_reload.py new file mode 100644 index 000000000000..fdd38a0cf3c9 --- /dev/null +++ b/components/src/dynamo/kv_dc_relay/tests/test_native_sources_reload.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native Relay source selection and HTTP state contract, without external services.""" + +import asyncio +import json +from contextlib import asynccontextmanager +from pathlib import Path + +import httpx +import pytest +import pytest_asyncio + +from tests.utils.constants import DynamoPortRange +from tests.utils.port_utils import allocate_port, deallocate_port + +pytestmark = [ + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.integration, + pytest.mark.asyncio, + # Keep the process-global native runtime out of the parent of later forked tests. + pytest.mark.forked, + pytest.mark.timeout(30), +] + + +def write_sources(path: Path, namespaces: list[str], **fields) -> None: + document = { + "version": 1, + "sources": [{"namespace": name} for name in namespaces], + **fields, + } + staged = path.with_suffix(".next") + staged.write_text(json.dumps(document)) + staged.replace(path) + + +@pytest.fixture +def sources_file(tmp_path): + path = tmp_path / "sources.json" + write_sources(path, []) + return path + + +@pytest.fixture +def system_port(request, monkeypatch): + port = allocate_port(DynamoPortRange.ROUTER) if request.param else -1 + monkeypatch.setenv("DYN_SYSTEM_PORT", str(port)) + monkeypatch.delenv("POD_UID", raising=False) + try: + yield port + finally: + if port != -1: + deallocate_port(port) + + +@pytest_asyncio.fixture +async def runtime(system_port): + pytest.importorskip("dynamo._core", reason="requires the native Dynamo extension") + from dynamo.runtime import DistributedRuntime + + runtime = DistributedRuntime( + asyncio.get_running_loop(), "mem", "tcp", event_plane="zmq" + ) + try: + yield runtime + finally: + runtime.shutdown() + + +@pytest.fixture +def relay_factory(runtime): + from dynamo.llm import KvDcRelay + + @asynccontextmanager + async def start(*, bind="127.0.0.1:0", **kwargs): + relay = KvDcRelay( + runtime.endpoint("test.relay.control"), + "test-dc", + bind=bind, + **kwargs, + ) + await relay.start() + try: + yield relay + finally: + await relay.shutdown() + + return start + + +async def wait_for_sources(relay, predicate): + last = None + for _ in range(150): + last = (await relay.health())["sources"] + if predicate(last): + return last + await asyncio.sleep(0.1) + pytest.fail(f"Relay sources did not converge: {last}") + + +@pytest.mark.parametrize("system_port", [True], indirect=True) +@pytest.mark.parametrize( + "mode,bind", + [("discovery", None), ("from-file", "127.0.0.1:0")], + ids=["discovery-wan-disabled", "from-file-wan-enabled"], +) +async def test_http_state( + relay_factory, sources_file, system_port, monkeypatch, mode, bind +): + monkeypatch.setenv("POD_UID", "test-pod") + options = {} + if mode == "from-file": + write_sources(sources_file, [], connectionRevision="revision-a") + options = { + "sources_file": str(sources_file), + "connection_revision": "revision-a", + } + + async with relay_factory(bind=bind, **options) as relay: + async with httpx.AsyncClient(timeout=1) as client: + url = f"http://127.0.0.1:{system_port}/engine/state" + for _ in range(100): + try: + response = await client.get(url) + except httpx.ConnectError: + await asyncio.sleep(0.1) + continue + response.raise_for_status() + state = response.json() + if state["ready"]: + break + await asyncio.sleep(0.1) + else: + pytest.fail("Relay HTTP state did not become ready") + + assert state["mode"] == mode + assert state["podUID"] == "test-pod" + assert state["connectionRevision"] == options.get("connection_revision") + health = await relay.health() + assert health["wan_enabled"] is (bind is not None) + assert state["sources"] == health["sources"] + assert state["sources"]["count"] == 0 + + +@pytest.mark.parametrize("system_port", [False], indirect=True) +async def test_file_sources_reload_on_same_relay(relay_factory, sources_file): + async with relay_factory(sources_file=str(sources_file)) as relay: + state = (await relay.health())["sources"] + assert state["count"] == 0 + assert state["appliedRevision"] is not None + assert state["desiredRevision"] == state["appliedRevision"] + assert state["lastError"] is None + previous = state["appliedRevision"] + for namespaces in (["a", "b"], ["b"], []): + write_sources(sources_file, namespaces) + state = await wait_for_sources( + relay, + lambda state: state["appliedRevision"] != previous + and state["count"] == len(namespaces), + ) + assert state["desiredRevision"] == state["appliedRevision"] + assert state["lastError"] is None + previous = state["appliedRevision"] + + +@pytest.mark.parametrize("system_port", [False], indirect=True) +async def test_invalid_update_retains_sources_and_recovers(relay_factory, sources_file): + write_sources(sources_file, ["a"]) + async with relay_factory(sources_file=str(sources_file)) as relay: + initial = (await relay.health())["sources"] + sources_file.write_text('{"private-invalid-file":') + + rejected = await wait_for_sources(relay, lambda state: state["lastError"]) + assert rejected["appliedRevision"] == initial["appliedRevision"] + assert rejected["desiredRevision"] is None + assert rejected["count"] == 1 + assert "private-invalid-file" not in rejected["lastError"] + + write_sources(sources_file, ["b", "c"]) + recovered = await wait_for_sources( + relay, lambda state: state["count"] == 2 and state["lastError"] is None + ) + assert recovered["appliedRevision"] != initial["appliedRevision"] + assert recovered["desiredRevision"] == recovered["appliedRevision"] + + +@pytest.mark.parametrize("system_port", [False], indirect=True) +@pytest.mark.parametrize( + ("document", "message"), + [ + pytest.param(None, "sources file is unavailable", id="missing"), + pytest.param('{"private-invalid-file":', "invalid sources JSON", id="invalid"), + pytest.param( + '{"version":1,"connectionRevision":"other","sources":[]}', + "sources connection revision does not match", + id="mismatched-connection", + ), + ], +) +async def test_invalid_initial_sources_fail_startup( + relay_factory, sources_file, document, message +): + if document is None: + sources_file.unlink() + else: + sources_file.write_text(document) + + with pytest.raises(Exception, match=message) as error: + async with relay_factory(sources_file=str(sources_file)): + pytest.fail("Relay accepted invalid initial sources") + assert "private-invalid-file" not in str(error.value) diff --git a/components/src/dynamo/mocker/args.py b/components/src/dynamo/mocker/args.py index 6974ccf1b4fa..ac1ca233c5d5 100644 --- a/components/src/dynamo/mocker/args.py +++ b/components/src/dynamo/mocker/args.py @@ -588,12 +588,11 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument( "--stagger-delay", type=float, - default=-1.0, + default=0.0, help=( - "Delay in seconds between launching each worker to avoid overwhelming " - "etcd/NATS/frontend with many workers. Set to 0 to disable staggering. " - "Use -1 for auto mode (0.1s for 32-128 workers, 0.2s for >128 workers, 0 otherwise). " - "Default: -1 (auto)" + "Delay in seconds between launching each worker. " + "Set to 0 to disable staggering (default). " + "Use -1 for auto mode (0.1s for 33-128 workers, 0.2s for >128 workers, 0 otherwise)." ), ) parser.add_argument( diff --git a/components/src/dynamo/planner/connectors/base.py b/components/src/dynamo/planner/connectors/base.py index 522848b0ccfd..479b5b979c7e 100644 --- a/components/src/dynamo/planner/connectors/base.py +++ b/components/src/dynamo/planner/connectors/base.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Optional, Protocol, TypeGuard, runtime_checkable from dynamo.planner.config.defaults import SubComponentType, TargetReplica +from dynamo.planner.core.types import WorkerCounts from dynamo.planner.monitoring.worker_info import WorkerInfo if TYPE_CHECKING: @@ -119,6 +120,33 @@ async def set_component_replicas( pass +@runtime_checkable +class StartupAwareConnector(Protocol): + """Optional inventory that distinguishes pending startup from other scaling.""" + + async def get_worker_inventory( + self, + prefill_component_name: Optional[str] = None, + decode_component_name: Optional[str] = None, + ) -> Optional[WorkerCounts]: + """Return serving counts and verified pending startup counts by role. + + Return None when optional startup reads are forbidden. The caller may + use legacy inventory, or power-aware inventory when power checks apply. + Names select the DGD components; an omitted name excludes that role. + ``WorkerCounts.pending_num_*`` may be positive only after ruling out + drain, rollout, and unobserved spec changes across the deployment. + A connector that also implements ``PowerAwareConnector`` must provide + the same rollout and terminating-Pod guarantees as its power-aware + counts, using the same snapshot for serving and pending inventory. + """ + ... + + +def is_startup_aware_connector(obj: object) -> TypeGuard[StartupAwareConnector]: + return callable(inspect.getattr_static(obj, "get_worker_inventory", None)) + + @runtime_checkable class PowerAwareConnector(Protocol): """Narrow read-only power capability โ€” Kubernetes-specific, not part of PlannerConnector. @@ -187,6 +215,8 @@ def is_power_aware_connector(obj: object) -> TypeGuard[PowerAwareConnector]: __all__ = [ "PlannerConnector", "PowerAwareConnector", + "StartupAwareConnector", "WorkerInfoProvider", "is_power_aware_connector", + "is_startup_aware_connector", ] diff --git a/components/src/dynamo/planner/connectors/clients/kubernetes_api.py b/components/src/dynamo/planner/connectors/clients/kubernetes_api.py index e94c51342399..f824431454e5 100644 --- a/components/src/dynamo/planner/connectors/clients/kubernetes_api.py +++ b/components/src/dynamo/planner/connectors/clients/kubernetes_api.py @@ -52,6 +52,7 @@ # Stable labels the operator stamps on every worker Pod. DYNAMO_DGD_NAME_LABEL = "nvidia.com/dynamo-graph-deployment-name" DYNAMO_COMPONENT_LABEL = "nvidia.com/dynamo-component" +GROVE_PCSG_REPLICA_INDEX_LABEL = "grove.io/podcliquescalinggroup-replica-index" def get_current_k8s_namespace() -> str: @@ -153,6 +154,26 @@ def update_service_replicas( else: raise + def get_service_replica_target( + self, graph_deployment_name: str, service_name: str + ) -> int: + """Read the authoritative scale target, including an unapplied DGDSA write.""" + try: + scale = self.custom_api.get_namespaced_custom_object_scale( + group=NVIDIA_API_GROUP, + version=DYNAMO_API_VERSION, + namespace=self.current_namespace, + plural=DGDSA_PLURAL, + name=f"{graph_deployment_name}-{service_name.lower()}", + ) + return int(scale["spec"]["replicas"]) + except client.ApiException as e: + if e.status != 404: + raise + deployment = self.get_graph_deployment(graph_deployment_name) + component = get_components_by_name(deployment)[service_name] + return Service(name=service_name, service=component).number_replicas() + def _update_dgd_replicas( self, graph_deployment_name: str, service_name: str, replicas: int ) -> None: @@ -325,6 +346,99 @@ def get_service_replica_status( return traffic_serving_replicas, is_stable + def pending_startup_replicas(self, deployment: dict, pods: list) -> dict[str, int]: + """Identify startup-only scaling; an empty result never authorizes a write. + + A Ready deficit alone is ambiguous: it also occurs during drain, rollout, + and stale status. Require observed spec, current worker revisions, and + no terminating/failed Pods before exposing pending startup capacity. + """ + if not self.is_spec_generation_observed(deployment): + return {} + phase = (deployment.get("status", {}).get("rollingUpdate") or {}).get("phase") + if phase not in (None, "", "Completed"): + return {} + pods = self.exclude_checkpoint_capture_pods(pods) + if not pods or any( + pod.metadata.deletion_timestamp is not None + or pod.status is None + or pod.status.phase not in ("Pending", "Running") + for pod in pods + ): + return {} + if not self.pcsg_pods_within_desired_replicas(deployment, pods): + return {} + pending: dict[str, int] = {} + statuses = deployment.get("status", {}).get("components", {}) + for name, spec in get_components_by_name(deployment).items(): + if get_component_type(spec) == "planner": + continue + status = statuses.get(name, {}) + desired = Service(name=name, service=spec).number_replicas() + ready, stable = self.get_service_replica_status(deployment, name) + replicas = status.get("replicas") + updated = status.get("updatedReplicas") + if ( + replicas is None + or updated is None + or not (0 <= ready <= replicas <= desired) + ): + return {} + # Grove PCSG counts a replica as updated only once it is available; + # PodClique/Deployment updated counts include unready new Pods. + is_pcsg = status.get("componentKind") == "PodCliqueScalingGroup" + if updated != (ready if is_pcsg else replicas): + return {} + if not stable: + if get_component_type(spec) not in ("prefill", "decode", "worker"): + return {} + if desired <= ready: + return {} + pending[name] = desired - ready + return pending + + @staticmethod + def exclude_checkpoint_capture_pods(pods: list) -> list: + """Capture Jobs inherit worker labels but do not serve inference traffic.""" + return [ + pod + for pod in pods + if not ( + (pod.metadata.labels or {}).get("nvidia.com/snapshot-job-uid") + and any( + owner.controller + and owner.api_version == "batch/v1" + and owner.kind == "Job" + and owner.name + == (pod.metadata.labels or {}).get("nvidia.com/snapshot-job") + for owner in (pod.metadata.owner_references or []) + ) + ) + ] + + def pcsg_pods_within_desired_replicas(self, deployment: dict, pods: list) -> bool: + """Reject excess PCSG groups hidden by its spec-derived replica count.""" + components = get_components_by_name(deployment) + pods_by_component = self.partition_pods_by_component(pods) + for name, status in deployment.get("status", {}).get("components", {}).items(): + if status.get("componentKind") != "PodCliqueScalingGroup": + continue + desired = Service( + name=name, service=components.get(name, {}) + ).number_replicas() + for pod in pods_by_component.get(name, []): + replica_index = (pod.metadata.labels or {}).get( + GROVE_PCSG_REPLICA_INDEX_LABEL + ) + if replica_index is None: + return False + try: + if not 0 <= int(replica_index) < desired: + return False + except (TypeError, ValueError): + return False + return True + def non_planner_components_stable(self, deployment: dict) -> tuple[bool, list[str]]: """Return ``(all_stable, unstable_names)`` for non-planner components.""" components = get_components_by_name(deployment) diff --git a/components/src/dynamo/planner/connectors/kubernetes.py b/components/src/dynamo/planner/connectors/kubernetes.py index 6560c8e60d6b..90ab1d59329c 100644 --- a/components/src/dynamo/planner/connectors/kubernetes.py +++ b/components/src/dynamo/planner/connectors/kubernetes.py @@ -17,8 +17,11 @@ import json import logging import os +from threading import Lock from typing import Optional +from kubernetes.client import ApiException + from dynamo.planner.config.defaults import SubComponentType, TargetReplica from dynamo.planner.connectors.base import PlannerConnector from dynamo.planner.connectors.clients.kubernetes_api import ( @@ -32,6 +35,7 @@ select_entry, worker_info_from_mdc, ) +from dynamo.planner.core.types import WorkerCounts from dynamo.planner.errors import ( DeploymentModelNameMismatchError, DeploymentValidationError, @@ -97,6 +101,12 @@ def __init__( # For backwards compatibility self.graph_deployment_name = self.parent_dgd_name self.raise_not_ready = raise_not_ready + # DGDSA application and DGD status are asynchronous. Hold subsequent + # writes until a startup reversal has actually finished, not just until + # the pre-write Ready condition is observed again. + self._startup_scale_down_lock = Lock() + self._startup_scale_down_targets: dict[str, int] = {} + self._startup_read_warnings: set[str] = set() async def async_init(self): """No-op asynchronous lifecycle hook.""" @@ -781,7 +791,140 @@ def _dgd_model_name() -> Optional[str]: ) return info - # todo -> how are we handling 3 active 2 more new workers pending? + def _list_startup_pods(self) -> Optional[list]: + """Return lifecycle Pods, or None when legacy RBAC forbids listing them.""" + try: + pods = self.kube_api.list_pods_for_graph(self.graph_deployment_name) + except ApiException as exc: + if exc.status != 403: + raise + with self._startup_scale_down_lock: + warn = "pods" not in self._startup_read_warnings + self._startup_read_warnings.add("pods") + if warn: + logger.warning( + "Pod list forbidden for %s; startup cancellation is disabled " + "until pods/list is granted", + self.graph_deployment_name, + ) + return None + with self._startup_scale_down_lock: + self._startup_read_warnings.discard("pods") + return KubernetesAPI.exclude_checkpoint_capture_pods(pods) + + def _startup_scale_down_in_progress(self, deployment: dict, pods: list) -> bool: + # Snapshot under the lock, then do Kubernetes I/O without holding it. + # Writers replace the dictionary so identity detects even a concurrent + # request for the same target, not only a different replica count. + with self._startup_scale_down_lock: + previous_targets = self._startup_scale_down_targets + targets = previous_targets.copy() + if not targets: + return False + if not self.kube_api.is_spec_generation_observed(deployment): + return True + if self.kube_api.has_terminating_pods(pods): + return True + if not self.kube_api.pcsg_pods_within_desired_replicas(deployment, pods): + return True + components = get_components_by_name(deployment) + startup = self.kube_api.pending_startup_replicas(deployment, pods) + remaining: dict[str, int] = {} + for name, target in targets.items(): + desired = Service( + name=name, service=components.get(name, {}) + ).number_replicas() + _, stable = self.kube_api.get_service_replica_status(deployment, name) + if desired != target: + try: + authoritative_target = self.kube_api.get_service_replica_target( + self.graph_deployment_name, name + ) + except ApiException as exc: + if exc.status != 403: + raise + # GET may be revoked after admission. Retain this request + # until its DGD target arrives; do not guess that an + # unobserved write was superseded or drop the latch. + with self._startup_scale_down_lock: + warn = "scale" not in self._startup_read_warnings + self._startup_read_warnings.add("scale") + if warn: + logger.warning( + "Scale get forbidden for %s; holding startup scale-down " + "until its DGD target is observed", + self.graph_deployment_name, + ) + remaining[name] = target + continue + with self._startup_scale_down_lock: + self._startup_read_warnings.discard("scale") + if authoritative_target == target or desired != authoritative_target: + remaining[name] = target + # A superseding DGDSA target has reached the observed DGD spec. + # Retire our old request; ordinary inventory guards its state. + elif not stable and name not in startup: + remaining[name] = target + # Once excess/terminating replicas are gone, a survivor becoming + # unready is startup capacity again, not an unfinished drain. + with self._startup_scale_down_lock: + if self._startup_scale_down_targets is previous_targets: + self._startup_scale_down_targets = remaining + return bool(self._startup_scale_down_targets) + + async def get_worker_inventory( + self, + prefill_component_name: Optional[str] = None, + decode_component_name: Optional[str] = None, + ) -> Optional[WorkerCounts]: + """Read selected roles' serving/pending inventory off the event loop. + + DGD status and Pods supply startup, drain, and rollout checks, including + the power-aware readiness guarantees. An omitted component name reports + zero counts for that role; pending counts never include draining workers. + Return None on a Pod-list 403 so callers can retain legacy behavior. + """ + return await asyncio.to_thread( + self._get_worker_inventory_sync, + prefill_component_name, + decode_component_name, + ) + + def _get_worker_inventory_sync( + self, + prefill_component_name: Optional[str], + decode_component_name: Optional[str], + ) -> Optional[WorkerCounts]: + deployment = self.kube_api.get_graph_deployment(self.graph_deployment_name) + pods = self._list_startup_pods() + if pods is None: + return None + p, d, stable = self._worker_counts_from_snapshot( + deployment, + prefill_component_name=prefill_component_name, + decode_component_name=decode_component_name, + pods_by_component=self.kube_api.partition_pods_by_component(pods), + power_aware=True, + ) + stable = ( + stable + and self.kube_api.is_spec_generation_observed(deployment) + and self.kube_api.pcsg_pods_within_desired_replicas(deployment, pods) + ) + pending = self.kube_api.pending_startup_replicas(deployment, pods) + if self._startup_scale_down_in_progress(deployment, pods): + stable, pending = False, {} + return WorkerCounts( + ready_num_prefill=p, + ready_num_decode=d, + expected_num_prefill=p if stable else None, + expected_num_decode=d if stable else None, + prefill_scaling_in_progress=not stable, + decode_scaling_in_progress=not stable, + pending_num_prefill=pending.get(prefill_component_name or "", 0), + pending_num_decode=pending.get(decode_component_name or "", 0), + ) + async def get_actual_worker_counts( self, prefill_component_name: Optional[str] = None, @@ -916,7 +1059,106 @@ async def set_component_replicas( deployment = self.kube_api.get_graph_deployment(self.graph_deployment_name) - if not self.kube_api.is_deployment_ready(deployment): + ready = self.kube_api.is_deployment_ready(deployment) + startup_reduction = False + reducing = any( + target.desired_replicas + < get_component_from_type_or_name( + deployment, + target.sub_component_type, + component_name=target.component_name, + ).number_replicas() + for target in target_replicas + ) + with self._startup_scale_down_lock: + scale_down_pending = bool(self._startup_scale_down_targets) + if not ready or scale_down_pending or reducing: + pods = self._list_startup_pods() + if pods is None: + # Keep ordinary Ready-deployment scaling working with old RBAC. + # Never release an accepted startup reversal without observing + # its drain, or use the startup exception to scale an unready DGD. + ready = ( + ready + and not scale_down_pending + and self.kube_api.non_planner_components_stable(deployment)[0] + ) + else: + if self._startup_scale_down_in_progress(deployment, pods): + logger.info("Startup scale-down still converging, ignoring scaling") + return + pending = self.kube_api.pending_startup_replicas(deployment, pods) + # Ready can still describe the state before a Pod deletion or the + # latest spec change. Recheck lifecycle before issuing a reduction. + phase = (deployment.get("status", {}).get("rollingUpdate") or {}).get( + "phase" + ) + ready = ( + ready + and self.kube_api.is_spec_generation_observed(deployment) + and not self.kube_api.has_terminating_pods(pods) + and self.kube_api.pcsg_pods_within_desired_replicas( + deployment, pods + ) + and self.kube_api.non_planner_components_stable(deployment)[0] + and phase in (None, "", "Completed") + and not pending + ) + startup_reduction = bool(pending) + any_reduction = False + for target in target_replicas: + if not startup_reduction: + break + service = get_component_from_type_or_name( + deployment, + target.sub_component_type, + component_name=target.component_name, + ) + serving, _ = self.kube_api.get_service_replica_status( + deployment, service.name + ) + desired = service.number_replicas() + if target.desired_replicas != desired: + startup_reduction &= ( + 0 <= target.desired_replicas <= serving + and target.desired_replicas < desired + ) + any_reduction = True + startup_reduction &= any_reduction + + if startup_reduction: + # Reconciliation needs Scale GET to recognize superseding writers. + # Check every changed component before the first PATCH, so partial + # RBAC cannot admit a cancellation we cannot subsequently track. + try: + for target in target_replicas: + service = get_component_from_type_or_name( + deployment, + target.sub_component_type, + component_name=target.component_name, + ) + if service.number_replicas() != target.desired_replicas: + self.kube_api.get_service_replica_target( + self.graph_deployment_name, service.name + ) + except ApiException as exc: + if exc.status != 403: + raise + startup_reduction = False + with self._startup_scale_down_lock: + warn = "scale" not in self._startup_read_warnings + self._startup_read_warnings.add("scale") + if warn: + logger.warning( + "Scale get forbidden for %s; startup cancellation is " + "disabled until Scale get is granted", + self.graph_deployment_name, + ) + else: + with self._startup_scale_down_lock: + self._startup_read_warnings.discard("scale") + + if not ready and not startup_reduction: if self.raise_not_ready: logger.warning( "Deployment %s is not ready, rejecting this scaling", @@ -948,6 +1190,12 @@ async def set_component_replicas( service.name, target_replica.desired_replicas, ) + if startup_reduction: + with self._startup_scale_down_lock: + self._startup_scale_down_targets = { + **self._startup_scale_down_targets, + service.name: target_replica.desired_replicas, + } else: logger.info( f"{target_replica.sub_component_type.value} component {service.name} already at desired replica count {target_replica.desired_replicas}, skipping" diff --git a/components/src/dynamo/planner/core/base.py b/components/src/dynamo/planner/core/base.py index f68d35ad664f..b5c5188ad4d6 100644 --- a/components/src/dynamo/planner/core/base.py +++ b/components/src/dynamo/planner/core/base.py @@ -631,6 +631,8 @@ async def _collect_worker_counts(self) -> WorkerCounts: decode_scaling_in_progress=( self.require_decode and state.decode.replicas.scaling ), + pending_num_prefill=state.prefill.replicas.pending_startup, + pending_num_decode=state.decode.replicas.pending_startup, ) async def _gather_tick_input(self, tick: ScheduledTick) -> TickInput: @@ -862,6 +864,9 @@ def _log_decision_summary(self, effects: PlannerEffects) -> None: diag = effects.diagnostics current_p, current_d = self._current_worker_counts() + if self._last_worker_counts is not None: + current_p += self._last_worker_counts.pending_num_prefill + current_d += self._last_worker_counts.pending_num_decode rec_p = decision.num_prefill if decision else None rec_d = decision.num_decode if decision else None diff --git a/components/src/dynamo/planner/core/load_scaling.py b/components/src/dynamo/planner/core/load_scaling.py index 5edc00d2ba36..7a7e15eba86d 100644 --- a/components/src/dynamo/planner/core/load_scaling.py +++ b/components/src/dynamo/planner/core/load_scaling.py @@ -59,7 +59,7 @@ def _advance_load(self, obs: FpmObservations) -> Optional[ScalingDecision]: def _advance_load_single( self, obs: FpmObservations, component: str ) -> Optional[ScalingDecision]: - if self._scaling_in_progress(component): + if self._load_scaling_blocked(component): logger.info(f"Scaling in progress for {component}, observing only") self._diag_load_reason = "scaling_in_progress" return None @@ -92,6 +92,8 @@ def _advance_load_single( if desired is None: return None + if (self._pending_num_p or self._pending_num_d) and desired > num_workers: + return None original_desired = desired if self._config.enable_throughput_scaling: bound = ( @@ -106,7 +108,7 @@ def _advance_load_single( component, ) - if desired < num_workers: + if desired < num_workers + self._pending_startup(component): if desired > original_desired: self._diag_load_reason = "scale_down_capped_by_throughput" else: @@ -133,7 +135,9 @@ def _advance_load_disagg(self, obs: FpmObservations) -> Optional[ScalingDecision self._diag_load_reason_prefill = "no_fpm_data" self._diag_load_reason_decode = "no_fpm_data" return None - if self._scaling_in_progress("prefill") or self._scaling_in_progress("decode"): + if self._load_scaling_blocked("prefill") or self._load_scaling_blocked( + "decode" + ): logger.info("Scaling in progress for disagg deployment, observing only") self._diag_load_reason = "scaling_in_progress" self._diag_load_reason_prefill = "scaling_in_progress" @@ -186,6 +190,14 @@ def _advance_load_disagg(self, obs: FpmObservations) -> Optional[ScalingDecision ) d_reason = self._diag_load_reason + if self._pending_num_p or self._pending_num_d: + # Preserve raw intent before budget clamps can turn an up request + # into a ready-equal value that looks like startup cancellation. + if p_desired is not None and p_desired > self._num_p_workers: + p_desired = None + if d_desired is not None and d_desired > self._num_d_workers: + d_desired = None + final_p = p_desired if p_desired is not None else self._num_p_workers final_d = d_desired if d_desired is not None else self._num_d_workers @@ -206,6 +218,12 @@ def _advance_load_disagg(self, obs: FpmObservations) -> Optional[ScalingDecision final_p = max(final_p, resolve_min_endpoint(self._config, "prefill")) final_d = max(final_d, resolve_min_endpoint(self._config, "decode")) + if self._pending_num_p or self._pending_num_d: + return self._startup_disagg_decision( + final_p if p_desired is not None else None, + final_d if d_desired is not None else None, + source="load", + ) throughput_lifted_proposal = self._config.enable_throughput_scaling and ( post_floor_p > original_p or post_floor_d > original_d ) @@ -230,10 +248,10 @@ def _reason(final: int, original: int, post_floor: int, current: int) -> str: return "scale_down_capped_by_throughput" if floor_capped else "no_change" self._diag_load_reason_prefill = _reason( - final_p, original_p, post_floor_p, self._num_p_workers + final_p, original_p, post_floor_p, self._num_p_workers + self._pending_num_p ) self._diag_load_reason_decode = _reason( - final_d, original_d, post_floor_d, self._num_d_workers + final_d, original_d, post_floor_d, self._num_d_workers + self._pending_num_d ) # Aggregate reason: prioritise "most interesting" across components. @@ -282,7 +300,7 @@ def _advance_load_agg(self, obs: FpmObservations) -> Optional[ScalingDecision]: return None num_workers = self._num_d_workers - if self._scaling_in_progress("decode"): + if self._load_scaling_blocked("decode"): logger.info( f"Scaling in progress ({num_workers} -> {self._expected_num_d}), observing only" ) @@ -300,6 +318,8 @@ def _advance_load_agg(self, obs: FpmObservations) -> Optional[ScalingDecision]: if desired is None: return None + if (self._pending_num_p or self._pending_num_d) and desired > num_workers: + return None original_desired = desired desired = max(desired, resolve_min_endpoint(self._config, "decode")) if self._config.enable_throughput_scaling: @@ -309,7 +329,7 @@ def _advance_load_agg(self, obs: FpmObservations) -> Optional[ScalingDecision]: "decode", ) - if desired < num_workers: + if desired < num_workers + self._pending_num_d: if desired > original_desired: self._diag_load_reason = "scale_down_capped_by_throughput" else: @@ -373,6 +393,8 @@ def _advance_load_agg(self, obs: FpmObservations) -> Optional[ScalingDecision]: # returning early. desired = num_workers + if (self._pending_num_p or self._pending_num_d) and desired > num_workers: + return None original_desired = desired desired = max(desired, resolve_min_endpoint(self._config, "decode")) if self._config.enable_throughput_scaling: @@ -386,7 +408,15 @@ def _advance_load_agg(self, obs: FpmObservations) -> Optional[ScalingDecision]: # distinct diagnostic reason even when the net result is no change. floor_capped = desired > original_desired and original_desired < num_workers - if desired == num_workers: + cancel_startup = ( + not any_refused + and self._pending_num_d > 0 + and p_desired is not None + and d_desired is not None + and p_desired <= num_workers + and d_desired <= num_workers + ) + if desired == num_workers and not cancel_startup: logger.info("Agg scaling: no scaling needed") if budget_reason is not None: self._diag_load_reason = budget_reason @@ -400,7 +430,7 @@ def _advance_load_agg(self, obs: FpmObservations) -> Optional[ScalingDecision]: ) return None - if desired < num_workers: + if desired < num_workers + self._pending_num_d: self._diag_load_reason = ( "scale_down_capped_by_throughput" if floor_capped else "scale_down" ) @@ -449,7 +479,7 @@ def _prefill_load_decision( # ``(SLA - T_own) * sensitivity``. The new request's own forward-pass # time (``T_own``) does not shrink with more workers, so it is # excluded from the safety-margin budget. - can_scale_down = num_workers > 1 + can_scale_down = num_workers > 1 or self._pending_startup("prefill") > 0 consolidation_refused = False for label, group in self._prefill_regression.query_groups(fpm_stats): queued = max(fpm.queued_requests.sum_prefill_tokens for fpm in group) @@ -542,7 +572,7 @@ def _decode_load_decision( # 2. SLA check: predicted ITL at the survivor's post-consolidation KV # must stay within ``SLA * sensitivity``. Decouples from cache # size -- engines often saturate latency well before cache. - can_scale_down = num_workers > 1 + can_scale_down = num_workers > 1 or self._pending_startup("decode") > 0 consolidation_refused = False for label, group in self._decode_regression.query_groups(fpm_stats): sched_kv = max(fpm.scheduled_requests.sum_decode_kv_tokens for fpm in group) @@ -612,7 +642,7 @@ def _agg_prefill_scaling( # applied only to the queue-induced portion of TTFT -- ``T_own`` (a # zero-queue prefill at the post-consolidation decode_kv) is treated # as the fixed cost the new request must pay regardless of N. - can_scale_down = num_workers > 1 + can_scale_down = num_workers > 1 or self._pending_startup("decode") > 0 consolidation_refused = False for _label, group in self._agg_regression.query_groups(fpm_stats): # Pre-consolidation prediction uses scheduled-only decode load. @@ -695,7 +725,7 @@ def _agg_decode_scaling( # can't model block eviction past the cache). # 2. SLA check via ``estimate_scheduled_decode_itl`` at the # post-consolidation combined kv. - can_scale_down = num_workers > 1 + can_scale_down = num_workers > 1 or self._pending_startup("decode") > 0 consolidation_refused = False for _label, group in self._agg_regression.query_groups(fpm_stats): sched_kv = max(fpm.scheduled_requests.sum_decode_kv_tokens for fpm in group) @@ -757,6 +787,17 @@ def _agg_decode_scaling( # Easy-mode decision methods (optimization_target != "sla") # ------------------------------------------------------------------ + def _consolidation_safe_easy_target( + self, desired: int, ready: int, values: list[float], up_threshold: float + ) -> int: + # Do not drain a serving worker if redistributing its load would + # immediately trigger scale-up. Keep this guard after startup settles + # too; otherwise cancelling pending capacity only delays oscillation. + if 0 < desired < ready: + if max(values) * ready / desired >= up_threshold: + return ready + return desired + def _prefill_easy_decision( self, fpm_stats: dict[tuple[str, int], ForwardPassMetrics], num_workers: int ) -> Optional[int]: @@ -828,7 +869,7 @@ def _prefill_easy_decision( return num_workers + 1 # Scale down if ALL engines below threshold - if num_workers > 1: + if num_workers > 1 or self._pending_startup("prefill"): if is_load: if all(v <= down_thresh for v in values): desired = max( @@ -839,7 +880,9 @@ def _prefill_easy_decision( f"Load prefill: all engines at or below scale-down " f"threshold ({down_thresh}), -> {desired}" ) - return desired + return self._consolidation_safe_easy_target( + desired, num_workers, values, up_thresh + ) elif is_latency: # For latency mode, scale down when ALL queues are empty if all(v <= down_thresh for v in values): @@ -850,7 +893,9 @@ def _prefill_easy_decision( logger.info( f"Easy prefill: all engines at zero queue, -> {desired}" ) - return desired + return self._consolidation_safe_easy_target( + desired, num_workers, values, up_thresh + ) else: if all(v < down_thresh for v in values): desired = max( @@ -861,7 +906,9 @@ def _prefill_easy_decision( f"Easy prefill: all engines below scale-down threshold " f"({down_thresh}), -> {desired}" ) - return desired + return self._consolidation_safe_easy_target( + desired, num_workers, values, up_thresh + ) self._diag_load_reason = "no_change" return None @@ -931,13 +978,15 @@ def _decode_easy_decision( if is_load else all(u < down_thresh for u in utils) ) - if num_workers > 1 and scale_down: + if (num_workers > 1 or self._pending_startup("decode")) and scale_down: desired = max(num_workers - 1, resolve_min_endpoint(self._config, "decode")) logger.info( f"Easy decode: all engines below scale-down threshold " f"({down_thresh}), -> {desired}" ) - return desired + return self._consolidation_safe_easy_target( + desired, num_workers, utils, up_thresh + ) self._diag_load_reason = "no_change" return None @@ -1009,13 +1058,15 @@ def _agg_easy_decision( if is_load else all(u < down_thresh for u in utils) ) - if num_workers > 1 and scale_down: + if (num_workers > 1 or self._pending_startup("decode")) and scale_down: desired = max(num_workers - 1, resolve_min_endpoint(self._config, "decode")) logger.info( f"Easy agg: all engines below scale-down threshold " f"({down_thresh}), -> {desired}" ) - return desired + return self._consolidation_safe_easy_target( + desired, num_workers, utils, up_thresh + ) self._diag_load_reason = "no_change" return None @@ -1066,7 +1117,7 @@ def _scale_decision( ) return num_workers + 1 - if num_workers > 1 and can_scale_down: + if can_scale_down: desired = max(num_workers - 1, min_endpoint) logger.info( f"Load-based {label}: post-consolidation prediction within " diff --git a/components/src/dynamo/planner/core/state_machine.py b/components/src/dynamo/planner/core/state_machine.py index 7df5b14510c6..d9a874f8c69a 100644 --- a/components/src/dynamo/planner/core/state_machine.py +++ b/components/src/dynamo/planner/core/state_machine.py @@ -101,6 +101,8 @@ def __init__( self._expected_num_d: Optional[int] = None self._prefill_scaling_in_progress: bool = False self._decode_scaling_in_progress: bool = False + self._pending_num_p = 0 + self._pending_num_d = 0 self._throughput_lower_bound_p: int = 1 self._throughput_lower_bound_d: int = 1 @@ -224,7 +226,92 @@ def advance_load( kv_hit_rate=predicted_kv_hit_rate, accept_length=predicted_accept_length, ) - return self._advance_load(obs) + decision = self._advance_load(obs) + if decision is not None and (self._pending_num_p or self._pending_num_d): + decision.num_prefill = self._startup_reduction( + decision.num_prefill, self._num_p_workers, self._pending_num_p + ) + decision.num_decode = self._startup_reduction( + decision.num_decode, self._num_d_workers, self._pending_num_d + ) + if decision.num_prefill is None and decision.num_decode is None: + return None + return decision + + @staticmethod + def _startup_reduction( + target: Optional[int], ready: int, pending: int + ) -> Optional[int]: + # Never turn a scale-up recommendation into a cancellation. Ready-equal + # targets are meaningful only when there are pending replicas to cancel. + if target is None or target > ready or target >= ready + pending: + return None + return target + + def _startup_disagg_decision( + self, + num_p: Optional[int], + num_d: Optional[int], + *, + source: Literal["load", "throughput"], + ) -> Optional[ScalingDecision]: + """Budget startup reductions against the unchanged peer's desired count. + + Serving counts remain the input to load/consolidation prediction. Only + allocation accounting includes pending replicas, and an omitted or up + recommendation preserves that role's entire allocation. + """ + targets = ( + self._startup_reduction(num_p, self._num_p_workers, self._pending_num_p), + self._startup_reduction(num_d, self._num_d_workers, self._pending_num_d), + ) + if targets == (None, None): + return None + p_gpu, d_gpu = self._resolve_disagg_gpu_costs() + if p_gpu is not None and d_gpu is not None: + # Adjustable roles may retain their serving replicas to respect the + # floor; fixed peers retain both serving and pending replicas. + current_p = self._num_p_workers + ( + self._pending_num_p if targets[0] is None else 0 + ) + current_d = self._num_d_workers + ( + self._pending_num_d if targets[1] is None else 0 + ) + fitted_p, fitted_d = fit_directional_budget_pair( + current_p, + current_d, + targets[0] if targets[0] is not None else current_p, + targets[1] if targets[1] is not None else current_d, + p_gpu, + d_gpu, + self._config.min_gpu_budget, + # Reductions cannot need a larger ceiling. The final startup + # projection checks the hard ceiling and power budget against + # the actual applied pair; using a floor-only fit here keeps + # its strict floor (no integer-step tolerance) intact. + -1, + self._min_endpoint_for("prefill"), + self._min_endpoint_for("decode"), + ) + targets = ( + self._startup_reduction( + fitted_p, self._num_p_workers, self._pending_num_p + ) + if targets[0] is not None + else None, + self._startup_reduction( + fitted_d, self._num_d_workers, self._pending_num_d + ) + if targets[1] is not None + else None, + ) + if targets == (None, None): + return None + if source == "load": + self._diag_load_reason = "scale_down" + else: + self._diag_throughput_reason = "scale" + return ScalingDecision(num_prefill=targets[0], num_decode=targets[1]) def advance_throughput_from_prediction( self, @@ -333,6 +420,18 @@ def _update_inventory(self, counts: WorkerCounts) -> None: self._expected_num_d = counts.expected_num_decode self._prefill_scaling_in_progress = counts.prefill_scaling_in_progress self._decode_scaling_in_progress = counts.decode_scaling_in_progress + self._pending_num_p = counts.pending_num_prefill + self._pending_num_d = counts.pending_num_decode + + def _pending_startup(self, component: str) -> int: + return self._pending_num_p if component == "prefill" else self._pending_num_d + + def _load_scaling_blocked(self, component: str) -> bool: + # Startup inventory is only supplied after deployment-wide drain and + # rollout checks. Decisions are still bounded to reductions below. + return self._scaling_in_progress(component) and not ( + self._pending_num_p or self._pending_num_d + ) def _scaling_in_progress(self, component: str) -> bool: if component == "prefill": diff --git a/components/src/dynamo/planner/core/throughput_scaling.py b/components/src/dynamo/planner/core/throughput_scaling.py index aca588f6f724..f0e14b6140f8 100644 --- a/components/src/dynamo/planner/core/throughput_scaling.py +++ b/components/src/dynamo/planner/core/throughput_scaling.py @@ -75,6 +75,7 @@ def _throughput_single( # Endpoint recovery is a hard invariant, not an ordinary throughput # movement, so it may exceed the per-observation delta cap. desired = max(desired, resolve_min_endpoint(self._config, component)) + requested = desired desired, _ceiling_reason = self._fit_single_throughput_ceiling( desired, component, @@ -99,11 +100,17 @@ def _throughput_single( ) return None + if (self._pending_num_p or self._pending_num_d) and ( + model_not_ready or requested > current + ): + if model_not_ready: + self._diag_throughput_reason = "model_not_ready" + return None desired, _budget_reason = self._apply_single_scaling_budget( desired, component, ) - if desired == current: + if desired == current and not self._pending_startup(component): self._diag_throughput_reason = ( "model_not_ready" if model_not_ready @@ -151,6 +158,15 @@ def _throughput_disagg( num_p = max(num_p, resolve_min_endpoint(self._config, "prefill")) num_d = max(num_d, resolve_min_endpoint(self._config, "decode")) bounded_p, bounded_d = num_p, num_d + if ( + self._pending_num_p or self._pending_num_d + ) and not self._config.enable_load_scaling: + if model_not_ready: + # A current-count fallback repairs floors after settlement; it + # is not demand evidence for cancelling pending capacity. + self._diag_throughput_reason = "model_not_ready" + return None + return self._startup_disagg_decision(num_p, num_d, source="throughput") num_p, num_d = self._fit_disagg_throughput_ceiling(num_p, num_d) budget_held = (num_p, num_d) == (self._num_p_workers, self._num_d_workers) and ( bounded_p, @@ -262,6 +278,7 @@ def _throughput_agg( desired, self._num_d_workers, "aggregated" ) desired = max(desired, resolve_min_endpoint(self._config, "decode")) + requested = desired desired, _ceiling_reason = self._fit_single_throughput_ceiling( desired, "decode", @@ -279,11 +296,15 @@ def _throughput_agg( ) return None + if self._pending_num_d and (model_not_ready or requested > self._num_d_workers): + if model_not_ready: + self._diag_throughput_reason = "model_not_ready" + return None desired, _budget_reason = self._apply_single_scaling_budget( desired, "decode", ) - if desired == self._num_d_workers: + if desired == self._num_d_workers and not self._pending_num_d: self._diag_throughput_reason = ( "model_not_ready" if model_not_ready diff --git a/components/src/dynamo/planner/core/types.py b/components/src/dynamo/planner/core/types.py index 951928464afd..a7676ac760f4 100644 --- a/components/src/dynamo/planner/core/types.py +++ b/components/src/dynamo/planner/core/types.py @@ -69,6 +69,15 @@ class WorkerCounts: expected_num_decode: Optional[int] = None prefill_scaling_in_progress: bool = False decode_scaling_in_progress: bool = False + # Positive only when the connector has verified a scale-up with no drain, + # rollout, or unobserved spec update anywhere in the deployment. + pending_num_prefill: int = 0 + pending_num_decode: int = 0 + + @property + def startup_in_progress(self) -> bool: + """True when verified startup-only inventory includes pending workers.""" + return self.pending_num_prefill > 0 or self.pending_num_decode > 0 @dataclass diff --git a/components/src/dynamo/planner/environment/base.py b/components/src/dynamo/planner/environment/base.py index e4891d2e43db..c42f5e34a18c 100644 --- a/components/src/dynamo/planner/environment/base.py +++ b/components/src/dynamo/planner/environment/base.py @@ -10,7 +10,12 @@ from dynamo.planner.config.backend_components import WORKER_COMPONENT_NAMES from dynamo.planner.config.defaults import SubComponentType, TargetReplica from dynamo.planner.config.planner_config import PlannerConfig -from dynamo.planner.connectors.base import PlannerConnector, is_power_aware_connector +from dynamo.planner.connectors.base import ( + PlannerConnector, + PowerAwareConnector, + is_power_aware_connector, + is_startup_aware_connector, +) from dynamo.planner.core.budget import minimum_power_footprint_fits from dynamo.planner.core.types import FpmObservations, TrafficObservation from dynamo.planner.environment.interface import ( @@ -45,7 +50,7 @@ class NoopTrafficMetricsProvider: async def collect_traffic(self) -> Optional[TrafficObservation]: return None - def collect_accept_length(self, interval_str: str) -> Optional[float]: + async def collect_accept_length(self, interval_str: str) -> Optional[float]: del interval_str return None @@ -157,8 +162,8 @@ def metrics_state(self) -> Metrics: async def collect_traffic(self) -> Optional[TrafficObservation]: return await self.traffic_provider.collect_traffic() - def collect_accept_length(self, interval_str: str) -> Optional[float]: - return self.traffic_provider.collect_accept_length(interval_str) + async def collect_accept_length(self, interval_str: str) -> Optional[float]: + return await self.traffic_provider.collect_accept_length(interval_str) async def collect_kv_hit_rate_observation( self, duration_s: float @@ -507,6 +512,7 @@ async def _refresh_replica_counts(self) -> None: if self.require_decode and self._state.decode.info is not None else None ) + power_controller: Optional[PowerAwareConnector] = None if self.config.enable_power_awareness: if not is_power_aware_connector(self.controller): raise DeploymentValidationError( @@ -519,7 +525,42 @@ async def _refresh_replica_counts(self) -> None: "this connector does not." ] ) - counts = await self.controller.get_power_aware_worker_counts( + power_controller = self.controller + if is_startup_aware_connector(self.controller): + inventory = await self.controller.get_worker_inventory( + prefill_component_name=prefill_name, + decode_component_name=decode_name, + ) + if inventory is not None: + for required, replicas, active, expected, scaling, pending in ( + ( + self.require_prefill, + self._state.prefill.replicas, + inventory.ready_num_prefill, + inventory.expected_num_prefill, + inventory.prefill_scaling_in_progress, + inventory.pending_num_prefill, + ), + ( + self.require_decode, + self._state.decode.replicas, + inventory.ready_num_decode, + inventory.expected_num_decode, + inventory.decode_scaling_in_progress, + inventory.pending_num_decode, + ), + ): + if required: + replicas.active = active or 0 + replicas.expected = expected + replicas.scaling = scaling + replicas.pending_startup = pending + return + # Never retain previously verified pending counts after losing access. + self._state.prefill.replicas.pending_startup = 0 + self._state.decode.replicas.pending_startup = 0 + if power_controller is not None: + counts = await power_controller.get_power_aware_worker_counts( prefill_component_name=prefill_name, decode_component_name=decode_name, ) diff --git a/components/src/dynamo/planner/environment/interface.py b/components/src/dynamo/planner/environment/interface.py index 52a87551a7e7..310b5b998536 100644 --- a/components/src/dynamo/planner/environment/interface.py +++ b/components/src/dynamo/planner/environment/interface.py @@ -51,7 +51,7 @@ def metrics_state(self) -> Metrics: async def collect_traffic(self) -> Optional[TrafficObservation]: pass - def collect_accept_length(self, interval_str: str) -> Optional[float]: + async def collect_accept_length(self, interval_str: str) -> Optional[float]: pass async def collect_kv_hit_rate_observation( diff --git a/components/src/dynamo/planner/environment/metrics_provider/interface.py b/components/src/dynamo/planner/environment/metrics_provider/interface.py index 1e91250fe3a4..08b92e39468b 100644 --- a/components/src/dynamo/planner/environment/metrics_provider/interface.py +++ b/components/src/dynamo/planner/environment/metrics_provider/interface.py @@ -13,7 +13,7 @@ class TrafficMetricsProvider(Protocol): async def collect_traffic(self) -> Optional[TrafficObservation]: pass - def collect_accept_length(self, interval_str: str) -> Optional[float]: + async def collect_accept_length(self, interval_str: str) -> Optional[float]: pass async def collect_kv_hit_rate_observation( diff --git a/components/src/dynamo/planner/environment/metrics_provider/prometheus_traffic_provider.py b/components/src/dynamo/planner/environment/metrics_provider/prometheus_traffic_provider.py index aef7eac1bd12..67eca07dfbd2 100644 --- a/components/src/dynamo/planner/environment/metrics_provider/prometheus_traffic_provider.py +++ b/components/src/dynamo/planner/environment/metrics_provider/prometheus_traffic_provider.py @@ -3,8 +3,10 @@ from __future__ import annotations +import asyncio import logging -from typing import Optional +from dataclasses import dataclass +from typing import Literal, Optional from dynamo.planner.config.planner_config import PlannerConfig from dynamo.planner.core.types import TrafficObservation @@ -13,11 +15,21 @@ RuntimeNamespaceSource, ) from dynamo.planner.environment.metrics_provider.interface import TrafficMetricsProvider +from dynamo.planner.environment.state import DeploymentState from dynamo.planner.monitoring.traffic_metrics import Metrics, PrometheusAPIClient logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class _QueryContext: + model_name: str + namespace: str + backend: str + component_name: Optional[str] + endpoint: Optional[str] + + class PrometheusTrafficProvider(TrafficMetricsProvider): def __init__( self, @@ -27,6 +39,8 @@ def __init__( metrics_state: Metrics, namespace_source: Optional[RuntimeNamespaceSource] = None, ) -> None: + self._collection_lock = asyncio.Lock() + self._collection_task: Optional[asyncio.Task[Metrics]] = None self.config = config self.state_source = state_source self.metrics_state = metrics_state @@ -48,37 +62,19 @@ def __init__( self.prometheus_traffic_client.warn_if_router_not_scraped() async def collect_traffic(self) -> Optional[TrafficObservation]: - model_name = self._model_name() - if model_name is None: - logger.info("Model name is not available, skipping traffic collection") + duration_s = self.config.throughput_adjustment_interval_seconds + result = await self._collect_metrics(f"{duration_s}s", scope="traffic") + if result is None: return None - - interval_str = f"{self.config.throughput_adjustment_interval_seconds}s" m = self.metrics_state - ttft = self.prometheus_traffic_client.get_avg_time_to_first_token( - interval_str, model_name - ) - m.ttft = ttft * 1000 if ttft is not None else None - itl = self.prometheus_traffic_client.get_avg_inter_token_latency( - interval_str, model_name - ) - m.itl = itl * 1000 if itl is not None else None - m.num_req = self.prometheus_traffic_client.get_avg_request_count( - interval_str, model_name - ) - m.request_duration = self.prometheus_traffic_client.get_avg_request_duration( - interval_str, model_name - ) - m.isl = self.prometheus_traffic_client.get_avg_input_sequence_tokens( - interval_str, model_name - ) - m.osl = self.prometheus_traffic_client.get_avg_output_sequence_tokens( - interval_str, model_name - ) - m.kv_hit_rate = self.prometheus_traffic_client.get_avg_kv_hit_rate( - interval_str, model_name, namespace=self._runtime_namespace() - ) - m.accept_length = self.collect_accept_length(interval_str) + m.ttft = result.ttft + m.itl = result.itl + m.num_req = result.num_req + m.request_duration = result.request_duration + m.isl = result.isl + m.osl = result.osl + m.kv_hit_rate = result.kv_hit_rate + m.accept_length = result.accept_length normalized_idle_metrics = m.normalize_idle_nans() if normalized_idle_metrics: @@ -87,7 +83,8 @@ async def collect_traffic(self) -> Optional[TrafficObservation]: ", ".join(normalized_idle_metrics), ) - if not m.is_valid(): + num_req, isl, osl = m.num_req, m.isl, m.osl + if num_req is None or isl is None or osl is None or not m.is_valid(): logger.info("Metrics contain None or NaN values, skipping") return None @@ -106,45 +103,129 @@ async def collect_traffic(self) -> Optional[TrafficObservation]: ) return TrafficObservation( - duration_s=self.config.throughput_adjustment_interval_seconds, - num_req=m.num_req, - isl=m.isl, - osl=m.osl, + duration_s=duration_s, + num_req=num_req, + isl=isl, + osl=osl, kv_hit_rate=m.kv_hit_rate, accept_length=m.accept_length, ) - def collect_accept_length(self, interval_str: str) -> Optional[float]: - if self.config.mode not in ("disagg", "decode", "agg"): - return None + async def _collect_metrics( + self, interval_str: str, *, scope: Literal["traffic", "kv", "accept_length"] + ) -> Optional[Metrics]: + async with self._collection_lock: + # Cancellation cannot stop Requests in a thread. Wait for the previous + # batch before reusing its session, even if its caller was cancelled. + if self._collection_task is not None: + await asyncio.wait({self._collection_task}) + context = self._query_context() + if context is None: + logger.info("Model name is not available, skipping traffic collection") + return None + task = asyncio.create_task( + asyncio.to_thread( + self._query_metrics, interval_str, context, scope=scope + ) + ) + self._collection_task = task + # Retrieve failures even when the caller abandons the shielded task. + task.add_done_callback(self._collection_done) + return await asyncio.shield(task) + + @staticmethod + def _collection_done(task: asyncio.Task[Metrics]) -> None: + if not task.cancelled(): + task.exception() + + def _query_context(self) -> Optional[_QueryContext]: state = self.state_source.deployment_state() - model_name = self._model_name() - decode_info = state.decode.info - if model_name is None or decode_info is None: + model_name = self._model_name(state) + if model_name is None: return None - if not decode_info.component_name or not decode_info.endpoint: + decode_info = state.decode.info + collect_accept = self.config.mode in ("disagg", "decode", "agg") + return _QueryContext( + model_name=model_name, + namespace=self._runtime_namespace(), + backend=self.config.backend, + component_name=( + decode_info.component_name if collect_accept and decode_info else None + ), + endpoint=decode_info.endpoint if collect_accept and decode_info else None, + ) + + def _query_metrics( + self, + interval_str: str, + context: _QueryContext, + *, + scope: Literal["traffic", "kv", "accept_length"], + ) -> Metrics: + # Only local results and the synchronous client are accessed off-loop. + m = Metrics() + if scope == "traffic": + ttft = self.prometheus_traffic_client.get_avg_time_to_first_token( + interval_str, context.model_name + ) + m.ttft = ttft * 1000 if ttft is not None else None + itl = self.prometheus_traffic_client.get_avg_inter_token_latency( + interval_str, context.model_name + ) + m.itl = itl * 1000 if itl is not None else None + m.num_req = self.prometheus_traffic_client.get_avg_request_count( + interval_str, context.model_name + ) + m.request_duration = ( + self.prometheus_traffic_client.get_avg_request_duration( + interval_str, context.model_name + ) + ) + m.isl = self.prometheus_traffic_client.get_avg_input_sequence_tokens( + interval_str, context.model_name + ) + m.osl = self.prometheus_traffic_client.get_avg_output_sequence_tokens( + interval_str, context.model_name + ) + if scope != "accept_length": + m.kv_hit_rate = self.prometheus_traffic_client.get_avg_kv_hit_rate( + interval_str, context.model_name, namespace=context.namespace + ) + m.accept_length = self._query_accept_length(interval_str, context) + + return m + + async def collect_accept_length(self, interval_str: str) -> Optional[float]: + """Collect accept length without overlapping other Prometheus queries.""" + result = await self._collect_metrics(interval_str, scope="accept_length") + return result.accept_length if result is not None else None + + def _query_accept_length( + self, interval_str: str, context: _QueryContext + ) -> Optional[float]: + if not context.component_name or not context.endpoint: return None return self.prometheus_traffic_client.get_avg_spec_decode_accept_length( interval_str, - self.config.backend, - decode_info.component_name, - model_name, - namespace=self._runtime_namespace(), - endpoint_name=decode_info.endpoint, + context.backend, + context.component_name, + context.model_name, + namespace=context.namespace, + endpoint_name=context.endpoint, ) async def collect_kv_hit_rate_observation( self, duration_s: float ) -> Optional[TrafficObservation]: - model_name = self._model_name() - if model_name is None or duration_s <= 0: + if duration_s <= 0: return None interval_str = f"{int(duration_s)}s" - hit_rate = self.prometheus_traffic_client.get_avg_kv_hit_rate( - interval_str, model_name, namespace=self._runtime_namespace() - ) - accept_length = self.collect_accept_length(interval_str) + result = await self._collect_metrics(interval_str, scope="kv") + if result is None: + return None + hit_rate = result.kv_hit_rate + accept_length = result.accept_length self.metrics_state.kv_hit_rate = hit_rate self.metrics_state.accept_length = accept_length hit_rate_str = f"{hit_rate:.3f}" if hit_rate is not None else "n/a" @@ -166,8 +247,8 @@ async def collect_kv_hit_rate_observation( accept_length=accept_length, ) - def _model_name(self) -> Optional[str]: - state = self.state_source.deployment_state() + @staticmethod + def _model_name(state: DeploymentState) -> Optional[str]: if state.model_name: return state.model_name if state.decode.info is not None and state.decode.info.model_name: diff --git a/components/src/dynamo/planner/environment/state.py b/components/src/dynamo/planner/environment/state.py index 63883c2ea16d..4b2e8facbee2 100644 --- a/components/src/dynamo/planner/environment/state.py +++ b/components/src/dynamo/planner/environment/state.py @@ -15,6 +15,7 @@ class ReplicaState: active: int = 0 expected: Optional[int] = None scaling: bool = False + pending_startup: int = 0 @dataclass diff --git a/components/src/dynamo/planner/plugins/builtins/observe.py b/components/src/dynamo/planner/plugins/builtins/observe.py index 2db4c7fd8f39..a850c472a548 100644 --- a/components/src/dynamo/planner/plugins/builtins/observe.py +++ b/components/src/dynamo/planner/plugins/builtins/observe.py @@ -100,6 +100,8 @@ def _collect_worker_counts(self) -> WorkerCounts: decode_scaling_in_progress=( self.require_decode and state.decode.replicas.scaling ), + pending_num_prefill=state.prefill.replicas.pending_startup, + pending_num_decode=state.decode.replicas.pending_startup, ) diff --git a/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py b/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py index 5a44bf688cc3..e4e995c19cf2 100644 --- a/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py +++ b/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py @@ -98,6 +98,7 @@ from dynamo.planner.plugins.clock import Clock, VirtualClock, WallClock from dynamo.planner.plugins.merge.types import ComponentKey from dynamo.planner.plugins.orchestrator.orchestrator import LocalPlannerOrchestrator +from dynamo.planner.plugins.orchestrator.pipeline import PipelineOutcome from dynamo.planner.plugins.registry.auth import AllowUnauthenticatedAuth from dynamo.planner.plugins.registry.circuit_breaker import CircuitBreaker from dynamo.planner.plugins.registry.config import build_auth_validator @@ -105,8 +106,11 @@ from dynamo.planner.plugins.scheduler import PluginScheduler from dynamo.planner.plugins.transport.config import make_transport_for_endpoint from dynamo.planner.plugins.types import ( + ComponentTarget, FpmData, ObservationData, + OverrideResult, + OverrideType, PipelineContext, TrafficMetrics, WorkerState, @@ -153,6 +157,14 @@ def __init__( # time) would leave plugins with execution_interval >> 10s # never re-firing after the first tick. self._clock: Clock = clock if clock is not None else WallClock() + self._startup_down_candidate: Optional[ + tuple[Optional[int], Optional[int]] + ] = None + self._startup_down_inventory: Optional[ + tuple[Optional[int], Optional[int], int, int] + ] = None + self._startup_down_since = 0.0 + self._startup_down_sources: frozenset[str] = frozenset() # Scale_interval cadence model โ€” pipeline fires once per # ``scale_interval_seconds`` regardless of individual plugin @@ -911,6 +923,207 @@ def _baseline_from_worker_counts( out[ComponentKey(sub_component_type="decode")] = counts.ready_num_decode return out + @staticmethod + def _startup_reduction_sources( + outcome: PipelineOutcome, role: str, ready: int + ) -> frozenset[str]: + # Hand-authored outcomes predating provenance metadata retain their + # explicit-component contract. Production pipelines always supply it. + if outcome.propose_results is None: + return ( + frozenset({"explicit_proposal"}) + if ComponentKey(sub_component_type=role) in outcome.proposed_components + else frozenset() + ) + sources: set[str] = set() + for merged, results in ( + (outcome.propose_outcome, outcome.propose_results), + (outcome.reconcile_outcome, outcome.reconcile_results or []), + ): + if merged is None or merged.proposal is None: + continue + merged_target = next( + ( + t.replicas + for t in merged.proposal.targets + if t.sub_component_type == role + ), + None, + ) + if merged_target is None or merged_target > ready: + continue + for result in results: + if not isinstance(result.result, OverrideResult): + continue + if any( + t.sub_component_type == role + and t.replicas is not None + and t.replicas <= ready + and t.type in (OverrideType.SET, OverrideType.AT_MOST) + for t in result.result.targets + ): + sources.add(result.plugin_id) + return frozenset(sources) + + def _project_startup_scale_down( + self, outcome: PipelineOutcome, counts: WorkerCounts + ) -> Optional[ScalingDecision]: + """Only explicit, sustained reductions may interrupt a verified startup. + + Keep all pending capacity in budget accounting. A merged ready-count + baseline is not a request to cancel that capacity. No up-proposal is + clamped into a down-proposal, and drain/rollout snapshots never enter + this path. The connector rechecks that invariant before writing. + """ + assert outcome.final_proposal is not None + proposed = { + t.sub_component_type: t.replicas for t in outcome.final_proposal.targets + } + targets: dict[str, int] = {} + sources: set[str] = set() + gpu_total = 0 + power_total = 0 + mode = self._config.mode + roles = ( + ("prefill", "decode") + if mode == "disagg" + else (("prefill",) if mode == "prefill" else ("decode",)) + ) + for role, ready, pending, caps in ( + ( + "prefill", + counts.ready_num_prefill, + counts.pending_num_prefill, + self._capabilities.prefill, + ), + ( + "decode", + counts.ready_num_decode, + counts.pending_num_decode, + self._capabilities.decode, + ), + ): + if role not in roles: + continue + if ready is None: + self._startup_down_candidate = None + return None + desired = ready + pending + target = proposed.get(role) + role_sources = self._startup_reduction_sources(outcome, role, ready) + if not role_sources: + target = None + if target is not None: + target = max( + target, + resolve_min_endpoint( + self._config, "prefill" if role == "prefill" else "decode" + ), + ) + if target <= ready and target < desired: + targets[role] = target + sources.update(role_sources) + effective = targets.get(role, desired) + gpu_cost = caps.resolved_gpu_cost_per_replica if caps else None + if gpu_cost is None and ( + self._config.min_gpu_budget >= 0 or self._config.max_gpu_budget >= 0 + ): + self._startup_down_candidate = None + return None + gpu_total += effective * (gpu_cost or 0) + watts = caps.power_watts_per_replica if caps else None + if self._config.enable_power_awareness and watts is None: + self._startup_down_candidate = None + return None + power_total += effective * (watts or 0) + + inventory = ( + counts.ready_num_prefill, + counts.ready_num_decode, + counts.pending_num_prefill, + counts.pending_num_decode, + ) + if not targets: + evaluated = outcome.evaluated_proposal_plugins + saved_targets = dict( + zip(("prefill", "decode"), self._startup_down_candidate or (None, None)) + ) + + def conflicts(target: ComponentTarget) -> bool: + saved = saved_targets.get(target.sub_component_type) + return ( + saved is not None + and target.replicas is not None + and target.replicas > saved + and target.type in (OverrideType.SET, OverrideType.AT_LEAST) + ) + + conflicts_with_candidate = any( + conflicts(target) + for result in (outcome.propose_results or []) + + (outcome.reconcile_results or []) + + (outcome.constrain_results or []) + if isinstance(result.result, OverrideResult) + for target in result.result.targets + ) + # A plugin throttled between observations has made no new claim. + # Preserve its timer, but never execute without a fresh proposal. + # A fresh ACCEPT (including missing FPM), error, or up signal + # invalidates the previous recommendation. + if ( + evaluated is None + or self._startup_down_sources.intersection(evaluated) + or conflicts_with_candidate + or inventory != self._startup_down_inventory + ): + self._startup_down_candidate = None + return None + + min_gpu, max_gpu = self._config.min_gpu_budget, self._config.max_gpu_budget + power_limit = self._config.total_gpu_power_limit + if ( + (min_gpu >= 0 and gpu_total < min_gpu) + or (max_gpu >= 0 and gpu_total > max_gpu) + or ( + self._config.enable_power_awareness + and power_limit is not None + and power_total > power_limit + ) + ): + self._startup_down_candidate = None + return None + candidate = (targets.get("prefill"), targets.get("decode")) + now = self._clock.monotonic() + # HOLD_LAST cache replay cannot count as a second observation. + if outcome.evaluated_proposal_plugins is not None and not sources.issubset( + outcome.evaluated_proposal_plugins + ): + if ( + candidate != self._startup_down_candidate + or inventory != self._startup_down_inventory + or sources != self._startup_down_sources + ): + self._startup_down_candidate = None + return None + if ( + candidate != self._startup_down_candidate + or inventory != self._startup_down_inventory + or sources != self._startup_down_sources + ): + self._startup_down_candidate = candidate + self._startup_down_inventory = inventory + self._startup_down_since = now + self._startup_down_sources = frozenset(sources) + return None + # Require a full load observation interval with the same recommendation + # before reversing a scale-up. Faster plugin ticks cannot bypass this. + if ( + now - self._startup_down_since + < self._config.load_adjustment_interval_seconds + ): + return None + return ScalingDecision(num_prefill=candidate[0], num_decode=candidate[1]) + def _project_scale_to(self, outcome, worker_counts: WorkerCounts): """Project the pipeline outcome onto ``PlannerEffects.scale_to`` with planner "no change -> None" detection. @@ -921,7 +1134,11 @@ def _project_scale_to(self, outcome, worker_counts: WorkerCounts): baseline peers without treating them as adjustable targets. """ if outcome.execute_action != "apply" or outcome.final_proposal is None: + self._startup_down_candidate = None return None + if worker_counts.startup_in_progress: + return self._project_startup_scale_down(outcome, worker_counts) + self._startup_down_candidate = None by_comp = { t.sub_component_type: t.replicas for t in outcome.final_proposal.targets diff --git a/components/src/dynamo/planner/plugins/orchestrator/pipeline.py b/components/src/dynamo/planner/plugins/orchestrator/pipeline.py index b71cb05c7a51..9a6b498d51b2 100644 --- a/components/src/dynamo/planner/plugins/orchestrator/pipeline.py +++ b/components/src/dynamo/planner/plugins/orchestrator/pipeline.py @@ -104,6 +104,12 @@ class PipelineOutcome: reconcile_outcome: Optional[MergeOutcome] = None constrain_outcome: Optional[MergeOutcome] = None proposed_components: frozenset[ComponentKey] = field(default_factory=frozenset) + # Startup reversal needs both intent and evaluation provenance: an idle + # plugin tick is not a fresh hold, and a CONSTRAIN ceiling is not demand. + propose_results: Optional[list[PluginResult]] = None + reconcile_results: Optional[list[PluginResult]] = None + constrain_results: Optional[list[PluginResult]] = None + evaluated_proposal_plugins: Optional[frozenset[str]] = None audit_events: list[str] = field(default_factory=list) @@ -390,6 +396,7 @@ async def _run_fanout_stage( clock: Clock, metrics: Optional[PluginFrameworkMetrics] = None, propose_results: Optional[list[ProposeResult]] = None, + evaluated_plugins: Optional[set[str]] = None, ) -> tuple[MergeOutcome, list[PluginResult]]: """PROPOSE / RECONCILE / CONSTRAIN fan-out-and-merge helper. @@ -414,6 +421,8 @@ async def _run_fanout_stage( """ active = scheduler.compute_active_set(tick_now, stage, ctx=ctx) plugins: list[RegisteredPlugin] = list(active.triggered) + if evaluated_plugins is not None: + evaluated_plugins.update(plugin.plugin_id for plugin in plugins) method = _STAGE_METHOD[stage] request = _stage_request(stage, ctx, proposals=propose_results) @@ -791,6 +800,7 @@ async def _body() -> PipelineOutcome: current_ctx = current_ctx.model_copy(update={"predictions": ca.prediction}) # ---- PROPOSE stage ---- + evaluated_proposal_plugins: set[str] = set() propose, propose_plugin_results = await _run_fanout_stage( stage="propose", scheduler=scheduler, @@ -801,6 +811,7 @@ async def _body() -> PipelineOutcome: set_allowed=True, clock=clock, metrics=metrics, + evaluated_plugins=evaluated_proposal_plugins, ) proposed_components = _proposed_component_mask(propose_plugin_results) if propose.short_circuited: @@ -823,7 +834,7 @@ async def _body() -> PipelineOutcome: # than only seeing the post-merge ctx.proposal). reconcile_baseline = _proposal_to_baseline(propose.proposal, baseline) propose_proposals = [_to_propose_result(pr) for pr in propose_plugin_results] - reconcile, _ = await _run_fanout_stage( + reconcile, reconcile_plugin_results = await _run_fanout_stage( stage="reconcile", scheduler=scheduler, circuit_breaker=circuit_breaker, @@ -834,6 +845,7 @@ async def _body() -> PipelineOutcome: clock=clock, metrics=metrics, propose_results=propose_proposals, + evaluated_plugins=evaluated_proposal_plugins, ) if reconcile.short_circuited: return PipelineOutcome( @@ -854,7 +866,7 @@ async def _body() -> PipelineOutcome: # ---- CONSTRAIN stage ---- # Baseline flows from RECONCILE's output. constrain_baseline = _proposal_to_baseline(reconcile.proposal, baseline) - constrain, _ = await _run_fanout_stage( + constrain, constrain_plugin_results = await _run_fanout_stage( stage="constrain", scheduler=scheduler, circuit_breaker=circuit_breaker, @@ -903,6 +915,10 @@ async def _body() -> PipelineOutcome: reconcile_outcome=reconcile, constrain_outcome=constrain, proposed_components=proposed_components, + propose_results=propose_plugin_results, + reconcile_results=reconcile_plugin_results, + constrain_results=constrain_plugin_results, + evaluated_proposal_plugins=frozenset(evaluated_proposal_plugins), audit_events=audit, ) diff --git a/components/src/dynamo/planner/tests/integration/test_virtual_connector.py b/components/src/dynamo/planner/tests/integration/test_virtual_connector.py index 6185210838e2..4db61b84d7a1 100644 --- a/components/src/dynamo/planner/tests/integration/test_virtual_connector.py +++ b/components/src/dynamo/planner/tests/integration/test_virtual_connector.py @@ -1,16 +1,24 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# -# This test requires etcd and nats to be running, and the bindings to be installed -# - import asyncio +import json import logging +import os +import re +import shutil +import subprocess +import sys +import time +from contextlib import contextmanager import pytest -from dynamo._core import DistributedRuntime, VirtualConnectorClient +from dynamo._core import ( + DistributedRuntime, + VirtualConnectorClient, + VirtualConnectorCoordinator, +) from dynamo.planner import SubComponentType, TargetReplica, VirtualConnector from dynamo.planner.monitoring.worker_info import build_worker_info_from_defaults @@ -24,6 +32,80 @@ logger = logging.getLogger(__name__) NAMESPACE = "test_virtual_connector" +ETCD_STARTUP_TIMEOUT = 5 +ETCD_SHUTDOWN_TIMEOUT = 5 + + +@contextmanager +def _isolated_etcd(tmp_path): + tmp_path.mkdir() + data_dir = tmp_path / "data" + log_path = tmp_path / "etcd.log" + command = [ + "etcd", + "--logger", + "zap", + "--data-dir", + str(data_dir), + "--listen-client-urls", + "http://localhost:0", + "--advertise-client-urls", + "http://localhost:0", + "--listen-peer-urls", + "http://localhost:0", + "--initial-advertise-peer-urls", + "http://localhost:0", + "--initial-cluster", + "default=http://localhost:0", + ] + + with log_path.open("w", encoding="utf-8") as log_writer: + etcd = subprocess.Popen( + command, + stdout=subprocess.DEVNULL, + stderr=log_writer, + ) + + try: + client_port = None + deadline = time.monotonic() + ETCD_STARTUP_TIMEOUT + with log_path.open(encoding="utf-8") as log_reader: + while time.monotonic() < deadline: + if etcd.poll() is not None: + break + + line = log_reader.readline() + if not line: + time.sleep(0.01) + continue + + try: + entry = json.loads(line) + except json.JSONDecodeError: + continue + + if "serving client" not in entry.get("msg", ""): + continue + + match = re.search(r":(\d+)$", entry.get("address", "")) + if match: + client_port = int(match.group(1)) + break + + if client_port is None: + details = log_path.read_text(encoding="utf-8") + raise RuntimeError(f"etcd failed to start:\n{details}") + + yield client_port + finally: + if etcd.poll() is None: + etcd.terminate() + try: + etcd.wait(timeout=ETCD_SHUTDOWN_TIMEOUT) + except subprocess.TimeoutExpired: + etcd.kill() + etcd.wait(timeout=ETCD_SHUTDOWN_TIMEOUT) + shutil.rmtree(tmp_path, ignore_errors=True) class DefaultWorkerInfoProvider: @@ -118,3 +200,64 @@ async def async_internal(distributed_runtime): assert event.num_prefill_workers == 0 assert event.num_decode_workers == 0 await client.complete(event) + + +@pytest.mark.timeout(15) +def test_wait_for_unacknowledged_decision(tmp_path): + # DistributedRuntime is process-global, so keep it out of the pytest worker. + with _isolated_etcd(tmp_path / "etcd") as etcd_port: + subprocess.run( + [sys.executable, __file__], + check=True, + env={**os.environ, "ETCD_ENDPOINTS": f"http://localhost:{etcd_port}"}, + timeout=30, + ) + + +async def _wait_for_unacknowledged_decision(): + runtime = DistributedRuntime( + asyncio.get_running_loop(), "etcd", "tcp", event_plane="zmq" + ) + try: + coord = VirtualConnectorCoordinator(runtime, NAMESPACE, 1, 30, 5) + await coord.async_init() + await coord.update_scaling_decision(1, 2) + + # A late consumer must find the decision even though its watch starts afterward. + client = VirtualConnectorClient(runtime, NAMESPACE) + await asyncio.wait_for(client.wait(), timeout=5) + first = await client.get() + assert ( + first.num_prefill_workers, + first.num_decode_workers, + first.decision_id, + ) == (1, 2, 0) + await client.complete(first) + await coord.wait_for_scaling_completion() + + waiter = asyncio.ensure_future(client.wait()) + try: + done, _ = await asyncio.wait([waiter], timeout=0.2) + assert not done, "An acknowledged decision must not wake the next wait" + await client.complete(first) + done, _ = await asyncio.wait([waiter], timeout=0.2) + assert not done, "An acknowledgement update is not a new decision" + + await coord.update_scaling_decision(0, 3) + await asyncio.wait_for(waiter, timeout=5) + second = await client.get() + assert ( + second.num_prefill_workers, + second.num_decode_workers, + second.decision_id, + ) == (0, 3, 1) + await client.complete(second) + finally: + waiter.cancel() + await asyncio.gather(waiter, return_exceptions=True) + finally: + runtime.shutdown() + + +if __name__ == "__main__": + asyncio.run(_wait_for_unacknowledged_decision()) diff --git a/components/src/dynamo/planner/tests/unit/test_kubernetes_connector.py b/components/src/dynamo/planner/tests/unit/test_kubernetes_connector.py index 54c37a93cdb7..cba968320335 100644 --- a/components/src/dynamo/planner/tests/unit/test_kubernetes_connector.py +++ b/components/src/dynamo/planner/tests/unit/test_kubernetes_connector.py @@ -53,6 +53,8 @@ def mock_kube_api(): mock_api.update_graph_replicas = AsyncMock() mock_api.wait_for_graph_deployment_ready = AsyncMock() mock_api.is_deployment_ready = Mock() + mock_api.pending_startup_replicas = Mock(return_value={}) + mock_api.non_planner_components_stable = Mock(return_value=(True, [])) # Default: no terminating pods; tests that want to simulate terminating pods # override this per-test. mock_api.has_terminating_pods = Mock(return_value=False) diff --git a/components/src/dynamo/planner/tests/unit/test_pending_startup_connector_review.py b/components/src/dynamo/planner/tests/unit/test_pending_startup_connector_review.py new file mode 100644 index 000000000000..563824121157 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_pending_startup_connector_review.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exercise provider status contracts and asynchronous startup cancellation.""" + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from datetime import datetime, timezone +from threading import Event, get_ident +from unittest.mock import Mock + +import pytest +from kubernetes import client + +from dynamo.planner.config.defaults import SubComponentType, TargetReplica +from dynamo.planner.connectors.clients.kubernetes_api import ( + GROVE_PCSG_REPLICA_INDEX_LABEL, + KubernetesAPI, +) +from dynamo.planner.tests.unit.test_pending_startup_scaling import ( + _connector, + _deployment, + _pods, +) + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + + +def _decode_target(replicas): + return [ + TargetReplica( + sub_component_type=SubComponentType.DECODE, desired_replicas=replicas + ) + ] + + +def _pcsg(deployment, pods): + status = deployment["status"]["components"]["d"] + status["componentKind"] = "PodCliqueScalingGroup" + status["availableReplicas"] = status.pop("readyReplicas") + status["updatedReplicas"] = status["availableReplicas"] + for index, pod in enumerate(pods[1:]): + pod.metadata.labels[GROVE_PCSG_REPLICA_INDEX_LABEL] = str(index) + + +def test_scale_target_reads_dgdsa_before_dgd_propagation(): + api = KubernetesAPI.__new__(KubernetesAPI) + api.current_namespace = "test" + api.custom_api = Mock() + api.custom_api.get_namespaced_custom_object_scale.return_value = { + "spec": {"replicas": 3} + } + assert api.get_service_replica_target("qwen", "Decode") == 3 + assert ( + api.custom_api.get_namespaced_custom_object_scale.call_args.kwargs["name"] + == "qwen-decode" + ) + + +@pytest.mark.parametrize("status", [404, 403]) +def test_scale_target_only_falls_back_on_missing_adapter(status): + api = KubernetesAPI.__new__(KubernetesAPI) + api.current_namespace = "test" + api.custom_api = Mock() + api.custom_api.get_namespaced_custom_object_scale.side_effect = client.ApiException( + status=status + ) + api.get_graph_deployment = Mock(return_value=_deployment()) + if status == 404: + assert api.get_service_replica_target("qwen", "d") == 2 + else: + with pytest.raises(client.ApiException): + api.get_service_replica_target("qwen", "d") + api.get_graph_deployment.assert_not_called() + + +@pytest.mark.asyncio +async def test_pcsg_unavailable_replica_can_be_cancelled(): + deployment, pods = _deployment(), _pods() + _pcsg(deployment, pods) + connector = _connector(deployment, pods) + inventory = await connector.get_worker_inventory("p", "d") + assert inventory.ready_num_decode == 1 + assert inventory.pending_num_decode == 1 + await connector.set_component_replicas(_decode_target(1), blocking=False) + connector.kube_api.update_graph_replicas.assert_called_once_with("qwen", "d", 1) + + +@pytest.mark.parametrize("unsafe", ["old_revision", "surplus_group", "missing_index"]) +def test_pcsg_startup_does_not_hide_old_revisions_or_extra_groups(unsafe): + deployment, pods = _deployment(), _pods() + _pcsg(deployment, pods) + if unsafe == "old_revision": + deployment["status"]["components"]["d"]["updatedReplicas"] = 0 + elif unsafe == "surplus_group": + pods[2].metadata.labels[GROVE_PCSG_REPLICA_INDEX_LABEL] = "2" + else: + pods[2].metadata.labels.pop(GROVE_PCSG_REPLICA_INDEX_LABEL) + connector = _connector(deployment, pods) + assert not connector._get_worker_inventory_sync("p", "d").startup_in_progress + + +@pytest.mark.asyncio +async def test_superseding_adapter_target_releases_old_latch_after_propagation(): + deployment, pods = _deployment(), _pods() + connector = _connector(deployment, pods) + await connector.set_component_replicas(_decode_target(1), blocking=False) + connector.kube_api.get_service_replica_target = Mock(return_value=3) + + # The new DGDSA target has not reached the DGD. Even an unrelated observed + # generation and a newly Ready second worker cannot release the latch yet. + deployment["metadata"]["generation"] = 3 + deployment["status"]["observedGeneration"] = 3 + deployment["status"]["components"]["d"]["readyReplicas"] = 2 + deployment["status"]["conditions"][0]["status"] = "True" + pods[2].status.phase = "Running" + assert (await connector.get_worker_inventory("p", "d")).decode_scaling_in_progress + await connector.set_component_replicas(_decode_target(1), blocking=False) + connector.kube_api.update_graph_replicas.assert_called_once() + + # The replacement target reaches both spec and status; no old request may + # keep a completely stable deployment in observing-only mode. + deployment["spec"]["components"][1]["replicas"] = 3 + deployment["status"]["components"]["d"] = { + "replicas": 3, + "updatedReplicas": 3, + "readyReplicas": 3, + } + pods.append(deepcopy(pods[2])) + pods[3].metadata.name = "d2" + assert not ( + await connector.get_worker_inventory("p", "d") + ).decode_scaling_in_progress + assert not connector._startup_scale_down_targets + + +@pytest.mark.asyncio +async def test_unrelated_generation_does_not_release_unapplied_scale_request(): + deployment, pods = _deployment(), _pods() + connector = _connector(deployment, pods) + await connector.set_component_replicas(_decode_target(1), blocking=False) + deployment["metadata"]["generation"] = 3 + deployment["status"]["observedGeneration"] = 3 + deployment["status"]["conditions"][0]["status"] = "True" + deployment["status"]["components"]["d"]["readyReplicas"] = 2 + pods[2].status.phase = "Running" + assert (await connector.get_worker_inventory("p", "d")).decode_scaling_in_progress + assert connector._startup_scale_down_targets == {"d": 1} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pcsg", [False, True]) +async def test_survivor_readiness_loss_after_drain_can_be_cancelled(pcsg): + deployment, pods = _deployment(), _pods() + deployment["spec"]["components"][1]["replicas"] = 3 + deployment["status"]["components"]["d"] = { + "replicas": 3, + "updatedReplicas": 3, + "readyReplicas": 2, + } + pods.append(deepcopy(pods[2])) + pods[3].metadata.name = "d2" + pods[2].status.phase = "Running" + if pcsg: + _pcsg(deployment, pods) + connector = _connector(deployment, pods) + await connector.set_component_replicas(_decode_target(2), blocking=False) + + deployment["spec"]["components"][1]["replicas"] = 2 + deployment["status"]["components"]["d"] = { + "replicas": 2, + "updatedReplicas": 2, + "readyReplicas": 1, + } + if pcsg: + _pcsg(deployment, pods) + # PCSG status.replicas alone must not hide the surplus group's live Pod. + if pcsg: + assert not (await connector.get_worker_inventory("p", "d")).startup_in_progress + pods[3].metadata.deletion_timestamp = datetime.now(timezone.utc) + assert not (await connector.get_worker_inventory("p", "d")).startup_in_progress + pods.pop() + inventory = await connector.get_worker_inventory("p", "d") + assert inventory.pending_num_decode == 1 + assert not connector._startup_scale_down_targets + await connector.set_component_replicas(_decode_target(1), blocking=False) + assert connector.kube_api.update_graph_replicas.call_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unsafe", ["drain", "unobserved", "rollout", "old_revision", "surplus_pcsg"] +) +async def test_stale_ready_condition_cannot_bypass_reduction_guards(unsafe): + deployment, pods = _deployment(), _pods() + deployment["status"]["conditions"][0]["status"] = "True" + if unsafe == "drain": + pods[2].metadata.deletion_timestamp = datetime.now(timezone.utc) + elif unsafe == "unobserved": + deployment["metadata"]["generation"] += 1 + elif unsafe == "rollout": + deployment["status"]["rollingUpdate"] = {"phase": "InProgress"} + elif unsafe == "old_revision": + deployment["status"]["components"]["d"]["updatedReplicas"] = 1 + else: + _pcsg(deployment, pods) + pods[2].metadata.labels[GROVE_PCSG_REPLICA_INDEX_LABEL] = "2" + connector = _connector(deployment, pods) + await connector.set_component_replicas(_decode_target(1), blocking=False) + connector.kube_api.update_graph_replicas.assert_not_called() + + +@pytest.mark.asyncio +async def test_startup_reduction_is_latched_even_with_stale_ready_condition(): + deployment, pods = _deployment(), _pods() + deployment["status"]["conditions"][0]["status"] = "True" + connector = _connector(deployment, pods) + await connector.set_component_replicas(_decode_target(1), blocking=False) + assert connector._startup_scale_down_targets == {"d": 1} + await connector.set_component_replicas(_decode_target(1), blocking=False) + connector.kube_api.update_graph_replicas.assert_called_once() + + +@pytest.mark.asyncio +async def test_checkpoint_capture_pod_does_not_block_serving_worker_cancellation(): + deployment, pods = _deployment(), _pods() + capture = deepcopy(pods[2]) + capture.metadata.name = "capture-abc" + capture.metadata.labels.update( + {"nvidia.com/snapshot-job": "capture", "nvidia.com/snapshot-job-uid": "uid"} + ) + capture.metadata.owner_references = [ + client.V1OwnerReference( + api_version="batch/v1", + kind="Job", + name="capture", + uid="job", + controller=True, + ) + ] + capture.status.phase = "Running" + capture.metadata.deletion_timestamp = datetime.now(timezone.utc) + pods.append(capture) + connector = _connector(deployment, pods) + assert (await connector.get_worker_inventory("p", "d")).pending_num_decode == 1 + await connector.set_component_replicas(_decode_target(1), blocking=False) + connector.kube_api.update_graph_replicas.assert_called_once_with("qwen", "d", 1) + + +def test_checkpoint_filter_does_not_hide_failed_worker_with_inherited_labels(): + deployment, pods = _deployment(), _pods() + pods[2].status.phase = "Failed" + pods[2].metadata.labels.update( + {"nvidia.com/snapshot-job": "capture", "nvidia.com/snapshot-job-uid": "uid"} + ) + pods[2].metadata.owner_references = [ + client.V1OwnerReference( + api_version="grove.io/v1alpha1", + kind="PodClique", + name="decode", + uid="clique", + controller=True, + ) + ] + connector = _connector(deployment, pods) + assert not connector._get_worker_inventory_sync("p", "d").startup_in_progress + + +@pytest.mark.asyncio +async def test_inventory_retirement_cannot_erase_concurrent_scale_down(): + deployment, pods = _deployment(), _pods() + connector = _connector(deployment, pods) + connector._startup_scale_down_targets = {"p": 1} + status = connector.kube_api.get_service_replica_status + inspecting_old_target, resume = Event(), Event() + main_thread = get_ident() + worker_calls = 0 + + def interleaved_status(snapshot, name): + nonlocal worker_calls + if get_ident() != main_thread: + worker_calls += 1 + # Two status reads classify startup; the third inspects the old + # completed p=1 target, just before it is retired. + if worker_calls == 3: + inspecting_old_target.set() + assert resume.wait(5) + return status(snapshot, name) + + connector.kube_api.get_service_replica_status = interleaved_status + with ThreadPoolExecutor(max_workers=1) as executor: + retirement = executor.submit( + connector._startup_scale_down_in_progress, deployment, pods + ) + try: + assert await asyncio.to_thread(inspecting_old_target.wait, 5) + await connector.set_component_replicas( + [ + TargetReplica( + sub_component_type=SubComponentType.DECODE, desired_replicas=1 + ) + ], + blocking=False, + ) + finally: + resume.set() + assert await asyncio.wrap_future(retirement) + + # The stale inventory observed p=1 settled, but d=1 was accepted meanwhile. + # It must remain held until the new reduction actually converges. + assert connector._startup_scale_down_targets == {"d": 1} + connector.kube_api.update_graph_replicas.assert_called_once_with("qwen", "d", 1) diff --git a/components/src/dynamo/planner/tests/unit/test_pending_startup_decision_review.py b/components/src/dynamo/planner/tests/unit/test_pending_startup_decision_review.py new file mode 100644 index 000000000000..84d9d2979b06 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_pending_startup_decision_review.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Startup cancellation across real plugin cadence, budget, and merge stages.""" + +from unittest.mock import Mock + +import pytest + +from dynamo.common.forward_pass_metrics import ( + ForwardPassMetrics, + ScheduledRequestMetrics, +) +from dynamo.planner.config.planner_config import PlannerConfig +from dynamo.planner.core.state_machine import PlannerScalingState +from dynamo.planner.core.types import ( + EngineCapabilities, + FpmObservations, + TickInput, + WorkerCapabilities, + WorkerCounts, +) +from dynamo.planner.plugins.clock import VirtualClock +from dynamo.planner.plugins.orchestrator.engine_adapter import OrchestratorEngineAdapter +from dynamo.planner.plugins.types import ( + ComponentTarget, + ConstrainStageResponse, + OverrideResult, + OverrideType, + ProposeStageResponse, + ReconcileStageResponse, +) + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + + +def _config(**kwargs): + return PlannerConfig( + mode="disagg", + optimization_target="load", + served_model_name="test", + prefill_scale_up_queue_tokens=1000, + prefill_scale_down_queue_tokens=100, + decode_scale_up_kv_rate=80, + decode_scale_down_kv_rate=20, + min_gpu_budget=-1, + max_gpu_budget=-1, + **kwargs, + ) + + +def _caps(): + return WorkerCapabilities( + prefill=EngineCapabilities( + num_gpu=1, max_num_batched_tokens=1024, max_kv_tokens=1000 + ), + decode=EngineCapabilities( + num_gpu=1, max_num_batched_tokens=1024, max_kv_tokens=1000 + ), + ) + + +def _counts(): + return WorkerCounts( + ready_num_prefill=1, + ready_num_decode=1, + expected_num_prefill=1, + expected_num_decode=2, + prefill_scaling_in_progress=True, + decode_scaling_in_progress=True, + pending_num_decode=1, + ) + + +def _obs(ready=1, kv=0): + return FpmObservations( + decode={ + (f"d{i}", 0): ForwardPassMetrics( + worker_id=f"d{i}", + scheduled_requests=ScheduledRequestMetrics(sum_decode_kv_tokens=kv), + ) + for i in range(ready) + }, + ) + + +@pytest.mark.asyncio +async def test_idle_plugin_ticks_preserve_startup_observation_window(): + config = _config( + load_adjustment_interval_seconds=30, scheduling={"scale_interval_seconds": 5} + ) + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + tick = adapter.initial_tick(0) + for _ in range(12): + effects = await adapter.tick( + tick, + TickInput( + now_s=tick.at_s, worker_counts=_counts(), fpm_observations=_obs() + ), + ) + if tick.at_s < 60: + assert effects.scale_to is None + else: + assert effects.scale_to is not None + assert effects.scale_to.num_decode == 1 + tick = effects.next_tick + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "role,floor,expected_at_60", + [("prefill", 1, 1), ("decode", 1, 1), ("decode", 2, None)], +) +async def test_interleaved_plugin_floor_only_resets_conflicting_candidate( + role, floor, expected_at_60 +): + config = _config( + load_adjustment_interval_seconds=30, scheduling={"scale_interval_seconds": 5} + ) + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + + class Floor: + async def Propose(self, req): + return ProposeStageResponse( + override=OverrideResult( + targets=[ + ComponentTarget( + sub_component_type=role, + replicas=floor, + type=OverrideType.AT_LEAST, + ) + ] + ) + ) + + adapter._orchestrator.register_internal("other_floor", "propose", 20, Floor()) + tick = adapter.initial_tick(0) + for _ in range(12): + effects = await adapter.tick( + tick, + TickInput( + now_s=tick.at_s, worker_counts=_counts(), fpm_observations=_obs() + ), + ) + if tick.at_s < 60: + assert effects.scale_to is None + tick = effects.next_tick + assert (effects.scale_to.num_decode if effects.scale_to else None) == expected_at_60 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("veto", ["hold", "missing_fpm"]) +async def test_fresh_evaluated_hold_or_missing_fpm_restarts_window(veto): + config = _config( + load_adjustment_interval_seconds=30, scheduling={"scale_interval_seconds": 5} + ) + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + tick = adapter.initial_tick(0) + for _ in range(24): + obs = _obs() + if tick.at_s == 60: + obs = _obs(kv=500) if veto == "hold" else FpmObservations() + effects = await adapter.tick( + tick, + TickInput(now_s=tick.at_s, worker_counts=_counts(), fpm_observations=obs), + ) + if tick.at_s < 120: + assert effects.scale_to is None + else: + assert effects.scale_to is not None + assert effects.scale_to.num_decode == 1 + tick = effects.next_tick + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "floor,ceiling,expected", [(5, -1, 2), (6, -1, None), (-1, 4, None)] +) +async def test_load_budget_preserves_pending_peer(floor, ceiling, expected): + config = _config() + config.min_gpu_budget, config.max_gpu_budget = floor, ceiling + counts = _counts() + counts.ready_num_decode = counts.expected_num_decode = 3 + counts.expected_num_prefill = 3 + counts.pending_num_prefill, counts.pending_num_decode = 2, 0 + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + tick = adapter.initial_tick(0) + for _ in range(2): + effects = await adapter.tick( + tick, + TickInput( + now_s=tick.at_s, worker_counts=counts, fpm_observations=_obs(ready=3) + ), + ) + tick = effects.next_tick + assert (effects.scale_to.num_decode if effects.scale_to else None) == expected + if effects.scale_to: + assert effects.scale_to.num_prefill is None + + +def test_throughput_floor_counts_preserved_pending_peer(): + config = _config() + config.enable_load_scaling = False + config.enable_throughput_scaling = True + config.min_gpu_budget = 5 + state = PlannerScalingState(config, _caps()) + counts = _counts() + counts.ready_num_decode = counts.expected_num_decode = 3 + counts.expected_num_prefill = 3 + counts.pending_num_prefill, counts.pending_num_decode = 2, 0 + state.observe_worker_counts(counts) + state._compute_prefill_replicas = Mock(return_value=3) + state._compute_decode_replicas = Mock(return_value=1) + decision = state._throughput_disagg(1, 100, 100) + assert decision is not None + # The floor limits D to 2 while the up-requested P keeps desired=3. + assert decision.num_prefill is None + assert decision.num_decode == 2 + + +def _target(replicas, kind=OverrideType.SET): + return OverrideResult( + targets=[ + ComponentTarget(sub_component_type="decode", replicas=replicas, type=kind) + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", ["capped_up", "floor_only", "reconcile_down"]) +async def test_external_pipeline_preserves_reduction_provenance(policy): + config = _config() + config.enable_load_scaling = False + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + + class Proposer: + async def Propose(self, req): + return ProposeStageResponse( + override=_target(1, OverrideType.AT_LEAST) + if policy == "floor_only" + else _target(3) + ) + + class Constraint: + async def Constrain(self, req): + return ConstrainStageResponse(override=_target(1, OverrideType.AT_MOST)) + + class Reconciler: + async def Reconcile(self, req): + return ReconcileStageResponse(override=_target(1)) + + adapter._orchestrator.register_internal("external", "propose", 0, Proposer()) + if policy == "capped_up": + adapter._orchestrator.register_internal("ceiling", "constrain", 0, Constraint()) + elif policy == "reconcile_down": + adapter._orchestrator.register_internal( + "down_policy", "reconcile", 0, Reconciler() + ) + tick = adapter.initial_tick(0) + for index in range(3): + effects = await adapter.tick( + tick, TickInput(now_s=tick.at_s, worker_counts=_counts()) + ) + if policy == "reconcile_down" and index > 0: + assert effects.scale_to is not None + assert effects.scale_to.num_decode == 1 + else: + assert effects.scale_to is None + tick = effects.next_tick diff --git a/components/src/dynamo/planner/tests/unit/test_pending_startup_scaling.py b/components/src/dynamo/planner/tests/unit/test_pending_startup_scaling.py new file mode 100644 index 000000000000..ce8128890cd8 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_pending_startup_scaling.py @@ -0,0 +1,519 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cancellation of startup capacity must not race a drain or manufacture demand.""" + +from copy import deepcopy +from datetime import datetime, timezone +from threading import Lock +from unittest.mock import Mock + +import pytest +from kubernetes import client + +from dynamo.common.forward_pass_metrics import ( + ForwardPassMetrics, + ScheduledRequestMetrics, +) +from dynamo.planner.config.defaults import SubComponentType, TargetReplica +from dynamo.planner.config.planner_config import PlannerConfig +from dynamo.planner.connectors.clients.kubernetes_api import KubernetesAPI +from dynamo.planner.connectors.kubernetes import KubernetesConnector +from dynamo.planner.core.state_machine import PlannerScalingState +from dynamo.planner.core.types import ( + EngineCapabilities, + FpmObservations, + TickInput, + WorkerCapabilities, + WorkerCounts, +) +from dynamo.planner.environment.base import PlannerEnvironmentImpl +from dynamo.planner.plugins.builtins.observe import EnvironmentObservePlugin +from dynamo.planner.plugins.clock import VirtualClock +from dynamo.planner.plugins.merge.types import ComponentKey +from dynamo.planner.plugins.orchestrator.engine_adapter import OrchestratorEngineAdapter +from dynamo.planner.plugins.orchestrator.pipeline import PipelineOutcome +from dynamo.planner.plugins.types import ComponentTarget, ScalingProposal + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + + +def _config(mode="disagg", **kwargs): + return PlannerConfig( + mode=mode, + optimization_target="load", + served_model_name="test", + prefill_scale_up_queue_tokens=1000, + prefill_scale_down_queue_tokens=100, + decode_scale_up_kv_rate=80, + decode_scale_down_kv_rate=20, + min_gpu_budget=-1, + max_gpu_budget=-1, + **kwargs, + ) + + +def _caps(): + return WorkerCapabilities( + prefill=EngineCapabilities( + num_gpu=1, max_num_batched_tokens=1024, max_kv_tokens=1000 + ), + decode=EngineCapabilities( + num_gpu=1, max_num_batched_tokens=1024, max_kv_tokens=1000 + ), + ) + + +def _counts(ready=1, pending=1): + return WorkerCounts( + ready_num_prefill=1, + ready_num_decode=ready, + prefill_scaling_in_progress=True, + decode_scaling_in_progress=True, + pending_num_decode=pending, + ) + + +def _obs(ready=1, kv=0): + return FpmObservations( + prefill={("p", 0): ForwardPassMetrics(worker_id="p")}, + decode={ + (f"d{i}", 0): ForwardPassMetrics( + worker_id=f"d{i}", + scheduled_requests=ScheduledRequestMetrics(sum_decode_kv_tokens=kv), + ) + for i in range(ready) + }, + ) + + +def _outcome(p=None, d=None, explicit=("prefill", "decode")): + return PipelineOutcome( + execute_action="apply", + final_proposal=ScalingProposal( + targets=[ + ComponentTarget(sub_component_type=role, replicas=value) + for role, value in (("prefill", p), ("decode", d)) + if value is not None + ] + ), + proposed_components=frozenset( + ComponentKey(sub_component_type=role) for role in explicit + ), + ) + + +@pytest.mark.parametrize("mode", ["disagg", "decode", "agg"]) +@pytest.mark.parametrize("ready", [1, 3]) +def test_load_can_cancel_startup_and_drain_serving_workers(mode, ready): + state = PlannerScalingState(_config(mode), _caps()) + state.observe_worker_counts(_counts(ready)) + decision = state.advance_load(_obs(ready)) + assert decision is not None + assert decision.num_decode == max(1, ready - 1) + assert decision.num_prefill is None + + +@pytest.mark.parametrize("kv", [500, 950]) +def test_hold_or_scale_up_signal_does_not_cancel_startup(kv): + state = PlannerScalingState(_config(), _caps()) + state.observe_worker_counts(_counts()) + assert state.advance_load(_obs(kv=kv)) is None + + +def test_unknown_scaling_or_drain_stays_blocked(): + state = PlannerScalingState(_config(), _caps()) + state.observe_worker_counts(_counts(ready=3, pending=0)) + assert state.advance_load(_obs(ready=3)) is None + + +def test_missing_fpm_does_not_cancel_pending_peer(): + state = PlannerScalingState(_config(), _caps()) + state.observe_worker_counts(_counts()) + obs = _obs() + obs.decode = None + assert state.advance_load(obs) is None + + +def test_startup_projection_requires_sustained_explicit_down_signal(): + config = _config() + clock = VirtualClock() + adapter = OrchestratorEngineAdapter(config, _caps(), clock=clock) + counts = _counts() + down = _outcome(d=1) + assert adapter._project_scale_to(down, counts) is None + # Repeated fast plugin ticks cannot accelerate the stabilization interval. + for _ in range(5): + assert adapter._project_scale_to(down, counts) is None + clock.advance(config.load_adjustment_interval_seconds) + # A renewed up signal resets stabilization; it must not become cancellation. + assert adapter._project_scale_to(_outcome(d=3), counts) is None + assert adapter._project_scale_to(down, counts) is None + clock.advance(config.load_adjustment_interval_seconds) + decision = adapter._project_scale_to(down, counts) + assert decision is not None and decision.num_decode == 1 + assert decision.num_prefill is None + # An omitted role's baseline echo is never a cancellation request. + assert adapter._project_scale_to(_outcome(p=1, d=1, explicit=()), counts) is None + + +@pytest.mark.parametrize( + "floor,ceiling,target,allowed", + [ + (3, -1, 1, False), # Ready-equal cancellation would cross GPU floor. + (-1, 2, 1, True), + (-1, 2, 5, False), # Up proposal must not be clamped into a down proposal. + ], +) +def test_startup_gpu_bounds(floor, ceiling, target, allowed): + config = _config() + config.min_gpu_budget, config.max_gpu_budget = floor, ceiling + clock = VirtualClock() + adapter = OrchestratorEngineAdapter(config, _caps(), clock=clock) + counts, outcome = _counts(), _outcome(d=target) + assert adapter._project_scale_to(outcome, counts) is None + clock.advance(config.load_adjustment_interval_seconds) + assert (adapter._project_scale_to(outcome, counts) is not None) is allowed + + +def _deployment(): + return { + "metadata": {"name": "qwen", "generation": 2}, + "spec": { + "components": [ + {"name": "p", "type": "prefill", "replicas": 1}, + {"name": "d", "type": "decode", "replicas": 2}, + ] + }, + "status": { + "observedGeneration": 2, + "conditions": [{"type": "Ready", "status": "False"}], + "components": { + "p": {"replicas": 1, "updatedReplicas": 1, "readyReplicas": 1}, + "d": {"replicas": 2, "updatedReplicas": 2, "readyReplicas": 1}, + }, + }, + } + + +def _pods(): + return [ + client.V1Pod( + metadata=client.V1ObjectMeta( + name=name, labels={"nvidia.com/dynamo-component": component} + ), + status=client.V1PodStatus(phase=phase), + ) + for name, component, phase in [ + ("p0", "p", "Running"), + ("d0", "d", "Running"), + ("d1", "d", "Pending"), + ] + ] + + +def _connector(deployment, pods): + api = KubernetesAPI.__new__(KubernetesAPI) + api.get_graph_deployment = Mock(side_effect=lambda _: deepcopy(deployment)) + api.list_pods_for_graph = Mock(return_value=pods) + # DGDSA writes precede DGD reconciliation; neither depends on Planner's latch. + scale_targets = { + component["name"]: component.get("replicas", 1) + for component in deployment["spec"]["components"] + } + api.update_graph_replicas = Mock( + side_effect=lambda _, name, target: scale_targets.__setitem__(name, target) + ) + api.get_service_replica_target = Mock( + side_effect=lambda _, name: scale_targets[name] + ) + connector = KubernetesConnector.__new__(KubernetesConnector) + connector.graph_deployment_name = "qwen" + connector.kube_api = api + connector.raise_not_ready = False + connector._startup_scale_down_lock = Lock() + connector._startup_scale_down_targets = {} + connector._startup_read_warnings = set() + return connector + + +@pytest.mark.parametrize( + "unsafe", + [ + "drain", + "scale_down", + "rollout", + "failed_rollout", + "stale", + "old_template", + "failed_pod", + ], +) +def test_inventory_refuses_to_label_unsafe_states_as_startup(unsafe): + deployment, pods = _deployment(), _pods() + if unsafe == "drain": + pods[0].metadata.deletion_timestamp = datetime.now(timezone.utc) + elif unsafe == "scale_down": + deployment["spec"]["components"][1]["replicas"] = 1 + elif unsafe in ("rollout", "failed_rollout"): + deployment["status"]["rollingUpdate"] = { + "phase": "Failed" if unsafe == "failed_rollout" else "InProgress" + } + elif unsafe == "stale": + deployment["metadata"]["generation"] += 1 + elif unsafe == "old_template": + deployment["status"]["components"]["d"]["updatedReplicas"] = 1 + else: + pods[2].status.phase = "Failed" + connector = _connector(deployment, pods) + inventory = connector._get_worker_inventory_sync("p", "d") + assert not inventory.startup_in_progress + + +@pytest.mark.asyncio +async def test_connector_cancels_startup_and_holds_until_drain_finishes(): + deployment, pods = _deployment(), _pods() + connector = _connector(deployment, pods) + inventory = await connector.get_worker_inventory("p", "d") + assert inventory.pending_num_decode == 1 + target = [ + TargetReplica(sub_component_type=SubComponentType.DECODE, desired_replicas=1) + ] + await connector.set_component_replicas(target, blocking=False) + connector.kube_api.update_graph_replicas.assert_called_once_with("qwen", "d", 1) + # Async DGDSA application has not reached the DGD yet: no repeat writes. + await connector.set_component_replicas(target, blocking=False) + assert not (await connector.get_worker_inventory("p", "d")).startup_in_progress + connector.kube_api.update_graph_replicas.assert_called_once() + deployment["spec"]["components"][1]["replicas"] = 1 + deployment["status"]["components"]["d"] = { + "replicas": 1, + "updatedReplicas": 1, + "readyReplicas": 1, + } + deployment["status"]["conditions"][0]["status"] = "True" + # Even if Ready/counts have converged, a deleting Pod still blocks reversal. + pods[2].metadata.deletion_timestamp = datetime.now(timezone.utc) + up = [TargetReplica(sub_component_type=SubComponentType.DECODE, desired_replicas=2)] + await connector.set_component_replicas(up, blocking=False) + connector.kube_api.update_graph_replicas.assert_called_once() + pods.pop() + await connector.set_component_replicas(up, blocking=False) + assert connector.kube_api.update_graph_replicas.call_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target", [-1, 3]) +async def test_unready_connector_rejects_non_reductions_and_negative_targets(target): + connector = _connector(_deployment(), _pods()) + await connector.set_component_replicas( + [ + TargetReplica( + sub_component_type=SubComponentType.DECODE, desired_replicas=target + ) + ], + blocking=False, + ) + connector.kube_api.update_graph_replicas.assert_not_called() + + +@pytest.mark.asyncio +async def test_builtin_pipeline_cancels_pending_decode_after_observation_window(): + config = _config() + clock = VirtualClock() + adapter = OrchestratorEngineAdapter(config, _caps(), clock=clock) + tick = adapter.initial_tick(start_s=0.0) + for index in range(2): + assert tick is not None + effects = await adapter.tick( + tick, + TickInput( + now_s=tick.at_s, + worker_counts=_counts(), + fpm_observations=_obs(), + ), + ) + if index == 0: + assert effects.scale_to is None + else: + assert effects.scale_to is not None + assert effects.scale_to.num_decode == 1 + assert effects.scale_to.num_prefill is None + tick = effects.next_tick + + +@pytest.mark.parametrize("mode", ["decode", "agg"]) +def test_startup_easy_consolidation_does_not_immediately_require_scale_up(mode): + config = _config(mode) + config.optimization_target = "throughput" + state = PlannerScalingState(config, _caps()) + state.observe_worker_counts(_counts(ready=2)) + # Each serving worker is under the 60% down threshold, but merging them + # would use 110% KV capacity. Only cancel the pending third worker. + decision = state.advance_load(_obs(ready=2, kv=550)) + assert decision is not None and decision.num_decode == 2 + # Once the pending worker is gone, unchanged load must continue to hold. + state.observe_worker_counts(WorkerCounts(ready_num_decode=2, expected_num_decode=2)) + for _ in range(5): + decision = state.advance_load(_obs(ready=2, kv=550)) + assert decision is None or decision.num_decode == 2 + + +@pytest.mark.parametrize("role", ["prefill", "decode"]) +def test_sla_one_ready_worker_can_cancel_startup_without_division_by_zero(role): + config = PlannerConfig( + mode=role, + optimization_target="sla", + enable_throughput_scaling=False, + enable_load_scaling=True, + served_model_name="test", + min_gpu_budget=-1, + max_gpu_budget=-1, + ) + state = PlannerScalingState(config, _caps()) + counts = _counts() + if role == "prefill": + counts.pending_num_prefill, counts.pending_num_decode = 1, 0 + state.observe_worker_counts(counts) + regression = Mock() + regression.has_sufficient_data.return_value = True + regression.avg_isl = 100 + obs = _obs() + regression.query_groups.return_value = [ + ("worker", list((obs.prefill if role == "prefill" else obs.decode).values())) + ] + regression.estimate_queued_prefill_time.side_effect = ( + lambda *args, **kwargs: 0.005 if kwargs.get("queue_scale") == 0 else 0.01 + ) + regression.estimate_scheduled_decode_itl.return_value = 0.001 + if role == "prefill": + state._prefill_regression = regression + else: + state._decode_regression = regression + decision = state.advance_load(obs) + assert decision is not None + assert (decision.num_prefill if role == "prefill" else decision.num_decode) == 1 + + +def test_pending_peer_capacity_is_preserved_and_charged_to_budget(): + config = _config() + config.min_gpu_budget = config.max_gpu_budget = 4 + clock = VirtualClock() + adapter = OrchestratorEngineAdapter(config, _caps(), clock=clock) + counts = _counts() + counts.pending_num_prefill = 2 + outcome = _outcome(p=1, d=1, explicit=("decode",)) + # p keeps desired=3, d cancels to 1: total=4, although only 2 are Ready. + assert adapter._project_scale_to(outcome, counts) is None + clock.advance(config.load_adjustment_interval_seconds) + decision = adapter._project_scale_to(outcome, counts) + assert ( + decision is not None + and decision.num_prefill is None + and decision.num_decode == 1 + ) + + +@pytest.mark.parametrize("minimum,ready,expected", [(2, 3, 2), (2, 1, None)]) +def test_startup_projection_preserves_component_minimum(minimum, ready, expected): + config = _config(decode_min_endpoint=minimum) + clock = VirtualClock() + adapter = OrchestratorEngineAdapter(config, _caps(), clock=clock) + counts, outcome = _counts(ready), _outcome(d=0) + adapter._project_scale_to(outcome, counts) + clock.advance(config.load_adjustment_interval_seconds) + decision = adapter._project_scale_to(outcome, counts) + assert (decision.num_decode if decision else None) == expected + + +@pytest.mark.asyncio +async def test_connector_rechecks_for_drain_between_observation_and_write(): + deployment, pods = _deployment(), _pods() + connector = _connector(deployment, pods) + assert (await connector.get_worker_inventory("p", "d")).startup_in_progress + pods[0].metadata.deletion_timestamp = datetime.now(timezone.utc) + await connector.set_component_replicas( + [TargetReplica(sub_component_type=SubComponentType.DECODE, desired_replicas=1)], + blocking=False, + ) + connector.kube_api.update_graph_replicas.assert_not_called() + + +@pytest.mark.asyncio +async def test_environment_passes_startup_inventory_to_observer_and_clears_it_on_drain(): + deployment, pods = _deployment(), _pods() + connector = _connector(deployment, pods) + environment = PlannerEnvironmentImpl( + config=_config(), + controller=connector, + require_prefill=True, + require_decode=True, + ) + environment.deployment_state().prefill.info = Mock(k8s_name="p") + environment.deployment_state().decode.info = Mock(k8s_name="d") + observer = EnvironmentObservePlugin( + environment, require_prefill=True, require_decode=True + ) + await environment._refresh_replica_counts() + assert observer._collect_worker_counts().pending_num_decode == 1 + pods[0].metadata.deletion_timestamp = datetime.now(timezone.utc) + await environment._refresh_replica_counts() + assert not observer._collect_worker_counts().startup_in_progress + + +@pytest.mark.parametrize("mode", ["disagg", "decode", "agg"]) +def test_budget_clamped_up_proposal_cannot_cancel_startup(mode): + config = _config(mode) + config.max_gpu_budget = 2 if mode == "disagg" else 1 + state = PlannerScalingState(config, _caps()) + state.observe_worker_counts(_counts()) + assert state.advance_load(_obs(kv=950)) is None + + +@pytest.mark.parametrize("requested,expected", [(1, 1), (4, None)]) +def test_throughput_cancels_startup_only_for_actual_down_intent(requested, expected): + config = _config() + config.enable_load_scaling = False + config.enable_throughput_scaling = True + config.max_gpu_budget = 2 + state = PlannerScalingState(config, _caps()) + state.observe_worker_counts(_counts()) + state._compute_prefill_replicas = Mock(return_value=1) + state._compute_decode_replicas = Mock(return_value=requested) + decision = state._throughput_disagg(1.0, 100.0, 100.0) + assert (decision.num_decode if decision else None) == expected + + +def test_worker_becoming_ready_restarts_startup_stabilization(): + config = _config() + clock = VirtualClock() + adapter = OrchestratorEngineAdapter(config, _caps(), clock=clock) + counts, outcome = _counts(ready=2, pending=2), _outcome(d=1) + adapter._project_scale_to(outcome, counts) + clock.advance(config.load_adjustment_interval_seconds) + counts.ready_num_decode, counts.pending_num_decode = 3, 1 + assert adapter._project_scale_to(outcome, counts) is None + clock.advance(config.load_adjustment_interval_seconds) + assert adapter._project_scale_to(outcome, counts).num_decode == 1 + + +@pytest.mark.parametrize("power_limit,allowed", [(500, False), (800, True)]) +def test_startup_power_budget_includes_pending_peers(power_limit, allowed): + config = _config() + config.enable_power_awareness = True + config.total_gpu_power_limit = power_limit + caps = _caps() + caps.prefill.power_watts_per_replica = caps.decode.power_watts_per_replica = 200 + counts = _counts() + counts.pending_num_prefill = 2 + clock = VirtualClock() + adapter = OrchestratorEngineAdapter(config, caps, clock=clock) + outcome = _outcome(d=1, explicit=("decode",)) + adapter._project_scale_to(outcome, counts) + clock.advance(config.load_adjustment_interval_seconds) + assert (adapter._project_scale_to(outcome, counts) is not None) is allowed diff --git a/components/src/dynamo/planner/tests/unit/test_power_environment.py b/components/src/dynamo/planner/tests/unit/test_power_environment.py index a47bb366658b..ec8baaa12ba6 100644 --- a/components/src/dynamo/planner/tests/unit/test_power_environment.py +++ b/components/src/dynamo/planner/tests/unit/test_power_environment.py @@ -17,6 +17,7 @@ from dynamo.planner.config.planner_config import PlannerConfig from dynamo.planner.connectors.base import is_power_aware_connector +from dynamo.planner.core.types import WorkerCounts from dynamo.planner.environment.base import PlannerEnvironmentImpl from dynamo.planner.errors import DeploymentValidationError, PowerAnnotationInvalidError from dynamo.planner.monitoring.dgd_services import ( @@ -449,3 +450,36 @@ async def wait_settled(*, include_planner=True): controller.wait_for_deployment_ready.assert_awaited_once_with(include_planner=False) assert backing_lookups == [] assert env.deployment_state().prefill.power_watts_per_replica is None + + +@pytest.mark.asyncio +async def test_startup_inventory_cannot_bypass_power_capability_validation(): + controller = Mock() + controller.get_worker_inventory = AsyncMock(return_value=WorkerCounts()) + env = _env(controller) + + with pytest.raises(DeploymentValidationError, match="PowerAwareConnector"): + await env._refresh_replica_counts() + controller.get_worker_inventory.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_power_capable_startup_inventory_preserves_pending_workers(): + controller = _controller() + controller.get_worker_inventory = AsyncMock( + return_value=WorkerCounts( + ready_num_prefill=1, + ready_num_decode=2, + prefill_scaling_in_progress=True, + decode_scaling_in_progress=True, + pending_num_decode=1, + ) + ) + env = _env(controller) + + await env._refresh_replica_counts() + + controller.get_worker_inventory.assert_awaited_once() + assert env.deployment_state().decode.replicas.active == 2 + assert env.deployment_state().decode.replicas.pending_startup == 1 + assert env.deployment_state().decode.replicas.scaling diff --git a/components/src/dynamo/planner/tests/unit/test_prometheus_traffic_provider.py b/components/src/dynamo/planner/tests/unit/test_prometheus_traffic_provider.py index a1f5a96ecdc0..ecba64c32917 100644 --- a/components/src/dynamo/planner/tests/unit/test_prometheus_traffic_provider.py +++ b/components/src/dynamo/planner/tests/unit/test_prometheus_traffic_provider.py @@ -1,11 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from unittest.mock import MagicMock, patch +import asyncio +import threading +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from aiohttp.test_utils import TestClient, TestServer from dynamo.planner.config.planner_config import PlannerConfig +from dynamo.planner.control_api import _build_app from dynamo.planner.environment.metrics_provider.prometheus_traffic_provider import ( PrometheusTrafficProvider, ) @@ -57,13 +61,14 @@ def _provider(namespace_source=None): return provider, client_class.return_value, client_patch -def test_accept_length_uses_current_runtime_namespace(): +@pytest.mark.asyncio +async def test_accept_length_uses_current_runtime_namespace(): namespace_source = MagicMock() namespace_source.runtime_namespace.return_value = "base-ns-workerhash" provider, client, client_patch = _provider(namespace_source) client.get_avg_spec_decode_accept_length.return_value = 2.5 try: - assert provider.collect_accept_length("30s") == 2.5 + assert await provider.collect_accept_length("30s") == 2.5 finally: client_patch.stop() @@ -93,11 +98,12 @@ def test_provider_passes_prometheus_request_timeout(): assert client_class.call_args.kwargs["request_timeout_seconds"] == 3.5 -def test_accept_length_falls_back_to_configured_namespace(): +@pytest.mark.asyncio +async def test_accept_length_falls_back_to_configured_namespace(): provider, client, client_patch = _provider() client.get_avg_spec_decode_accept_length.return_value = 2.5 try: - provider.collect_accept_length("30s") + await provider.collect_accept_length("30s") finally: client_patch.stop() @@ -114,6 +120,8 @@ def test_accept_length_falls_back_to_configured_namespace(): "get_avg_time_to_first_token", "get_avg_inter_token_latency", "get_avg_request_count", + "get_avg_input_sequence_tokens", + "get_avg_output_sequence_tokens", ], ) async def test_invalid_required_metric_is_checked_before_float_logging(missing_method): @@ -133,3 +141,267 @@ async def test_invalid_required_metric_is_checked_before_float_logging(missing_m client_patch.stop() assert observation is None + + +@pytest.fixture +async def traffic_provider(): + provider, client, client_patch = _provider() + client.get_avg_time_to_first_token.return_value = 0.1 + client.get_avg_inter_token_latency.return_value = 0.01 + client.get_avg_request_count.return_value = 2.0 + client.get_avg_request_duration.return_value = 0.2 + client.get_avg_input_sequence_tokens.return_value = 100.0 + client.get_avg_output_sequence_tokens.return_value = 20.0 + client.get_avg_kv_hit_rate.return_value = 0.5 + client.get_avg_spec_decode_accept_length.return_value = 2.5 + try: + yield provider, client + finally: + if provider._collection_task is not None: + await asyncio.gather(provider._collection_task, return_exceptions=True) + client_patch.stop() + + +@pytest.fixture +async def blocked_query(traffic_provider): + _, client = traffic_provider + started = asyncio.Event() + release = threading.Event() + loop = asyncio.get_running_loop() + + def query(*args, **kwargs): + loop.call_soon_threadsafe(started.set) + assert release.wait(timeout=5), "Test did not release Prometheus query" + return 0.5 + + client.get_avg_kv_hit_rate.side_effect = query + try: + yield started, release + finally: + release.set() + + +@pytest.mark.asyncio +@pytest.mark.timeout(10) +@pytest.mark.parametrize("full", [True, False]) +async def test_control_api_responds_during_collection( + traffic_provider, blocked_query, full +): + provider, client = traffic_provider + started, release = blocked_query + controller = MagicMock() + controller.get_min_endpoints = AsyncMock(return_value={"decode_min_endpoint": 1}) + controller.patch_min_endpoints = AsyncMock(return_value={"decode_min_endpoint": 2}) + async with TestClient(TestServer(_build_app(controller))) as http: + collection = asyncio.create_task( + provider.collect_traffic() + if full + else provider.collect_kv_hit_rate_observation(30) + ) + try: + await asyncio.wait_for(started.wait(), timeout=2) + response = await asyncio.wait_for(http.get("/v1/min-endpoints"), timeout=2) + assert response.status == 200 + assert await response.json() == {"decode_min_endpoint": 1} + response = await asyncio.wait_for( + http.patch("/v1/min-endpoints", json={"decode_min_endpoint": 2}), + timeout=2, + ) + assert response.status == 200 + assert await response.json() == {"decode_min_endpoint": 2} + assert not collection.done() + assert provider.metrics_state == Metrics() + # Later queries must use the original snapshot, not live state. + provider.config.namespace = "changed-ns" + provider.state_source.deployment_state.return_value.decode.info.endpoint = ( + "changed-endpoint" + ) + provider.config.throughput_adjustment_interval_seconds = 60 + finally: + release.set() + observation = await collection + + assert observation.duration_s == 30 + assert observation.kv_hit_rate == 0.5 + assert observation.accept_length == 2.5 + assert provider.metrics_state.kv_hit_rate == 0.5 + if full: + assert provider.metrics_state.ttft == 100 + assert provider.metrics_state.itl == 10 + assert observation.num_req == 2 + client.get_avg_spec_decode_accept_length.assert_called_once_with( + "30s", + "vllm", + "backend", + "Qwen/Qwen3", + namespace="base-ns", + endpoint_name="generate", + ) + + +@pytest.mark.asyncio +@pytest.mark.timeout(10) +@pytest.mark.parametrize("full", [True, False]) +async def test_cancelled_collection_does_not_publish_or_overlap( + traffic_provider, blocked_query, full +): + provider, client = traffic_provider + started, release = blocked_query + collection = asyncio.create_task( + provider.collect_traffic() + if full + else provider.collect_kv_hit_rate_observation(30) + ) + await asyncio.wait_for(started.wait(), timeout=2) + collection.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(collection, timeout=2) + + # A new caller must wait without starting a second query. Cancelling that + # waiter must not cancel the thread-backed task either. + waiter = asyncio.create_task(provider.collect_kv_hit_rate_observation(30)) + await asyncio.sleep(0) + assert not waiter.done() + assert client.get_avg_kv_hit_rate.call_count == 1 + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + release.set() + await provider._collection_task + assert provider.metrics_state == Metrics() + + observation = await provider.collect_kv_hit_rate_observation(30) + assert observation.kv_hit_rate == 0.5 + assert client.get_avg_kv_hit_rate.call_count == 2 + + +@pytest.mark.asyncio +async def test_query_failure_preserves_metrics_and_allows_retry(traffic_provider): + provider, client = traffic_provider + client.get_avg_request_duration.side_effect = RuntimeError("query failed") + with pytest.raises(RuntimeError, match="query failed"): + await provider.collect_traffic() + assert provider.metrics_state == Metrics() + client.get_avg_request_duration.side_effect = None + assert await provider.collect_traffic() is not None + + +@pytest.mark.asyncio +async def test_idle_normalization_preserves_load_metrics(traffic_provider): + provider, client = traffic_provider + provider.metrics_state.p_load = 0.4 + provider.metrics_state.d_load = 0.6 + client.get_avg_request_count.return_value = 0.0 + client.get_avg_time_to_first_token.return_value = float("nan") + client.get_avg_inter_token_latency.return_value = float("nan") + client.get_avg_input_sequence_tokens.return_value = float("nan") + client.get_avg_output_sequence_tokens.return_value = float("nan") + client.get_avg_request_duration.return_value = float("nan") + observation = await provider.collect_traffic() + assert observation.num_req == observation.isl == observation.osl == 0 + assert provider.metrics_state.ttft == provider.metrics_state.itl == 0 + assert provider.metrics_state.p_load == 0.4 + assert provider.metrics_state.d_load == 0.6 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_source", ["global", "decode", "prefill"]) +async def test_query_context_uses_one_deployment_snapshot( + traffic_provider, model_source +): + provider, client = traffic_provider + state = _state_source().deployment_state() + if model_source != "global": + state.model_name = None + if model_source == "decode": + state.decode.info.model_name = "Qwen/Qwen3" + else: + state.prefill.info = WorkerInfo(model_name="Qwen/Qwen3") + other = DeploymentState(model_name="other-model") + other.decode.info = WorkerInfo(component_name="other", endpoint="other-endpoint") + provider.state_source.deployment_state.side_effect = [state, other] + + assert await provider.collect_accept_length("30s") == 2.5 + + provider.state_source.deployment_state.assert_called_once_with() + client.get_avg_spec_decode_accept_length.assert_called_once_with( + "30s", + "vllm", + "backend", + "Qwen/Qwen3", + namespace="base-ns", + endpoint_name="generate", + ) + client.get_avg_kv_hit_rate.assert_not_called() + client.get_avg_time_to_first_token.assert_not_called() + assert provider.metrics_state == Metrics() + + +@pytest.mark.asyncio +@pytest.mark.timeout(10) +@pytest.mark.parametrize("cancel_waiter", [False, True]) +async def test_accept_length_waits_for_collection( + traffic_provider, blocked_query, cancel_waiter +): + provider, client = traffic_provider + started, release = blocked_query + collection = asyncio.create_task(provider.collect_traffic()) + await asyncio.wait_for(started.wait(), timeout=2) + waiter = asyncio.create_task(provider.collect_accept_length("10s")) + try: + await asyncio.sleep(0) + assert not waiter.done() + client.get_avg_spec_decode_accept_length.assert_not_called() + if cancel_waiter: + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + finally: + release.set() + await collection + await asyncio.gather(waiter, return_exceptions=True) + + if cancel_waiter: + assert await provider.collect_accept_length("10s") == 2.5 + else: + assert waiter.result() == 2.5 + assert client.get_avg_spec_decode_accept_length.call_count == 2 + assert client.get_avg_kv_hit_rate.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.timeout(10) +async def test_direct_accept_length_is_responsive_and_cancellation_safe( + traffic_provider, blocked_query +): + provider, client = traffic_provider + started, release = blocked_query + client.get_avg_spec_decode_accept_length.side_effect = ( + client.get_avg_kv_hit_rate.side_effect + ) + controller = MagicMock() + controller.get_min_endpoints = AsyncMock(return_value={"decode_min_endpoint": 1}) + async with TestClient(TestServer(_build_app(controller))) as http: + collection = asyncio.create_task(provider.collect_accept_length("30s")) + waiter = None + try: + await asyncio.wait_for(started.wait(), timeout=2) + response = await asyncio.wait_for(http.get("/v1/min-endpoints"), timeout=2) + assert response.status == 200 + assert not collection.done() + collection.cancel() + with pytest.raises(asyncio.CancelledError): + await collection + waiter = asyncio.create_task(provider.collect_accept_length("10s")) + await asyncio.sleep(0) + assert not waiter.done() + assert client.get_avg_spec_decode_accept_length.call_count == 1 + finally: + release.set() + await asyncio.gather(collection, return_exceptions=True) + if waiter is not None: + await asyncio.gather(waiter, return_exceptions=True) + assert waiter.result() == 0.5 + assert client.get_avg_spec_decode_accept_length.call_count == 2 + client.get_avg_kv_hit_rate.assert_not_called() + assert provider.metrics_state == Metrics() diff --git a/components/src/dynamo/planner/tests/unit/test_startup_rbac_compatibility.py b/components/src/dynamo/planner/tests/unit/test_startup_rbac_compatibility.py new file mode 100644 index 000000000000..37ff28103a9f --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_startup_rbac_compatibility.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Older service accounts keep ordinary scaling without granting startup reads.""" + +from copy import deepcopy +from unittest.mock import AsyncMock, Mock, call + +import pytest +from kubernetes.client import ApiException + +from dynamo.planner.config.defaults import SubComponentType, TargetReplica +from dynamo.planner.environment.base import PlannerEnvironmentImpl +from dynamo.planner.errors import DynamoGraphDeploymentNotReadyError +from dynamo.planner.monitoring.worker_info import WorkerInfo +from dynamo.planner.tests.unit.test_pending_startup_scaling import ( + _config, + _connector, + _deployment, + _pods, +) + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + + +def _settle(deployment, replicas=3): + deployment["spec"]["components"][1]["replicas"] = replicas + deployment["status"]["conditions"][0]["status"] = "True" + deployment["status"]["components"]["d"] = { + "replicas": replicas, + "updatedReplicas": replicas, + "readyReplicas": replicas, + } + + +def _env(connector, *, power=False): + config = _config() + config.enable_power_awareness = power + environment = PlannerEnvironmentImpl( + config=config, + controller=connector, + require_prefill=True, + require_decode=True, + ) + environment.deployment_state().prefill.info = WorkerInfo(k8s_name="p") + environment.deployment_state().decode.info = WorkerInfo(k8s_name="d") + return environment + + +def _target(replicas): + return [ + TargetReplica( + sub_component_type=SubComponentType.DECODE, desired_replicas=replicas + ) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target", [2, 4]) +async def test_old_rbac_preserves_replica_refresh_and_ready_scaling(target): + deployment = _deployment() + _settle(deployment) + connector = _connector(deployment, _pods()) + connector.kube_api.list_pods_for_graph.side_effect = ApiException(status=403) + environment = _env(connector) + + await environment._refresh_replica_counts() + replicas = environment.deployment_state().decode.replicas + assert ( + replicas.active, + replicas.expected, + replicas.scaling, + replicas.pending_startup, + ) == (3, 3, False, 0) + await connector.set_component_replicas(_target(target), blocking=False) + connector.kube_api.update_graph_replicas.assert_called_once_with( + "qwen", "d", target + ) + + +@pytest.mark.asyncio +async def test_pod_permission_changes_enable_and_disable_startup_inventory(): + connector = _connector(_deployment(), _pods()) + environment = _env(connector) + connector.kube_api.list_pods_for_graph.side_effect = ApiException(status=403) + await environment._refresh_replica_counts() + replicas = environment.deployment_state().decode.replicas + assert replicas.active == 1 and replicas.scaling and replicas.pending_startup == 0 + + connector.kube_api.list_pods_for_graph.side_effect = None + await environment._refresh_replica_counts() + assert replicas.pending_startup == 1 + + connector.kube_api.list_pods_for_graph.side_effect = ApiException(status=403) + await environment._refresh_replica_counts() + assert replicas.pending_startup == 0 + await connector.set_component_replicas(_target(1), blocking=False) + connector.kube_api.update_graph_replicas.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raise_not_ready", [False, True]) +async def test_old_rbac_keeps_unready_rejection(raise_not_ready): + connector = _connector(_deployment(), _pods()) + connector.raise_not_ready = raise_not_ready + connector.kube_api.list_pods_for_graph.side_effect = ApiException(status=403) + if raise_not_ready: + with pytest.raises(DynamoGraphDeploymentNotReadyError): + await connector.set_component_replicas(_target(1), blocking=False) + else: + await connector.set_component_replicas(_target(1), blocking=False) + connector.kube_api.update_graph_replicas.assert_not_called() + + +@pytest.mark.asyncio +async def test_pod_read_denial_does_not_release_accepted_downscale(): + deployment = _deployment() + _settle(deployment, replicas=1) + connector = _connector(deployment, _pods()) + connector._startup_scale_down_targets = {"d": 1} + connector.kube_api.list_pods_for_graph.side_effect = ApiException(status=403) + assert await connector.get_worker_inventory("p", "d") is None + await connector.set_component_replicas(_target(2), blocking=False) + assert connector._startup_scale_down_targets == {"d": 1} + connector.kube_api.update_graph_replicas.assert_not_called() + + +@pytest.mark.asyncio +async def test_power_awareness_does_not_fall_back_to_unverified_counts(): + connector = _connector(_deployment(), _pods()) + connector.kube_api.list_pods_for_graph.side_effect = ApiException(status=403) + connector.get_actual_worker_counts = AsyncMock() + connector.get_power_aware_worker_counts = AsyncMock( + wraps=connector.get_power_aware_worker_counts + ) + with pytest.raises(ApiException) as error: + await _env(connector, power=True)._refresh_replica_counts() + assert error.value.status == 403 + connector.get_power_aware_worker_counts.assert_awaited_once() + connector.get_actual_worker_counts.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [401, 500]) +async def test_non_permission_errors_still_propagate(status): + connector = _connector(_deployment(), _pods()) + connector.kube_api.list_pods_for_graph.side_effect = ApiException(status=status) + connector.get_actual_worker_counts = AsyncMock() + with pytest.raises(ApiException) as error: + await _env(connector)._refresh_replica_counts() + assert error.value.status == status + connector.get_actual_worker_counts.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_forbidden_dgd_read_is_not_treated_as_optional_pod_access(): + connector = _connector(_deployment(), _pods()) + connector.kube_api.get_graph_deployment.side_effect = ApiException(status=403) + connector.get_actual_worker_counts = AsyncMock() + with pytest.raises(ApiException): + await _env(connector)._refresh_replica_counts() + connector.get_actual_worker_counts.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raise_not_ready", [False, True]) +@pytest.mark.parametrize("ready_target", [2, 4]) +async def test_scale_get_denial_preserves_scaling_after_external_update( + raise_not_ready, ready_target +): + deployment, pods = _deployment(), _pods() + connector = _connector(deployment, pods) + connector.raise_not_ready = raise_not_ready + connector.kube_api.get_service_replica_target.side_effect = ApiException(status=403) + environment = _env(connector, power=True) + await environment._refresh_replica_counts() + assert environment.deployment_state().decode.replicas.pending_startup == 1 + + if raise_not_ready: + with pytest.raises(DynamoGraphDeploymentNotReadyError): + await connector.set_component_replicas(_target(1), blocking=False) + else: + await connector.set_component_replicas(_target(1), blocking=False) + connector.kube_api.update_graph_replicas.assert_not_called() + assert connector._startup_scale_down_targets == {} + + # An external writer scales to three workers while Scale GET stays forbidden. + _settle(deployment) + deployment["metadata"]["generation"] = 3 + deployment["status"]["observedGeneration"] = 3 + pods[-1].status.phase = "Running" + pods.append(deepcopy(pods[-1])) + pods[-1].metadata.name = "d2" + await environment._refresh_replica_counts() + replicas = environment.deployment_state().decode.replicas + assert (replicas.active, replicas.scaling, replicas.pending_startup) == ( + 3, + False, + 0, + ) + await connector.set_component_replicas(_target(ready_target), blocking=False) + connector.kube_api.update_graph_replicas.assert_called_once_with( + "qwen", "d", ready_target + ) + assert connector._startup_scale_down_targets == {} + connector.kube_api.get_service_replica_target.assert_called_once_with("qwen", "d") + + +@pytest.mark.asyncio +async def test_all_startup_targets_are_checked_before_writes_and_retried(): + deployment, pods = _deployment(), _pods() + deployment["spec"]["components"][0]["replicas"] = 2 + deployment["status"]["components"]["p"] = { + "replicas": 2, + "updatedReplicas": 2, + "readyReplicas": 2, + } + pods.append(deepcopy(pods[0])) + pods[-1].metadata.name = "p1" + connector = _connector(deployment, pods) + read_target = connector.kube_api.get_service_replica_target.side_effect + connector.kube_api.get_service_replica_target.side_effect = [ + 2, + ApiException(status=403), + ] + targets = [ + TargetReplica(sub_component_type=SubComponentType.PREFILL, desired_replicas=1), + *_target(1), + ] + await connector.set_component_replicas(targets, blocking=False) + connector.kube_api.update_graph_replicas.assert_not_called() + assert connector._startup_scale_down_targets == {} + assert connector.kube_api.get_service_replica_target.call_args_list == [ + call("qwen", "p"), + call("qwen", "d"), + ] + + # Granting GET takes effect on the next attempt, without restarting Planner. + connector.kube_api.get_service_replica_target.side_effect = read_target + await connector.set_component_replicas(targets, blocking=False) + assert connector.kube_api.update_graph_replicas.call_args_list == [ + call("qwen", "p", 1), + call("qwen", "d", 1), + ] + assert connector._startup_scale_down_targets == {"p": 1, "d": 1} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [401, 500]) +async def test_startup_scale_get_non_permission_errors_propagate_before_writes(status): + connector = _connector(_deployment(), _pods()) + connector.kube_api.get_service_replica_target.side_effect = ApiException( + status=status + ) + with pytest.raises(ApiException) as error: + await connector.set_component_replicas(_target(1), blocking=False) + assert error.value.status == status + connector.kube_api.update_graph_replicas.assert_not_called() + assert connector._startup_scale_down_targets == {} + + +@pytest.mark.asyncio +async def test_scale_get_denial_holds_until_the_accepted_target_is_observed(): + deployment = _deployment() + connector = _connector(deployment, _pods()) + await connector.set_component_replicas(_target(1), blocking=False) + # Permission can still be revoked after admission; keep guarding that write. + connector.kube_api.get_service_replica_target = Mock( + side_effect=ApiException(status=403) + ) + inventory = await connector.get_worker_inventory("p", "d") + assert inventory is not None and inventory.decode_scaling_in_progress + assert not inventory.startup_in_progress + assert connector._startup_scale_down_targets == {"d": 1} + + _settle(deployment, replicas=1) + inventory = await connector.get_worker_inventory("p", "d") + assert inventory is not None and not inventory.decode_scaling_in_progress + assert connector._startup_scale_down_targets == {} + connector.kube_api.get_service_replica_target.assert_called_once() + + +@pytest.mark.asyncio +async def test_forbidden_replica_write_is_not_silently_ignored(): + deployment = _deployment() + _settle(deployment) + connector = _connector(deployment, _pods()) + connector.kube_api.list_pods_for_graph.side_effect = ApiException(status=403) + connector.kube_api.update_graph_replicas.side_effect = ApiException(status=403) + with pytest.raises(ApiException) as error: + await connector.set_component_replicas(_target(2), blocking=False) + assert error.value.status == 403 + assert connector._startup_scale_down_targets == {} + + +@pytest.mark.asyncio +async def test_forbidden_pod_read_does_not_admit_startup_via_stale_ready(): + deployment = _deployment() + deployment["status"]["conditions"][0]["status"] = "True" + connector = _connector(deployment, _pods()) + connector.kube_api.list_pods_for_graph.side_effect = ApiException(status=403) + await connector.set_component_replicas(_target(1), blocking=False) + connector.kube_api.update_graph_replicas.assert_not_called() diff --git a/components/src/dynamo/planner/tests/unit/test_throughput_scaling.py b/components/src/dynamo/planner/tests/unit/test_throughput_scaling.py index 9ca3c0c4c7df..70a2b7e92237 100644 --- a/components/src/dynamo/planner/tests/unit/test_throughput_scaling.py +++ b/components/src/dynamo/planner/tests/unit/test_throughput_scaling.py @@ -544,3 +544,57 @@ def test_engine_rps_recommendation_is_independent_of_sidecar_cost( ) assert replicas == 3 + + +@pytest.mark.parametrize("mode", ["prefill", "decode", "disagg", "agg"]) +@pytest.mark.parametrize("model_ready", [False, True]) +def test_startup_cancellation_requires_a_throughput_estimate(mode, model_ready): + estimate = 1 if model_ready else None + state = _disagg_state( + 1, + 1, + estimate, + estimate, + enable_load_scaling=False, + min_gpus=-1, + max_gpus=-1, + ) + state._config.mode = mode + state.observe_worker_counts( + WorkerCounts( + ready_num_prefill=1, + ready_num_decode=1, + prefill_scaling_in_progress=mode in ("prefill", "disagg"), + decode_scaling_in_progress=mode != "prefill", + pending_num_prefill=1 if mode in ("prefill", "disagg") else 0, + pending_num_decode=1 if mode != "prefill" else 0, + ) + ) + if mode == "agg": + state.update_capabilities( + WorkerCapabilities( + decode=EngineCapabilities( + gpu_cost_per_replica=1, + max_num_batched_tokens=4096, + ) + ) + ) + state._agg_regression = SimpleNamespace( + find_engine_capacity_rps=lambda **_kwargs: ( + SimpleNamespace(rps=1.0, ttft_ms=1.0, itl_ms=1.0, eligible=True) + if model_ready + else None + ) + ) + decision = state._throughput_agg(1.0, 1.0, 1.0) + elif mode == "disagg": + decision = state._throughput_disagg(1.0, 1.0, 1.0) + else: + decision = state._throughput_single(1.0, 1.0, 1.0, mode) + + if model_ready: + assert decision is not None + assert (decision.num_prefill if mode == "prefill" else decision.num_decode) == 1 + else: + assert decision is None + assert state.diagnostics().throughput_decision_reason == "model_not_ready" diff --git a/components/src/dynamo/profiler/tests/unit/test_profiler_protocol.py b/components/src/dynamo/profiler/tests/unit/test_profiler_protocol.py index 28dba9564514..04d7a1691861 100644 --- a/components/src/dynamo/profiler/tests/unit/test_profiler_protocol.py +++ b/components/src/dynamo/profiler/tests/unit/test_profiler_protocol.py @@ -4,6 +4,8 @@ """Unit tests for profiler config_modifiers/protocol helpers.""" import copy +import errno +from pathlib import Path from unittest.mock import AsyncMock, patch import pytest @@ -32,10 +34,17 @@ EngineType, SearchStrategy, ) + from dynamo.profiler.utils.dgd_materialization import ( + DGDMaterializationPurpose, + materialize_dgd, + ) from dynamo.profiler.utils.dgdr_v1beta1_types import ( DynamoGraphDeploymentRequestSpec, OverridesSpec, ) + from dynamo.profiler.utils.model_info import ( + model_ref_allows_implicit_trust_remote_code, + ) from dynamo.profiler.utils.profile_common import ProfilerOperationalConfig except ImportError: pytest.skip("dynamo.llm bindings not available", allow_module_level=True) @@ -1735,6 +1744,82 @@ def test_materialize_dgd_shell_form_preserves_syntax() -> None: assert result_args[0] == original_cmd + " --trust-remote-code" +@pytest.mark.parametrize( + ("path_kind", "expected"), + [ + ("directory", True), + ("symlink", True), + ("file", False), + ("missing", False), + ("child_of_file", False), + ], +) +def test_implicit_trust_requires_local_directory(tmp_path, path_kind, expected): + model_path = tmp_path / "model" + if path_kind == "directory": + model_path.mkdir() + elif path_kind == "symlink": + target = tmp_path / "snapshot" + target.mkdir() + model_path.symlink_to(target, target_is_directory=True) + elif path_kind in ("file", "child_of_file"): + model_path.touch() + if path_kind == "child_of_file": + model_path /= "child" + + assert model_ref_allows_implicit_trust_remote_code(model_path) is expected + + +@pytest.mark.parametrize("explicit_trust", [False, True]) +def test_materialize_dgd_inaccessible_model_path( + tmp_path, monkeypatch, caplog, explicit_trust +): + model_path = tmp_path / "model" + model_path.mkdir() + (model_path / "config.json").write_text("{}") + error = PermissionError(errno.EACCES, "Cannot inspect model", str(model_path)) + original_stat = Path.stat + + def stat(self, *args, **kwargs): + if self == model_path: + raise error + return original_stat(self, *args, **kwargs) + + monkeypatch.setattr(Path, "stat", stat) + config = _make_dgd_with_workers("decode") + if explicit_trust: + _main_container(_components_by_name(config)["decode"])["args"].append( + "--trust-remote-code" + ) + original_config = copy.deepcopy(config) + + if explicit_trust: + result = materialize_dgd( + config, + purpose=DGDMaterializationPurpose.FINAL_OUTPUT, + runtime_backend="vllm", + model_name_or_path=str(model_path), + ) + assert result == original_config + else: + with pytest.raises(RuntimeError, match="Cannot inspect model path") as exc: + materialize_dgd( + config, + purpose=DGDMaterializationPurpose.FINAL_OUTPUT, + runtime_backend="vllm", + model_name_or_path=str(model_path), + ) + assert exc.value.__cause__ is error + assert str(model_path) in str(exc.value) + assert "symlink ownership" in str(exc.value) + assert "modelCache.pvcModelPath" in str(exc.value) + assert "mutable remote" not in str(exc.value) + + assert config == original_config + assert "auto_map detection is inconclusive" in caplog.text + assert "injecting --trust-remote-code" not in caplog.text + + def test_model_has_auto_map_returns_true_on_unexpected_error() -> None: """Unexpected errors (network, auth) must return True (conservative default) rather than silently returning False and risking a missed injection.""" diff --git a/components/src/dynamo/profiler/tests/unit/test_resolve_model_path.py b/components/src/dynamo/profiler/tests/unit/test_resolve_model_path.py index 20739d37445a..9301ec36fca6 100644 --- a/components/src/dynamo/profiler/tests/unit/test_resolve_model_path.py +++ b/components/src/dynamo/profiler/tests/unit/test_resolve_model_path.py @@ -4,11 +4,15 @@ """Unit tests for resolve_model_path() and the rapid.py / thorough.py call sites that feed its result into aiconfigurator.""" +import asyncio import copy +import errno +import os from unittest.mock import AsyncMock, MagicMock, patch import pandas as pd import pytest +import yaml pytestmark = [ pytest.mark.pre_merge, @@ -18,6 +22,7 @@ ] try: + from dynamo.profiler.profile_sla import run_profile from dynamo.profiler.rapid import ( _generate_dgd_from_pick, _run_autoscale_sim, @@ -28,7 +33,10 @@ run_thorough, ) from dynamo.profiler.utils.config_modifiers import CONFIG_MODIFIERS - from dynamo.profiler.utils.dgd_materialization import DGDMaterializationPurpose + from dynamo.profiler.utils.dgd_materialization import ( + DGDMaterializationPurpose, + materialize_dgd, + ) from dynamo.profiler.utils.dgdr_v1beta1_types import ( DynamoGraphDeploymentRequestSpec, HardwareSpec, @@ -167,6 +175,85 @@ def test_returns_hf_id_when_dir_has_no_config_json(self, tmp_path): dgdr = _make_dgdr(modelCache=_pvc_model_cache(str(tmp_path), "model")) assert resolve_model_path(dgdr) == _HF_ID + def test_inaccessible_pvc_config_during_materialization( + self, tmp_path, monkeypatch + ): + local_dir = tmp_path / "model" + _make_model_dir(local_dir) + config_path = local_dir / "config.json" + dgdr = _make_dgdr( + backend="vllm", modelCache=_pvc_model_cache(str(tmp_path), "model") + ) + error = PermissionError(errno.EACCES, "Permission denied", str(config_path)) + original_stat = os.stat + + def stat(path, *args, **kwargs): + if os.fspath(path) == str(config_path): + raise error + return original_stat(path, *args, **kwargs) + + monkeypatch.setattr(os, "stat", stat) + blueprint = { + "spec": { + "components": [ + { + "name": "worker", + "type": "worker", + "podTemplate": { + "spec": {"containers": [{"name": "main", "args": []}]} + }, + } + ] + } + } + with ( + patch("dynamo.profiler.utils.model_info.hf_hub_download") as download, + pytest.raises(RuntimeError, match="Cannot inspect PVC model config") as exc, + ): + materialize_dgd( + blueprint, + purpose=DGDMaterializationPurpose.FINAL_OUTPUT, + runtime_backend=dgdr.backend, + model_name_or_path=resolve_model_path(dgdr), + ) + assert exc.value.__cause__ is error + assert str(config_path) in str(exc.value) + assert "symlink ownership" in str(exc.value) + assert "modelCache.pvcModelPath" in str(exc.value) + download.assert_not_called() + + def test_run_profile_records_pvc_inspection_failure(self, tmp_path, monkeypatch): + local_dir = tmp_path / "model" + _make_model_dir(local_dir) + config_path = local_dir / "config.json" + dgdr = _make_dgdr(modelCache=_pvc_model_cache(str(tmp_path), "model")) + error = PermissionError(errno.EACCES, "Permission denied", str(config_path)) + original_stat = os.stat + + def stat(path, *args, **kwargs): + if os.fspath(path) == str(config_path): + raise error + return original_stat(path, *args, **kwargs) + + monkeypatch.setattr(os, "stat", stat) + output_dir = tmp_path / "output" + with ( + patch( + "dynamo.profiler.profile_sla.check_model_hardware_support", + side_effect=AssertionError("Must not fall back to the Hub model"), + ), + pytest.raises(RuntimeError, match="Cannot inspect PVC model config") as exc, + ): + asyncio.run( + run_profile(dgdr, ProfilerOperationalConfig(output_dir=str(output_dir))) + ) + status = yaml.safe_load((output_dir / "profiler_status.yaml").read_text()) + assert status["status"] == "failed" + assert status["error"] == str(exc.value) + assert str(config_path) in status["error"] + assert "modelCache.pvcModelPath" in status["error"] + assert exc.value.__cause__ is error + # --------------------------------------------------------------------------- # rapid.py call sites diff --git a/components/src/dynamo/profiler/utils/dgdr_v1beta1_types.py b/components/src/dynamo/profiler/utils/dgdr_v1beta1_types.py index 68f6d952cd1e..952b472cde8a 100644 --- a/components/src/dynamo/profiler/utils/dgdr_v1beta1_types.py +++ b/components/src/dynamo/profiler/utils/dgdr_v1beta1_types.py @@ -174,7 +174,7 @@ class OverridesSpec(BaseModel): profilingJob: Optional[Dict[str, Any]] = Field( default=None, - description="ProfilingJob allows overriding the profiling Job specification. Fields set here are merged into the controller-generated Job spec.", + description="ProfilingJob allows overriding the profiling Job specification. Fields set here are merged into the controller-generated Job spec. Security: creating a DGDR is workload-creation authority in its namespace โ€” these overrides carry the same blast radius as creating a Job or Pod directly there, by design. Pod security is enforced centrally by Kubernetes Pod Security Admission on the resulting Pods once the namespace is labeled (see pod-security.kubernetes.io/enforce): it applies the full Pod Security Standards โ€” covering privileged, host namespaces, and hostPath, not only securityContext โ€” not this API. ServiceAccount identity is a separate layer: these overrides can set serviceAccountName and automountServiceAccountToken, which are bounded by RBAC and namespace membership rather than PSA โ€” the same authority any Pod author in the namespace already holds. Grant create/update on DGDRs only to principals trusted to create Pods in the namespace. The profiling Job always runs in the DGDR's own namespace and overrides cannot change that โ€” but namespace containment is not node or cross-tenant isolation. Dynamo's workloads, including this Job, satisfy the baseline standard, so enforce baseline (non-exempt) on every resulting Pod to close the privileged, host-namespace, host-device, and hostPath paths.", ) trustRemoteCode: bool = Field( default=False, diff --git a/components/src/dynamo/profiler/utils/model_info.py b/components/src/dynamo/profiler/utils/model_info.py index 5ed7730f0fbe..f669cfcfa743 100644 --- a/components/src/dynamo/profiler/utils/model_info.py +++ b/components/src/dynamo/profiler/utils/model_info.py @@ -4,6 +4,7 @@ import json import logging from pathlib import Path +from stat import S_ISDIR from typing import Optional, Union from huggingface_hub import hf_hub_download, model_info @@ -144,8 +145,9 @@ def model_has_auto_map( Works for both local directories and Hub model IDs. Reads ``config.json`` directly (no ``AutoConfig`` load) so it works even for architectures - that ``transformers`` doesn't know about. Returns False on any read - error so callers can treat detection as best-effort. + that ``transformers`` doesn't know about. Returns False for missing or + malformed configs. Unexpected read errors return True so callers apply + their trust policy; this does not itself authorize custom code execution. """ path = Path(model_name_or_path) try: @@ -184,11 +186,12 @@ def model_has_auto_map( return False except Exception as e: # Unexpected failure (network, auth, I/O). We cannot determine whether - # the model needs trust_remote_code, so conservatively return True to - # avoid workers crashing at load time. + # the model needs trust_remote_code, so return True and let the caller + # decide whether custom code execution is allowed. logger.warning( "model_has_auto_map: unexpected error reading config.json for %s: %s " - "โ€” defaulting to True (injecting --trust-remote-code).", + "โ€” auto_map detection is inconclusive; returning True for the " + "caller's trust policy check.", model_name_or_path, e, ) @@ -205,9 +208,23 @@ def model_ref_allows_implicit_trust_remote_code( remote HF model IDs are treated as mutable and must opt in explicitly. Only local directories (including PVC-resolved snapshots) qualify for implicit ``--trust-remote-code`` injection. + + Raises RuntimeError when the path cannot be inspected, rather than + treating an inaccessible local path as a remote model ID. """ path = Path(model_name_or_path) - return path.exists() and path.is_dir() + try: + return S_ISDIR(path.stat().st_mode) + except (FileNotFoundError, NotADirectoryError): + return False + except OSError as e: + raise RuntimeError( + f"Cannot inspect model path {str(path)!r} to determine whether " + f"--trust-remote-code may be enabled automatically: {e}. " + "Check directory permissions and symlink ownership for the profiler " + "user. For PVC models, set modelCache.pvcModelPath to the actual " + "snapshot directory or create the symlink with the profiler's UID." + ) from e class ModelInfo(BaseModel): diff --git a/components/src/dynamo/profiler/utils/profile_common.py b/components/src/dynamo/profiler/utils/profile_common.py index c45b06840e04..ecacef158698 100644 --- a/components/src/dynamo/profiler/utils/profile_common.py +++ b/components/src/dynamo/profiler/utils/profile_common.py @@ -19,6 +19,7 @@ import logging import os from dataclasses import dataclass, field +from stat import S_ISREG import pandas as pd @@ -164,8 +165,19 @@ def resolve_model_path(dgdr: DynamoGraphDeploymentRequestSpec) -> str: dgdr.modelCache.pvcMountPath, dgdr.modelCache.pvcModelPath, ) - if os.path.isfile(os.path.join(local_path, "config.json")): - return local_path + config_path = os.path.join(local_path, "config.json") + try: + if S_ISREG(os.stat(config_path).st_mode): + return local_path + except (FileNotFoundError, NotADirectoryError): + pass + except OSError as e: + raise RuntimeError( + f"Cannot inspect PVC model config {config_path!r}: {e}. " + "Check directory permissions and symlink ownership for the " + "profiler user, or set modelCache.pvcModelPath to an accessible " + "snapshot directory." + ) from e return dgdr.model diff --git a/components/src/dynamo/sglang/AGENTS.md b/components/src/dynamo/sglang/AGENTS.md index 0d9e87816945..47439a27ac31 100644 --- a/components/src/dynamo/sglang/AGENTS.md +++ b/components/src/dynamo/sglang/AGENTS.md @@ -30,6 +30,56 @@ support the current version plus 1 version back (N and N-1). The pattern: component files. Do not version-check with `sglang.__version__` -- import probing is more reliable since SGLang's internal layout doesn't always match the version string. +## Multi-process gateway (`--gateway-workers N`) + +`gateway.py`. One `dynamo.sglang` process fronts every DP rank of its engine: SGLang's +`TokenizerManager` intake and the Dynamo handler's token relay run on one GIL. With +`--tokenizer-worker-num N > 1` SGLang puts a `MultiTokenizerRouter` in the engine process; it +has no `generate_request`, so the leader (node rank 0) does not serve. Instead it publishes +the launch data with SGLang's shared-memory contract (`write_data_for_multi_tokenizer`) and +spawns N children, `python -m dynamo.sglang ` with `DYN_SGLANG_GATEWAY_PARENT_PID` +and `DYN_SGLANG_GATEWAY_CHILD_INDEX` set. A child builds a `TokenizerWorker` registered with +the router (`build_gateway_engine`), wraps it in `GatewayEngine` (what the handlers use of an +`sgl.Engine`: `tokenizer_manager`, `server_args`, `port_args`, `async_generate`, scheduler +info) and runs the ordinary `init_decode`/`init_prefill` path as its own endpoint instance +(`attached_engine=`, distinct from `snapshot_engine=`), so the router sees N instances per +engine. If SGLang exposes `Engine.attach_tokenizer_worker`, children use it instead of the +facade. + +Configuration (`effective_gateway_workers`, `validate_gateway_mode`, applied in `main.py` +before snapshot preparation and runtime creation): `--gateway-workers N` sets +`tokenizer_worker_num` to N; `--tokenizer-worker-num N` alone means N gateways; a tokenizer count +above 1 that differs from `--gateway-workers` is an error. Gateway mode is rejected for the direct-engine workers +(embedding, rerank, multimodal, diffusion), with `--enable-lora` (dynamic LoRA state would +live in one child), with `--enable-forward-pass-metrics` (the schedulers stamp FPM with the +non-serving leader's identity) and in snapshot mode. + +Ports and metrics: the leader gives `DYN_SYSTEM_PORT` to child 0 and runs without a system +status server; the other children bind a random system port (`DYN_SYSTEM_PORT=0`, logged by the +runtime), because any fixed offset can collide with another worker group's configured port. +Engine-target `--engine-routes` need `Engine.attach_tokenizer_worker`; the facade only exposes +the tokenizer manager. +The schedulers push KV metrics to one PULL socket: child 0 (`owns_engine_metrics()`) binds it +and re-publishes every `KvMetrics` on an ipc PUB (`metrics_fanout_endpoint()`) that the other +children subscribe to, so every gateway identity reports the engine's real KV usage. Every +child republishes KV events under its own worker id, otherwise the router would see prefixes +on instance 0 only; the cost is N copies of each KV event. On multi-node engines the +non-leader nodes resolve all N leader instances by worker group id +(`Client.wait_for_instances_by_runtime_data`) and attribute their remote-rank KV events to +each. Only the metrics owner publishes the engine-level gauges (total blocks, cache +usage, load time), so a scrape across children counts the engine once; every child still publishes +its own routing usage. Each child's model card carries `dynamo.sglang.gateway_engine` +(`host:leader_pid`) and `dynamo.sglang.gateway_workers`, so anything that counts workers or sums +per-worker capacity from discovery can collapse the N instances of one engine. Only child 0's +system port is fixed; sibling health lives on random ports, so fixed probes see one of N processes. + +Lifecycle: the leader owns the engine subprocesses and the shared memory, runs the deferred +shutdown handlers, and terminates and reaps the children with the worker's own shutdown budget +(`DYN_GRACEFUL_SHUTDOWN_GRACE_PERIOD_SECS` + drain + cleanup); a child exiting during shutdown is +not an error. Children run a parent watchdog (`PR_SET_PDEATHSIG` plus a liveness poll) and +SIGTERM themselves when the leader dies, since without the schedulers they would stay registered +and fail every request. + ## Entry Point `__main__.py` -> `main.py:main()` -> `main.py:worker()` diff --git a/components/src/dynamo/sglang/_compat.py b/components/src/dynamo/sglang/_compat.py index 090e90ed5937..fec412cb7f5a 100644 --- a/components/src/dynamo/sglang/_compat.py +++ b/components/src/dynamo/sglang/_compat.py @@ -62,6 +62,24 @@ # portion when minimum supported SGLang is 0.5.19+. _sglang_publish = None +try: + from sglang.srt.observability.req_time_stats import APIServerReqTimeStats +except ImportError: + # Fail closed for downstream SGLang builds that omit request-time + # statistics or dispatch timestamps. + APIServerReqTimeStats = None + + +def supports_disagg_prefill_cancel_anytime(engine: Any) -> bool: + """Return whether aborts can be ordered after scheduler dispatch.""" + tokenizer_manager = getattr(engine, "tokenizer_manager", None) + if not isinstance(getattr(tokenizer_manager, "rid_to_state", None), Mapping): + return False + if APIServerReqTimeStats is None: + return False + fields = getattr(APIServerReqTimeStats, "__dataclass_fields__", {}) + return "api_server_dispatch_finish_time" in fields + def get_sglang_model_config(server_args: Any) -> Any: """Return the resolved model config across SGLang ServerArgs APIs. @@ -336,6 +354,33 @@ def filter_supported_async_generate_kwargs( return {key: value for key, value in kwargs.items() if key in supported_kwarg_names} +def cache_salt_kwargs(engine: Any, cache_salt: str | None) -> dict[str, Any]: + """Preserve cache isolation, rejecting salts an older engine cannot honor.""" + if not cache_salt: + return {} + # SGLang 0.5.11 in the XPU image lacks cache_salt. Remove this check when + # the XPU pin supports the explicit cache_salt argument (0.5.18+). + kwargs = filter_supported_async_generate_kwargs(engine, {"cache_salt": cache_salt}) + if "cache_salt" not in kwargs: + raise ValueError("cache_salt is not supported by the installed SGLang engine") + return kwargs + + +def prefill_dp_rank_kwargs(engine: Any, prefill_dp_rank: Any) -> dict[str, Any]: + """Hand SGLang's decode the prefill DP rank the router already chose. + + Without ``disagg_prefill_dp_rank`` the decode scheduler parks the request + and resolves the rank over HTTP against the prefill bootstrap server, which + only learns the room once the prefill scheduler has created its KV sender; + the prefill forward cannot start before that round trip completes. + """ + if prefill_dp_rank is None: + return {} + return filter_supported_async_generate_kwargs( + engine, {"disagg_prefill_dp_rank": int(prefill_dp_rank)} + ) + + def require_reasoning_kwargs(engine: Any, request: Mapping[str, Any]) -> dict[str, Any]: """Build the optional SGLang per-request reasoning-gate argument.""" require_reasoning = bool(request.get("require_reasoning", False)) @@ -350,6 +395,7 @@ def require_reasoning_kwargs(engine: Any, request: Mapping[str, Any]) -> dict[st __all__ = [ "ConfigArgumentMerger", + "cache_salt_kwargs", "ensure_sglang_tensor_image_size", "filter_supported_async_generate_kwargs", "get_encoder_preprocessor_modules", diff --git a/components/src/dynamo/sglang/args.py b/components/src/dynamo/sglang/args.py index 769c3188e746..faeec86d271f 100644 --- a/components/src/dynamo/sglang/args.py +++ b/components/src/dynamo/sglang/args.py @@ -38,12 +38,24 @@ ensure_sglang_tensor_image_size, get_sglang_model_config, resolved_server_args, + sglang_uses_mla_backend, ) from dynamo.sglang.backend_args import DynamoSGLangArgGroup, DynamoSGLangConfig +from dynamo.sglang.elastic_ep_preflight import check_elastic_ep_backend configure_dynamo_logging() PREFILL_DECODE_DISAGGREGATION_MODE = "pd" +# Non-MLA attention backends that read the DCP parallel state in their forward +# path. `aiter` is here because SGLang itself allows dcp_size > 1 on ROCm. +DCP_CAPABLE_ATTENTION_BACKENDS = frozenset({"triton", "aiter"}) + +ATTENTION_BACKEND_CLI_FIELDS = ( + "attention_backend", + "prefill_attention_backend", + "decode_attention_backend", +) + class DynamoConfig(DynamoRuntimeConfig, DynamoSGLangConfig): """Combined configuration container for SGLang server and Dynamo args.""" @@ -66,9 +78,19 @@ def validate(self) -> None: class Config: """Combined configuration container for SGLang server and Dynamo args.""" - def __init__(self, server_args: ServerArgs, dynamo_args: DynamoConfig) -> None: + def __init__( + self, + server_args: ServerArgs, + dynamo_args: DynamoConfig, + *, + attention_backend_from_cli: bool = False, + ) -> None: self.server_args = server_args self.dynamo_args = dynamo_args + # Whether the launch named an attention backend itself. Only the CLI + # knows this; the resolved configuration reports the same field whether + # SGLang chose the value or the user did. + self.attention_backend_from_cli = attention_backend_from_cli self.serving_mode = self._set_serving_strategy() def _set_serving_strategy(self): @@ -81,8 +103,19 @@ def _set_serving_strategy(self): else: return DisaggregationMode.AGGREGATED + def validate_engine_server_args(self, server_args: Any) -> None: + """Re-run the launch checks a built engine's configuration can answer. + + Call this on ``engine.server_args`` before the engine takes a request + of any kind, including a warmup one. + """ + _validate_dcp_attention_backend( + server_args, backend_from_cli=self.attention_backend_from_cli + ) + def use_resolved_server_args(self, server_args: Any) -> Any: """Switch post-runtime Dynamo code to SGLang's resolved configuration.""" + self.validate_engine_server_args(server_args) self.server_args = resolved_server_args(server_args) return self.server_args @@ -112,6 +145,8 @@ def _unsupported_fpm_trace_role(dynamo_config: DynamoConfig) -> Optional[str]: """Return the worker role when the selected path does not create an FPM relay.""" if is_snapshot_enabled(): return "snapshot" + if dynamo_config.rerank_worker: + return "rerank" if dynamo_config.embedding_worker: return "embedding" if ( @@ -190,6 +225,84 @@ def _validate_parser_flags( sys.exit(1) +def _attention_backend_from_cli(parsed_args: Namespace) -> bool: + """Return whether the launch named an attention backend on the command line. + + The phase-specific flags take priority over --attention-backend, so any one + of the three means the backend in force was chosen by the user. + """ + return any( + getattr(parsed_args, field, None) for field in ATTENTION_BACKEND_CLI_FIELDS + ) + + +def _validate_dcp_attention_backend( + server_args: Any, *, backend_from_cli: bool +) -> None: + """Reject --dcp-size > 1 on an attention backend that never reads it. + + SGLang sizes the KV cache pool the DCP way: DCP ranks replicate the KV + heads instead of splitting them, and partition the context dimension. Only + some attention backends honor that in their forward path; the rest still do + a plain tensor-parallel head split, so the two disagree on the KV row width + and the scheduler dies on the first real request. Fail here instead, before + the worker registers and starts taking traffic. + + Call this twice: once on the arguments the CLI produced, which catches a + backend the user named, and once on the engine's own configuration, which + is the only place a backend SGLang chose for itself can be read. SGLang + 0.5.19 keeps ``ServerArgs`` at what the caller asked for and resolves in a + separate pass, so at CLI time an automatic backend is still ``None``. + """ + # Diffusion/video argument stubs and older SGLang releases omit dcp_size. + dcp_size = int(getattr(server_args, "dcp_size", 1) or 1) + if dcp_size <= 1: + return + + # MLA KV pools have no per-rank KV head split to disagree about, and every + # MLA attention backend consumes the DCP parallel state. + if sglang_uses_mla_backend(server_args): + return + + # Read the effective backend, not the flag the user typed: fa3 is the + # automatic choice for an MHA model on Hopper with no --attention-backend. + resolved = resolved_server_args(server_args) + base_backend = getattr(resolved, "attention_backend", None) + phase_backends = { + "prefill": getattr(resolved, "prefill_attention_backend", None) or base_backend, + "decode": getattr(resolved, "decode_attention_backend", None) or base_backend, + } + + unsupported = sorted( + (phase, backend) + for phase, backend in phase_backends.items() + # A backend SGLang has not decided yet is unknown, not unsupported. At + # CLI time that is every automatic backend; the engine's configuration + # carries the decision, and this runs again there. + if backend is not None and backend not in DCP_CAPABLE_ATTENTION_BACKENDS + ) + if not unsupported: + return + + named = ", ".join( + f"{phase} attention backend '{backend}'" for phase, backend in unsupported + ) + automatic = ( + "" + if backend_from_cli + else " SGLang selected it automatically because no attention backend was passed." + ) + raise ValueError( + f"--dcp-size {dcp_size} is not supported with the {named}.{automatic} " + "Decode context parallel replicates the KV heads across DCP ranks when " + "it sizes the KV cache pool, but this backend still splits them across " + "tensor-parallel ranks, so the worker crashes on its first request. " + "Use --dcp-size 1, or an attention backend that implements decode " + "context parallel (--attention-backend triton), or an MLA model with " + "one of SGLang's MLA attention backends." + ) + + def _has_cli_flag(args: list[str], flag: str) -> bool: """Return True when a CLI flag is present in '--flag val' or '--flag=val' form.""" return any(arg == flag or arg.startswith(f"{flag}=") for arg in args) @@ -441,6 +554,16 @@ async def parse_args(args: list[str]) -> Config: args_dict["disaggregation_bootstrap_port"] = bootstrap_port parsed_args = Namespace(**args_dict) + # Read off the parsed flags rather than ServerArgs: ServerArgs.from_cli_args + # downloads the model, and the diffusion and video paths build a stub that + # carries neither flag, so both would leave this unchecked. parsed_args + # comes from ServerArgs.add_cli_args, which declares both options across + # the supported SGLang releases. + check_elastic_ep_backend( + parsed_args.elastic_ep_backend, + parsed_args.enable_dp_attention, + ) + # Dynamo argument processing # If an endpoint is provided, validate and use it # otherwise fall back to default endpoints @@ -450,9 +573,14 @@ async def parse_args(args: list[str]) -> Config: if dynamo_config.enable_multimodal: parsed_args.enable_multimodal = True - # If --embedding-worker is set, also set SGLang's --is-embedding flag - if dynamo_config.embedding_worker: + # Both dedicated pooling modes use SGLang's embedding engine. + if dynamo_config.embedding_worker or dynamo_config.rerank_worker: parsed_args.is_embedding = True + if dynamo_config.rerank_worker and ( + parsed_args.disaggregation_mode != "null" + or getattr(parsed_args, "dllm_algorithm", None) + ): + raise ValueError("--rerank-worker requires aggregated cross-encoder serving") # Enable encoder_only mode for multimodal encode workers to load only vision encoder # This significantly reduces memory usage by avoiding loading the full LLM weights @@ -461,7 +589,9 @@ async def parse_args(args: list[str]) -> Config: endpoint = dynamo_config.endpoint if endpoint is None: - if dynamo_config.embedding_worker: + if dynamo_config.rerank_worker: + endpoint = f"dyn://{namespace}.rerank.generate" + elif dynamo_config.embedding_worker: endpoint = f"dyn://{namespace}.backend.generate" elif dynamo_config.image_diffusion_worker: endpoint = f"dyn://{namespace}.backend.generate" @@ -644,6 +774,9 @@ async def parse_args(args: list[str]) -> Config: "values are always higher priority at the API layer." ) + backend_from_cli = _attention_backend_from_cli(parsed_args) + _validate_dcp_attention_backend(server_args, backend_from_cli=backend_from_cli) + if dynamo_config.use_sglang_tokenizer: warnings.warn( "--use-sglang-tokenizer is deprecated and will be removed in a future " @@ -684,7 +817,9 @@ async def parse_args(args: list[str]) -> Config: logging.debug(f"Dynamo configs: {dynamo_config}") - return Config(server_args, dynamo_config) + return Config( + server_args, dynamo_config, attention_backend_from_cli=backend_from_cli + ) @contextlib.contextmanager diff --git a/components/src/dynamo/sglang/backend_args.py b/components/src/dynamo/sglang/backend_args.py index 8c7c051dcda6..f5e269428617 100644 --- a/components/src/dynamo/sglang/backend_args.py +++ b/components/src/dynamo/sglang/backend_args.py @@ -67,6 +67,22 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: "the same SGLang-native pre/post processing with KV router support.", ) + add_argument( + g, + flag_name="--gateway-workers", + env_var="DYN_SGL_GATEWAY_WORKERS", + default=None, + arg_type=int, + help="Run N gateway processes in front of this engine, each a Dynamo " + "endpoint instance with its own SGLang request gateway (SGLang calls " + "that process a tokenizer worker; it also handles request intake and " + "output relay, which is the work being spread). Sets SGLang's " + "--tokenizer-worker-num to N; a --tokenizer-worker-num above 1 that " + "differs from N is an error, and --tokenizer-worker-num N alone runs N " + "gateways as well. Decode and prefill LLM workers only; not with " + "--enable-lora, --enable-forward-pass-metrics or snapshot mode. Child 0 " + "takes DYN_SYSTEM_PORT, the other children bind a random system port.", + ) add_negatable_bool_argument( g, flag_name="--enable-multimodal", @@ -109,6 +125,14 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: help="Run as embedding worker component (Dynamo flag, also sets SGLang's --is-embedding).", ) + add_negatable_bool_argument( + g, + flag_name="--rerank-worker", + env_var="DYN_SGL_RERANK_WORKER", + default=False, + help="Run a dedicated text cross-encoder rerank worker. Requires a rerank-capable frontend; also sets SGLang's --is-embedding.", + ) + add_negatable_bool_argument( g, flag_name="--image-diffusion-worker", @@ -179,6 +203,7 @@ class DynamoSGLangConfig(ConfigBase): """Configuration for Dynamo SGLang wrapper (SGLang-specific only).""" use_sglang_tokenizer: bool + gateway_workers: Optional[int] = None # Internal roles derived from the canonical multimodal arguments in args.py. multimodal_encode_worker: bool = False multimodal_worker: bool = False @@ -186,6 +211,7 @@ class DynamoSGLangConfig(ConfigBase): dedicated_mm_encoder: bool = False embedding_transfer_mode: EmbeddingTransferMode embedding_worker: bool + rerank_worker: bool = False image_diffusion_worker: bool disagg_config: Optional[str] = None @@ -203,6 +229,17 @@ class DynamoSGLangConfig(ConfigBase): def validate(self) -> None: _reject_removed_multimodal_env_vars() + if self.rerank_worker and ( + self.embedding_worker + or self.image_diffusion_worker + or self.video_generation_worker + or self.enable_multimodal + or self.dedicated_mm_encoder + or self.frontend_decoding + ): + raise ValueError( + "--rerank-worker cannot be combined with other worker modes" + ) if not isinstance(self.embedding_transfer_mode, EmbeddingTransferMode): self.embedding_transfer_mode = EmbeddingTransferMode( @@ -214,6 +251,8 @@ def validate(self) -> None: "Both 'disagg_config' and 'disagg_config_key' must be provided together." ) + if self.gateway_workers is not None and self.gateway_workers < 1: + raise ValueError("--gateway-workers must be a positive integer") self.validate_multimodal_topology() self.validate_dedicated_mm_encoder() diff --git a/components/src/dynamo/sglang/elastic_ep_preflight.py b/components/src/dynamo/sglang/elastic_ep_preflight.py new file mode 100644 index 000000000000..215d60acccd2 --- /dev/null +++ b/components/src/dynamo/sglang/elastic_ep_preflight.py @@ -0,0 +1,290 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validate SGLang's mooncake elastic-EP backend before model loading. + +The check avoids importing ``sglang`` or eagerly importing ``torch`` so it can +diagnose a broken engine image. +""" + +from __future__ import annotations + +import importlib +import importlib.util +import os +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as _distribution_version +from typing import Dict, List, Optional, Tuple + +MOONCAKE_BACKEND = "mooncake" + +# The torch backend names SGLang registers its elastic-EP process groups under +# (sglang/srt/distributed/parallel_state.py): the CUDA device group and the CPU +# group it uses for metadata collectives. The extension must register both with +# their expected device capabilities, or engine startup still fails later. +_REQUIRED_TORCH_BACKENDS = {"mooncake": "cuda", "mooncake-cpu": "cpu"} + +# mooncake renamed this extension ``mooncake.ep`` -> ``mooncake.pg``; SGLang +# v0.5.16 imports the old name and v0.5.18 the new one. +_PROCESS_GROUP_MODULES = ("mooncake.pg", "mooncake.ep") + +# Where SGLang names that extension. Read as text: importing them to ask would +# need the working engine this check exists to doubt, and would pull in torch +# and CUDA before it can answer. +_SGLANG_ELASTIC_EP_SOURCES = ( + "srt/elastic_ep/elastic_ep.py", + "srt/distributed/parallel_state.py", +) + +# Both distributions ship the same extension; only one is normally installed. +_MOONCAKE_DISTRIBUTIONS = ( + "mooncake-transfer-engine-cuda13", + "mooncake-transfer-engine", +) + +# Compiled ProcessGroup extensions are named for the exact torch release they +# were built against, e.g. ``pg_2_11_0`` / ``ep_2_9_1``. +_EXTENSION_PREFIXES = ("pg_", "ep_") + + +def _required_process_group_modules() -> Tuple[str, ...]: + """The extension names the installed SGLang actually imports. + + An image can hold the wrong half of the ``mooncake.ep`` -> ``mooncake.pg`` + rename: the wheel imports cleanly and the engine still cannot start, + because it asks for the other name. Reading the engine's own import line + pins which name has to work. Falls back to accepting either when those + sources cannot be read, since refusing a worker over an upstream file move + would be worse than the late failure this check replaces. + """ + # find_spec raises ImportError when a parent package fails to import and + # ValueError when the module is already in sys.modules without a spec. + # Both mean "cannot tell", which is what the widened fallback answers. Any + # other error is not something this fallback was reasoned about, so it + # propagates rather than silently widening the check. + try: + spec = importlib.util.find_spec("sglang") + except (ImportError, ValueError): + return _PROCESS_GROUP_MODULES + if spec is None or not spec.submodule_search_locations: + return _PROCESS_GROUP_MODULES + + required: List[str] = [] + for location in spec.submodule_search_locations: + for relative in _SGLANG_ELASTIC_EP_SOURCES: + try: + with open(os.path.join(location, relative), encoding="utf-8") as handle: + source = handle.read() + except OSError: + continue + for module_name in _PROCESS_GROUP_MODULES: + if module_name in source and module_name not in required: + required.append(module_name) + return tuple(required) or _PROCESS_GROUP_MODULES + + +def _import_process_group_extension() -> Optional[str]: + """Return ``None`` when the extension imports, else the collected failures. + + Importing it is what SGLang itself does later; doing it here surfaces the + torch-version skew (``mooncake.pg`` raises ``ImportError`` when no + ``pg_`` module matches) before any GPU work has started. + """ + failures: List[str] = [] + for module_name in _required_process_group_modules(): + try: + importlib.import_module(module_name) + return None + # Deliberately broad. This is the one place that runs the compiled + # ProcessGroup extension's own module initialization, and the question + # being asked is "does SGLang's import of this succeed", not "does it + # fail in a way we predicted". A mismatched build surfaces as + # ImportError, but the same import also runs C++ static initializers + # and dlopens CUDA, which report as OSError, RuntimeError, or a + # torch-defined exception depending on the wheel. Narrowing here would + # let the unlisted ones escape as a raw traceback from argument parsing + # instead of the diagnostic below, which is the failure mode this whole + # module exists to remove. The exception type is kept verbatim in the + # message so nothing about the real cause is lost. + except Exception as exc: # noqa: BLE001 + failures.append(f"{module_name}: {type(exc).__name__}: {exc}") + return "; ".join(failures) if failures else "no attempt was made" + + +def _installed_mooncake_versions() -> Dict[str, str]: + """Map each installed mooncake distribution to its version.""" + found: Dict[str, str] = {} + for distribution in _MOONCAKE_DISTRIBUTIONS: + try: + found[distribution] = _distribution_version(distribution) + except PackageNotFoundError: + continue + return found + + +def _available_extension_modules() -> List[str]: + """Names of the compiled ProcessGroup extensions the wheel actually ships. + + Resolved from the package directory rather than by import, so this stays + cheap and side-effect free when the extension is the thing that is broken. + """ + try: + spec = importlib.util.find_spec("mooncake") + except (ImportError, ValueError): + # Same two failure modes as the SGLang resolver above. Either way the + # wheel's layout cannot be listed, which this reports as "none found". + return [] + if spec is None or not spec.submodule_search_locations: + return [] + + modules = set() + for location in spec.submodule_search_locations: + try: + entries = os.listdir(location) + except OSError: + continue + for entry in entries: + if not entry.startswith(_EXTENSION_PREFIXES): + continue + modules.add(entry.split(".", 1)[0]) + return sorted(modules) + + +def _registered_torch_backends() -> Optional[Dict[str, Tuple[str, ...]]]: + """torch's backend -> supported-device map, or ``None`` when unreadable. + + ``Backend.backend_capability`` is a torch internal. Reporting "could not + read the registry" as "mooncake is not registered" would let a rename in + somebody else's private attribute refuse every worker that asked for this + transport, which is more fatal than the failure being guarded against. + ``None`` keeps the two apart so the caller can skip only this assertion. + """ + # ImportError covers a torch that is absent or does not import; + # AttributeError covers the rename this docstring is about. Both are + # "unreadable". torch types the mapping's values as lists, so copy them + # into tuples rather than asserting the annotation over torch's. + try: + import torch.distributed as torch_distributed + + return { + name: tuple(devices) + for name, devices in torch_distributed.Backend.backend_capability.items() + } + except (ImportError, AttributeError): + return None + + +def _torch_version() -> str: + """The running torch version, or why it could not be determined. + + Only a genuinely absent distribution may be reported as "not installed": a + torch that is present but fails to load is a different problem, and an + operator reading this diagnostic needs to see which one they have. + """ + try: + import torch + + return str(torch.__version__) + except ModuleNotFoundError as exc: + if exc.name == "torch": + return "not installed" + return f"installed but not importable: {type(exc).__name__}: {exc}" + # A torch that is installed but unusable raises from its own loader, most + # often ImportError or OSError for a missing CUDA runtime library. + except (ImportError, OSError) as exc: + return f"installed but not importable: {type(exc).__name__}: {exc}" + + +def _format_versions(versions: Dict[str, str]) -> str: + if not versions: + return "none installed" + return ", ".join(f"{name} {value}" for name, value in sorted(versions.items())) + + +def _build_diagnostic( + import_failure: Optional[str], + registered_backends: Dict[str, Tuple[str, ...]], + missing_backends: Tuple[str, ...], + enable_dp_attention: bool, +) -> str: + mooncake_backends = { + name: devices + for name, devices in registered_backends.items() + if MOONCAKE_BACKEND in name + } + lines = [ + "--elastic-ep-backend mooncake was requested, but the mooncake torch " + "ProcessGroup backend is not usable in this image. SGLang builds its " + "elastic-EP process groups from that backend after the model is " + "loaded, so leaving this unchecked fails the worker late and from " + "inside engine code.", + f" mooncake ProcessGroup extension: {import_failure or 'imported, but registered no usable backend'}", + f" installed mooncake distributions: {_format_versions(_installed_mooncake_versions())}", + f" torch: {_torch_version()}", + f" ProcessGroup extensions this mooncake ships: {', '.join(_available_extension_modules()) or 'none found'}", + f" mooncake backends registered with torch.distributed: {mooncake_backends or 'none'}", + ] + if missing_backends: + lines.append( + " required by SGLang but missing or unavailable for its expected device: " + f"{', '.join(missing_backends)}" + ) + if enable_dp_attention: + lines.append( + " --enable-dp-attention is also set: DP attention synchronizes its " + "MLP batch metadata with an all-gather over this group on the very " + "first forward pass, so this combination reaches the backend " + "immediately at startup." + ) + lines.append( + "Install a mooncake-transfer-engine build whose ProcessGroup extension " + "matches the running torch, or drop --elastic-ep-backend." + ) + return "\n".join(lines) + + +def check_elastic_ep_backend( + elastic_ep_backend: Optional[str], + enable_dp_attention: bool = False, +) -> None: + """Reject ``--elastic-ep-backend mooncake`` when the image cannot serve it. + + Silent for every other configuration, including an unset backend: this must + never turn into an unconditional startup failure for workers that asked for + nothing of the sort. + + Raises: + ValueError: The mooncake backend was requested but its torch + ProcessGroup extension does not import, or imports without + registering both of the backends SGLang asks torch for with their + required device capabilities. + """ + if not elastic_ep_backend: + return + if str(elastic_ep_backend).strip().lower() != MOONCAKE_BACKEND: + return + + import_failure = _import_process_group_extension() + registered_backends: Dict[str, Tuple[str, ...]] = {} + missing_backends: Tuple[str, ...] = () + if import_failure is None: + readable_backends = _registered_torch_backends() + if readable_backends is None: + return + missing_backends = tuple( + name + for name, required_device in _REQUIRED_TORCH_BACKENDS.items() + if required_device not in readable_backends.get(name, ()) + ) + if not missing_backends: + return + registered_backends = readable_backends + + raise ValueError( + _build_diagnostic( + import_failure, + registered_backends, + missing_backends, + enable_dp_attention, + ) + ) diff --git a/components/src/dynamo/sglang/engine_generate.py b/components/src/dynamo/sglang/engine_generate.py index daa33f3dad42..a3bff59b3b8d 100644 --- a/components/src/dynamo/sglang/engine_generate.py +++ b/components/src/dynamo/sglang/engine_generate.py @@ -7,6 +7,7 @@ from collections.abc import AsyncIterator, Mapping from typing import Any +from uuid import uuid4 from pydantic import TypeAdapter from sglang.srt.managers.io_struct import GenerateReqInput @@ -14,10 +15,16 @@ from dynamo.common.backend import logprobs as _shared_logprobs SGLANG_GENERATE_CAPABILITY = "sglang_generate" +DISAGG_PREFILL_CANCEL_ANYTIME_V1 = "disagg_prefill_cancel_anytime_v1" _PAYLOAD_KEY = "sglang_tito" _GENERATE_REQUEST_ADAPTER = TypeAdapter(GenerateReqInput) +def new_sglang_request_id() -> str: + """Return an unguessable fixed-width ID used only inside SGLang.""" + return uuid4().hex + + def native_generate_payload( request: Mapping[str, Any], ) -> Mapping[str, Any] | None: @@ -33,7 +40,7 @@ def build_native_generate_request( native_payload: Mapping[str, Any], *, input_ids: list[int], - fallback_rid: str, + request_id: str, priority: int | None, sampling_overrides: Mapping[str, Any] | None = None, bootstrap_host: str | None = None, @@ -41,7 +48,9 @@ def build_native_generate_request( bootstrap_room: int | None = None, external_trace_header: dict[str, str] | None = None, routed_dp_rank: int | None = None, + prefill_dp_rank: int | None = None, lora_path: str | None = None, + cache_salt: str | None = None, ) -> GenerateReqInput: """Reconstruct the installed SGLang version native request. @@ -49,10 +58,27 @@ def build_native_generate_request( ``extra_args.sglang_tito``. Dynamo replaces only canonical input, routing state, and fields supplied by the selected worker. SGLang owns all remaining validation. + + ``request_id`` is router-owned and always replaces a caller-supplied + ``rid`` so cancellation cannot target another request. + Native session requests pass their caller-visible node ID as the fallback + because SGLang keys session continuations by the submitted top-level + request ID. """ payload = dict(native_payload) + salt = cache_salt if cache_salt is not None else payload.get("cache_salt") + if salt is None or salt == "": + payload.pop("cache_salt", None) + else: + # The separately pinned XPU release (0.5.11) silently ignores unknown + # dataclass fields. Remove this check when that pin supports cache_salt. + if "cache_salt" not in GenerateReqInput.__dataclass_fields__: + raise ValueError( + "cache_salt is not supported by the installed SGLang engine" + ) + payload["cache_salt"] = salt payload["input_ids"] = input_ids - payload["rid"] = payload.get("rid") or fallback_rid + payload["rid"] = request_id payload["stream"] = True if priority is None: payload.pop("priority", None) @@ -80,6 +106,11 @@ def build_native_generate_request( ): if value is not None: payload[name] = value + if ( + prefill_dp_rank is not None + and "disagg_prefill_dp_rank" in GenerateReqInput.__dataclass_fields__ + ): + payload["disagg_prefill_dp_rank"] = prefill_dp_rank native_request = _GENERATE_REQUEST_ADAPTER.validate_python(payload) _shared_logprobs.validate_sglang_top_logprobs( diff --git a/components/src/dynamo/sglang/gateway.py b/components/src/dynamo/sglang/gateway.py new file mode 100644 index 000000000000..a0610b06377d --- /dev/null +++ b/components/src/dynamo/sglang/gateway.py @@ -0,0 +1,370 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multi-process request gateway for the SGLang worker: the leader keeps the engine +and spawns N ``dynamo.sglang`` children that each serve requests through their own +SGLang ``TokenizerWorker``. Design notes: AGENTS.md, "Multi-process gateway". +""" + +from __future__ import annotations + +import asyncio +import ctypes +import logging +import os +import signal +import socket +import subprocess +import sys +import tempfile +import threading +import types +from typing import Awaitable, Callable, Optional + +import sglang as sgl + +from dynamo.common.snapshot.constants import SNAPSHOT_CONTROL_DIR_ENV +from dynamo.common.utils.graceful_shutdown import get_grace_period_seconds + +ENV_PARENT_PID = "DYN_SGLANG_GATEWAY_PARENT_PID" +ENV_CHILD_INDEX = "DYN_SGLANG_GATEWAY_CHILD_INDEX" +ENV_SYSTEM_PORT = "DYN_SYSTEM_PORT" +ENV_SYSTEM_PORT_BASE = "DYN_SGLANG_GATEWAY_SYSTEM_PORT" +ENV_LOAD_TIME = "DYN_SGLANG_GATEWAY_LOAD_TIME_S" +# Runtime-data keys on every gateway child's model card: consumers that count +# workers or sum per-worker capacity can collapse the N instances of one engine. +GATEWAY_ENGINE_ID_KEY = "dynamo.sglang.gateway_engine" +GATEWAY_WORKERS_KEY = "dynamo.sglang.gateway_workers" +# Children drain like any worker (grace + drain + cleanup); give them that budget. +CHILD_DRAIN_AND_CLEANUP_SECS = 60.0 + +DIRECT_ENGINE_WORKER_FLAGS = ( + "image_diffusion_worker", + "video_generation_worker", + "rerank_worker", + "embedding_worker", + "multimodal_encode_worker", + "multimodal_worker", + "diffusion_worker", +) + + +class GatewayEngine: + """What the request handlers need from an ``sgl.Engine``. + + A child never owns scheduler processes; it only holds a ``TokenizerWorker`` bound + to the parent's router. ``Engine.async_generate`` only touches + ``self.tokenizer_manager``, so it can be reused unchanged. + """ + + def __init__(self, tokenizer_manager, server_args, port_args, scheduler_info): + self.tokenizer_manager = tokenizer_manager + self.server_args = server_args + self.port_args = port_args + self._scheduler_init_result = types.SimpleNamespace( + scheduler_infos=[scheduler_info] + ) + + async_generate = sgl.Engine.async_generate + _resolve_routed_dp_rank = sgl.Engine._resolve_routed_dp_rank + + def shutdown(self): + pass + + +def is_gateway_child() -> bool: + return ENV_PARENT_PID in os.environ + + +def gateway_child_index() -> int: + return int(os.environ.get(ENV_CHILD_INDEX, "0")) + + +def owns_engine_metrics() -> bool: + """SGLang's schedulers push KV metrics to one PULL socket, so exactly one gateway + process may consume them: child 0 (or the single worker when the mode is off).""" + return gateway_child_index() == 0 + + +def effective_gateway_workers(server_args, dynamo_args) -> int: + """Gateway count implied by ``--gateway-workers`` and ``--tokenizer-worker-num``. + + The engine runs one SGLang tokenizer worker per gateway process. A tokenizer + count above 1 that differs from ``--gateway-workers`` is an error rather than a + silent choice; the default of 1 lets ``--gateway-workers N`` stand alone.""" + requested = getattr(dynamo_args, "gateway_workers", None) + tokenizer_workers = getattr(server_args, "tokenizer_worker_num", 1) or 1 + if requested is None: + return tokenizer_workers + if tokenizer_workers > 1 and tokenizer_workers != requested: + raise ValueError( + f"--gateway-workers {requested} conflicts with --tokenizer-worker-num " + f"{tokenizer_workers}: each gateway process is one SGLang tokenizer " + "worker, so drop --tokenizer-worker-num or give both the same value" + ) + return requested + + +def gateway_worker_count(server_args, dynamo_args) -> int: + if (getattr(server_args, "node_rank", 0) or 0) != 0 or is_gateway_child(): + return 1 + return effective_gateway_workers(server_args, dynamo_args) + + +def validate_gateway_mode(server_args, dynamo_args, count: int) -> None: + if count <= 1: + return + direct = [f for f in DIRECT_ENGINE_WORKER_FLAGS if getattr(dynamo_args, f, False)] + if direct: + raise ValueError( + "gateway mode (--gateway-workers / --tokenizer-worker-num > 1) is only " + "supported by the decode and prefill LLM workers, not with " + f"--{direct[0].replace('_', '-')}" + ) + if getattr(server_args, "enable_lora", False): + raise ValueError( + "gateway mode is not supported with --enable-lora: dynamic LoRA state " + "lives in each gateway process" + ) + if getattr(server_args, "enable_forward_pass_metrics", False): + raise ValueError( + "gateway mode is not supported with --enable-forward-pass-metrics: the " + "schedulers stamp forward-pass metrics with the identity of the process " + "that created the engine, which serves no requests in gateway mode" + ) + if os.environ.get(SNAPSHOT_CONTROL_DIR_ENV): + raise ValueError( + "gateway mode is not supported in snapshot mode " + f"({SNAPSHOT_CONTROL_DIR_ENV} is set): snapshot warmup needs a single " + "tokenizer manager" + ) + routes = getattr(dynamo_args, "engine_routes", None) + if routes and getattr(sgl.Engine, "attach_tokenizer_worker", None) is None: + from dynamo.sglang.engine_routes import parse_engine_route_descriptors + + engine_targets = [ + d.path + for d in parse_engine_route_descriptors(routes) + if d.target == "engine" + ] + if engine_targets: + raise ValueError( + "gateway mode with engine-target --engine-routes " + f"({', '.join(engine_targets)}) needs an SGLang with " + "Engine.attach_tokenizer_worker; on this SGLang the gateway children " + "only expose the tokenizer manager" + ) + + +def reserve_system_port_for_children() -> None: + """Leader side, before the runtime starts. The leader serves no requests, so the + configured ``DYN_SYSTEM_PORT`` goes to child 0; the leader runs without a system + status server and the other children bind a random port (0), because any fixed + offset can collide with another worker group's configured port.""" + raw = os.environ.get(ENV_SYSTEM_PORT) + try: + port = int(raw) if raw is not None else -1 + except ValueError: + return + if port <= 0: + return + os.environ[ENV_SYSTEM_PORT_BASE] = str(port) + os.environ[ENV_SYSTEM_PORT] = "-1" + + +def child_environment(index: int, load_time: Optional[float] = None) -> dict[str, str]: + env = { + **os.environ, + ENV_PARENT_PID: str(os.getpid()), + ENV_CHILD_INDEX: str(index), + } + base = env.pop(ENV_SYSTEM_PORT_BASE, None) + if base is not None: + env[ENV_SYSTEM_PORT] = base if index == 0 else "0" + env.pop(ENV_LOAD_TIME, None) + if index == 0 and load_time is not None: + # The leader measured the load; the metrics owner publishes it once. + env[ENV_LOAD_TIME] = repr(load_time) + return env + + +def attached_engine_load_time() -> Optional[float]: + raw = os.environ.get(ENV_LOAD_TIME) + return float(raw) if raw else None + + +def follow_pause_broadcasts( + tokenizer_manager, on_change: Callable[[], Awaitable[None]] +) -> bool: + """A pause or resume issued through any sibling reaches this process only as the + router's broadcast into its ``TokenizerWorker``; run ``on_change`` after each one so + this child's discovery registration follows the shared state.""" + apply = getattr(tokenizer_manager, "_apply_pause_continue_broadcast", None) + if apply is None: + return False + + async def apply_and_follow(obj): + await apply(obj) + await on_change() + + tokenizer_manager._apply_pause_continue_broadcast = apply_and_follow + return True + + +def gateway_engine_id() -> Optional[str]: + pid = os.environ.get(ENV_PARENT_PID) + if pid is None: + return None + return f"{socket.gethostname()}:{pid}" + + +def child_shutdown_timeout() -> float: + return get_grace_period_seconds() + CHILD_DRAIN_AND_CLEANUP_SECS + + +def start_parent_watchdog( + parent_pid: int, + on_parent_death: Optional[Callable[[], None]] = None, + poll_seconds: float = 2.0, +) -> threading.Thread: + """The leader owns the schedulers; a child that outlives it would stay registered + and fail every request. Ask the kernel for SIGTERM on parent death and poll as a + fallback for the window before the flag is set and for non-Linux hosts.""" + if on_parent_death is None: + + def on_parent_death() -> None: + logging.error( + "gateway parent pid=%d is gone; shutting down child pid=%d", + parent_pid, + os.getpid(), + ) + os.kill(os.getpid(), signal.SIGTERM) + + try: + libc = ctypes.CDLL("libc.so.6", use_errno=True) + libc.prctl(1, signal.SIGTERM) # PR_SET_PDEATHSIG + except (OSError, AttributeError): + pass + + def poll() -> None: + while True: + try: + os.kill(parent_pid, 0) + except ProcessLookupError: + on_parent_death() + return + except PermissionError: + pass + if os.getppid() != parent_pid: + on_parent_death() + return + threading.Event().wait(poll_seconds) + + thread = threading.Thread(target=poll, name="gateway-parent-watchdog", daemon=True) + thread.start() + return thread + + +def metrics_fanout_endpoint() -> Optional[str]: + """Where child 0 re-publishes the schedulers' KV metrics for its siblings.""" + pid = os.environ.get(ENV_PARENT_PID) + if pid is None: + return None + return f"ipc://{tempfile.gettempdir()}/dynamo_sglang_gateway_metrics_{pid}" + + +def build_gateway_engine(): + """Child side: join the parent's engine through a ``TokenizerWorker``.""" + parent_pid = int(os.environ[ENV_PARENT_PID]) + start_parent_watchdog(parent_pid) + attach = getattr(sgl.Engine, "attach_tokenizer_worker", None) + if attach is not None: + engine = attach(parent_pid) + logging.info( + "gateway child pid=%d attached via Engine.attach_tokenizer_worker", + os.getpid(), + ) + return engine + + from sglang.srt.managers.multi_tokenizer_mixin import ( + get_tokenizer_worker_class, + read_from_shared_memory, + ) + from sglang.srt.runtime_context import publish + + port_args, server_args, scheduler_info = read_from_shared_memory( + f"multi_tokenizer_args_{parent_pid}" + ) + publish(server_args, role="tokenizer") + port_args.tokenizer_ipc_name = ( + f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}" + ) + tm = get_tokenizer_worker_class(server_args)(server_args, port_args) + tm.max_req_input_len = scheduler_info["max_req_input_len"] + tm.set_startup_time(scheduler_info["startup_time"]) + logging.info( + "gateway child pid=%d registered TokenizerWorker ipc=%s", + os.getpid(), + port_args.tokenizer_ipc_name, + ) + return GatewayEngine(tm, server_args, port_args, scheduler_info) + + +def _reap(proc: subprocess.Popen, timeout: Optional[float] = None) -> None: + try: + proc.wait(timeout=child_shutdown_timeout() if timeout is None else timeout) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + + +async def serve_via_gateway_children( + engine, + count: int, + shutdown_event: asyncio.Event, + load_time: Optional[float] = None, +) -> None: + from sglang.srt.managers.multi_tokenizer_mixin import write_data_for_multi_tokenizer + + shm = getattr(engine, "_multi_tokenizer_shm", None) + owns_shm = shm is None + if shm is None: + scheduler_info = { + **engine._scheduler_init_result.scheduler_infos[0], + "startup_time": engine.tokenizer_manager.startup_time, + } + shm = write_data_for_multi_tokenizer( + engine.port_args, engine.server_args, scheduler_info + ) + argv = sys.argv[1:] + procs: list[subprocess.Popen] = [] + try: + for index in range(count): + procs.append( + subprocess.Popen( + [sys.executable, "-m", "dynamo.sglang", *argv], + env=child_environment(index, load_time), + ) + ) + logging.info( + "gateway parent pid=%d spawned %d children: %s", + os.getpid(), + count, + [p.pid for p in procs], + ) + while not shutdown_event.is_set(): + await asyncio.sleep(2) + dead = [p for p in procs if p.poll() is not None] + if dead and not shutdown_event.is_set(): + raise RuntimeError( + f"gateway child pid={dead[0].pid} exited rc={dead[0].returncode}" + ) + finally: + for p in procs: + if p.poll() is None: + p.terminate() + await asyncio.gather(*(asyncio.to_thread(_reap, p) for p in procs)) + if owns_shm: + try: + shm.unlink() + except FileNotFoundError: + pass diff --git a/components/src/dynamo/sglang/health_check.py b/components/src/dynamo/sglang/health_check.py index 1b1296dd2095..df1ad3717a19 100644 --- a/components/src/dynamo/sglang/health_check.py +++ b/components/src/dynamo/sglang/health_check.py @@ -87,6 +87,36 @@ def __init__( super().__init__() +class SglangEmbeddingHealthCheckPayload(HealthCheckPayload): + """Send an embedding request through the worker's normal encode path.""" + + def __init__( + self, + model_name: str, + engine: Optional[sgl.Engine] = None, + use_text_input: bool = False, + ) -> None: + self.default_payload = { + "model": model_name, + "input": ( + "Test" if use_text_input else [_get_bos_token_id_from_engine(engine)] + ), + } + super().__init__() + + +class SglangRerankHealthCheckPayload(HealthCheckPayload): + """Probe the cross-encoder path with a single query/document pair.""" + + def __init__(self, model: str) -> None: + self.default_payload = { + "model": model, + "query": "health check", + "documents": ["health check"], + } + super().__init__() + + class SglangDisaggHealthCheckPayload(HealthCheckPayload): """SGLang-specific health check payload for PD-disaggregated mode. diff --git a/components/src/dynamo/sglang/init_embedding.py b/components/src/dynamo/sglang/init_embedding.py index 41f3e0ae9155..a423fe607f0e 100644 --- a/components/src/dynamo/sglang/init_embedding.py +++ b/components/src/dynamo/sglang/init_embedding.py @@ -12,13 +12,16 @@ from dynamo.llm import ModelInput, ModelType, WorkerType from dynamo.runtime import DistributedRuntime from dynamo.sglang.args import Config -from dynamo.sglang.health_check import SglangHealthCheckPayload +from dynamo.sglang.health_check import ( + SglangEmbeddingHealthCheckPayload, + SglangRerankHealthCheckPayload, +) from dynamo.sglang.publisher import ( set_forward_pass_metrics_worker_id, setup_sgl_metrics, ) from dynamo.sglang.register import register_model_with_readiness_gate -from dynamo.sglang.request_handlers import EmbeddingWorkerHandler +from dynamo.sglang.request_handlers import EmbeddingWorkerHandler, RerankWorkerHandler from dynamo.sglang.request_handlers.embedding.metrics import init_embedding_metrics @@ -30,6 +33,25 @@ async def init_embedding( run_deferred_handlers: Callable[[], Awaitable[None]] | None = None, ) -> None: """Initialize embedding worker component""" + await _init_pooling( + runtime, + config, + shutdown_event, + shutdown_endpoints, + run_deferred_handlers, + rerank=False, + ) + + +async def _init_pooling( + runtime: DistributedRuntime, + config: Config, + shutdown_event: asyncio.Event, + shutdown_endpoints: list, + run_deferred_handlers: Callable[[], Awaitable[None]] | None = None, + *, + rerank: bool, +) -> None: server_args, dynamo_args = config.server_args, config.dynamo_args generate_endpoint = runtime.endpoint( @@ -46,35 +68,43 @@ async def init_embedding( engine, config, generate_endpoint ) - # Wire ``dynamo_embedding_*`` histograms to the worker's /metrics - # endpoint via a dedicated CollectorRegistry โ€” the default Prometheus - # global registry is not exposed by the Dynamo SGLang publisher. - # Done AFTER engine init (and after setup_sgl_metrics) so the lazy - # ``prometheus_client`` import in metrics.py doesn't interfere with - # SGLang's multiprocess Prometheus setup. - from prometheus_client import CollectorRegistry + if not rerank: + # Wire ``dynamo_embedding_*`` histograms to the worker's /metrics + # endpoint via a dedicated CollectorRegistry โ€” the default Prometheus + # global registry is not exposed by the Dynamo SGLang publisher. + # Done AFTER engine init (and after setup_sgl_metrics) so the lazy + # ``prometheus_client`` import in metrics.py doesn't interfere with + # SGLang's multiprocess Prometheus setup. + from prometheus_client import CollectorRegistry - embedding_metrics_registry = CollectorRegistry() - register_engine_metrics_callback( - endpoint=generate_endpoint, - registry=embedding_metrics_registry, - metric_prefix_filters=["dynamo_embedding_"], - namespace_name=dynamo_args.namespace, - component_name=dynamo_args.component, - endpoint_name=dynamo_args.endpoint, - model_name=server_args.served_model_name, - ) - init_embedding_metrics(embedding_metrics_registry) + embedding_metrics_registry = CollectorRegistry() + register_engine_metrics_callback( + endpoint=generate_endpoint, + registry=embedding_metrics_registry, + metric_prefix_filters=["dynamo_embedding_"], + namespace_name=dynamo_args.namespace, + component_name=dynamo_args.component, + endpoint_name=dynamo_args.endpoint, + model_name=server_args.served_model_name, + ) + init_embedding_metrics(embedding_metrics_registry) ready_event = asyncio.Event() - - handler = EmbeddingWorkerHandler(engine, config, publisher, shutdown_event) - health_check_payload = SglangHealthCheckPayload( - engine, use_text_input=dynamo_args.use_sglang_tokenizer - ).to_dict() - - register_model_taint_route(runtime, generate_endpoint) + handler: EmbeddingWorkerHandler | RerankWorkerHandler | None = None try: + if rerank: + handler = RerankWorkerHandler(engine, config, publisher, shutdown_event) + health_check_payload = SglangRerankHealthCheckPayload( + server_args.served_model_name or server_args.model_path + ).to_dict() + else: + handler = EmbeddingWorkerHandler(engine, config, publisher, shutdown_event) + health_check_payload = SglangEmbeddingHealthCheckPayload( + server_args.served_model_name or server_args.model_path, + engine, + use_text_input=dynamo_args.use_sglang_tokenizer, + ).to_dict() + register_model_taint_route(runtime, generate_endpoint) await asyncio.gather( generate_endpoint.serve_endpoint( handler.generate, @@ -88,14 +118,14 @@ async def init_embedding( server_args, dynamo_args, input_type=ModelInput.Text, - output_type=ModelType.Embedding, + output_type=ModelType.Rerank if rerank else ModelType.Embedding, readiness_gate=ready_event, worker_type=WorkerType.Aggregated, needs=[], ), ) except Exception as e: - logging.error(f"Failed to serve embedding endpoints: {e}") + logging.error("Failed to serve pooling endpoint: %s", e) raise finally: metrics_task.cancel() @@ -104,7 +134,10 @@ async def init_embedding( except asyncio.CancelledError: logging.info("Metrics task successfully cancelled") pass - handler.cleanup() + if handler is not None: + handler.cleanup() + else: + engine.shutdown() if run_deferred_handlers is not None: logging.info("Running deferred handlers") await run_deferred_handlers() diff --git a/components/src/dynamo/sglang/init_llm.py b/components/src/dynamo/sglang/init_llm.py index 10236f0e897c..ccff8406e2be 100644 --- a/components/src/dynamo/sglang/init_llm.py +++ b/components/src/dynamo/sglang/init_llm.py @@ -15,6 +15,11 @@ from dynamo.llm import ModelInput, ModelType, WorkerType from dynamo.runtime import DistributedRuntime from dynamo.sglang.args import Config +from dynamo.sglang.gateway import ( + attached_engine_load_time, + gateway_worker_count, + serve_via_gateway_children, +) from dynamo.sglang.health_check import ( SglangDisaggHealthCheckPayload, SglangHealthCheckPayload, @@ -48,6 +53,7 @@ async def init_decode( shutdown_endpoints: list, run_deferred_handlers: Callable[[], Awaitable[None]] | None = None, snapshot_engine: Optional[sgl.Engine] = None, + attached_engine: Optional[object] = None, ) -> None: server_args, dynamo_args = config.server_args, config.dynamo_args @@ -62,6 +68,7 @@ async def init_decode( ) # Use pre-created engine if provided (snapshot mode) + load_time: Optional[float] if snapshot_engine is not None: engine = snapshot_engine load_time = 0.0 @@ -70,6 +77,11 @@ async def init_decode( "Snapshot ServerArgs must disable forward-pass metrics before " "engine creation" ) + elif attached_engine is not None: + # Gateway child: the parent owns the engine, this process only holds a + # TokenizerWorker registered with its router. + engine = attached_engine + load_time = attached_engine_load_time() else: set_forward_pass_metrics_worker_id(server_args, generate_endpoint) start_time = time.time() @@ -77,6 +89,19 @@ async def init_decode( load_time = time.time() - start_time server_args = config.use_resolved_server_args(engine.server_args) + gateway_count = gateway_worker_count(server_args, dynamo_args) + if gateway_count > 1: + # engine.tokenizer_manager is SGLang's MultiTokenizerRouter here and cannot + # serve requests; gateway children do, this process keeps the engine alive. + try: + await serve_via_gateway_children( + engine, gateway_count, shutdown_event, load_time=load_time + ) + finally: + engine.shutdown() + if run_deferred_handlers is not None: + await run_deferred_handlers() + return if server_args.enable_trace: set_global_trace_level(dynamo_args.sglang_trace_level) @@ -100,8 +125,9 @@ async def init_decode( # which take a different init path entirely. Narrow for mypy. assert publisher is not None, "setup_sgl_metrics returned None on chat path" - publisher.component_gauges.set_model_load_time(load_time) - logging.debug(f"SGLang model load time: {load_time:.2f}s") + if load_time is not None: + publisher.component_gauges.set_model_load_time(load_time) + logging.debug(f"SGLang model load time: {load_time:.2f}s") if server_args.node_rank >= 1: await handle_non_leader_node(engine, publisher, metrics_task) @@ -129,6 +155,8 @@ async def init_decode( first_token_source=first_token_source, ) handler.register_engine_routes(runtime) + if attached_engine is not None: + handler.follow_shared_pause_state() if config.serving_mode == DisaggregationMode.DECODE: health_check_payload = SglangDisaggHealthCheckPayload( @@ -207,6 +235,7 @@ async def init_prefill( shutdown_endpoints: list, run_deferred_handlers: Callable[[], Awaitable[None]] | None = None, snapshot_engine: Optional[sgl.Engine] = None, + attached_engine: Optional[object] = None, ) -> None: server_args, dynamo_args = config.server_args, config.dynamo_args @@ -221,6 +250,7 @@ async def init_prefill( ) # Use pre-created engine if provided (snapshot mode) + load_time: Optional[float] if snapshot_engine is not None: engine = snapshot_engine load_time = 0.0 @@ -229,6 +259,11 @@ async def init_prefill( "Snapshot ServerArgs must disable forward-pass metrics before " "engine creation" ) + elif attached_engine is not None: + # Gateway child: the parent owns the engine, this process only holds a + # TokenizerWorker registered with its router. + engine = attached_engine + load_time = attached_engine_load_time() else: set_forward_pass_metrics_worker_id(server_args, generate_endpoint) start_time = time.time() @@ -236,6 +271,19 @@ async def init_prefill( load_time = time.time() - start_time server_args = config.use_resolved_server_args(engine.server_args) + gateway_count = gateway_worker_count(server_args, dynamo_args) + if gateway_count > 1: + # engine.tokenizer_manager is SGLang's MultiTokenizerRouter here and cannot + # serve requests; gateway children do, this process keeps the engine alive. + try: + await serve_via_gateway_children( + engine, gateway_count, shutdown_event, load_time=load_time + ) + finally: + engine.shutdown() + if run_deferred_handlers is not None: + await run_deferred_handlers() + return if server_args.enable_trace: set_global_trace_level(dynamo_args.sglang_trace_level) @@ -259,7 +307,8 @@ async def init_prefill( # which take a different init path entirely. Narrow for mypy. assert publisher is not None, "setup_sgl_metrics returned None on chat path" - publisher.component_gauges.set_model_load_time(load_time) + if load_time is not None: + publisher.component_gauges.set_model_load_time(load_time) if server_args.node_rank >= 1: await handle_non_leader_node(engine, publisher, metrics_task) @@ -280,6 +329,8 @@ async def init_prefill( engine, config, publisher, generate_endpoint, shutdown_event ) handler.register_engine_routes(runtime) + if attached_engine is not None: + handler.follow_shared_pause_state() health_check_payload = SglangPrefillHealthCheckPayload(engine).to_dict() diff --git a/components/src/dynamo/sglang/init_rerank.py b/components/src/dynamo/sglang/init_rerank.py new file mode 100644 index 000000000000..347ab5e691e1 --- /dev/null +++ b/components/src/dynamo/sglang/init_rerank.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from typing import Awaitable, Callable + +from dynamo.runtime import DistributedRuntime +from dynamo.sglang.args import Config +from dynamo.sglang.init_embedding import _init_pooling + + +async def init_rerank( + runtime: DistributedRuntime, + config: Config, + shutdown_event: asyncio.Event, + shutdown_endpoints: list, + run_deferred_handlers: Callable[[], Awaitable[None]] | None = None, +) -> None: + """Serve and advertise a dedicated cross-encoder worker.""" + await _init_pooling( + runtime, + config, + shutdown_event, + shutdown_endpoints, + run_deferred_handlers, + rerank=True, + ) diff --git a/components/src/dynamo/sglang/main.py b/components/src/dynamo/sglang/main.py index a17ea0f42d35..503ddb2ab6c6 100644 --- a/components/src/dynamo/sglang/main.py +++ b/components/src/dynamo/sglang/main.py @@ -19,6 +19,14 @@ from dynamo.runtime.logging import configure_dynamo_logging from dynamo.sglang._compat import override_server_args from dynamo.sglang.args import parse_args +from dynamo.sglang.gateway import ( + build_gateway_engine, + effective_gateway_workers, + gateway_worker_count, + is_gateway_child, + reserve_system_port_for_children, + validate_gateway_mode, +) from dynamo.sglang.init_diffusion import ( init_image_diffusion, init_llm_diffusion, @@ -31,6 +39,7 @@ init_multimodal_prefill_worker, init_multimodal_worker, ) +from dynamo.sglang.init_rerank import init_rerank from dynamo.sglang.nixl_telemetry import install_per_rank_nixl_prometheus_ports from dynamo.sglang.shutdown import install_graceful_shutdown from dynamo.sglang.snapshot import prepare_snapshot_engine @@ -61,12 +70,30 @@ async def worker(argv: list[str] | None = None): load_format=setup_gms(config.server_args), ) + gateway_count = effective_gateway_workers(config.server_args, config.dynamo_args) + validate_gateway_mode(config.server_args, config.dynamo_args, gateway_count) + if gateway_count > 1 and config.server_args.tokenizer_worker_num != gateway_count: + # SGLang launches its MultiTokenizerRouter (the shared-memory handoff the + # gateway children join) with exactly tokenizer_worker_num workers. + override_server_args( + config.server_args, + "dynamo.gateway", + tokenizer_worker_num=gateway_count, + ) + if gateway_worker_count(config.server_args, config.dynamo_args) > 1: + reserve_system_port_for_children() + # Snapshot mode: engine must be created before runtime so CRIU captures no # NATS/etcd connections. - snapshot_controller = await prepare_snapshot_engine(config.server_args) + snapshot_controller = await prepare_snapshot_engine(config) dynamo_args = config.dynamo_args snapshot_engine = None + attached_engine = None + if snapshot_controller is None and is_gateway_child(): + # Spawned by gateway.serve_via_gateway_children: share the parent's engine + # through a TokenizerWorker instead of launching schedulers. + attached_engine = build_gateway_engine() if snapshot_controller is not None: snapshot_engine = snapshot_controller.engine dynamo_args = await refresh_snapshot_restore_config( @@ -105,6 +132,10 @@ async def worker(argv: list[str] | None = None): await init_video_diffusion( runtime, config, shutdown_endpoints, run_deferred_handlers ) + elif config.dynamo_args.rerank_worker: + await init_rerank( + runtime, config, shutdown_event, shutdown_endpoints, run_deferred_handlers + ) elif config.dynamo_args.embedding_worker: await init_embedding( runtime, @@ -154,6 +185,7 @@ async def worker(argv: list[str] | None = None): shutdown_endpoints, run_deferred_handlers, snapshot_engine=snapshot_engine, + attached_engine=attached_engine, ) else: await init_prefill( @@ -163,6 +195,7 @@ async def worker(argv: list[str] | None = None): shutdown_endpoints, run_deferred_handlers, snapshot_engine=snapshot_engine, + attached_engine=attached_engine, ) diff --git a/components/src/dynamo/sglang/nixl_telemetry.py b/components/src/dynamo/sglang/nixl_telemetry.py index 4a28b0651633..3d9c6153eea0 100644 --- a/components/src/dynamo/sglang/nixl_telemetry.py +++ b/components/src/dynamo/sglang/nixl_telemetry.py @@ -101,7 +101,11 @@ def _node_local_rank( def _assign_nixl_prometheus_port(target: Any, args: tuple, kwargs: dict) -> None: - """Rewrite this process's exporter port before the NIXL agent is built.""" + """Rewrite this process's exporter port before the NIXL agent is built. + + Raises ``ValueError`` when the NIXL environment is invalid or the rank's + fixed port cannot fit in Dynamo's declared scrape range. + """ base_port = nixl_prometheus_base_port() if base_port is None: return @@ -163,8 +167,10 @@ def install_per_rank_nixl_prometheus_ports() -> None: """Point SGLang's scheduler launches at the wrapper, when telemetry is on. A no-op when NIXL Prometheus telemetry is disabled, so a deployment that - does not scrape NIXL keeps SGLang's own entry point. Raises ``RuntimeError`` - when telemetry is on but this SGLang offers no override point. + does not scrape NIXL keeps SGLang's own entry point. Raises ``ValueError`` + for invalid NIXL environment values or an unusable fixed port, and + ``RuntimeError`` when telemetry is on but this SGLang offers no override + point. """ if nixl_prometheus_base_port() is None: return diff --git a/components/src/dynamo/sglang/pause.py b/components/src/dynamo/sglang/pause.py deleted file mode 100644 index e75f6d7fcee6..000000000000 --- a/components/src/dynamo/sglang/pause.py +++ /dev/null @@ -1,77 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import logging -from typing import Any - -from sglang.srt.managers.io_struct import ( - ContinueGenerationReqInput, - PauseGenerationReqInput, - ReleaseMemoryOccupationReqInput, - ResumeMemoryOccupationReqInput, -) - -logger = logging.getLogger(__name__) - - -class SGLangEnginePauseController: - def __init__(self, engine: Any): - self._engine = engine - self._is_paused = False - self._generation_paused = False - - @property - def is_paused(self) -> bool: - return self._is_paused - - @property - def needs_resume_recovery(self) -> bool: - return self._generation_paused - - async def pause(self, tags: list[str] | None = None) -> bool: - if self._is_paused or self._generation_paused: - return False - - await self._engine.tokenizer_manager.pause_generation(PauseGenerationReqInput()) - self._generation_paused = True - try: - await self._engine.tokenizer_manager.release_memory_occupation( - ReleaseMemoryOccupationReqInput(tags=tags), - None, - ) - except Exception: - try: - await self._engine.tokenizer_manager.continue_generation( - ContinueGenerationReqInput() - ) - self._generation_paused = False - except Exception: - logger.exception( - "failed to resume generation after memory release failed" - ) - raise - - self._is_paused = True - return True - - async def resume(self, tags: list[str] | None = None) -> bool: - if not self._is_paused and not self._generation_paused: - return False - - if self._is_paused: - await self._engine.tokenizer_manager.resume_memory_occupation( - ResumeMemoryOccupationReqInput(tags=tags), - None, - ) - if self._generation_paused: - await self._engine.tokenizer_manager.continue_generation( - ContinueGenerationReqInput() - ) - self._generation_paused = False - return True - - def mark_resumed(self) -> None: - self._is_paused = False - self._generation_paused = False diff --git a/components/src/dynamo/sglang/protocol.py b/components/src/dynamo/sglang/protocol.py index 54a7a393231e..0ddffb92d085 100644 --- a/components/src/dynamo/sglang/protocol.py +++ b/components/src/dynamo/sglang/protocol.py @@ -63,6 +63,16 @@ class EmbeddingRequest(BaseModel): encoding_format: Literal["float", "base64"] = "float" +class RerankRequest(BaseModel): + """SGLang-compatible text-only cross-encoder rerank request.""" + + model: str + query: str + documents: List[str] + top_n: Optional[int] = None + return_documents: bool = True + + class DisaggPreprocessedRequest(BaseModel): request: Union[PreprocessedRequest, ChatCompletionRequest] sampling_params: dict diff --git a/components/src/dynamo/sglang/publisher.py b/components/src/dynamo/sglang/publisher.py index 38e839529227..867e3ffa9797 100644 --- a/components/src/dynamo/sglang/publisher.py +++ b/components/src/dynamo/sglang/publisher.py @@ -33,6 +33,11 @@ local_dp_rank_bounds, publishes_kv_events, ) +from dynamo.sglang.gateway import ( + effective_gateway_workers, + metrics_fanout_endpoint, + owns_engine_metrics, +) def get_local_dp_rank_range(server_args) -> range: @@ -59,15 +64,20 @@ def set_forward_pass_metrics_worker_id( ) -async def _resolve_multinode_leader_worker_id( +async def _resolve_multinode_leader_worker_ids( generate_endpoint: Endpoint, server_args, -) -> Optional[int]: - """Return the routable leader worker id for SGLang non-leader nodes.""" + expected: int = 1, +) -> list[int]: + """Return the routable leader worker ids for SGLang non-leader nodes. + + In gateway mode the leader node registers ``expected`` endpoint instances for + the same engine (one per gateway child); remote-rank KV events must be + attributed to every one of them or the router only sees them on one.""" node_rank = getattr(server_args, "node_rank", 0) or 0 nnodes = getattr(server_args, "nnodes", 1) or 1 if node_rank <= 0 or nnodes <= 1: - return None + return [] worker_group_id = get_sglang_worker_group_id(server_args) client = await generate_endpoint.client() @@ -75,11 +85,21 @@ async def _resolve_multinode_leader_worker_id( timeout_s = getattr(server_args, "dist_timeout", None) timeout_s = None if timeout_s is None else float(timeout_s) try: - worker_id = await client.wait_for_instance_by_runtime_data( - SGLANG_WORKER_GROUP_ID_KEY, - worker_group_id, - timeout_s=timeout_s, - ) + if expected > 1: + worker_ids = await client.wait_for_instances_by_runtime_data( + SGLANG_WORKER_GROUP_ID_KEY, + worker_group_id, + expected, + timeout_s=timeout_s, + ) + else: + worker_ids = [ + await client.wait_for_instance_by_runtime_data( + SGLANG_WORKER_GROUP_ID_KEY, + worker_group_id, + timeout_s=timeout_s, + ) + ] except Exception as e: raise RuntimeError( "Failed to resolve SGLang leader worker_id for non-leader " @@ -88,13 +108,18 @@ async def _resolve_multinode_leader_worker_id( ) from e logging.info( - "Using SGLang leader worker_id=%s for non-leader KV event " + "Using SGLang leader worker_ids=%s for non-leader KV event " "publishing via worker_group_id=%s", - worker_id, + worker_ids, worker_group_id, ) - return int(worker_id) + return [int(w) for w in worker_ids] + if expected > 1: + raise RuntimeError( + "gateway mode on a multi-node engine needs dist_init_addr so the " + "non-leader nodes can find every gateway instance of their engine" + ) instances = await client.wait_for_instances() if len(instances) == 1: worker_id = int(instances[0]) @@ -102,14 +127,14 @@ async def _resolve_multinode_leader_worker_id( "Using SGLang leader worker_id=%s for non-leader KV event publishing", worker_id, ) - return worker_id + return [worker_id] logging.warning( "Expected exactly one SGLang leader endpoint instance for non-leader " "KV event attribution, got %d; skipping non-leader KV event publishing", len(instances), ) - return None + return [] def format_zmq_endpoint(endpoint_template: str, ip_address: str) -> str: @@ -143,6 +168,32 @@ def format_zmq_endpoint(endpoint_template: str, ip_address: str) -> str: # This is the same pattern used by dynamo+vLLM. +def _open_metrics_sockets( + ctx: zmq.asyncio.Context, + metrics_ipc_name: str, + fanout_endpoint: Optional[str], + owner: bool, +) -> tuple[zmq.asyncio.Socket, Optional[zmq.asyncio.Socket]]: + """The schedulers push KvMetrics to one PULL socket. Its owner (the single + worker, or gateway child 0) re-publishes every message on ``fanout_endpoint`` + so sibling gateways can report the same usage for their own identities.""" + # SGLang's get_zmq_socket only configures PUSH/PULL/DEALER/REQ/REP/PAIR and + # raises for PUB/SUB, so the fan-out pair is created directly. + if owner: + sock = get_zmq_socket(ctx, zmq.PULL, metrics_ipc_name, True) + fanout = None + if fanout_endpoint is not None: + fanout = ctx.socket(zmq.PUB) + fanout.bind(fanout_endpoint) + return sock, fanout + if fanout_endpoint is None: + raise ValueError("a sibling gateway needs the metrics fan-out endpoint") + sock = ctx.socket(zmq.SUB) + sock.setsockopt(zmq.SUBSCRIBE, b"") + sock.connect(fanout_endpoint) + return sock, None + + class DynamoSglangPublisher: """ Handles SGLang kv events and metrics reception and publishing. @@ -189,13 +240,19 @@ def __init__( # need KV event publishing which is set up separately in init_kv_event_publish() node_rank = getattr(self.server_args, "node_rank", 0) or 0 self._ctx: zmq.asyncio.Context | None = None + self._sock: zmq.asyncio.Socket | None = None + self._fanout: zmq.asyncio.Socket | None = None + # Engine-level gauges (total blocks, cache usage) describe one engine; only + # the process that consumes the schedulers' metrics publishes them, so a + # scrape across gateway children does not count the engine N times. + self._publishes_engine_gauges = owns_engine_metrics() if node_rank == 0: self._ctx = zmq.asyncio.Context() - self._sock = get_zmq_socket( + self._sock, self._fanout = _open_metrics_sockets( self._ctx, - zmq.PULL, self.engine.port_args.metrics_ipc_name, - True, + metrics_fanout_endpoint(), + self._publishes_engine_gauges, ) else: self._ctx = None @@ -223,6 +280,8 @@ async def run(self) -> None: # Receive KvMetrics object from SGLang scheduler via ZMQ # KvMetrics class: sglang/srt/observability/scheduler_metrics_mixin.py kv_metrics = await self._sock.recv_pyobj() + if self._fanout is not None: + await self._fanout.send_pyobj(kv_metrics) dp_rank = ( kv_metrics.data_parallel_rank if kv_metrics.data_parallel_rank is not None @@ -236,13 +295,12 @@ async def run(self) -> None: self.metrics_publisher.publish( dp_rank, kv_used_blocks=active_decode_blocks ) - dp_rank_str = str(dp_rank) - # Publish total blocks (always available in KvMetrics) - self.component_gauges.set_total_blocks(dp_rank_str, total_blocks) - # Publish GPU cache usage percentage (always available in KvMetrics) - self.component_gauges.set_gpu_cache_usage( - dp_rank_str, kv_metrics.gpu_cache_usage_perc - ) + if self._publishes_engine_gauges: + dp_rank_str = str(dp_rank) + self.component_gauges.set_total_blocks(dp_rank_str, total_blocks) + self.component_gauges.set_gpu_cache_usage( + dp_rank_str, kv_metrics.gpu_cache_usage_perc + ) except Exception: if self._running: logging.exception( @@ -254,9 +312,11 @@ def cleanup(self) -> None: self._running = False # Close ZMQ socket and context - if self._sock is not None: + for sock in (self._sock, self._fanout): + if sock is None: + continue try: - self._sock.close(linger=0) + sock.close(linger=0) except Exception as e: logging.warning(f"Failed to close ZMQ socket: {e}") @@ -286,9 +346,10 @@ def init_engine_metrics_publish(self) -> None: """Publish initial dummy metrics to bootstrap the metrics endpoint.""" logging.info("Sending dummy metrics to initialize") self.metrics_publisher.publish(self.dp_rank, kv_used_blocks=0) - dp_rank_str = str(self.dp_rank) - self.component_gauges.set_total_blocks(dp_rank_str, 0) - self.component_gauges.set_gpu_cache_usage(dp_rank_str, 0.0) + if self._publishes_engine_gauges: + dp_rank_str = str(self.dp_rank) + self.component_gauges.set_total_blocks(dp_rank_str, 0) + self.component_gauges.set_gpu_cache_usage(dp_rank_str, 0.0) def init_kv_event_publish(self) -> List[KvEventPublisher]: """Initialize KV event publisher(s) if configured. @@ -480,7 +541,7 @@ async def setup_sgl_metrics( and starts a ``DynamoSglangPublisher`` that pulls scheduler metrics over ZMQ and (optionally) forwards KV events / FPM stats. - For **embedding workers** (``config.dynamo_args.embedding_worker``), + For **embedding and rerank workers**, the chat-shaped pipeline is **skipped entirely**: pooling engines have no KV cache, no prefill/decode phase, and no scheduler metrics worth collecting, so every metric in that pipeline would emit zeros @@ -501,9 +562,11 @@ async def setup_sgl_metrics( """ metrics_labels = [("model", engine.server_args.served_model_name)] - if getattr(config.dynamo_args, "embedding_worker", False): + if getattr(config.dynamo_args, "embedding_worker", False) or getattr( + config.dynamo_args, "rerank_worker", False + ): logging.info( - "Embedding worker: skipping chat-shaped Prometheus + KV-event " + "Pooling worker: skipping chat-shaped Prometheus + KV-event " "wiring (no KV cache, no prefill/decode, no scheduler metrics). " "Embedding-shaped metrics are registered separately." ) @@ -556,6 +619,9 @@ async def _idle() -> None: publisher.init_engine_metrics_publish() node_rank = getattr(config.server_args, "node_rank", 0) or 0 + # Every gateway child republishes KV events under its own worker id so the + # router can match prefixes on any instance of the engine; scheduler metrics + # and forward-pass metrics are consumed once, by the process that owns them. if node_rank <= 0 and config.dynamo_args.use_kv_events: publisher.init_kv_event_publish() publisher.init_fpm_relay() @@ -588,11 +654,12 @@ async def handle_non_leader_node( if publisher.dynamo_args.use_kv_events and publishes_kv_events( publisher.server_args ): - kv_worker_id = await _resolve_multinode_leader_worker_id( + kv_worker_ids = await _resolve_multinode_leader_worker_ids( publisher.generate_endpoint, publisher.server_args, + effective_gateway_workers(publisher.server_args, publisher.dynamo_args), ) - if kv_worker_id is not None: + for kv_worker_id in kv_worker_ids: publisher.kv_worker_id = kv_worker_id publisher.init_kv_event_publish() diff --git a/components/src/dynamo/sglang/register.py b/components/src/dynamo/sglang/register.py index 42e6c020ce55..dcaf9e38a4ae 100644 --- a/components/src/dynamo/sglang/register.py +++ b/components/src/dynamo/sglang/register.py @@ -30,7 +30,10 @@ WorkerType, register_model, ) -from dynamo.sglang._compat import sglang_uses_mla_backend +from dynamo.sglang._compat import ( + sglang_uses_mla_backend, + supports_disagg_prefill_cancel_anytime, +) from dynamo.sglang._disagg import SGLANG_WORKER_GROUP_ID_KEY, get_sglang_worker_group_id from dynamo.sglang.args import DynamoConfig, use_modelexpress_remote_instance from dynamo.sglang.capacity import ( @@ -40,7 +43,16 @@ model_card_dp_rank_bounds, runtime_capacity, ) -from dynamo.sglang.engine_generate import SGLANG_GENERATE_CAPABILITY +from dynamo.sglang.engine_generate import ( + DISAGG_PREFILL_CANCEL_ANYTIME_V1, + SGLANG_GENERATE_CAPABILITY, +) +from dynamo.sglang.gateway import ( + GATEWAY_ENGINE_ID_KEY, + GATEWAY_WORKERS_KEY, + effective_gateway_workers, + gateway_engine_id, +) SGLANG_HICACHE_MOONCAKE_RUNTIME_KEY = "sglang_hicache_mooncake" SPEC_DECODE_RUNTIME_KEY = "spec_decode" @@ -146,14 +158,14 @@ async def _register_model_with_runtime_config( """ runtime_config = await get_runtime_config(engine, server_args, dynamo_args) - if dynamo_args.use_sglang_tokenizer: + if dynamo_args.use_sglang_tokenizer and not ( + output_type.supports_embedding() or output_type.supports_rerank() + ): logging.warning( "Using the sglang tokenizer/detokenizer instead. The dynamo tokenizer/detokenizer will not be used and only v1/chat/completions will be available" ) input_type = ModelInput.Text - # Only override output_type for chat models, not for embeddings - if output_type != ModelType.Embedding: - output_type = ModelType.Chat + output_type = ModelType.Chat if runtime_config is not None and _supports_engine_generate( input_type, output_type, worker_type @@ -401,6 +413,22 @@ async def get_runtime_config( runtime_config = ModelRuntimeConfig() runtime_config.kv_state_endpoint = dynamo_args.kv_state_endpoint runtime_config.context_length = server_args.context_length + llm_handler = not dynamo_args.enable_multimodal + # Both disaggregated legs implement exact-RID cancellation before output, + # including waiting until the scheduler has accepted the request. + if ( + server_args.disaggregation_mode + in { + "prefill", + "decode", + } + and llm_handler + and supports_disagg_prefill_cancel_anytime(engine) + ): + runtime_config.set_engine_specific( + DISAGG_PREFILL_CANCEL_ANYTIME_V1, + json.dumps(True), + ) # Multimodal encode workers have no tokenizer manager and delegate # generation overflow handling to their downstream backend. if engine is not None: @@ -456,6 +484,16 @@ async def get_runtime_config( e, ) + gateway_engine = gateway_engine_id() + if gateway_engine is not None: + runtime_config.set_engine_specific( + GATEWAY_ENGINE_ID_KEY, json.dumps(gateway_engine) + ) + runtime_config.set_engine_specific( + GATEWAY_WORKERS_KEY, + json.dumps(effective_gateway_workers(server_args, dynamo_args)), + ) + # Set topology and KV transfer policy for topology-aware routing apply_topology_config(runtime_config) diff --git a/components/src/dynamo/sglang/request_handlers/__init__.py b/components/src/dynamo/sglang/request_handlers/__init__.py index 0c1fde3e2bc8..3799365ca0a9 100644 --- a/components/src/dynamo/sglang/request_handlers/__init__.py +++ b/components/src/dynamo/sglang/request_handlers/__init__.py @@ -14,6 +14,7 @@ "PrefillWorkerHandler": ".llm", # Embedding handlers "EmbeddingWorkerHandler": ".embedding", + "RerankWorkerHandler": ".rerank", # Image diffusion handlers "ImageDiffusionWorkerHandler": ".image_diffusion", # Video generation handlers @@ -36,21 +37,16 @@ def __getattr__(name: str) -> Any: __all__ = [ - # Base handlers "BaseGenerativeHandler", "BaseWorkerHandler", - # LLM handlers "DecodeWorkerHandler", "DiffusionWorkerHandler", - "PrefillWorkerHandler", - # Embedding handlers "EmbeddingWorkerHandler", - # Image diffusion handlers "ImageDiffusionWorkerHandler", - # Video generation handlers - "VideoGenerationWorkerHandler", - # Multimodal handlers "MultimodalEncodeWorkerHandler", "MultimodalPrefillWorkerHandler", "MultimodalWorkerHandler", + "PrefillWorkerHandler", + "RerankWorkerHandler", + "VideoGenerationWorkerHandler", ] diff --git a/components/src/dynamo/sglang/request_handlers/cancellation.py b/components/src/dynamo/sglang/request_handlers/cancellation.py new file mode 100644 index 000000000000..dab569828f2c --- /dev/null +++ b/components/src/dynamo/sglang/request_handlers/cancellation.py @@ -0,0 +1,458 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SGLang request cancellation and shutdown coordination.""" + +import asyncio +import logging +from collections.abc import Mapping +from contextlib import asynccontextmanager +from typing import Any, AsyncGenerator, AsyncIterator + +from dynamo._core import Context +from dynamo.llm.exceptions import EngineShutdown +from dynamo.sglang._compat import resolved_server_args + +_CANCELLATION_POLL_MAX_DELAY_S = 0.05 +_CANCELLATION_ABORT_RETRY_LIMIT = 8 +_CANCELLATION_DRAIN_TIMEOUT_S = 1.0 +_CANCELLATION_REGISTRATION_WAIT_TIMEOUT_S = 30.0 +# Leave most of the drain window available after an ordered abort is submitted. +_CANCELLATION_DISPATCH_WAIT_TIMEOUT_S = 0.25 + + +async def _next_stream_item(iterator: AsyncIterator[Any]) -> Any: + return await anext(iterator) + + +def _consume_detached_task(task: asyncio.Task[Any]) -> None: + """Consume a detached cancellation task result without surfacing closure.""" + try: + task.result() + except (asyncio.CancelledError, StopAsyncIteration): + pass + except Exception: + logging.exception("Detached SGLang task failed during cancellation") + + +def _cancel_and_detach(task: asyncio.Task[Any]) -> None: + """Request cancellation without allowing resistant iterators to block cleanup.""" + task.cancel() + task.add_done_callback(_consume_detached_task) + + +class CancellationMixin: + """Coordinate Dynamo cancellation with SGLang request aborts.""" + + engine: Any + config: Any + shutdown_event: asyncio.Event | None + _abort_tasks: set[asyncio.Task[Any]] + + def _track_abort_task(self, task: asyncio.Task[Any]) -> None: + self._abort_tasks.add(task) + task.add_done_callback(self._abort_tasks.discard) + task.add_done_callback(_consume_detached_task) + + def _start_ordered_abort( + self, + tokenizer_manager: Any, + request_id_future: asyncio.Future, + submitted_request_id: str, + registry: Mapping[str, Any], + context_id: str, + ) -> asyncio.Task[Any]: + task = asyncio.create_task( + self._abort_after_registration( + tokenizer_manager, + request_id_future, + submitted_request_id, + registry, + context_id, + ) + ) + self._track_abort_task(task) + return task + + def _start_abort_retry( + self, + tokenizer_manager: Any, + request_id: str, + registry: Mapping[str, Any], + state: Any, + ) -> None: + retry_task = asyncio.create_task( + self._retry_abort(tokenizer_manager, request_id, registry, state) + ) + self._track_abort_task(retry_task) + + def _cancel_abort_tasks(self) -> None: + for task in tuple(self._abort_tasks): + task.cancel() + + def _abort_requests(self, request_ids: set[str], context: Context) -> None: + if not request_ids: + return + tokenizer_manager = getattr(self.engine, "tokenizer_manager", None) + if tokenizer_manager is None: + logging.error("SGLang tokenizer_manager not found for abort requests") + return + for request_id in request_ids: + tokenizer_manager.abort_request(rid=request_id, abort_all=False) + request_ids.clear() + logging.info("Aborted Request ID: %s", context.id()) + + async def _stream_until_cancelled( + self, + stream_source: AsyncIterator[Any], + cancellation_task: asyncio.Task[Any], + ) -> AsyncGenerator[Any, None]: + """Drain briefly after abort, then stop a stream that did not terminate.""" + iterator = aiter(stream_source) + drain_deadline: float | None = None + while True: + next_item = asyncio.create_task(_next_stream_item(iterator)) + try: + if drain_deadline is None: + done, _ = await asyncio.wait( + (next_item, cancellation_task), + return_when=asyncio.FIRST_COMPLETED, + ) + if cancellation_task in done: + cancellation_task.result() + drain_deadline = ( + asyncio.get_running_loop().time() + + _CANCELLATION_DRAIN_TIMEOUT_S + ) + if next_item in done: + try: + yield next_item.result() + except StopAsyncIteration: + return + continue + + assert drain_deadline is not None + remaining = drain_deadline - asyncio.get_running_loop().time() + if remaining <= 0: + logging.warning( + "Timed out draining SGLang stream after cancellation" + ) + return + done, _ = await asyncio.wait((next_item,), timeout=remaining) + if next_item not in done: + logging.warning( + "Timed out draining SGLang stream after cancellation" + ) + return + try: + yield next_item.result() + except StopAsyncIteration: + return + finally: + if not next_item.done(): + _cancel_and_detach(next_item) + + @asynccontextmanager + async def _wait_for_signal(self, context: Context) -> AsyncGenerator[bool, None]: + cancellation_future = context.async_killed_or_stopped() + shutdown_task = ( + asyncio.create_task(self.shutdown_event.wait()) + if self.shutdown_event + else None + ) + wait_for: list[asyncio.Future[Any]] = [cancellation_future] + if shutdown_task is not None: + wait_for.append(shutdown_task) + try: + done, _ = await asyncio.wait(wait_for, return_when=asyncio.FIRST_COMPLETED) + yield shutdown_task is not None and shutdown_task in done + finally: + pending = [awaitable for awaitable in wait_for if not awaitable.done()] + for awaitable in pending: + awaitable.cancel() + await asyncio.gather(*pending, return_exceptions=True) + + @staticmethod + def _request_registry( + tokenizer_manager: Any, submitted_request_id: str | None + ) -> Mapping[str, Any] | None: + if tokenizer_manager is None or submitted_request_id is None: + return None + candidate = getattr(tokenizer_manager, "rid_to_state", None) + return candidate if isinstance(candidate, Mapping) else None + + async def _wait_for_registration( + self, + submitted_request_id: str, + request_id_future: asyncio.Future, + registry: Mapping[str, Any], + ) -> str | None: + delay = 0.001 + while submitted_request_id not in registry: + if request_id_future.cancelled(): + return None + if request_id_future.done(): + return request_id_future.result() + await asyncio.sleep(delay) + delay = min(delay * 2, _CANCELLATION_POLL_MAX_DELAY_S) + return submitted_request_id + + async def _abort_after_registration( + self, + tokenizer_manager: Any, + request_id_future: asyncio.Future, + submitted_request_id: str, + registry: Mapping[str, Any], + context_id: str, + ) -> None: + try: + request_id = await asyncio.wait_for( + self._wait_for_registration( + submitted_request_id, + request_id_future, + registry, + ), + timeout=_CANCELLATION_REGISTRATION_WAIT_TIMEOUT_S, + ) + except asyncio.TimeoutError: + logging.warning( + "Timed out waiting for SGLang Request ID %s to register", + submitted_request_id, + ) + return + if request_id is None: + logging.debug( + "Abandoning SGLang abort for Context %s; request never registered", + context_id, + ) + return + await self._abort_sglang_request( + tokenizer_manager, + request_id, + registry, + context_id, + ) + + @staticmethod + def _resolved_request_id(request_id_future: asyncio.Future) -> str: + if request_id_future.done() and not request_id_future.cancelled(): + try: + return request_id_future.result() + except Exception: + pass + return "unknown" + + async def _handle_cancellation( + self, + request_id_future: asyncio.Future, + context: Context, + submitted_request_id: str | None = None, + request_ids: set[str] | None = None, + ) -> asyncio.Task[Any] | None: + """Wait for cancellation, then order an exact SGLang abort.""" + logging.debug("Cancellation monitor started for Context: %s", context.id()) + ordered_abort_task = None + request_id = submitted_request_id + try: + if request_id is None: + request_id = await request_id_future + async with self._wait_for_signal(context) as shutdown_requested: + tokenizer_manager = getattr(self.engine, "tokenizer_manager", None) + registry = self._request_registry( + tokenizer_manager, submitted_request_id + ) + logging.info( + "Cancellation or shutdown signal received for SGLang Request ID %s, Context: %s", + request_id, + context.id(), + ) + if tokenizer_manager is None: + logging.error( + "SGLang tokenizer_manager not found for abort request: %s", + context.id(), + ) + elif submitted_request_id is not None and registry is not None: + ordered_abort_task = self._start_ordered_abort( + tokenizer_manager, + request_id_future, + submitted_request_id, + registry, + context.id(), + ) + elif request_ids is not None: + self._abort_requests(request_ids, context) + else: + try: + await self._abort_sglang_request( + tokenizer_manager, request_id, registry, context.id() + ) + except Exception: + logging.exception( + "Failed to abort SGLang Request ID %s, Context: %s", + request_id, + context.id(), + ) + if shutdown_requested: + if ordered_abort_task is not None: + try: + await asyncio.wait_for( + asyncio.shield(ordered_abort_task), + timeout=_CANCELLATION_DRAIN_TIMEOUT_S, + ) + except asyncio.TimeoutError: + logging.warning( + "Timed out waiting for ordered SGLang abort during shutdown" + ) + except asyncio.CancelledError: + if not ordered_abort_task.cancelled(): + raise + except Exception: + pass + raise EngineShutdown("Engine was shut down during token generation") + return ordered_abort_task + except asyncio.CancelledError: + logging.debug( + "Cancellation monitor task cancelled for SGLang Request ID %s, Context: %s", + self._resolved_request_id(request_id_future), + context.id(), + ) + raise + + async def _wait_for_dispatch( + self, + registry: Mapping[str, Any], + request_id: str, + state: Any, + ) -> bool: + time_stats = getattr(state, "time_stats", None) + if time_stats is None or not hasattr( + time_stats, "api_server_dispatch_finish_time" + ): + return False + delay = 0.001 + deadline = ( + asyncio.get_running_loop().time() + _CANCELLATION_DISPATCH_WAIT_TIMEOUT_S + ) + while registry.get(request_id) is state: + if getattr(time_stats, "api_server_dispatch_finish_time", None): + return True + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + logging.warning( + "Timed out waiting for SGLang Request ID %s to dispatch", + request_id, + ) + return False + await asyncio.sleep(min(delay, remaining)) + delay = min(delay * 2, _CANCELLATION_POLL_MAX_DELAY_S) + return False + + @staticmethod + def _submit_abort(tokenizer_manager: Any, request_id: str, context_id: str) -> None: + tokenizer_manager.abort_request(rid=request_id, abort_all=False) + logging.info("Aborted Request ID: %s", context_id) + + @staticmethod + def _requires_abort_retry(server_args: Any, dispatch_observed: bool) -> bool: + return ( + not dispatch_observed + or getattr(server_args, "pp_size", 1) > 1 + or ( + getattr(server_args, "enable_dp_attention", False) + and not getattr( + server_args, "enable_dp_attention_local_control_broadcast", False + ) + ) + ) + + async def _retry_abort( + self, + tokenizer_manager: Any, + request_id: str, + registry: Mapping[str, Any], + state: Any, + ) -> None: + delay = 0.05 + for retry in range(_CANCELLATION_ABORT_RETRY_LIMIT): + await asyncio.sleep(delay) + if registry.get(request_id) is not state: + return + logging.debug("Retrying SGLang abort_request for Request ID %s", request_id) + tokenizer_manager.abort_request(rid=request_id, abort_all=False) + delay = min(delay * 2, 1.0) + if registry.get(request_id) is state: + logging.warning( + "SGLang request %s remained registered after %d abort retries", + request_id, + retry + 1, + ) + + async def _abort_sglang_request( + self, + tokenizer_manager: Any, + request_id: str, + registry: Mapping[str, Any] | None, + context_id: str, + ) -> None: + state = registry.get(request_id) if registry is not None else None + if registry is None or state is None: + self._submit_abort(tokenizer_manager, request_id, context_id) + return + dispatch_observed = await self._wait_for_dispatch(registry, request_id, state) + if registry.get(request_id) is not state: + return + self._submit_abort(tokenizer_manager, request_id, context_id) + server_args = resolved_server_args(self.config.server_args) + if self._requires_abort_retry(server_args, dispatch_observed): + self._start_abort_retry(tokenizer_manager, request_id, registry, state) + + @asynccontextmanager + async def _cancellation_monitor( + self, + request_id_future: asyncio.Future, + context: Context, + submitted_request_id: str | None = None, + request_ids: set[str] | None = None, + ) -> AsyncGenerator[asyncio.Task, None]: + """Own the cancellation monitor task for one response stream.""" + logging.debug( + "Creating cancellation monitor task for Context: %s", context.id() + ) + cancellation_task = asyncio.create_task( + self._handle_cancellation( + request_id_future, + context, + submitted_request_id, + request_ids, + ) + ) + + try: + yield cancellation_task + finally: + request_id = self._resolved_request_id(request_id_future) + ordered_abort_task = None + try: + if not cancellation_task.done(): + logging.debug( + "Cancelling cancellation monitor task for SGLang Request ID %s, Context: %s", + request_id, + context.id(), + ) + cancellation_task.cancel() + try: + await cancellation_task + except asyncio.CancelledError: + pass + else: + ordered_abort_task = cancellation_task.result() + finally: + if not request_id_future.done() and ordered_abort_task is None: + request_id_future.cancel() + if request_ids is not None: + if ordered_abort_task is not None and submitted_request_id is not None: + request_ids.discard(submitted_request_id) + self._abort_requests(request_ids, context) + + if self.shutdown_event and self.shutdown_event.is_set(): + raise EngineShutdown("Engine was shut down during token generation") diff --git a/components/src/dynamo/sglang/request_handlers/handler_base.py b/components/src/dynamo/sglang/request_handlers/handler_base.py index 26cf4e1fd0e8..0c851937e264 100644 --- a/components/src/dynamo/sglang/request_handlers/handler_base.py +++ b/components/src/dynamo/sglang/request_handlers/handler_base.py @@ -9,17 +9,7 @@ import re import threading from abc import ABC, abstractmethod -from contextlib import asynccontextmanager -from typing import ( - Any, - AsyncGenerator, - AsyncIterator, - Dict, - Generic, - Optional, - Tuple, - TypeVar, -) +from typing import Any, AsyncIterator, Dict, Generic, Optional, Tuple, TypeVar import sglang as sgl from sglang.srt.managers.io_struct import ProfileReq @@ -44,17 +34,26 @@ register_llm, unregister_llm, ) -from dynamo.llm.exceptions import EngineShutdown from dynamo.runtime import DistributedRuntime +from dynamo.sglang._compat import supports_disagg_prefill_cancel_anytime from dynamo.sglang.args import Config from dynamo.sglang.capacity import kv_event_block_size from dynamo.sglang.engine_routes import resolve_configured_engine_routes -from dynamo.sglang.pause import SGLangEnginePauseController +from dynamo.sglang.gateway import follow_pause_broadcasts from dynamo.sglang.publisher import DynamoSglangPublisher +from dynamo.sglang.request_handlers.cancellation import CancellationMixin logger = logging.getLogger(__name__) +_DEFAULT_NATIVE_ENGINE_ROUTES = ( + ("pause_generation", "pause_generation:tm"), + ("continue_generation", "continue_generation:tm"), + ("release_memory_occupation", "release_memory_occupation:tm"), + ("resume_memory_occupation", "resume_memory_occupation:tm"), +) + + RequestT = TypeVar("RequestT") ResponseT = TypeVar("ResponseT") @@ -549,7 +548,9 @@ async def list_loras(self, _request: Optional[Dict[str, Any]] = None): yield {"status": "error", "message": str(e)} -class BaseWorkerHandler(LoraMixin, BaseGenerativeHandler[RequestT, ResponseT]): +class BaseWorkerHandler( + CancellationMixin, LoraMixin, BaseGenerativeHandler[RequestT, ResponseT] +): """Abstract base class for SGLang LLM worker handlers. Extends BaseGenerativeHandler with LLM-specific functionality: @@ -584,6 +585,14 @@ def __init__( self.generate_endpoint = generate_endpoint self.publisher = publisher self.shutdown_event = shutdown_event + self._abort_tasks: set[asyncio.Task[Any]] = set() + self._supports_ordered_cancellation = ( + engine is not None and supports_disagg_prefill_cancel_anytime(engine) + ) + logger.info( + "SGLang ordered cancellation support: %s", + self._supports_ordered_cancellation, + ) if publisher is not None: self.metrics_publisher = publisher.metrics_publisher self.kv_publisher = publisher.kv_publisher @@ -607,10 +616,7 @@ def __init__( # have an sgl.Engine. self.input_param_manager = InputParamManager(None) self._engine_supports_priority = False - self._pause_controller = ( - SGLangEnginePauseController(engine) if engine is not None else None - ) - self._pause_lock = asyncio.Lock() + self._engine_route_lock = asyncio.Lock() # Serializes elastic-EP scaling: SGLang tracks a single in-flight scale # phase, so concurrent scale_elastic_ep calls must not overlap. @@ -629,122 +635,6 @@ def _priority_kwargs(self, priority: Any) -> Dict[str, Any]: return {"priority": normalized} return {} - async def release_memory_occupation(self, body: dict) -> dict: - """Release GPU memory occupation and unregister from discovery. - - Args: - body: Optional dict with "tags" to target specific memory regions. - - Order of operations: - 1. Unregister from discovery - stop accepting new requests - 2. Pause generation - drain in-flight requests - 3. Release memory - safe now that no requests are active - """ - if self._pause_controller is None: - return { - "status": "error", - "message": "memory control not supported on this worker", - } - - body = body or {} - tags = body.get("tags") - async with self._pause_lock: - if self._pause_controller.is_paused: - return { - "status": "ok", - "message": "Memory already released", - } - if self._pause_controller.needs_resume_recovery: - return { - "status": "error", - "message": "resume_memory_occupation required before retrying release", - } - - unregistered = False - try: - # Stop new requests and drain in-flight work before releasing memory. - if self.generate_endpoint is not None: - await self.generate_endpoint.unregister_endpoint_instance() - unregistered = True - - await self._pause_controller.pause(tags) - - return { - "status": "ok", - "message": ( - f"Memory released for tags: {tags}" - if tags is not None - else "Memory released" - ), - } - except Exception as e: - logging.error(f"Failed to release memory occupation: {e}") - # If pause rolled back cleanly the engine is serving-safe again, - # but discovery still shows us unregistered and resume will - # early-return. Re-register so the worker rejoins the routing pool. - if ( - unregistered - and not self._pause_controller.is_paused - and not self._pause_controller.needs_resume_recovery - and self.generate_endpoint is not None - ): - try: - await self.generate_endpoint.register_endpoint_instance() - logging.info( - "Re-registered endpoint after failed memory release rollback" - ) - except Exception as reg_err: - logging.error( - f"Failed to re-register endpoint after release failure: {reg_err}" - ) - return {"status": "error", "message": str(e)} - - async def resume_memory_occupation(self, body: dict) -> dict: - """Resume GPU memory occupation and re-register to discovery. - - Args: - body: Optional dict with "tags" to target specific memory regions. - - Order of operations: - 1. Resume memory - restore GPU allocations - 2. Continue generation - ready to serve requests - 3. Re-register to discovery - allow frontend to route here - """ - if self._pause_controller is None: - return { - "status": "error", - "message": "memory control not supported on this worker", - } - - body = body or {} - tags = body.get("tags") - async with self._pause_lock: - needs_recovery = self._pause_controller.needs_resume_recovery - if not self._pause_controller.is_paused and not needs_recovery: - return { - "status": "ok", - "message": "Memory already resumed", - } - - try: - await self._pause_controller.resume(tags) - - if self.generate_endpoint is not None: - await self.generate_endpoint.register_endpoint_instance() - self._pause_controller.mark_resumed() - - return { - "status": "ok", - "message": ( - f"Memory resumed for tags: {tags}" - if tags is not None - else "Memory resumed" - ), - } - except Exception as e: - logging.error(f"Failed to resume memory occupation: {e}") - return {"status": "error", "message": str(e)} - async def clear_kv_blocks(self, request: Optional[Dict[str, Any]] = None): """Flush SGLang's local cache when no requests are active.""" tokenizer_manager = ( @@ -760,7 +650,7 @@ async def clear_kv_blocks(self, request: Optional[Dict[str, Any]] = None): return try: - async with self._pause_lock: + async with self._engine_route_lock: if getattr(tokenizer_manager, "rid_to_state", None): yield { "status": "error", @@ -978,6 +868,48 @@ async def is_scaling_elastic_ep(self, body: dict) -> dict: return backend_error return dict(self.engine.tokenizer_manager.get_elastic_ep_state()) + async def _sync_discovery_with_sglang_pause_state(self) -> None: + """Make discovery match SGLang's authoritative generation pause state.""" + tokenizer_manager = ( + getattr(self.engine, "tokenizer_manager", None) + if self.engine is not None + else None + ) + if ( + self.generate_endpoint is None + or tokenizer_manager is None + or not hasattr(tokenizer_manager, "is_pause") + ): + return + + if tokenizer_manager.is_pause: + await self.generate_endpoint.unregister_endpoint_instance() + else: + await self.generate_endpoint.register_endpoint_instance() + + def follow_shared_pause_state(self) -> None: + async def resync() -> None: + async with self._engine_route_lock: + await self._sync_discovery_with_sglang_pause_state() + + follow_pause_broadcasts(getattr(self.engine, "tokenizer_manager", None), resync) + + async def _invoke_engine_route(self, route_handler, body: dict) -> dict: + """Invoke one engine route and then synchronize worker discovery.""" + async with self._engine_route_lock: + try: + return await route_handler(body) + finally: + await self._sync_discovery_with_sglang_pause_state() + + def _wrap_engine_route(self, route_handler): + """Add discovery synchronization to an SGLang route.""" + + async def synchronized_handler(body: dict) -> dict: + return await self._invoke_engine_route(route_handler, body) + + return synchronized_handler + def register_engine_routes(self, runtime: DistributedRuntime) -> None: """Register all engine routes for this handler. @@ -991,8 +923,6 @@ def register_engine_routes(self, runtime: DistributedRuntime) -> None: built_in_routes = { "control/start_profile": self.start_profile, "control/stop_profile": self.stop_profile, - "control/release_memory_occupation": self.release_memory_occupation, - "control/resume_memory_occupation": self.resume_memory_occupation, "control/update_weights_from_disk": self.update_weights_from_disk, "control/update_weights_from_tensor": self.update_weights_from_tensor, "control/update_weights_from_distributed": ( @@ -1016,11 +946,23 @@ def register_engine_routes(self, runtime: DistributedRuntime) -> None: "with a built-in route" ) + configured_paths = {path for path, _ in configured_routes} + # Expose the SGLang lifecycle methods without requiring --engine-route. + # An explicit route with the same public path overrides the default. + default_native_routes = resolve_configured_engine_routes( + self.engine, + [ + descriptor + for path, descriptor in _DEFAULT_NATIVE_ENGINE_ROUTES + if path not in configured_paths + ], + ) + register_model_taint_route(runtime, self.generate_endpoint) for path, handler in built_in_routes.items(): runtime.register_engine_route(path, handler) - for path, configured_handler in configured_routes: - runtime.register_engine_route(path, configured_handler) + for path, route_handler in [*default_native_routes, *configured_routes]: + runtime.register_engine_route(path, self._wrap_engine_route(route_handler)) @abstractmethod def generate(self, request: RequestT, context: Context) -> AsyncIterator[ResponseT]: @@ -1037,6 +979,7 @@ def generate(self, request: RequestT, context: Context) -> AsyncIterator[Respons def cleanup(self) -> None: """Cleanup resources. Override in subclasses as needed.""" + self._cancel_abort_tasks() if self.publisher is not None: self.publisher.cleanup() @@ -1233,152 +1176,3 @@ def _get_bootstrap_info(engine: sgl.Engine) -> Tuple[str, int]: ) return bootstrap_host, bootstrap_port - - async def _handle_cancellation( - self, request_id_future: asyncio.Future, context: Context - ): - """Background task to handle cancellation and shutdown by monitoring both signals. - - Args: - request_id_future: Future that will be set with the SGLang request ID - when the first response arrives. - context: Context object for cancellation handling. - - Raises: - EngineShutdown: If shutdown event was triggered. - """ - cancellation_future: asyncio.Future[Any] | None = None - shutdown_task: asyncio.Task[Any] | None = None - try: - logging.debug(f"Cancellation monitor started for Context: {context.id()}") - - # Always wait for the request ID to ensure we can abort the request - sglang_request_id = await request_id_future - logging.debug( - f"Cancellation monitor received SGLang Request ID {sglang_request_id} for Context: {context.id()}" - ) - logging.debug(f"Request ID future cancelled for Context: {context.id()}") - - # Get the cancellation future - cancellation_future = context.async_killed_or_stopped() - - # Build list of futures/tasks to wait for - wait_for: list[asyncio.Future[Any]] = [cancellation_future] - - if self.shutdown_event: - # Create task for shutdown monitoring and add to wait list - shutdown_task = asyncio.create_task(self.shutdown_event.wait()) - wait_for.append(shutdown_task) - - # Wait for whichever happens first - done, pending = await asyncio.wait( - wait_for, - return_when=asyncio.FIRST_COMPLETED, - ) - - # Cancel the pending task/future - for task in pending: - task.cancel() - try: - await task - except asyncio.CancelledError: - pass - - logging.info( - f"Cancellation or shutdown signal received for SGLang Request ID {sglang_request_id}, Context: {context.id()}" - ) - - # Call abort_request on the tokenizer_manager through the engine - if ( - hasattr(self.engine, "tokenizer_manager") - and self.engine.tokenizer_manager - ): - logging.info( - f"Calling SGLang abort_request for Request ID {sglang_request_id}" - ) - self.engine.tokenizer_manager.abort_request( - rid=sglang_request_id, abort_all=False - ) - logging.info(f"Aborted Request ID: {context.id()}") - else: - logging.error( - f"SGLang tokenizer_manager not found for abort request: {context.id()}" - ) - - # Check which event triggered and raise EngineShutdown if shutdown - if shutdown_task and shutdown_task in done: - raise EngineShutdown("Engine was shut down during token generation") - - except asyncio.CancelledError: - # Task was cancelled, which is expected when generation completes - request_id = "unknown" - if request_id_future.done() and not request_id_future.cancelled(): - try: - request_id = request_id_future.result() - except Exception: - pass - logging.debug( - f"Cancellation monitor task cancelled for SGLang Request ID {request_id}, Context: {context.id()}" - ) - raise - finally: - for awaitable in (cancellation_future, shutdown_task): - if awaitable is None or awaitable.done(): - continue - awaitable.cancel() - try: - await awaitable - except (asyncio.CancelledError, Exception): - pass - - @asynccontextmanager - async def _cancellation_monitor( - self, request_id_future: asyncio.Future, context: Context - ) -> AsyncGenerator[asyncio.Task, None]: - """ - Context manager for monitoring request cancellation and shutdown. - Automatically creates a background task to monitor for cancellation and - shutdown events, cleaning it up when the context exits. - - If shutdown event was triggered, raises EngineShutdown on exit. - - Args: - request_id_future: Future that will be set with the SGLang request ID - when the first response arrives. - context: Context object for cancellation handling - - Yields: - asyncio.Task: The cancellation monitoring task being managed - """ - logging.debug(f"Creating cancellation monitor task for Context: {context.id()}") - - # Start the cancellation monitoring task - cancellation_task = asyncio.create_task( - self._handle_cancellation(request_id_future, context) - ) - - try: - yield cancellation_task - finally: - # Clean up the background cancellation task - request_id = "unknown" - if request_id_future.done() and not request_id_future.cancelled(): - try: - request_id = request_id_future.result() - except Exception: - pass - - if not cancellation_task.done(): - logging.debug( - f"Cancelling cancellation monitor task for SGLang Request ID {request_id}, Context: {context.id()}" - ) - cancellation_task.cancel() - try: - await cancellation_task - except asyncio.CancelledError: - pass - else: - cancellation_task.result() - - if self.shutdown_event and self.shutdown_event.is_set(): - raise EngineShutdown("Engine was shut down during token generation") diff --git a/components/src/dynamo/sglang/request_handlers/image_diffusion/image_diffusion_handler.py b/components/src/dynamo/sglang/request_handlers/image_diffusion/image_diffusion_handler.py index 5e78da2ea36e..497d1e48840b 100644 --- a/components/src/dynamo/sglang/request_handlers/image_diffusion/image_diffusion_handler.py +++ b/components/src/dynamo/sglang/request_handlers/image_diffusion/image_diffusion_handler.py @@ -8,12 +8,16 @@ import random import time import uuid +from contextlib import AsyncExitStack from typing import Any, AsyncGenerator, Optional import torch from PIL import Image from dynamo._core import Context +from dynamo.common.http.base import HttpStatusError +from dynamo.common.http.media_reference import local_media_reference +from dynamo.common.http.url_validator import UrlValidationError, UrlValidationPolicy from dynamo.common.protocols.image_protocol import ImageNvExt from dynamo.common.storage import upload_to_fs from dynamo.llm.exceptions import InvalidArgument @@ -170,16 +174,43 @@ async def _generate_images( "seed": seed if seed is not None else random.randint(0, 1000000), } - # Add image_path for I2I/TI2I if provided - if input_reference is not None: - if not input_reference.strip(): - raise InvalidArgument("input_reference must be a non-empty string") - args["image_path"] = input_reference - - result = await asyncio.to_thread( - self.generator.generate, - sampling_params_kwargs=args, - ) + # Add image_path for I2I/TI2I if provided. A URL is fetched through the + # SSRF-safe client (revalidating each redirect hop) into a temp file, so + # the generator only ever opens a trusted local path โ€” never a URL it + # would fetch and redirect itself. + async with AsyncExitStack() as stack: + if input_reference is not None: + if not input_reference.strip(): + raise InvalidArgument("input_reference must be a non-empty string") + try: + args["image_path"] = await stack.enter_async_context( + local_media_reference( + input_reference, UrlValidationPolicy.from_env() + ) + ) + except UrlValidationError as exc: + # A policy verdict on a client-supplied reference is a bad + # request. UrlValidationError is a ValueError, which the + # binding already maps to InvalidArgument + # (backend.rs py_err_to_dynamo) โ€” this says so at the call + # site rather than relying on that fallback. + raise InvalidArgument(str(exc)) from exc + except HttpStatusError as exc: + if not 400 <= exc.status < 500: + raise + # The origin the client picked answered 4xx. image_loader + # takes the same line for an unreachable user-supplied URL + # ("a client error (400) rather than an internal server + # fault"). The URL is not echoed back: the client sent it, + # and it has no length limit. + raise InvalidArgument( + f"input_reference could not be fetched (HTTP {exc.status})" + ) from exc + + result = await asyncio.to_thread( + self.generator.generate, + sampling_params_kwargs=args, + ) # DiffGenerator.generate() returns GenerationResult | list[GenerationResult] | None if result is None: diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index f92ee7493b39..63a3d6871c50 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -16,13 +16,18 @@ from dynamo.common.backend import logprobs as _shared_logprobs from dynamo.common.constants import DisaggregationMode from dynamo.common.metadata_upload import MetadataUploader -from dynamo.common.multimodal.image_loader import ImageLoader +from dynamo.common.multimodal.image_loader import ( + ImageLoader, + image_cache_scope_from_request, +) from dynamo.common.multimodal.video_loader import VideoLoader from dynamo.common.utils.engine_response import normalize_finish_reason from dynamo.llm import HttpError -from dynamo.llm.exceptions import EngineShutdown +from dynamo.llm.exceptions import EngineShutdown, InvalidArgument from dynamo.sglang._compat import ( + cache_salt_kwargs, filter_supported_async_generate_kwargs, + prefill_dp_rank_kwargs, require_reasoning_kwargs, ) from dynamo.sglang._disagg import validate_disagg_parallel_sampling @@ -32,6 +37,7 @@ build_native_generate_request, native_generate_payload, native_generate_stream, + new_sglang_request_id, ) from dynamo.sglang.publisher import DynamoSglangPublisher from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler @@ -43,6 +49,7 @@ extract_media_urls, raise_if_unextracted_multimodal, ) +from dynamo.sglang.request_utils import request_cache_salt _SAMPLING_OPTION_FIELDS = ( "presence_penalty", @@ -54,6 +61,38 @@ "min_p", ) BYPASS_REMOTE_PREFILL_ANNOTATION = "x-bypass-remote-prefill" +_MAX_ABORT_MESSAGE_LENGTH = 8192 + + +def _raise_if_sglang_error(finish_reason: dict[str, Any]) -> None: + """Propagate error-bearing aborts before their placeholder output is emitted.""" + if finish_reason.get("type") != "abort": + return + + status_code = finish_reason.get("status_code") + if status_code is None and finish_reason.get("err_type") is None: + # SGLang also uses FINISH_ABORT for ordinary cancellation. Its default + # message is "Aborted", but both error metadata fields are absent/null. + return + + if ( + isinstance(status_code, bool) + or not isinstance(status_code, int) + or not 400 <= status_code < 600 + ): + status_code = 500 + + message = finish_reason.get("message") + if not isinstance(message, str) or not message: + message = "SGLang aborted the request with an error" + message = message[:_MAX_ABORT_MESSAGE_LENGTH] + + if status_code == 400: + # Explicit InvalidArgument keeps a client-visible validation message + # through Dynamo's Python/Rust boundary. Generic HttpError messages are + # diagnostic-only in the semantic error protocol. + raise InvalidArgument(message) + raise HttpError(status_code, message) def _raise_if_conditional_disagg_bypass(request: Dict[str, Any]) -> None: @@ -142,6 +181,95 @@ def _sampling_option_params(values: Dict[str, Any]) -> Dict[str, Any]: return params +def _ordered_cancellation_request_id( + request_id: str, + sampling_params: Any, + *, + supported: bool, + batched: bool = False, +) -> str | None: + """Use pre-output cancellation only when SGLang preserves the submitted ID.""" + if not supported or batched: + return None + if isinstance(sampling_params, list): + sampling_params = sampling_params[0] if sampling_params else {} + if isinstance(sampling_params, Mapping): + sample_count = sampling_params.get("n") or 1 + beam_width = sampling_params.get("beam_width") or 1 + # SGLang parallel sampling replaces normalized ``rid_`` values + # with unrelated UUIDs before scheduler dispatch. Until SGLang exposes + # a group abort API, the submitted ID cannot order a safe exact abort. + if sample_count > 1 and beam_width <= 1: + return None + return request_id + + +def _native_payload_is_batched(native_payload: Mapping[str, Any]) -> bool: + for field in ("prompt", "text"): + if isinstance(native_payload.get(field), list): + return True + input_ids = native_payload.get("input_ids") + if isinstance(input_ids, list) and input_ids and isinstance(input_ids[0], list): + return True + input_embeds = native_payload.get("input_embeds") + return bool( + isinstance(input_embeds, list) + and input_embeds + and isinstance(input_embeds[0], list) + and input_embeds[0] + and isinstance(input_embeds[0][0], list) + ) + + +def _native_session_request_id( + native_payload: Mapping[str, Any], context_id: str +) -> str | None: + """Preserve the public node ID used by SGLang session continuations.""" + if native_payload.get("session_params") is None: + return None + request_id = native_payload.get("rid") + if request_id is None or request_id == "": + return context_id + if not isinstance(request_id, str): + raise ValueError("native SGLang session requests require a scalar rid") + return request_id + + +def _public_native_response_id( + engine_response_id: Any, + response_index: Any, + internal_request_id: str | None, + response_request_id: str | list[str] | None, +) -> Any: + """Map only response IDs derived from Dynamo's submitted SGLang ID.""" + if not isinstance(engine_response_id, str) or not isinstance( + internal_request_id, str + ): + return engine_response_id + if engine_response_id == internal_request_id: + return ( + response_request_id + if isinstance(response_request_id, str) + else engine_response_id + ) + if ( + not isinstance(response_index, int) + or isinstance(response_index, bool) + or response_index < 0 + or engine_response_id != f"{internal_request_id}_{response_index}" + ): + return engine_response_id + if isinstance(response_request_id, str): + return f"{response_request_id}_{response_index}" + if ( + isinstance(response_request_id, list) + and response_index < len(response_request_id) + and isinstance(response_request_id[response_index], str) + ): + return response_request_id[response_index] + return engine_response_id + + def _user_stop_token_ids(request: Dict[str, Any]) -> set[int]: stop_conditions = request.get("stop_conditions") if isinstance(stop_conditions, dict): @@ -362,12 +490,25 @@ def _build_sampling_params(self, request: Dict[str, Any]) -> Dict[str, Any]: _plain = stop_conditions.get("stop_token_ids") or [] _merged = list(set(_hidden).union(_plain)) stop_token_ids = _merged if _merged else None + # A tokenizer-free SGLang rejects positive min_new_tokens. Let + # Dynamo's decoder enforce the floor and keep the engine running so + # it cannot stop before Dynamo reaches it. + min_tokens = stop_conditions.get("min_tokens") + hold_engine_open = bool( + min_tokens + and min_tokens > 0 + and self.config.server_args.skip_tokenizer_init + ) + if hold_engine_open: + stop_token_ids = None param_mapping = { "n": sampling_opts.get("n"), "max_new_tokens": stop_conditions.get("max_tokens"), - "min_new_tokens": stop_conditions.get("min_tokens"), - "ignore_eos": stop_conditions.get("ignore_eos"), + "min_new_tokens": None if hold_engine_open else min_tokens, + "ignore_eos": True + if hold_engine_open + else stop_conditions.get("ignore_eos"), "stop_token_ids": stop_token_ids, **_sampling_option_params(sampling_opts), **self._get_guided_decoding_params( @@ -416,6 +557,7 @@ def _native_generate_stream( native_payload: Mapping[str, Any], input_param: Dict[str, Any], context: Context, + sglang_request_id: str, priority: int | None, ) -> AsyncIterator[Dict[str, Any]]: """Build and dispatch one native SGLang request.""" @@ -436,16 +578,18 @@ def _native_generate_stream( native_request = build_native_generate_request( native_payload, input_ids=input_ids, - fallback_rid=context.trace_id or context.id(), + request_id=sglang_request_id, priority=self._priority_kwargs(priority).get("priority"), bootstrap_host=bootstrap_info.get("bootstrap_host"), bootstrap_port=bootstrap_info.get("bootstrap_port"), bootstrap_room=bootstrap_info.get("bootstrap_room"), - external_trace_header=context.trace_headers() - if self.enable_trace - else None, + external_trace_header=( + context.trace_headers() if self.enable_trace else None + ), routed_dp_rank=routing.get("dp_rank"), + prefill_dp_rank=routing.get("prefill_dp_rank"), lora_path=self._resolve_lora(request), + cache_salt=request_cache_salt(request), ) return native_generate_stream(self.engine, native_request) @@ -471,24 +615,60 @@ async def generate( if self._first_token_source is not None: self._first_token_source.bind(context, routing.get("dp_rank")) _raise_if_conditional_disagg_bypass(request) - trace_id = context.trace_id + sglang_request_id = new_sglang_request_id() input_param = self._get_input_param(request) priority = (request.get("routing") or {}).get("priority") native_payload = native_generate_payload(request) + native_session_request_id = None if native_payload is not None: + native_session_request_id = _native_session_request_id( + native_payload, context.id() + ) + if native_session_request_id is not None: + sglang_request_id = native_session_request_id + logging.info( + "Submitted SGLang Request ID: %s, Context: %s", + sglang_request_id, + context.id(), + ) + if native_payload is not None: + submitted_request_id = _ordered_cancellation_request_id( + sglang_request_id, + native_payload.get("sampling_params"), + supported=( + getattr(self, "_supports_ordered_cancellation", False) + # Session node IDs are caller-visible. Waiting for SGLang's + # response avoids aborting an unrelated request if a caller + # reuses an active node ID. + and native_session_request_id is None + ), + batched=_native_payload_is_batched(native_payload), + ) stream = self._native_generate_stream( request, native_payload, input_param, context, + sglang_request_id, priority, ) - async for output in self._process_native_generate_stream(stream, context): + async for output in self._process_native_generate_stream( + stream, + context, + submitted_request_id=submitted_request_id, + internal_request_id=sglang_request_id, + response_request_id=native_payload.get("rid") or context.id(), + ): yield output return priority_kwargs = self._priority_kwargs(priority) sampling_params = self._build_sampling_params(request) + submitted_request_id = _ordered_cancellation_request_id( + sglang_request_id, + sampling_params, + supported=getattr(self, "_supports_ordered_cancellation", False), + ) logprob_kwargs = self._build_logprob_kwargs(request) metadata_uploader = self._metadata_uploader_from_request(request) @@ -533,6 +713,7 @@ async def generate( decode = await self.engine.async_generate( **input_param, **decode_mm_kwargs, + **cache_salt_kwargs(self.engine, request_cache_salt(request)), sampling_params=sampling_params, stream=True, **require_reasoning_kwargs(self.engine, request), @@ -540,8 +721,9 @@ async def generate( bootstrap_host=bootstrap_info["bootstrap_host"], bootstrap_port=bootstrap_info["bootstrap_port"], bootstrap_room=bootstrap_info["bootstrap_room"], + **prefill_dp_rank_kwargs(self.engine, routing.get("prefill_dp_rank")), external_trace_header=trace_header, - rid=trace_id, + rid=sglang_request_id, data_parallel_rank=dp_rank, lora_path=lora_path, **logprob_kwargs, @@ -555,6 +737,7 @@ async def generate( return_tokens_as_token_ids, user_stop_token_ids=user_stop_token_ids, metadata_uploader=metadata_uploader, + submitted_request_id=submitted_request_id, ): yield out else: @@ -564,6 +747,7 @@ async def generate( request=request, user_stop_token_ids=user_stop_token_ids, metadata_uploader=metadata_uploader, + submitted_request_id=submitted_request_id, ): yield out else: @@ -582,7 +766,10 @@ async def generate( assert self._image_loader is not None image_items = mm_data.get(IMAGE_URL_KEY) or [] if image_items: - image_data = await self._image_loader.load_image_batch(image_items) + image_data = await self._image_loader.load_image_batch( + image_items, + cache_scope=image_cache_scope_from_request(request), + ) else: image_data = None @@ -616,6 +803,7 @@ async def generate( agg = await self.engine.async_generate( **input_param, + **cache_salt_kwargs(self.engine, request_cache_salt(request)), image_data=image_data, audio_data=audio_data, video_data=video_data, @@ -625,7 +813,7 @@ async def generate( **self._routed_experts_kwargs, **mm_hashes_kwargs, external_trace_header=trace_header, - rid=trace_id, + rid=sglang_request_id, data_parallel_rank=dp_rank, lora_path=lora_path, **logprob_kwargs, @@ -639,6 +827,7 @@ async def generate( return_tokens_as_token_ids, user_stop_token_ids=user_stop_token_ids, metadata_uploader=metadata_uploader, + submitted_request_id=submitted_request_id, ): yield out else: @@ -648,6 +837,7 @@ async def generate( request=request, user_stop_token_ids=user_stop_token_ids, metadata_uploader=metadata_uploader, + submitted_request_id=submitted_request_id, ): yield out @@ -655,25 +845,65 @@ async def _process_native_generate_stream( self, stream_source: AsyncIterator[Dict[str, Any]], context: Context, + submitted_request_id: str | None = None, + internal_request_id: str | None = None, + response_request_id: str | list[str] | None = None, ) -> AsyncGenerator[Dict[str, Any], None]: """Forward opaque SGLang chunks while retaining engine cancellation.""" request_id_future: asyncio.Future[str] = asyncio.Future() + request_ids: set[str] = set() first_output_seen = False - async with self._cancellation_monitor(request_id_future, context): - async for chunk in stream_source: + async with self._cancellation_monitor( + request_id_future, + context, + submitted_request_id, + request_ids=request_ids, + ) as cancellation_task: + async for chunk in self._stream_until_cancelled( + stream_source, cancellation_task + ): native_response = chunk["engine_data"]["sglang_response"] - if not request_id_future.done(): - sglang_request_id = native_response.get("meta_info", {}).get("id") - if sglang_request_id: + output = chunk + meta_info = native_response.get("meta_info", {}) + sglang_request_id = meta_info.get("id") + if sglang_request_id: + request_ids.add(sglang_request_id) + if not request_id_future.done(): request_id_future.set_result(sglang_request_id) logging.debug(f"New SGLang Request ID: {sglang_request_id}") + if meta_info.get("finish_reason"): + request_ids.discard(sglang_request_id) + if context.is_stopped(): + if submitted_request_id is None: + self._abort_requests(request_ids, context) + continue if not first_output_seen and ( native_response.get("output_ids") or native_response.get("text") ): first_output_seen = True context.notify_first_token() + if response_request_id is not None and isinstance(meta_info, dict): + engine_response_id = meta_info.get("id") + public_response_id = _public_native_response_id( + engine_response_id, + native_response.get("index"), + internal_request_id, + response_request_id, + ) + if public_response_id != engine_response_id: + public_response = { + **native_response, + "meta_info": {**meta_info, "id": public_response_id}, + } + output = { + **chunk, + "engine_data": { + **chunk["engine_data"], + "sglang_response": public_response, + }, + } if not context.is_stopped(): - yield chunk + yield output async def _process_token_stream( self, @@ -682,6 +912,7 @@ async def _process_token_stream( return_tokens_as_token_ids: bool = False, user_stop_token_ids: set[int] | None = None, metadata_uploader: MetadataUploader | None = None, + submitted_request_id: str | None = None, ) -> AsyncGenerator[Dict[str, Any], None]: """Process token-based stream output. @@ -697,21 +928,32 @@ async def _process_token_stream( """ # Use Future pattern for request ID - will be set when first response arrives request_id_future: asyncio.Future[str] = asyncio.Future() + request_ids: set[str] = set() first_output_seen = False - async with self._cancellation_monitor(request_id_future, context): - async for res in stream_source: + async with self._cancellation_monitor( + request_id_future, + context, + submitted_request_id, + request_ids=request_ids, + ) as cancellation_task: + async for res in self._stream_until_cancelled( + stream_source, cancellation_task + ): meta_info = res.get("meta_info", {}) - # Extract SGLang request ID from the first response and set the future - if not request_id_future.done(): - sglang_request_id = meta_info.get("id") - if sglang_request_id: + sglang_request_id = meta_info.get("id") + if sglang_request_id: + request_ids.add(sglang_request_id) + if not request_id_future.done(): request_id_future.set_result(sglang_request_id) logging.debug(f"New SGLang Request ID: {sglang_request_id}") + if meta_info.get("finish_reason"): + request_ids.discard(sglang_request_id) + if context.is_stopped(): + if submitted_request_id is None: + # A choice's first chunk can arrive after the monitor fired. + self._abort_requests(request_ids, context) + continue - # Check cancellation before yielding to allow proper cleanup. - # This lets SGLang proceed to the second token generation, which will - # async context switch and allow the abort monitor to signal cancellation. - # The loop should exit by itself when context.is_stopped() returns True. # SGLang omits index for non-n/legacy chunks; treat those as # choice 0 while preserving explicit indices for n>1. output_idx = res.get("index") or 0 @@ -728,6 +970,7 @@ async def _process_token_stream( raise EngineShutdown( "Engine was shut down during token generation" ) + _raise_if_sglang_error(finish_reason) out["finish_reason"] = normalize_finish_reason( finish_reason["type"] ) @@ -739,11 +982,9 @@ async def _process_token_stream( # With stream_output=True, output_ids contains only new tokens (disjoint) output_ids = res.get("output_ids", []) - # Empty, non-final chunks can happen during scheduler idle ticks. - # Keep waiting for the next chunk unless cancellation was requested. + # Keep draining empty, non-final chunks after cancellation so the + # abort monitor can run and SGLang can finish request cleanup. if not output_ids and not finish_reason: - if context.is_stopped(): - break continue if output_ids and not first_output_seen: @@ -821,6 +1062,7 @@ async def _process_text_stream( request: Dict[str, Any] | None = None, user_stop_token_ids: set[int] | None = None, metadata_uploader: MetadataUploader | None = None, + submitted_request_id: str | None = None, ) -> AsyncGenerator[Dict[str, Any], None]: """Process text-based stream output in OpenAI format. @@ -835,16 +1077,30 @@ async def _process_text_stream( # Use Future pattern for request ID - will be set when first response arrives request_id_future: asyncio.Future[str] = asyncio.Future() + request_ids: set[str] = set() first_output_seen = False - async with self._cancellation_monitor(request_id_future, context): - async for res in stream_source: + async with self._cancellation_monitor( + request_id_future, + context, + submitted_request_id, + request_ids=request_ids, + ) as cancellation_task: + async for res in self._stream_until_cancelled( + stream_source, cancellation_task + ): meta_info = res.get("meta_info", {}) - # Extract SGLang request ID from the first response and set the future - if not request_id_future.done(): - sglang_request_id = meta_info.get("id") - if sglang_request_id: + sglang_request_id = meta_info.get("id") + if sglang_request_id: + request_ids.add(sglang_request_id) + if not request_id_future.done(): request_id_future.set_result(sglang_request_id) logging.debug(f"New SGLang Request ID: {sglang_request_id}") + if meta_info.get("finish_reason"): + request_ids.discard(sglang_request_id) + if context.is_stopped(): + if submitted_request_id is None: + self._abort_requests(request_ids, context) + continue # Check cancellation before yielding to allow proper cleanup. # This lets SGLang proceed to the second token generation, which will @@ -871,6 +1127,7 @@ async def _process_text_stream( raise EngineShutdown( "Engine was shut down during token generation" ) + _raise_if_sglang_error(finish_reason) finish_reason_type = normalize_finish_reason(finish_reason["type"]) else: finish_reason_type = None @@ -888,7 +1145,7 @@ async def _process_text_stream( ) response = { - "id": meta_info["id"], + "id": context.id(), "created": int(time.time()), "choices": [choice_data], "model": self.config.server_args.served_model_name, diff --git a/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py b/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py index c188e6d639bb..8220c9250b7f 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py @@ -10,7 +10,7 @@ from dynamo._core import Context from dynamo.health_check import HEALTH_CHECK_KEY -from dynamo.sglang._compat import require_reasoning_kwargs +from dynamo.sglang._compat import cache_salt_kwargs, require_reasoning_kwargs from dynamo.sglang._disagg import validate_disagg_parallel_sampling from dynamo.sglang.agent_session import agent_session_kwargs from dynamo.sglang.args import Config @@ -18,14 +18,20 @@ build_native_generate_request, native_generate_payload, native_generate_stream, + new_sglang_request_id, ) from dynamo.sglang.publisher import DynamoSglangPublisher from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler -from dynamo.sglang.request_handlers.llm.decode_handler import _sampling_option_params +from dynamo.sglang.request_handlers.llm.decode_handler import ( + _native_payload_is_batched, + _ordered_cancellation_request_id, + _sampling_option_params, +) from dynamo.sglang.request_handlers.llm.mm_disagg_utils import ( build_disagg_mm_kwargs, raise_if_unextracted_multimodal, ) +from dynamo.sglang.request_utils import request_cache_salt # Sentinel value matching u32::MAX from the C/Go prefill-routing ABI. # This remains as a compatibility fallback for older callers that still encode @@ -87,7 +93,12 @@ async def generate( """ validate_disagg_parallel_sampling(request) logging.debug(f"New Request ID: {context.id()}") - trace_id = context.trace_id + sglang_request_id = new_sglang_request_id() + logging.debug( + "Submitted SGLang Request ID: %s, Context: %s", + sglang_request_id, + context.id(), + ) if "request" in request: # DisaggPreprocessedRequest format @@ -171,12 +182,15 @@ async def generate( priority_kwargs = self._priority_kwargs(priority) if native_payload is not None: input_ids = input_param.get("input_ids") + native_payload_is_batched = _native_payload_is_batched( + {"input_ids": input_ids} + ) if not isinstance(input_ids, list): raise ValueError("native SGLang Generate requires token input") native_request = build_native_generate_request( native_payload, input_ids=input_ids, - fallback_rid=trace_id or context.id(), + request_id=sglang_request_id, priority=priority_kwargs.get("priority"), sampling_overrides={"n": 1, "max_new_tokens": 1}, bootstrap_host=bootstrap_host, @@ -185,12 +199,26 @@ async def generate( external_trace_header=trace_header, routed_dp_rank=dp_rank, lora_path=lora_path, + cache_salt=request_cache_salt(inner_request), + ) + submitted_request_id = _ordered_cancellation_request_id( + sglang_request_id, + native_request.sampling_params, + supported=getattr(self, "_supports_ordered_cancellation", False), + batched=native_payload_is_batched, ) results = native_generate_stream(self.engine, native_request) else: + submitted_request_id = _ordered_cancellation_request_id( + sglang_request_id, + sampling_params, + supported=getattr(self, "_supports_ordered_cancellation", False), + batched=False, + ) results = await self.engine.async_generate( **input_param, **mm_kwargs, + **cache_salt_kwargs(self.engine, request_cache_salt(inner_request)), sampling_params=sampling_params, stream=True, **require_reasoning_kwargs(self.engine, inner_request), @@ -198,7 +226,7 @@ async def generate( bootstrap_port=bootstrap_port, bootstrap_room=bootstrap_room, external_trace_header=trace_header, - rid=trace_id, + rid=sglang_request_id, data_parallel_rank=dp_rank, lora_path=lora_path, **priority_kwargs, @@ -220,33 +248,43 @@ async def generate( "disaggregated_params": bootstrap_info, } - task = asyncio.create_task(self._consume_results(results, context)) + task = asyncio.create_task( + self._consume_results(results, submitted_request_id, context) + ) self._consume_tasks.add(task) task.add_done_callback(self._consume_tasks.discard) await task async def _consume_results( - self, results: AsyncIterator[Any], context: Context + self, + results: AsyncIterator[Any], + submitted_request_id: str | None, + context: Context, ) -> None: """Consume async generator results without processing. Args: results: Async generator from engine.async_generate. + submitted_request_id: Exact engine ID known before output, when supported. context: Context object for cancellation handling. """ - # Use Future pattern for request ID - will be set when first response arrives + # Preserve the response ID as a fallback if SGLang replaces the submitted ID. request_id_future: asyncio.Future[str] = asyncio.Future() - async with self._cancellation_monitor(request_id_future, context): - async for res in results: - # Extract SGLang request ID from the first response and set the future + async with self._cancellation_monitor( + request_id_future, context, submitted_request_id + ) as cancellation_task: + async for res in self._stream_until_cancelled(results, cancellation_task): if not request_id_future.done(): - meta_info = res.get("meta_info", {}) + meta_info = res.get("meta_info") or ( + res.get("engine_data", {}) + .get("sglang_response", {}) + .get("meta_info", {}) + ) sglang_request_id = meta_info.get("id") if sglang_request_id: request_id_future.set_result(sglang_request_id) logging.debug(f"New Prefill Request ID: {sglang_request_id}") - # Note: No explicit cancellation checks needed here. - # When abort_request is called by the cancellation monitor, - # SGLang will terminate this async generator automatically. + # The shared iterator briefly drains after abort so SGLang can + # clean up, then closes a stream that does not terminate. diff --git a/components/src/dynamo/sglang/request_handlers/multimodal/encode_worker_handler.py b/components/src/dynamo/sglang/request_handlers/multimodal/encode_worker_handler.py index 42209b9edd90..7ff645abcce2 100644 --- a/components/src/dynamo/sglang/request_handlers/multimodal/encode_worker_handler.py +++ b/components/src/dynamo/sglang/request_handlers/multimodal/encode_worker_handler.py @@ -24,6 +24,7 @@ from dynamo._core import Client, Context from dynamo.common.http import fetch_bytes +from dynamo.common.http.media_reference import max_media_bytes from dynamo.common.http.url_validator import UrlValidationPolicy, validate_media_url from dynamo.common.memory.multimodal_embedding_cache_manager import ( CachedEmbedding, @@ -35,7 +36,13 @@ MissingMediaDecoderError, video_decoder_missing, ) -from dynamo.common.multimodal.image_loader import DECODED_VARIANT_KEY, URL_VARIANT_KEY +from dynamo.common.multimodal.image_loader import ( + DECODED_VARIANT_KEY, + URL_VARIANT_KEY, + image_cache_scope_from_request, + image_cache_session_scoped_from_env, + scope_image_cache_key, +) from dynamo.common.multimodal.media_descriptor import decoded_content_hash_key from dynamo.common.multimodal.media_source import ( is_local_media_url, @@ -263,6 +270,8 @@ class MultimodalEncodeWorkerHandler(BaseWorkerHandler[SglangMultimodalRequest, s tokens, transfers embeddings, and forwards to the PD worker. """ + _session_scoped_cache = False + def __init__( self, config: Config, @@ -276,6 +285,7 @@ def __init__( self.model = config.server_args.model_path self._missing_video_cache_key_config_warned = False self._decoded_content_hash_warning_emitted = False + self._session_scoped_cache = image_cache_session_scoped_from_env() self._image_loader: Optional[ImageLoader] = ( ImageLoader(enable_frontend_decoding=True) if config.dynamo_args.frontend_decoding @@ -614,7 +624,12 @@ async def _maybe_nvdec_decoder(self, url: str) -> Optional[Any]: normalized = await validate_media_url(url, self._url_policy) scheme = urlparse(normalized).scheme if scheme in ("http", "https"): - content = await fetch_bytes(normalized, 30.0, policy=self._url_policy) + content = await fetch_bytes( + normalized, + 30.0, + policy=self._url_policy, + max_bytes=max_media_bytes(), + ) elif is_local_media_url(normalized): content = await read_local_media_bytes(normalized, self._url_policy) else: @@ -680,8 +695,9 @@ async def _build_encode_inputs( so SGLang does not re-download what we already validated and hold), or the original URL string when nothing was fetched. Non-video modalities and decoded inputs are returned unchanged. When - NVDEC is disabled or ineligible the URLs are returned policy-validated - and normalized, since SGLang fetches them with its own session. + NVDEC is disabled or ineligible, remote URLs are fetched under Dynamo's + policy and size limit and passed to SGLang as bytes; other URLs remain + policy-validated and normalized. Called from both the cached and uncached encode paths. The embedding cache is disabled by default, so routing this only through the cached @@ -691,9 +707,8 @@ async def _build_encode_inputs( return media_inputs if not self._nvdec_video_enabled(): # NVDEC off (CPU image, DYN_DISABLE_NVDEC, or a gated model type): - # these URLs go straight to SGLang's software path, which fetches - # them with its own session and never consults our url policy. Run - # the policy here so a source we would refuse is refused before + # these inputs go straight to SGLang's software path. Validate them + # here so a source we would refuse is refused before # SGLang can reach it -- and before we answer with anything about # this deployment, since a request we reject is not the place to # report which decoders are installed. @@ -713,7 +728,20 @@ async def _build_encode_inputs( and not _software_video_decoder_imports() ): raise video_decoder_missing("sglang", "decord2", "decord", None) - return validated + fetched_inputs: list[Any] = [] + for media_input in validated: + if isinstance(media_input, str) and urlparse(media_input).scheme in ( + "http", + "https", + ): + media_input = await fetch_bytes( + media_input, + 30.0, + policy=self._url_policy, + max_bytes=max_media_bytes(), + ) + fetched_inputs.append(media_input) + return fetched_inputs encode_inputs: list[Any] = [] for media_input in media_inputs: if not isinstance(media_input, str): @@ -918,14 +946,17 @@ def _parse_media_item(item: Any, media_name: str) -> tuple[str, Any]: return variant, item[variant] async def _prepare_image_inputs( - self, image_items: list[Any] + self, + image_items: list[Any], + cache_scope: str | None = None, ) -> tuple[list[Any], list[Optional[str]], dict[int, Optional[CachedEmbedding]],]: """Prepare MMEncoder inputs and aligned embedding-cache keys. URL variants stay as strings so the existing SGLang loading path is unchanged. Decoded variants are read from NIXL and become PIL Images. Their cache keys come from the canonical content hash serialized by the - Rust media decoder. + Rust media decoder. When session scoping is enabled, a missing scope + leaves both URL and decoded inputs unkeyed. """ encoder_inputs: list[Any] = [None] * len(image_items) cache_keys: list[Optional[str]] = [None] * len(image_items) @@ -943,7 +974,11 @@ async def _prepare_image_inputs( raise ValueError(f"Unsupported image data variant: {item}") encoder_inputs[index] = url if cache is not None: - cache_keys[index] = self._url_hash(url) + cache_keys[index] = scope_image_cache_key( + self._url_hash(url), + cache_scope, + session_scoped_cache=self._session_scoped_cache, + ) continue if not isinstance(value, dict): @@ -955,9 +990,8 @@ async def _prepare_image_inputs( ) if cache is not None: - cache_key = decoded_content_hash_key(value) - cache_keys[index] = cache_key - if cache_key is None: + content_hash = decoded_content_hash_key(value) + if content_hash is None: if not self._decoded_content_hash_warning_emitted: logger.warning( "Frontend-decoded image descriptor has a missing or invalid " @@ -967,10 +1001,17 @@ async def _prepare_image_inputs( ) self._decoded_content_hash_warning_emitted = True else: - cached_entry = cache.get(cache_key) - prechecked_entries[index] = cached_entry - if cached_entry is not None: - continue + cache_key = scope_image_cache_key( + content_hash, + cache_scope, + session_scoped_cache=self._session_scoped_cache, + ) + cache_keys[index] = cache_key + if cache_key is not None: + cached_entry = cache.get(cache_key) + prechecked_entries[index] = cached_entry + if cached_entry is not None: + continue decoded_items.append({DECODED_VARIANT_KEY: value}) decoded_indices.append(index) @@ -1030,7 +1071,10 @@ async def generate( image_inputs, image_cache_keys, image_prechecked_entries, - ) = await self._prepare_image_inputs(image_items) + ) = await self._prepare_image_inputs( + image_items, + cache_scope=image_cache_scope_from_request(raw_request), + ) video_cache_keys: list[Optional[str]] if self._embedding_cache is None: video_cache_keys = [None] * len(video_urls) diff --git a/components/src/dynamo/sglang/request_handlers/rerank/__init__.py b/components/src/dynamo/sglang/request_handlers/rerank/__init__.py new file mode 100644 index 000000000000..d379f644f529 --- /dev/null +++ b/components/src/dynamo/sglang/request_handlers/rerank/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dynamo.sglang.request_handlers.rerank.rerank_handler import RerankWorkerHandler + +__all__ = ["RerankWorkerHandler"] diff --git a/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py b/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py new file mode 100644 index 000000000000..0b2f8165b189 --- /dev/null +++ b/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import heapq +from collections.abc import AsyncGenerator +from typing import Any + +import sglang as sgl +from sglang.srt.managers.io_struct import EmbeddingReqInput + +from dynamo._core import Context +from dynamo.sglang.args import Config +from dynamo.sglang.protocol import RerankRequest +from dynamo.sglang.publisher import DynamoSglangPublisher +from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler + + +def _looks_like_decoder_reranker(tokenizer_manager: Any) -> bool: + """Detect decoder-only Qwen3 rerankers that need SGLang's generation path.""" + tokenizer = getattr(tokenizer_manager, "tokenizer", None) + chat_template = getattr(tokenizer, "chat_template", "") + template = chat_template.lower() if isinstance(chat_template, str) else "" + model_config = getattr(tokenizer_manager, "model_config", None) + model_path = str(getattr(model_config, "model_path", "")).lower() + has_yes_no_template = ( + "answer can only be" in template and "yes" in template and "no" in template + ) + return ( + has_yes_no_template + or "qwen3-reranker" in model_path + or "qwen3-vl-reranker" in model_path + ) + + +class RerankWorkerHandler(BaseWorkerHandler): + """Runs stable text-only cross-encoder reranking through SGLang pooling.""" + + def __init__( + self, + engine: sgl.Engine, + config: Config, + publisher: DynamoSglangPublisher | None = None, + shutdown_event: asyncio.Event | None = None, + ) -> None: + self.validate_engine(engine) + super().__init__(engine, config, publisher, None, shutdown_event) + + @staticmethod + def validate_engine(engine: sgl.Engine) -> None: + if _looks_like_decoder_reranker(engine.tokenizer_manager): + raise ValueError( + "Dynamo's SGLang /v1/rerank integration currently supports " + "cross-encoder rerankers only; decoder-only Qwen3 and Qwen3-VL " + "rerankers are not supported" + ) + + def cleanup(self) -> None: + super().cleanup() + self.engine.shutdown() + + @staticmethod + def _validate(request: RerankRequest) -> None: + if not request.query.strip(): + raise ValueError("Query cannot be empty or whitespace only") + if not request.documents: + raise ValueError("Documents cannot be empty") + if any(not document.strip() for document in request.documents): + raise ValueError("Each document cannot be empty or whitespace only") + if request.top_n is not None and request.top_n < 1: + raise ValueError("Parameter 'top_n' must be larger than 0") + + @staticmethod + def _score(item: Any, index: int) -> float: + if not isinstance(item, dict) or "embedding" not in item: + raise ValueError(f"Missing embedding score for rerank at index {index}") + score = item["embedding"] + if isinstance(score, list): + if not score or not isinstance(score[0], (int, float)): + raise ValueError( + f"Invalid embedding score for rerank at index {index}: {score!r}" + ) + score = score[0] + if not isinstance(score, (int, float)): + raise ValueError( + f"Invalid embedding score for rerank at index {index}: {score!r}" + ) + return float(score) + + async def generate( + self, request: dict[str, Any], context: Context + ) -> AsyncGenerator[list[dict[str, Any]], None]: + rerank_request = RerankRequest(**request) + self._validate(rerank_request) + + tokenizer_manager = self.engine.tokenizer_manager + pairs = [ + [rerank_request.query, document] for document in rerank_request.documents + ] + trace_id = context.trace_id + request_ids = ( + [f"{trace_id}-{index}" for index in range(len(pairs))] + if trace_id is not None + else None + ) + internal_request = EmbeddingReqInput( + text=pairs, + is_cross_encoder_request=True, + external_trace_header=( + context.trace_headers() if self.enable_trace else None + ), + rid=request_ids, + ) + responses = tokenizer_manager.generate_request(internal_request, None) + result = await anext(responses) + if not isinstance(result, list): + result = [result] + if len(result) != len(rerank_request.documents): + raise ValueError( + "SGLang returned an unexpected number of rerank scores: " + f"expected {len(rerank_request.documents)}, got {len(result)}" + ) + + ranked = [] + for index, item in enumerate(result): + response: dict[str, Any] = { + "score": self._score(item, index), + "index": index, + } + if rerank_request.return_documents: + response["document"] = rerank_request.documents[index] + meta_info = item.get("meta_info") if isinstance(item, dict) else None + if meta_info is not None: + response["meta_info"] = meta_info + ranked.append(response) + + if rerank_request.top_n is not None: + ranked = heapq.nlargest( + rerank_request.top_n, ranked, key=lambda item: item["score"] + ) + else: + ranked.sort(key=lambda item: item["score"], reverse=True) + yield ranked diff --git a/components/src/dynamo/sglang/request_handlers/video_generation/video_generation_handler.py b/components/src/dynamo/sglang/request_handlers/video_generation/video_generation_handler.py index f94c83dd8335..7ab995502fd1 100644 --- a/components/src/dynamo/sglang/request_handlers/video_generation/video_generation_handler.py +++ b/components/src/dynamo/sglang/request_handlers/video_generation/video_generation_handler.py @@ -7,11 +7,14 @@ import logging import random import time +from contextlib import AsyncExitStack from typing import Any, AsyncGenerator, Optional import torch from dynamo._core import Context +from dynamo.common.http.media_reference import local_media_reference +from dynamo.common.http.url_validator import UrlValidationPolicy from dynamo.common.storage import upload_to_fs from dynamo.sglang.args import Config from dynamo.sglang.protocol import ( @@ -227,24 +230,31 @@ async def _generate_video( "seed": seed if seed is not None else random.randint(0, 1000000), } - # Add image_path for I2V if provided - if input_reference: - args["image_path"] = input_reference - - logger.info( - f"Generating video with {num_frames} frames at {width}x{height}, " - f"{num_inference_steps} steps, request_id={request_id}" - ) + # Add image_path for I2V if provided. A URL is fetched through the + # SSRF-safe client (revalidating each redirect hop) into a temp file, so + # the generator only ever opens a trusted local path. + async with AsyncExitStack() as stack: + if input_reference: + args["image_path"] = await stack.enter_async_context( + local_media_reference( + input_reference, UrlValidationPolicy.from_env() + ) + ) - # Serialize access -- DiffGenerator has mutable state (CUDA graph - # caches, shared config objects) and is not thread-safe. - async with self._generate_lock: - # Run in thread pool to avoid blocking event loop - result = await asyncio.to_thread( - self.generator.generate, - sampling_params_kwargs=args, + logger.info( + f"Generating video with {num_frames} frames at {width}x{height}, " + f"{num_inference_steps} steps, request_id={request_id}" ) + # Serialize access -- DiffGenerator has mutable state (CUDA graph + # caches, shared config objects) and is not thread-safe. + async with self._generate_lock: + # Run in thread pool to avoid blocking event loop + result = await asyncio.to_thread( + self.generator.generate, + sampling_params_kwargs=args, + ) + # DiffGenerator.generate() returns GenerationResult | list[GenerationResult] | None if result is None: raise RuntimeError("DiffGenerator returned None") diff --git a/components/src/dynamo/sglang/request_utils.py b/components/src/dynamo/sglang/request_utils.py new file mode 100644 index 000000000000..8d59a60ecf8f --- /dev/null +++ b/components/src/dynamo/sglang/request_utils.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Mapping +from typing import Any + + +def request_cache_salt(request: Mapping[str, Any]) -> str | None: + """Return the first non-empty cache salt, preferring resolved routing hints.""" + extra_args = request.get("extra_args") or {} + for source in ( + request.get("routing"), + extra_args.get("nvext") if isinstance(extra_args, Mapping) else None, + request.get("nvext"), + request, + ): + if isinstance(source, Mapping): + salt = source.get("cache_salt") + if salt: + return salt + return None diff --git a/components/src/dynamo/sglang/snapshot.py b/components/src/dynamo/sglang/snapshot.py index 8279a8fa28ed..2a4086726cec 100644 --- a/components/src/dynamo/sglang/snapshot.py +++ b/components/src/dynamo/sglang/snapshot.py @@ -8,9 +8,15 @@ import logging import os import time -from typing import Any +from typing import TYPE_CHECKING, Any import sglang as sgl +from sglang.srt.managers.io_struct import ( + ContinueGenerationReqInput, + PauseGenerationReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, +) from dynamo.common.snapshot.lifecycle import ( EngineSnapshotController, @@ -19,11 +25,34 @@ ) from dynamo.sglang._compat import override_server_args, resolved_server_args -from .pause import SGLangEnginePauseController +if TYPE_CHECKING: + from dynamo.sglang.args import Config logger = logging.getLogger(__name__) +class _SGLangSnapshotLifecycle: + """Adapt SGLang's native controls to the snapshot lifecycle interface.""" + + def __init__(self, engine: sgl.Engine): + self._tokenizer_manager = engine.tokenizer_manager + + async def pause(self) -> None: + await self._tokenizer_manager.pause_generation(PauseGenerationReqInput()) + await self._tokenizer_manager.release_memory_occupation( + ReleaseMemoryOccupationReqInput(), None + ) + + async def resume(self) -> None: + await self._tokenizer_manager.resume_memory_occupation( + ResumeMemoryOccupationReqInput(), None + ) + await self._tokenizer_manager.continue_generation(ContinueGenerationReqInput()) + + def mark_resumed(self) -> None: + pass + + async def warmup_engine(engine: sgl.Engine, server_args: Any) -> None: """Warm up the direct SGLang Engine generation path before snapshot capture.""" DUMMY_WARMUP_MAX_NEW_TOKENS = 8 @@ -115,7 +144,7 @@ async def warmup_engine(engine: sgl.Engine, server_args: Any) -> None: async def prepare_snapshot_engine( - server_args, + config: "Config", ) -> EngineSnapshotController[sgl.Engine] | None: """Single entry point for Dynamo Snapshot integration. @@ -134,6 +163,8 @@ async def prepare_snapshot_engine( if snapshot_config is None: return None + server_args = config.server_args + configure_snapshot_capture_env() logger.info("Snapshot mode enabled (watcher-driven signals)") @@ -164,6 +195,7 @@ async def prepare_snapshot_engine( logger.info( f"SGLang engine loaded in {time.time() - start_time:.2f}s (snapshot mode)" ) + config.validate_engine_server_args(engine.server_args) runtime_server_args = resolved_server_args(engine.server_args) await warmup_engine(engine, runtime_server_args) @@ -171,7 +203,7 @@ async def prepare_snapshot_engine( snapshot_controller = EngineSnapshotController( engine=engine, - pause_controller=SGLangEnginePauseController(engine), + pause_controller=_SGLangSnapshotLifecycle(engine), snapshot_config=snapshot_config, ) if not await snapshot_controller.wait_for_restore(): diff --git a/components/src/dynamo/sglang/tests/test_elastic_ep_preflight.py b/components/src/dynamo/sglang/tests/test_elastic_ep_preflight.py new file mode 100644 index 000000000000..328b62cb4ab1 --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_elastic_ep_preflight.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Behaviour of the ``--elastic-ep-backend mooncake`` startup check. + +Imports neither ``sglang`` nor any ``dynamo.sglang`` module that imports it, +so these run in an image where the engine is not installed or not importable +-- which is the only reason the check lives in its own module rather than in +``args.py``. The environment probes are the seams; the decision the check +makes over their results is what is exercised here. +""" + +import builtins +import importlib.util +import sys + +import pytest + +from dynamo.sglang import elastic_ep_preflight +from dynamo.sglang.elastic_ep_preflight import check_elastic_ep_backend + +pytestmark = [ + pytest.mark.unit, + pytest.mark.gpu_0, + pytest.mark.pre_merge, +] + +_IMPORT_FAILURE = ( + "mooncake.pg: ImportError: Mooncake PG was not built against torch==2.11.0; " + "mooncake.ep: ModuleNotFoundError: No module named 'mooncake.ep'" +) + + +def _simulate_environment( + monkeypatch, + *, + import_failure, + registered_backends, + mooncake_versions, + extension_modules, + torch_version="2.11.0+cu130", +): + """Pin every probe so the check runs against a known image shape.""" + monkeypatch.setattr( + elastic_ep_preflight, + "_import_process_group_extension", + lambda: import_failure, + ) + monkeypatch.setattr( + elastic_ep_preflight, + "_registered_torch_backends", + lambda: dict(registered_backends), + ) + monkeypatch.setattr( + elastic_ep_preflight, + "_installed_mooncake_versions", + lambda: dict(mooncake_versions), + ) + monkeypatch.setattr( + elastic_ep_preflight, + "_available_extension_modules", + lambda: list(extension_modules), + ) + monkeypatch.setattr(elastic_ep_preflight, "_torch_version", lambda: torch_version) + + +def _simulate_broken_mooncake(monkeypatch, **overrides): + """An image whose mooncake ProcessGroup extension does not load.""" + kwargs = { + "import_failure": _IMPORT_FAILURE, + "registered_backends": {"gloo": ("cpu",), "nccl": ("cuda",)}, + "mooncake_versions": {"mooncake-transfer-engine-cuda13": "0.3.11.post1"}, + "extension_modules": ["pg_2_9_1", "pg_2_10_0"], + } + kwargs.update(overrides) + _simulate_environment(monkeypatch, **kwargs) + + +def _simulate_healthy_mooncake(monkeypatch): + """An image whose mooncake ProcessGroup extension loads and registers.""" + _simulate_environment( + monkeypatch, + import_failure=None, + registered_backends={ + "gloo": ("cpu",), + "nccl": ("cuda",), + "mooncake": ("cuda",), + "mooncake-cpu": ("cpu",), + }, + mooncake_versions={"mooncake-transfer-engine-cuda13": "0.3.11.post1"}, + extension_modules=["pg_2_11_0"], + ) + + +def test_rejects_mooncake_backend_the_image_cannot_serve(monkeypatch): + """The worker fails at argument parsing, naming what the operator needs.""" + _simulate_broken_mooncake(monkeypatch) + + with pytest.raises(ValueError) as excinfo: + check_elastic_ep_backend("mooncake", True) + + message = str(excinfo.value) + assert "--elastic-ep-backend" in message + # The three facts that separate a bad wheel from a torch-version skew, so + # the failure is diagnosable without reproducing the crash. + assert "0.3.11.post1" in message + assert "2.11.0+cu130" in message + assert "pg_2_9_1" in message + # --enable-dp-attention is what drives the collective on the first forward + # pass, so it has to be called out when it is part of the configuration. + assert "--enable-dp-attention" in message + + +def test_reports_missing_mooncake_install_explicitly(monkeypatch): + """No mooncake at all reads as an absence, not a blank version field.""" + _simulate_broken_mooncake(monkeypatch, mooncake_versions={}, extension_modules=[]) + + with pytest.raises(ValueError) as excinfo: + check_elastic_ep_backend("mooncake") + + assert "none installed" in str(excinfo.value) + + +@pytest.mark.parametrize("requested_backend", [None, "nccl"]) +def test_ignores_backends_other_than_mooncake(monkeypatch, requested_backend): + """Negative control: the same broken image is fine for everyone else. + + Guards against the check becoming an unconditional startup failure for + workers that never asked for the mooncake transport. + """ + _simulate_broken_mooncake(monkeypatch) + + check_elastic_ep_backend(requested_backend, True) + + +def test_accepts_mooncake_when_the_image_can_serve_it(monkeypatch): + """Negative control: a working image is not blocked.""" + _simulate_healthy_mooncake(monkeypatch) + + check_elastic_ep_backend("mooncake", True) + + +def test_rejects_a_wheel_that_registers_only_half_the_process_groups(monkeypatch): + """A device backend without its CPU peer is not a usable image. + + SGLang builds both an elastic-EP device group and a CPU group for the + metadata collectives, so accepting on ``mooncake`` alone lets a wheel that + registered only that one pass here and still fail during engine startup. + """ + _simulate_environment( + monkeypatch, + import_failure=None, + registered_backends={ + "gloo": ("cpu",), + "nccl": ("cuda",), + "mooncake": ("cuda",), + }, + mooncake_versions={"mooncake-transfer-engine-cuda13": "0.3.11.post1"}, + extension_modules=["pg_2_11_0"], + ) + + with pytest.raises(ValueError) as excinfo: + check_elastic_ep_backend("mooncake") + + # The operator has to be told which half is absent; the registry line alone + # shows a mooncake backend present and reads as a working image. + assert "mooncake-cpu" in str(excinfo.value) + + +def test_unreadable_torch_backend_registry_does_not_block_startup(monkeypatch): + """A registry that cannot be read is not evidence that mooncake is absent. + + ``Backend.backend_capability`` is a torch internal. Reading it and failing + must not be worth more than the extension importing cleanly, or a torch-side + rename refuses every worker on this transport. + """ + _simulate_broken_mooncake(monkeypatch, import_failure=None) + monkeypatch.setattr( + elastic_ep_preflight, "_registered_torch_backends", lambda: None + ) + + check_elastic_ep_backend("mooncake", True) + + +def test_import_probe_reports_every_module_name_it_tried(monkeypatch): + """The real probe, run against an import system with no mooncake in it. + + Blocking at ``sys.meta_path`` rather than by uninstalling makes the case + behave the same in an image that does ship a working wheel. + """ + monkeypatch.setattr( + elastic_ep_preflight, + "_required_process_group_modules", + lambda: ("mooncake.pg", "mooncake.ep"), + ) + + class _RefuseMooncake: + def find_spec(self, name, path=None, target=None): + if name == "mooncake" or name.startswith("mooncake."): + raise ImportError(f"blocked for this test: {name}") + return None + + monkeypatch.setattr(sys, "meta_path", [_RefuseMooncake(), *sys.meta_path]) + for cached in [ + name + for name in sys.modules + if name == "mooncake" or name.startswith("mooncake.") + ]: + monkeypatch.delitem(sys.modules, cached) + + failure = elastic_ep_preflight._import_process_group_extension() + + assert failure is not None + # Both names have to appear: which one an image ships is the first thing + # an operator needs to compare against the engine version. + assert "mooncake.pg" in failure + assert "mooncake.ep" in failure + + +def _install_fake_sglang_sources(monkeypatch, tmp_path, import_line): + """Point the resolver at a source tree naming one ProcessGroup module.""" + package = tmp_path / "sglang" + elastic_ep = package / "srt" / "elastic_ep" + elastic_ep.mkdir(parents=True) + (elastic_ep / "elastic_ep.py").write_text(import_line, encoding="utf-8") + + spec = importlib.util.spec_from_file_location( + "sglang", package / "__init__.py", submodule_search_locations=[str(package)] + ) + monkeypatch.setattr( + elastic_ep_preflight.importlib.util, + "find_spec", + lambda name: spec if name == "sglang" else None, + ) + + +@pytest.mark.parametrize( + "import_line, expected", + [ + ("from mooncake.pg import MooncakeBackendOptions\n", ("mooncake.pg",)), + ("from mooncake.ep import MooncakeBackendOptions\n", ("mooncake.ep",)), + ], +) +def test_probes_only_the_module_the_installed_engine_imports( + monkeypatch, tmp_path, import_line, expected +): + """The half of the rename the engine does not use must not satisfy the check. + + Accepting either name lets an image whose wheel ships only the other one + pass here and still fail during engine startup, which is the failure this + module exists to move earlier. + """ + _install_fake_sglang_sources(monkeypatch, tmp_path, import_line) + + assert elastic_ep_preflight._required_process_group_modules() == expected + + +def test_unreadable_engine_sources_probe_both_module_names(monkeypatch, tmp_path): + """Sources that are present but unreadable are not evidence for either name. + + This is the case the resolver is actually written for: an editable install + whose engine tree the worker's uid cannot read, which is the same condition + the repository ``conftest.py`` reports at collection time. Narrowing on a + guess would refuse workers over a permissions problem, so the resolver + widens back to both names instead. + """ + _install_fake_sglang_sources( + monkeypatch, tmp_path, "from mooncake.pg import MooncakeBackendOptions\n" + ) + + probed = [] + real_open = builtins.open + + def _refuse_engine_sources(file, *args, **kwargs): + if str(file).endswith(".py") and "sglang" in str(file): + probed.append(str(file)) + raise PermissionError(13, "Permission denied", str(file)) + return real_open(file, *args, **kwargs) + + monkeypatch.setattr(builtins, "open", _refuse_engine_sources) + + assert elastic_ep_preflight._required_process_group_modules() == ( + "mooncake.pg", + "mooncake.ep", + ) + # The fallback must come from the read failing, not from the resolver + # skipping the sources it was pointed at. + assert probed diff --git a/components/src/dynamo/sglang/tests/test_runtime_metadata.py b/components/src/dynamo/sglang/tests/test_runtime_metadata.py index 76ff8c13ee6f..0d1b94cbc181 100644 --- a/components/src/dynamo/sglang/tests/test_runtime_metadata.py +++ b/components/src/dynamo/sglang/tests/test_runtime_metadata.py @@ -135,13 +135,31 @@ def test_token_budget_matches_sglang_policy( assert _get_token_budget(engine, server_args) == expected -def test_runtime_config_without_engine_omits_token_budget(monkeypatch, caplog): +@pytest.mark.parametrize( + "disaggregation_mode, runtime_supported, enable_multimodal, capability_expected", + [ + (None, True, False, False), + ("null", True, False, False), + ("prefill", False, False, False), + ("prefill", True, True, False), + ("prefill", True, False, True), + ("decode", True, False, True), + ], +) +def test_runtime_config_publishes_supported_disagg_capabilities( + monkeypatch, + caplog, + disaggregation_mode, + runtime_supported, + enable_multimodal, + capability_expected, +): from dynamo.sglang import register server_args = SimpleNamespace( allow_auto_truncate=False, context_length=4096, - disaggregation_mode=None, + disaggregation_mode=disaggregation_mode, max_prefill_tokens=None, page_size=16, speculative_algorithm="NONE", @@ -149,6 +167,7 @@ def test_runtime_config_without_engine_omits_token_budget(monkeypatch, caplog): ) dynamo_args = register.DynamoConfig() dynamo_args.enable_local_indexer = False + dynamo_args.enable_multimodal = enable_multimodal capacity = SimpleNamespace( max_num_seqs=None, max_num_batched_tokens=None, @@ -164,12 +183,32 @@ def test_runtime_config_without_engine_omits_token_budget(monkeypatch, caplog): monkeypatch.setattr(register, "get_spec_decode_runtime_data", lambda _: None) monkeypatch.setattr(register, "_get_mooncake_runtime_data", lambda _: None) monkeypatch.setattr(register, "runtime_capacity", lambda *_: capacity) + monkeypatch.setattr( + register, "supports_disagg_prefill_cancel_anytime", lambda _: runtime_supported + ) + engine = None + if runtime_supported: + engine = SimpleNamespace( + tokenizer_manager=SimpleNamespace( + context_len=4096, + validate_total_tokens=True, + num_reserved_tokens=0, + rid_to_state={}, + ), + _scheduler_init_result=SimpleNamespace(scheduler_infos=[{}]), + ) runtime_config = asyncio.run( - register.get_runtime_config(None, server_args, dynamo_args) + register.get_runtime_config(engine, server_args, dynamo_args) ) - assert TOKEN_BUDGET_RUNTIME_KEY not in runtime_config.runtime_data + assert ( + TOKEN_BUDGET_RUNTIME_KEY in runtime_config.runtime_data + ) is runtime_supported + capability = register.DISAGG_PREFILL_CANCEL_ANYTIME_V1 + assert (capability in runtime_config.runtime_data) is capability_expected + if capability_expected: + assert json.loads(runtime_config.runtime_data[capability]) is True assert "Failed to get runtime config" not in caplog.text diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_cancellation.py b/components/src/dynamo/sglang/tests/test_sglang_decode_cancellation.py new file mode 100644 index 000000000000..c2024a44333b --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_cancellation.py @@ -0,0 +1,1190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import logging +from types import SimpleNamespace + +import pytest + +from dynamo.common.constants import DisaggregationMode +from dynamo.llm.exceptions import EngineShutdown +from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler +from dynamo.sglang.request_handlers.llm.decode_handler import DecodeWorkerHandler + +pytestmark = [ + pytest.mark.unit, + pytest.mark.sglang, + pytest.mark.fault_tolerance, + pytest.mark.gpu_0, + pytest.mark.profiled_vram_gib(0), + pytest.mark.pre_merge, +] + + +async def _collect(stream): + return [item async for item in stream] + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancellation_monitor_rechecks_shutdown_after_cleanup(): + handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler) + handler.shutdown_event = asyncio.Event() + handler._abort_tasks = set() + + async def set_shutdown_when_cancelled(*_args): + try: + await asyncio.Future() + except asyncio.CancelledError: + handler.shutdown_event.set() + raise + + handler._handle_cancellation = set_shutdown_when_cancelled + request_id_future = asyncio.get_running_loop().create_future() + request_id_future.set_result("sglang-request-id") + context = SimpleNamespace(id=lambda: "request-id") + + with pytest.raises(EngineShutdown, match="shut down during token generation"): + async with handler._cancellation_monitor(request_id_future, context): + await asyncio.sleep(0) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancellation_monitor_keeps_pending_ordered_abort_alive( + decode_cancellation_case, +): + case = decode_cancellation_case + request_id_future = asyncio.get_running_loop().create_future() + + async with case.handler._cancellation_monitor( + request_id_future, + case.context, + submitted_request_id="internal-request-id", + ) as cancellation_task: + case.cancelled.set() + await asyncio.wait_for(cancellation_task, timeout=1) + + assert not request_id_future.cancelled() + case.registry["internal-request-id"] = SimpleNamespace( + time_stats=SimpleNamespace(api_server_dispatch_finish_time=1.0) + ) + case.dispatched.set() + await asyncio.wait_for(case.aborted.wait(), timeout=1) + await asyncio.gather(*tuple(case.handler._abort_tasks)) + + assert case.abort_calls == [("internal-request-id", False)] + case.registry.clear() + request_id_future.cancel() + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_pending_ordered_abort_has_bounded_registration_wait( + decode_cancellation_case, monkeypatch, caplog +): + caplog.set_level(logging.WARNING) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation._CANCELLATION_REGISTRATION_WAIT_TIMEOUT_S", + 0, + ) + case = decode_cancellation_case + request_id_future = asyncio.get_running_loop().create_future() + + async with case.handler._cancellation_monitor( + request_id_future, + case.context, + submitted_request_id="internal-request-id", + ) as cancellation_task: + case.cancelled.set() + await asyncio.wait_for(cancellation_task, timeout=1) + + await asyncio.gather(*tuple(case.handler._abort_tasks)) + assert not case.abort_calls + assert ( + "Timed out waiting for SGLang Request ID internal-request-id to register" + in caplog.messages + ) + request_id_future.cancel() + + +@pytest.fixture +def decode_cancellation_case(monkeypatch): + handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler) + handler._abort_tasks = set() + handler.shutdown_event = None + handler.use_sglang_tokenizer = False + handler.config = SimpleNamespace( + server_args=SimpleNamespace(served_model_name="test-model"), + dynamo_args=SimpleNamespace(enable_rl=False), + ) + handler._first_token_source = None + handler._supports_ordered_cancellation = True + monkeypatch.setattr( + "dynamo.sglang.request_handlers.llm.decode_handler.new_sglang_request_id", + lambda: "internal-request-id", + ) + handler.shutdown_event = asyncio.Event() + handler.serving_mode = DisaggregationMode.DECODE + handler.enable_trace = False + handler._routed_experts_kwargs = {} + handler._enable_frontend_decoding = False + handler._mm_hashes_supported = False + handler._get_input_param = lambda request: {"input_ids": [1]} + handler._build_sampling_params = lambda request: {"max_new_tokens": 1} + handler._build_logprob_kwargs = lambda request: {} + handler._resolve_lora = lambda request: None + handler._priority_kwargs = lambda priority: {} + monkeypatch.setattr( + "dynamo.sglang.request_handlers.llm.decode_handler.require_reasoning_kwargs", + lambda *args: {}, + ) + + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation.resolved_server_args", + lambda server_args: server_args, + ) + polling = asyncio.Queue() + + async def observed_sleep(delay): + polling.put_nowait(delay) + await asyncio.sleep(delay) + + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation.asyncio", + SimpleNamespace(**{**vars(asyncio), "sleep": observed_sleep}), + ) + started = asyncio.Event() + allow_registration = asyncio.Event() + registered = asyncio.Event() + allow_dispatch = asyncio.Event() + allow_dispatch.set() + dispatched = asyncio.Event() + aborted = asyncio.Event() + cancelled = asyncio.Event() + drained = asyncio.Event() + abort_calls = [] + registry = {} + + def cancellation_future(): + return asyncio.create_task(cancelled.wait()) + + context = SimpleNamespace( + id=lambda: "dynamo-request", + trace_id="trace-request", + async_killed_or_stopped=cancellation_future, + is_stopped=cancelled.is_set, + notify_first_token=lambda: None, + ) + case = SimpleNamespace( + handler=handler, + context=context, + started=started, + allow_registration=allow_registration, + registered=registered, + allow_dispatch=allow_dispatch, + dispatched=dispatched, + aborted=aborted, + cancelled=cancelled, + drained=drained, + abort_calls=abort_calls, + registry=registry, + polling=polling, + finish_without_cancel=False, + first_response=False, + cancel_before_response=False, + fail_registration=False, + fail_dispatch=False, + first_response_consumed=asyncio.Event(), + native_request=None, + request={ + "token_ids": [1], + "bootstrap_info": { + "bootstrap_host": "localhost", + "bootstrap_port": 0, + "bootstrap_room": 1, + }, + }, + ) + + async def stream(rid): + started.set() + try: + await allow_registration.wait() + if case.fail_registration: + raise ValueError("registration failed") + state = SimpleNamespace( + time_stats=SimpleNamespace(api_server_dispatch_finish_time=0.0) + ) + registry[rid] = state + registered.set() + await allow_dispatch.wait() + if case.fail_dispatch: + raise ValueError("dispatch failed") + state.time_stats.api_server_dispatch_finish_time = 1.0 + dispatched.set() + if case.finish_without_cancel: + return + if case.first_response: + if case.cancel_before_response: + cancelled.set() + yield { + "output_ids": [], + "meta_info": {"id": rid, "finish_reason": None}, + } + case.first_response_consumed.set() + await aborted.wait() + finally: + registry.pop(rid, None) + drained.set() + + def abort_request(*, rid, abort_all): + abort_calls.append((rid, abort_all)) + # Registration alone does not make a request visible to the scheduler. + if rid in registry and dispatched.is_set(): + aborted.set() + + async def async_generate(**kwargs): + return stream(kwargs["rid"]) + + def generate_request(request, context): + case.native_request = request + return stream(request.rid) + + handler.engine = SimpleNamespace( + async_generate=async_generate, + tokenizer_manager=SimpleNamespace( + rid_to_state=registry, + abort_request=abort_request, + generate_request=generate_request, + ), + ) + return case + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize("disaggregated", [False, True]) +@pytest.mark.parametrize("output_mode", ["tokens", "text", "native"]) +@pytest.mark.parametrize( + "phase", ["before_registration", "before_dispatch", "after_dispatch"] +) +@pytest.mark.parametrize("signal", ["cancel", "shutdown"]) +async def test_decode_cancels_before_first_response( + decode_cancellation_case, output_mode, phase, signal, disaggregated, caplog +): + caplog.set_level(logging.INFO) + case = decode_cancellation_case + abort_message = f"Aborted Request ID: {case.context.id()}" + if phase == "before_dispatch": + case.allow_dispatch.clear() + expected_id = "internal-request-id" + if not disaggregated: + case.handler.serving_mode = DisaggregationMode.AGGREGATED + case.handler.use_sglang_tokenizer = output_mode == "text" + if output_mode == "native": + case.request["extra_args"] = {"sglang_tito": {"rid": "caller-supplied-id"}} + + consumer = asyncio.create_task( + _collect(case.handler.generate(case.request, case.context)) + ) + try: + await asyncio.wait_for(case.started.wait(), timeout=1) + if phase != "before_registration": + case.allow_registration.set() + await asyncio.wait_for(case.registered.wait(), timeout=1) + if phase == "after_dispatch": + await asyncio.wait_for(case.dispatched.wait(), timeout=1) + if signal == "cancel": + case.cancelled.set() + else: + case.handler.shutdown_event.set() + if phase == "before_registration": + await asyncio.wait_for(case.polling.get(), timeout=1) + assert not case.abort_calls + assert abort_message not in caplog.messages + case.allow_registration.set() + elif phase == "before_dispatch": + await asyncio.wait_for(case.polling.get(), timeout=1) + assert not case.abort_calls + assert abort_message not in caplog.messages + assert not case.aborted.is_set() + assert case.registry + case.allow_dispatch.set() + + if signal == "shutdown": + with pytest.raises(EngineShutdown): + await asyncio.wait_for(consumer, timeout=1) + else: + assert await asyncio.wait_for(consumer, timeout=1) == [] + abort_results = await asyncio.gather( + *tuple(case.handler._abort_tasks), return_exceptions=True + ) + assert all( + result is None or isinstance(result, asyncio.CancelledError) + for result in abort_results + ) + assert "Detached SGLang task failed during cancellation" not in caplog.messages + assert case.abort_calls == [(expected_id, False)] + assert caplog.messages.count(abort_message) == 1 + assert case.dispatched.is_set() + assert case.aborted.is_set() + assert case.drained.is_set() + assert not case.registry + if output_mode == "native": + assert case.native_request is not None + expected_bootstrap = ( + ("localhost", 0, 1) if disaggregated else (None, None, None) + ) + assert ( + case.native_request.bootstrap_host, + case.native_request.bootstrap_port, + case.native_request.bootstrap_room, + ) == expected_bootstrap + else: + assert case.native_request is None + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize( + "outcome", ["complete", "registration_error", "cancel_after_response"] +) +async def test_decode_cancellation_preserves_stream_lifetime( + decode_cancellation_case, outcome +): + case = decode_cancellation_case + case.finish_without_cancel = outcome == "complete" + case.fail_registration = outcome == "registration_error" + case.first_response = outcome == "cancel_after_response" + case.allow_registration.set() + consumer = asyncio.create_task( + _collect(case.handler.generate(case.request, case.context)) + ) + try: + if outcome == "registration_error": + with pytest.raises(ValueError, match="registration failed"): + await asyncio.wait_for(consumer, timeout=1) + else: + if case.first_response: + await asyncio.wait_for(case.first_response_consumed.wait(), timeout=1) + case.cancelled.set() + assert await asyncio.wait_for(consumer, timeout=1) == [] + expected = [("internal-request-id", False)] if case.first_response else [] + assert case.abort_calls == expected + assert case.drained.is_set() + assert not case.registry + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize("native", [False, True]) +async def test_unsupported_runtime_uses_response_id_cancellation( + decode_cancellation_case, native +): + case = decode_cancellation_case + case.handler._supports_ordered_cancellation = False + case.first_response = True + case.allow_registration.set() + if native: + case.request["extra_args"] = {"sglang_tito": {"rid": "caller-supplied-id"}} + + consumer = asyncio.create_task( + _collect(case.handler.generate(case.request, case.context)) + ) + try: + await asyncio.wait_for(case.first_response_consumed.wait(), timeout=1) + assert not case.abort_calls + + case.cancelled.set() + + outputs = await asyncio.wait_for(consumer, timeout=1) + assert case.abort_calls == [("internal-request-id", False)] + if native: + assert case.native_request.rid == "internal-request-id" + assert ( + outputs[0]["engine_data"]["sglang_response"]["meta_info"]["id"] + == "caller-supplied-id" + ) + else: + assert outputs == [] + assert case.drained.is_set() + assert not case.registry + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize( + ("public_request_id", "expected_request_id"), + [("session-request-2", "session-request-2"), (None, "dynamo-request")], +) +async def test_native_session_request_preserves_session_tree_ids( + decode_cancellation_case, public_request_id, expected_request_id +): + case = decode_cancellation_case + native_payload = {"session_params": {"id": "session-1", "rid": "session-request-1"}} + if public_request_id is not None: + native_payload["rid"] = public_request_id + case.request["extra_args"] = {"sglang_tito": native_payload} + case.finish_without_cancel = True + case.allow_registration.set() + + outputs = await asyncio.wait_for( + _collect(case.handler.generate(case.request, case.context)), timeout=1 + ) + + assert outputs == [] + assert case.native_request.rid == expected_request_id + assert case.native_request.session_params["rid"] == "session-request-1" + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_native_session_cancellation_uses_confirmed_response_id( + decode_cancellation_case, +): + case = decode_cancellation_case + case.request["extra_args"] = { + "sglang_tito": { + "rid": "session-request-2", + "session_params": {"id": "session-1", "rid": "session-request-1"}, + } + } + case.first_response = True + case.allow_registration.set() + + consumer = asyncio.create_task( + _collect(case.handler.generate(case.request, case.context)) + ) + try: + await asyncio.wait_for(case.first_response_consumed.wait(), timeout=1) + assert not case.abort_calls + + case.cancelled.set() + + outputs = await asyncio.wait_for(consumer, timeout=1) + assert case.abort_calls == [("session-request-2", False)] + assert ( + outputs[0]["engine_data"]["sglang_response"]["meta_info"]["id"] + == "session-request-2" + ) + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_decode_logs_internal_and_context_request_ids( + decode_cancellation_case, caplog +): + caplog.set_level(logging.INFO) + case = decode_cancellation_case + case.finish_without_cancel = True + case.allow_registration.set() + + await asyncio.wait_for( + _collect(case.handler.generate(case.request, case.context)), timeout=1 + ) + + assert ( + "Submitted SGLang Request ID: internal-request-id, Context: dynamo-request" + in caplog.messages + ) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_parallel_sampling_does_not_guess_sglang_child_request_ids( + decode_cancellation_case, +): + case = decode_cancellation_case + case.handler.serving_mode = DisaggregationMode.AGGREGATED + case.handler._build_sampling_params = lambda request: { + "max_new_tokens": 1, + "n": 3, + } + case.first_response = True + case.allow_registration.set() + + async def async_generate(**kwargs): + assert kwargs["rid"] == "internal-request-id" + return case.handler.engine.tokenizer_manager.generate_request( + SimpleNamespace(rid="actual-child-uuid"), case.context + ) + + case.handler.engine.async_generate = async_generate + consumer = asyncio.create_task( + _collect(case.handler.generate(case.request, case.context)) + ) + try: + await asyncio.wait_for(case.first_response_consumed.wait(), timeout=1) + assert not case.abort_calls + + case.cancelled.set() + + assert await asyncio.wait_for(consumer, timeout=1) == [] + assert case.abort_calls == [("actual-child-uuid", False)] + assert case.drained.is_set() + assert not case.registry + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancelled_stream_drain_is_bounded( + decode_cancellation_case, monkeypatch, caplog +): + caplog.set_level(logging.WARNING) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation._CANCELLATION_DRAIN_TIMEOUT_S", + 0, + ) + + async def stalled_stream(): + await asyncio.Future() + yield {} + + cancellation_task = asyncio.create_task(asyncio.sleep(0)) + await cancellation_task + outputs = await _collect( + decode_cancellation_case.handler._stream_until_cancelled( + stalled_stream(), cancellation_task + ) + ) + + assert outputs == [] + assert "Timed out draining SGLang stream after cancellation" in caplog.messages + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancelled_stream_starts_deadline_when_chunk_is_already_ready( + decode_cancellation_case, monkeypatch, caplog +): + caplog.set_level(logging.WARNING) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation._CANCELLATION_DRAIN_TIMEOUT_S", + 0, + ) + + async def simultaneous_wait(awaitables, **_kwargs): + await asyncio.gather(*awaitables) + return set(awaitables), set() + + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation.asyncio.wait", + simultaneous_wait, + ) + + async def buffered_stream(): + yield {"chunk": 1} + yield {"chunk": 2} + + cancellation_task = asyncio.create_task(asyncio.sleep(0)) + await cancellation_task + stream = decode_cancellation_case.handler._stream_until_cancelled( + buffered_stream(), cancellation_task + ) + + assert await anext(stream) == {"chunk": 1} + with pytest.raises(StopAsyncIteration): + await anext(stream) + assert "Timed out draining SGLang stream after cancellation" in caplog.messages + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancelled_stream_drain_does_not_wait_for_resistant_iterator( + decode_cancellation_case, monkeypatch, caplog +): + caplog.set_level(logging.WARNING) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation._CANCELLATION_DRAIN_TIMEOUT_S", + 0.01, + ) + cancellation_seen = asyncio.Event() + iterator_finished = asyncio.Event() + + async def resistant_stream(): + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + cancellation_seen.set() + await asyncio.sleep(0.2) + iterator_finished.set() + yield {} + + cancellation_task = asyncio.create_task(asyncio.sleep(0)) + await cancellation_task + started = asyncio.get_running_loop().time() + outputs = await _collect( + decode_cancellation_case.handler._stream_until_cancelled( + resistant_stream(), cancellation_task + ) + ) + elapsed = asyncio.get_running_loop().time() - started + + assert outputs == [] + assert elapsed < 0.1 + assert "Timed out draining SGLang stream after cancellation" in caplog.messages + await asyncio.wait_for(iterator_finished.wait(), timeout=1) + assert cancellation_seen.is_set() + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_decode_cancellation_drains_buffered_empty_chunk( + decode_cancellation_case, +): + case = decode_cancellation_case + case.first_response = True + case.cancel_before_response = True + case.allow_registration.set() + + outputs = await asyncio.wait_for( + _collect(case.handler.generate(case.request, case.context)), timeout=1 + ) + + assert outputs == [] + assert case.abort_calls == [("internal-request-id", False)] + assert case.first_response_consumed.is_set() + assert case.drained.is_set() + assert not case.registry + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize("submitted_request_id", [None, "internal-request-id_0"]) +async def test_cancellation_monitor_aborts_parallel_ids_on_stream_close( + decode_cancellation_case, submitted_request_id +): + case = decode_cancellation_case + request_id_future = asyncio.get_running_loop().create_future() + request_id_future.set_result("internal-request-id_0") + request_ids = {"internal-request-id_0", "internal-request-id_1"} + + async with case.handler._cancellation_monitor( + request_id_future, + case.context, + submitted_request_id=submitted_request_id, + request_ids=request_ids, + ): + pass + + assert set(case.abort_calls) == { + ("internal-request-id_0", False), + ("internal-request-id_1", False), + } + assert not request_ids + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize("state_kind", ["no_registry", "missing", "null_stats"]) +@pytest.mark.parametrize("abort_fails", [False, True]) +async def test_cancellation_monitor_logs_only_submitted_abort( + decode_cancellation_case, caplog, state_kind, abort_fails +): + caplog.set_level(logging.INFO) + case = decode_cancellation_case + rid = case.context.trace_id + registry = None if state_kind == "no_registry" else case.registry + if state_kind == "null_stats": + case.registry[rid] = SimpleNamespace(time_stats=None) + + def abort_request(*, rid, abort_all): + assert f"Aborted Request ID: {case.context.id()}" not in caplog.messages + if abort_fails: + raise RuntimeError("abort failed") + case.abort_calls.append((rid, abort_all)) + case.registry.pop(rid, None) + + manager = case.handler.engine.tokenizer_manager + manager.abort_request = abort_request + abort = case.handler._abort_sglang_request( + manager, rid, registry, case.context.id() + ) + if abort_fails: + with pytest.raises(RuntimeError, match="abort failed"): + await abort + else: + await abort + await asyncio.gather(*case.handler._abort_tasks) + assert case.abort_calls == ([] if abort_fails else [(rid, False)]) + assert caplog.messages.count(f"Aborted Request ID: {case.context.id()}") == ( + 0 if abort_fails else 1 + ) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize("reuse_rid", [False, True]) +@pytest.mark.parametrize("waiting_for", ["dispatch", "retry"]) +async def test_cancellation_monitor_stops_when_request_finishes( + decode_cancellation_case, monkeypatch, reuse_rid, waiting_for, caplog +): + """Neither a dispatch wait nor a retry may abort a replacement request.""" + caplog.set_level(logging.INFO) + case = decode_cancellation_case + rid = case.context.trace_id + case.registry[rid] = ( + SimpleNamespace(time_stats=SimpleNamespace(api_server_dispatch_finish_time=0.0)) + if waiting_for == "dispatch" + else object() + ) + sleeping = asyncio.Event() + resume = asyncio.Event() + + async def controlled_sleep(delay): + sleeping.set() + await resume.wait() + + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation.asyncio", + SimpleNamespace(**{**vars(asyncio), "sleep": controlled_sleep}), + ) + request_id_future = asyncio.get_running_loop().create_future() + monitor = asyncio.create_task( + case.handler._handle_cancellation(request_id_future, case.context, rid) + ) + try: + case.cancelled.set() + await asyncio.wait_for(sleeping.wait(), timeout=1) + case.registry.pop(rid) + if reuse_rid: + case.registry[rid] = object() + resume.set() + await asyncio.wait_for(monitor, timeout=1) + await asyncio.gather(*case.handler._abort_tasks) + expected = [(rid, False)] if waiting_for == "retry" else [] + assert case.abort_calls == expected + assert caplog.messages.count(f"Aborted Request ID: {case.context.id()}") == len( + expected + ) + assert not request_id_future.done() + finally: + monitor.cancel() + await asyncio.gather(monitor, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize("outcome", ["complete", "error"]) +async def test_decode_stream_exit_retains_pending_abort_wait( + decode_cancellation_case, monkeypatch, outcome +): + case = decode_cancellation_case + case.allow_registration.set() + case.allow_dispatch.clear() + case.finish_without_cancel = True + sleeping = asyncio.Event() + wait_cancelled = asyncio.Event() + + async def controlled_sleep(delay): + sleeping.set() + try: + await asyncio.Future() + finally: + wait_cancelled.set() + + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation.asyncio", + SimpleNamespace(**{**vars(asyncio), "sleep": controlled_sleep}), + ) + case.fail_dispatch = outcome == "error" + consumer = asyncio.create_task( + _collect(case.handler.generate(case.request, case.context)) + ) + try: + await asyncio.wait_for(case.registered.wait(), timeout=1) + case.cancelled.set() + await asyncio.wait_for(sleeping.wait(), timeout=1) + case.allow_dispatch.set() + if outcome == "error": + with pytest.raises(ValueError, match="dispatch failed"): + await asyncio.wait_for(consumer, timeout=1) + else: + assert await asyncio.wait_for(consumer, timeout=1) == [] + assert not case.abort_calls + assert case.drained.is_set() + assert not case.registry + abort_tasks = tuple(case.handler._abort_tasks) + assert abort_tasks + case.handler._cancel_abort_tasks() + await asyncio.gather(*abort_tasks, return_exceptions=True) + assert wait_cancelled.is_set() + finally: + consumer.cancel() + await asyncio.gather(consumer, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize("signal", ["cancel", "shutdown"]) +async def test_cancellation_monitor_without_tokenizer_manager( + decode_cancellation_case, signal +): + case = decode_cancellation_case + case.handler.engine = None + if signal == "cancel": + case.cancelled.set() + else: + case.handler.shutdown_event.set() + monitor = case.handler._handle_cancellation( + asyncio.get_running_loop().create_future(), case.context, case.context.trace_id + ) + if signal == "shutdown": + with pytest.raises(EngineShutdown): + await asyncio.wait_for(monitor, timeout=1) + else: + await asyncio.wait_for(monitor, timeout=1) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancellation_monitor_without_request_registry(decode_cancellation_case): + case = decode_cancellation_case + abort_calls = [] + + def abort_request(*, rid, abort_all): + abort_calls.append((rid, abort_all)) + + tokenizer_manager = SimpleNamespace(abort_request=abort_request) + tokenizer_manager.rid_to_state = [] + case.handler.engine = SimpleNamespace(tokenizer_manager=tokenizer_manager) + case.cancelled.set() + + await asyncio.wait_for( + case.handler._handle_cancellation( + asyncio.get_running_loop().create_future(), + case.context, + case.context.trace_id, + ), + timeout=1, + ) + + assert abort_calls == [(case.context.trace_id, False)] + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize( + "server_args, has_dispatch_time, expect_retry", + [ + ({}, True, False), + ({"dp_size": 2}, True, False), + ({"enable_dp_attention": True}, True, True), + ( + { + "enable_dp_attention": True, + "enable_dp_attention_local_control_broadcast": True, + }, + True, + False, + ), + ({"pp_size": 2}, True, True), + ({}, False, True), + ], + ids=["tp", "dp", "dp-attention-global", "dp-attention-local", "pp", "legacy"], +) +async def test_cancellation_monitor_retries_only_without_ordered_dispatch( + decode_cancellation_case, monkeypatch, server_args, has_dispatch_time, expect_retry +): + case = decode_cancellation_case + # SGLang may keep raw and effective server arguments separately. + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation.resolved_server_args", + lambda raw_args: SimpleNamespace(**server_args), + ) + rid = case.context.trace_id + case.registry[rid] = ( + SimpleNamespace(time_stats=SimpleNamespace(api_server_dispatch_finish_time=1.0)) + if has_dispatch_time + else object() + ) + + def abort_request(*, rid, abort_all): + case.abort_calls.append((rid, abort_all)) + # Leave the state registered after the first abort. A single-abort path + # must exit on its own, independently of the consumer draining the stream. + if len(case.abort_calls) == 2: + case.registry.pop(rid) + + case.handler.engine.tokenizer_manager.abort_request = abort_request + case.cancelled.set() + await asyncio.wait_for( + case.handler._handle_cancellation( + asyncio.get_running_loop().create_future(), case.context, rid + ), + timeout=1, + ) + await asyncio.gather(*case.handler._abort_tasks) + assert case.abort_calls == [(rid, False)] * (2 if expect_retry else 1) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancellation_monitor_bounds_abort_retries( + decode_cancellation_case, monkeypatch, caplog +): + caplog.set_level(logging.WARNING) + case = decode_cancellation_case + rid = "internal-request-id" + case.registry[rid] = SimpleNamespace( + time_stats=SimpleNamespace(api_server_dispatch_finish_time=0.0) + ) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation._CANCELLATION_DISPATCH_WAIT_TIMEOUT_S", + 0, + ) + + async def immediate_sleep(_delay): + await asyncio.sleep(0) + + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation.asyncio", + SimpleNamespace(**{**vars(asyncio), "sleep": immediate_sleep}), + ) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation.resolved_server_args", + lambda raw_args: SimpleNamespace(), + ) + + await asyncio.wait_for( + case.handler._abort_sglang_request( + case.handler.engine.tokenizer_manager, + rid, + case.registry, + case.context.id(), + ), + timeout=1, + ) + assert ( + "Timed out waiting for SGLang Request ID internal-request-id to dispatch" + in caplog.messages + ) + + await asyncio.gather(*case.handler._abort_tasks) + assert case.abort_calls == [(rid, False)] * 9 + assert ( + "SGLang request internal-request-id remained registered after 8 abort retries" + in caplog.messages + ) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_stream_drain_starts_before_abort_retries_complete( + decode_cancellation_case, monkeypatch +): + case = decode_cancellation_case + rid = "internal-request-id" + state = SimpleNamespace( + time_stats=SimpleNamespace(api_server_dispatch_finish_time=1.0) + ) + case.registry[rid] = state + retry_started = asyncio.Event() + release_retry = asyncio.Event() + + async def blocked_retry(*_args): + retry_started.set() + await release_retry.wait() + + async def stalled_stream(): + await asyncio.Future() + yield {} + + monkeypatch.setattr(case.handler, "_retry_abort", blocked_retry) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation.resolved_server_args", + lambda raw_args: SimpleNamespace(pp_size=2), + ) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation._CANCELLATION_DRAIN_TIMEOUT_S", + 0, + ) + + case.cancelled.set() + cancellation_task = asyncio.create_task( + case.handler._handle_cancellation( + asyncio.get_running_loop().create_future(), case.context, rid + ) + ) + stream_task = asyncio.create_task( + _collect( + case.handler._stream_until_cancelled(stalled_stream(), cancellation_task) + ) + ) + try: + await asyncio.wait_for(retry_started.wait(), timeout=1) + assert await asyncio.wait_for(stream_task, timeout=0.1) == [] + assert cancellation_task.done() + assert case.abort_calls == [(rid, False)] + finally: + retry_tasks = tuple(case.handler._abort_tasks) + release_retry.set() + await asyncio.gather( + cancellation_task, stream_task, *retry_tasks, return_exceptions=True + ) + assert not case.handler._abort_tasks + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancel_abort_tasks_cancels_and_releases_owned_tasks( + decode_cancellation_case, monkeypatch +): + case = decode_cancellation_case + retry_started = asyncio.Event() + + async def blocked_retry(*_args): + retry_started.set() + await asyncio.Future() + + monkeypatch.setattr(case.handler, "_retry_abort", blocked_retry) + case.handler._start_abort_retry(object(), "request-id", {}, object()) + await asyncio.wait_for(retry_started.wait(), timeout=1) + retry_task = next(iter(case.handler._abort_tasks)) + + case.handler._cancel_abort_tasks() + with pytest.raises(asyncio.CancelledError): + await retry_task + await asyncio.sleep(0) + + assert not case.handler._abort_tasks + + +def test_base_worker_cleanup_cancels_abort_tasks(): + handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler) + handler.publisher = None + cleanup_calls = [] + handler._cancel_abort_tasks = lambda: cleanup_calls.append("cancel") + + BaseWorkerHandler.cleanup(handler) + + assert cleanup_calls == ["cancel"] + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize("phase", ["registration", "dispatch"]) +async def test_cancellation_monitor_preserves_abort_ordering( + decode_cancellation_case, phase +): + case = decode_cancellation_case + rid = "internal-request-id" + state = SimpleNamespace( + time_stats=SimpleNamespace(api_server_dispatch_finish_time=0.0) + ) + + if phase == "dispatch": + case.registry[rid] = state + + case.cancelled.set() + operation = asyncio.create_task( + case.handler._handle_cancellation( + asyncio.get_running_loop().create_future(), + case.context, + rid, + ) + ) + + await asyncio.wait_for(case.polling.get(), timeout=1) + await asyncio.wait_for(operation, timeout=1) + assert not case.abort_calls + case.registry[rid] = state + state.time_stats.api_server_dispatch_finish_time = 1.0 + + await asyncio.gather(*case.handler._abort_tasks) + assert case.abort_calls == [(rid, False)] + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_ordered_abort_stops_when_stream_ends_before_registration( + decode_cancellation_case, caplog +): + caplog.set_level(logging.DEBUG) + case = decode_cancellation_case + request_id_future = asyncio.get_running_loop().create_future() + request_id_future.cancel() + + await case.handler._abort_after_registration( + case.handler.engine.tokenizer_manager, + request_id_future, + "internal-request-id", + case.registry, + case.context.id(), + ) + + assert not case.abort_calls + assert ( + "Abandoning SGLang abort for Context dynamo-request; request never registered" + in caplog.messages + ) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancellation_monitor_logs_fallback_abort_failure( + decode_cancellation_case, caplog +): + caplog.set_level(logging.ERROR) + case = decode_cancellation_case + case.handler.engine.tokenizer_manager.rid_to_state = [] + + def abort_request(*, rid, abort_all): + raise RuntimeError("abort failed") + + case.handler.engine.tokenizer_manager.abort_request = abort_request + case.cancelled.set() + + await case.handler._handle_cancellation( + asyncio.get_running_loop().create_future(), + case.context, + "internal-request-id", + ) + + assert ( + "Failed to abort SGLang Request ID internal-request-id, Context: dynamo-request" + in caplog.messages + ) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_shutdown_survives_ordered_abort_cleanup(decode_cancellation_case): + case = decode_cancellation_case + rid = "internal-request-id" + case.registry[rid] = SimpleNamespace( + time_stats=SimpleNamespace(api_server_dispatch_finish_time=0.0) + ) + case.handler.shutdown_event = asyncio.Event() + case.handler.shutdown_event.set() + operation = asyncio.create_task( + case.handler._handle_cancellation( + asyncio.get_running_loop().create_future(), + case.context, + rid, + ) + ) + + await asyncio.wait_for(case.polling.get(), timeout=1) + abort_task = next(iter(case.handler._abort_tasks)) + abort_task.cancel() + + with pytest.raises(EngineShutdown): + await operation diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index 529d4066d4cc..676c71976359 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, Mock import pytest +from sglang.srt.managers.io_struct import GenerateReqInput from dynamo.common.constants import DisaggregationMode from dynamo.common.metadata_upload import MetadataUploader @@ -26,8 +27,11 @@ from dynamo.sglang.request_handlers.llm.decode_handler import ( DecodeWorkerHandler, _extract_sglang_stop_reason, + _native_payload_is_batched, _nvext_extra_field_requested, _openai_stop_sampling_params, + _ordered_cancellation_request_id, + _public_native_response_id, _user_stop_token_ids, ) from dynamo.sglang.request_handlers.llm.mm_disagg_utils import ( @@ -37,6 +41,7 @@ ) from dynamo.sglang.request_handlers.llm.prefill_handler import PrefillWorkerHandler from dynamo.sglang.request_handlers.multimodal.worker_handler import SglangUtils +from dynamo.sglang.request_utils import request_cache_salt pytestmark = [ pytest.mark.unit, @@ -48,28 +53,6 @@ ] -@pytest.mark.asyncio -async def test_cancellation_monitor_rechecks_shutdown_after_cleanup(): - handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler) - handler.shutdown_event = asyncio.Event() - - async def set_shutdown_when_cancelled(*_args): - try: - await asyncio.Future() - except asyncio.CancelledError: - handler.shutdown_event.set() - raise - - handler._handle_cancellation = set_shutdown_when_cancelled - request_id_future = asyncio.get_running_loop().create_future() - request_id_future.set_result("sglang-request-id") - context = SimpleNamespace(id=lambda: "request-id") - - with pytest.raises(EngineShutdown, match="shut down during token generation"): - async with handler._cancellation_monitor(request_id_future, context): - await asyncio.sleep(0) - - def _read_zstd_payload(path): import zstandard as zstd @@ -242,19 +225,35 @@ def test_openai_stop_sampling_params_maps_token_id_stop_array(): } -def _new_decode_handler(*, use_sglang_tokenizer: bool = False, enable_rl: bool = False): +def _new_decode_handler( + *, + use_sglang_tokenizer: bool = False, + skip_tokenizer_init: bool = False, + enable_rl: bool = False, +): handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler) handler.shutdown_event = None handler.use_sglang_tokenizer = use_sglang_tokenizer handler.config = SimpleNamespace( - server_args=SimpleNamespace(served_model_name="test-model"), + server_args=SimpleNamespace( + served_model_name="test-model", + skip_tokenizer_init=skip_tokenizer_init, + ), dynamo_args=SimpleNamespace(enable_rl=enable_rl), ) handler._first_token_source = None @asynccontextmanager async def no_cancellation_monitor(*args, **kwargs): - yield None + async def wait_forever(): + await asyncio.Future() + + task = asyncio.create_task(wait_forever()) + try: + yield task + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) handler._cancellation_monitor = no_cancellation_monitor return handler @@ -264,19 +263,38 @@ async def no_cancellation_monitor(*args, **kwargs): @pytest.mark.parametrize( "processor_name", ["_process_token_stream", "_process_text_stream"] ) -async def test_shutdown_abort_chunk_raises_engine_shutdown(processor_name): +@pytest.mark.parametrize( + "finish_reason", + [ + pytest.param({"type": "abort"}, id="cancellation"), + pytest.param( + { + "type": "abort", + "message": "Failed to compile json grammar", + "status_code": 400, + "err_type": "BadRequestError", + }, + id="validation-error", + ), + ], +) +async def test_shutdown_abort_chunk_raises_engine_shutdown( + processor_name, finish_reason +): + """Shutdown remains retryable even when its abort carries validation details.""" handler = _new_decode_handler() handler.shutdown_event = asyncio.Event() handler.shutdown_event.set() - context = SimpleNamespace(id=lambda: "request-id") + context = SimpleNamespace(id=lambda: "request-id", is_stopped=lambda: False) async def stream(): + """Yield the terminal abort from the shutting-down backend.""" yield { "text": "", "output_ids": [], "meta_info": { "id": "sglang-request-id", - "finish_reason": {"type": "abort"}, + "finish_reason": finish_reason, }, } @@ -285,6 +303,208 @@ async def stream(): pass +@pytest.fixture +def abort_context(): + """Track first-token notification without cancelling the test request.""" + return SimpleNamespace( + id=lambda: "request-id", + is_stopped=lambda: False, + notify_first_token=Mock(), + ) + + +def _abort_chunk(finish_reason): + """Model the placeholder token and text SGLang can emit after grammar failure.""" + return { + "output_ids": [101], + "text": "_color", + "meta_info": { + "id": "sglang-request-id", + "finish_reason": finish_reason, + "prompt_tokens": 1, + "completion_tokens": 1, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "processor_name", ["_process_token_stream", "_process_text_stream"] +) +@pytest.mark.parametrize( + ("status_code", "message", "err_type"), + [ + pytest.param( + 400, + "Failed to compile json grammar: [04:33:48] " + '/project/cpp/json_schema_converter.cc:2999: Unsupported type "invalid_type"\n', + "BadRequestError", + id="invalid-schema-type", + ), + pytest.param( + 422, + "Request validation failed", + "UnprocessableEntityError", + id="other-client-error", + ), + pytest.param( + 500, + "Grammar compiler failed internally", + "InternalServerError", + id="backend-server-error", + ), + ], +) +async def test_error_abort_preserves_http_error_before_output( + processor_name, status_code, message, err_type, abort_context +): + """Raise the backend error before exposing placeholder output or token timing.""" + handler = _new_decode_handler() + finish_reason = { + "type": "abort", + "message": message, + "status_code": status_code, + "err_type": err_type, + } + stream = getattr(handler, processor_name)( + _stream([_abort_chunk(finish_reason)]), abort_context + ) + + expected_error = InvalidArgument if status_code == 400 else HttpError + with pytest.raises(expected_error) as error: + await anext(stream) + + if status_code == 400: + assert str(error.value) == message + else: + assert error.value.code == status_code + assert error.value.message == message + abort_context.notify_first_token.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "processor_name", ["_process_token_stream", "_process_text_stream"] +) +async def test_error_abort_after_output_preserves_error_and_stops( + processor_name, abort_context +): + """Preserve an earlier delta, then terminate with the original validation error.""" + handler = _new_decode_handler() + normal_chunk = _abort_chunk(None) + normal_chunk["output_ids"] = [42] + normal_chunk["text"] = "hello" + message = "Failed to compile json grammar: unsupported schema type" + error_chunk = _abort_chunk( + { + "type": "abort", + "message": message, + "status_code": 400, + "err_type": "BadRequestError", + } + ) + stream = getattr(handler, processor_name)( + _stream([normal_chunk, error_chunk]), abort_context + ) + + first = await anext(stream) + if processor_name == "_process_token_stream": + assert first["token_ids"] == [42] + assert first.get("finish_reason") is None + else: + assert first["choices"][0]["delta"]["content"] == "hello" + assert first["choices"][0]["finish_reason"] is None + + with pytest.raises(InvalidArgument) as error: + await anext(stream) + assert str(error.value) == message + + # Neither the abort's placeholder nor a successful terminal chunk can follow. + with pytest.raises(StopAsyncIteration): + await anext(stream) + abort_context.notify_first_token.assert_called_once_with() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "processor_name", ["_process_token_stream", "_process_text_stream"] +) +@pytest.mark.parametrize( + "status_fields", + [ + pytest.param({}, id="missing"), + pytest.param({"status_code": None}, id="null"), + pytest.param({"status_code": "400"}, id="string"), + pytest.param({"status_code": True}, id="boolean"), + pytest.param({"status_code": 400.0}, id="float"), + pytest.param({"status_code": 200}, id="success-status"), + pytest.param({"status_code": 302}, id="redirect-status"), + pytest.param({"status_code": 600}, id="out-of-range"), + ], +) +async def test_error_abort_with_missing_or_invalid_status_uses_server_error( + processor_name, status_fields, abort_context +): + """Malformed error metadata fails as a server error instead of a completion.""" + handler = _new_decode_handler() + message = "Failed to compile json grammar" + finish_reason = { + "type": "abort", + "message": message, + "err_type": "BadRequestError", + **status_fields, + } + stream = getattr(handler, processor_name)( + _stream([_abort_chunk(finish_reason)]), abort_context + ) + + with pytest.raises(HttpError) as error: + await anext(stream) + + assert error.value.code == 500 + assert error.value.message == message + abort_context.notify_first_token.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "processor_name", ["_process_token_stream", "_process_text_stream"] +) +@pytest.mark.parametrize( + "finish_fields", + [ + pytest.param({}, id="missing"), + pytest.param({"message": None}, id="null"), + pytest.param({"message": ""}, id="empty"), + pytest.param( + {"message": "Aborted", "status_code": None, "err_type": None}, + id="sglang-cancellation", + ), + ], +) +async def test_abort_without_error_details_remains_cancelled( + processor_name, finish_fields, abort_context +): + """An ordinary SGLang cancellation keeps its existing finish mapping.""" + handler = _new_decode_handler() + chunk = _abort_chunk({"type": "abort", **finish_fields}) + chunk["output_ids"] = [] + chunk["text"] = "" + + chunks = await _collect( + getattr(handler, processor_name)(_stream([chunk]), abort_context) + ) + + assert len(chunks) == 1 + if processor_name == "_process_token_stream": + assert chunks[0]["finish_reason"] == "cancelled" + assert chunks[0]["token_ids"] == [] + else: + assert chunks[0]["choices"][0]["finish_reason"] == "cancelled" + assert chunks[0]["choices"][0]["delta"]["content"] == "" + abort_context.notify_first_token.assert_not_called() + + @pytest.mark.asyncio @pytest.mark.parametrize( "processor_name", ["_process_token_stream", "_process_text_stream"] @@ -321,8 +541,67 @@ async def stream(): pass -def test_engine_generate_preserves_native_fields_and_overrides_worker_state(): +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize( + "processor_name", ["_process_token_stream", "_process_text_stream"] +) +async def test_ordered_cancellation_skips_stopped_chunk_processing(processor_name): + handler = _new_decode_handler() + notified = False + + def notify_first_token(): + nonlocal notified + notified = True + + context = SimpleNamespace( + id=lambda: "request-id", + is_stopped=lambda: True, + notify_first_token=notify_first_token, + ) + + class UnexpectedUploader: + async def upload_choice(self, *_args): + raise AssertionError("stopped chunks must not upload metadata") + + async def stream(): + yield { + "text": "ignored", + "output_ids": [1], + "meta_info": { + "id": "internal-request-id", + "finish_reason": {"type": "stop"}, + }, + } + + outputs = await _collect( + getattr(handler, processor_name)( + stream(), + context, + metadata_uploader=UnexpectedUploader(), + submitted_request_id="internal-request-id", + ) + ) + + assert outputs == [] + assert not notified + + +@pytest.mark.parametrize( + "body_salt,routing_salt,expected_salt", + [ + (None, None, None), + ("", None, None), + ("body-salt", None, "body-salt"), + ("body-salt", "", None), + ("body-salt", "tenant-a", "tenant-a"), + ], +) +def test_engine_generate_preserves_native_fields_and_overrides_worker_state( + body_salt, routing_salt, expected_salt +): request = { + "cache_salt": body_salt, "rid": "resolved-request", "sampling_params": { "max_new_tokens": 32, @@ -343,17 +622,19 @@ def test_engine_generate_preserves_native_fields_and_overrides_worker_state(): native = build_native_generate_request( request, input_ids=[7, 8], - fallback_rid="fallback-request", + request_id="internal-request-id", priority=9, sampling_overrides={"n": 1, "max_new_tokens": 1}, bootstrap_host="prefill.internal", routed_dp_rank=3, + cache_salt=routing_salt, ) - assert native.rid == "resolved-request" + assert native.rid == "internal-request-id" assert native.input_ids == [7, 8] assert native.stream is True assert native.priority == 9 + assert native.cache_salt == expected_salt assert native.session_id == "session-1" assert native.return_logprob is True assert native.return_text_in_logprobs is True @@ -370,6 +651,37 @@ def test_engine_generate_preserves_native_fields_and_overrides_worker_state(): } +def test_native_generate_rejects_salt_without_engine_support(monkeypatch): + monkeypatch.delitem(GenerateReqInput.__dataclass_fields__, "cache_salt") + kwargs = {"input_ids": [1], "request_id": "request", "priority": None} + + with pytest.raises(ValueError, match="cache_salt is not supported"): + build_native_generate_request({"cache_salt": "tenant-a"}, **kwargs) + + native = build_native_generate_request({"cache_salt": ""}, **kwargs) + assert native.cache_salt is None + + +def test_request_cache_salt_precedence(): + request = { + "routing": {"cache_salt": "routing"}, + "extra_args": {"nvext": {"cache_salt": "extra"}}, + "nvext": {"cache_salt": "nvext"}, + "cache_salt": "body", + } + for source in [ + request["routing"], + request["extra_args"]["nvext"], + request["nvext"], + request, + ]: + assert request_cache_salt(request) == source["cache_salt"] + source["cache_salt"] = "" + assert request_cache_salt(request) is None + assert request_cache_salt({}) is None + assert request_cache_salt({"extra_args": [1], "cache_salt": "body"}) == "body" + + def test_engine_generate_requires_object_sampling_params_for_prefill_override(): request = {"sampling_params": [1, 2]} @@ -377,7 +689,7 @@ def test_engine_generate_requires_object_sampling_params_for_prefill_override(): build_native_generate_request( request, input_ids=[1], - fallback_rid="prefill-request", + request_id="prefill-request", priority=None, sampling_overrides={"max_new_tokens": 1}, ) @@ -390,7 +702,7 @@ def test_engine_generate_rejects_top_logprobs_by_default(monkeypatch): build_native_generate_request( {"return_logprob": True, "top_logprobs_num": 1}, input_ids=[1], - fallback_rid="request", + request_id="request", priority=None, ) @@ -401,7 +713,7 @@ def test_engine_generate_allows_top_logprobs_with_escape_hatch(monkeypatch): native = build_native_generate_request( {"return_logprob": True, "top_logprobs_num": 2}, input_ids=[1], - fallback_rid="request", + request_id="request", priority=None, ) @@ -439,6 +751,93 @@ async def generate_request(self, request, request_context): assert chunks[0]["engine_data"]["sglang_response"] is native_response +@pytest.mark.asyncio +async def test_native_generate_stream_maps_internal_id_without_mutation(): + native_response = { + "output_ids": [101], + "meta_info": {"id": "internal-request-id"}, + } + + async def stream(): + yield { + "token_ids": [], + "engine_data": {"sglang_response": native_response}, + } + + handler = _new_decode_handler() + chunks = await _collect( + handler._process_native_generate_stream( + stream(), + _Context(), + submitted_request_id="internal-request-id", + internal_request_id="internal-request-id", + response_request_id="caller-visible-id", + ) + ) + + mapped_response = chunks[0]["engine_data"]["sglang_response"] + assert mapped_response["meta_info"]["id"] == "caller-visible-id" + assert native_response["meta_info"]["id"] == "internal-request-id" + assert mapped_response is not native_response + assert mapped_response["meta_info"] is not native_response["meta_info"] + + +@pytest.mark.asyncio +async def test_native_generate_stream_maps_batched_ids_without_collapsing(): + native_responses = [ + { + "output_ids": [101 + index], + "index": index, + "meta_info": {"id": f"internal-request-id_{index}"}, + } + for index in range(2) + ] + + async def stream(): + for native_response in native_responses: + yield { + "token_ids": [], + "engine_data": {"sglang_response": native_response}, + } + + handler = _new_decode_handler() + chunks = await _collect( + handler._process_native_generate_stream( + stream(), + _Context(), + submitted_request_id=None, + internal_request_id="internal-request-id", + response_request_id=["caller-request-0", "caller-request-1"], + ) + ) + + assert [ + chunk["engine_data"]["sglang_response"]["meta_info"]["id"] for chunk in chunks + ] == ["caller-request-0", "caller-request-1"] + assert [response["meta_info"]["id"] for response in native_responses] == [ + "internal-request-id_0", + "internal-request-id_1", + ] + + +@pytest.mark.parametrize( + ("engine_id", "index", "public_id", "expected"), + [ + ("internal_0", 0, "caller", "caller_0"), + ("actual-child-a", 0, "caller", "actual-child-a"), + ("internal_2", 2, ["caller-0", "caller-1"], "internal_2"), + ("internal_0", True, "caller", "internal_0"), + ("internal_0", 0, [123], "internal_0"), + ], +) +def test_public_native_response_id_maps_only_derived_string_ids( + engine_id, index, public_id, expected +): + assert ( + _public_native_response_id(engine_id, index, "internal", public_id) == expected + ) + + def _new_token_input_handler(maximum_input_token_id: int = 151935): handler = _new_decode_handler() handler._max_input_token_id = maximum_input_token_id @@ -663,6 +1062,9 @@ async def _stream(items): class _Context: + def id(self): + return "public-request-id" + def is_stopped(self): return False @@ -685,6 +1087,55 @@ def test_build_sampling_params_passes_n_for_token_requests(): assert sampling_params["max_new_tokens"] == 8 +@pytest.mark.parametrize( + ("sampling_params", "supported", "expected"), + [ + ({}, True, "request-id"), + ({"n": 1}, True, "request-id"), + ([{"n": 1}], True, "request-id"), + ({"n": 3}, True, None), + ([{"n": 3}], True, None), + ({"n": 3, "beam_width": 2}, True, "request-id"), + ({"n": 1}, False, None), + ], +) +def test_ordered_cancellation_requires_stable_sglang_request_id( + sampling_params, supported, expected +): + assert ( + _ordered_cancellation_request_id( + "request-id", sampling_params, supported=supported + ) + == expected + ) + + +@pytest.mark.parametrize( + ("native_payload", "expected"), + [ + ({}, False), + ({"prompt": "hello"}, False), + ({"prompt": ["hello", "world"]}, True), + ({"text": ["hello", "world"]}, True), + ({"input_ids": [1, 2]}, False), + ({"input_ids": [[1], [2]]}, True), + ({"input_embeds": [[0.1], [0.2]]}, False), + ({"input_embeds": [[[0.1]], [[0.2]]]}, True), + ], +) +def test_native_payload_batch_detection(native_payload, expected): + assert _native_payload_is_batched(native_payload) is expected + + +def test_ordered_cancellation_rejects_native_batch(): + assert ( + _ordered_cancellation_request_id( + "request-id", {"n": 1}, supported=True, batched=True + ) + is None + ) + + def test_build_sampling_params_forwards_repetition_controls_for_token_requests(): handler = _new_decode_handler(use_sglang_tokenizer=False) @@ -813,23 +1264,55 @@ def test_guided_decoding_params_are_accepted_by_sglang(guided_decoding): SamplingParams(**sampling_params) -def test_build_sampling_params_maps_min_tokens_for_token_requests(): - handler = _new_decode_handler(use_sglang_tokenizer=False) +def test_build_sampling_params_holds_tokenizer_free_engine_open_for_min_tokens(): + handler = _new_decode_handler( + use_sglang_tokenizer=False, + skip_tokenizer_init=True, + ) sampling_params = handler._build_sampling_params( { "sampling_options": {}, "stop_conditions": { "max_tokens": 64, - "min_tokens": 64, - "ignore_eos": True, + "min_tokens": 8, + "stop_token_ids": [7], + "stop_token_ids_hidden": [151643], }, } ) - assert sampling_params["min_new_tokens"] == 64 - assert sampling_params["max_new_tokens"] == 64 + assert "min_new_tokens" not in sampling_params + assert "stop_token_ids" not in sampling_params assert sampling_params["ignore_eos"] is True + assert sampling_params["max_new_tokens"] == 64 + + from sglang.srt.sampling.sampling_params import SamplingParams + + SamplingParams(**sampling_params).normalize(tokenizer=None) + + +def test_build_sampling_params_forwards_min_tokens_when_engine_has_tokenizer(): + handler = _new_decode_handler( + use_sglang_tokenizer=False, + skip_tokenizer_init=False, + ) + + sampling_params = handler._build_sampling_params( + { + "sampling_options": {}, + "stop_conditions": { + "max_tokens": 64, + "min_tokens": 8, + "stop_token_ids": [7], + }, + } + ) + + assert sampling_params["min_new_tokens"] == 8 + assert sampling_params["stop_token_ids"] == [7] + assert "ignore_eos" not in sampling_params + assert sampling_params["max_new_tokens"] == 64 def test_build_sampling_params_maps_min_tokens_for_sglang_tokenizer_requests(): @@ -844,7 +1327,10 @@ def test_build_sampling_params_maps_min_tokens_for_sglang_tokenizer_requests(): def test_build_sampling_params_omits_min_new_tokens_when_min_tokens_absent(): - handler = _new_decode_handler(use_sglang_tokenizer=False) + handler = _new_decode_handler( + use_sglang_tokenizer=False, + skip_tokenizer_init=True, + ) sampling_params = handler._build_sampling_params( {"sampling_options": {}, "stop_conditions": {"max_tokens": 8}} @@ -854,7 +1340,10 @@ def test_build_sampling_params_omits_min_new_tokens_when_min_tokens_absent(): def test_build_sampling_params_forwards_explicit_zero_min_tokens(): - handler = _new_decode_handler(use_sglang_tokenizer=False) + handler = _new_decode_handler( + use_sglang_tokenizer=False, + skip_tokenizer_init=True, + ) sampling_params = handler._build_sampling_params( {"sampling_options": {}, "stop_conditions": {"max_tokens": 8, "min_tokens": 0}} @@ -1492,6 +1981,7 @@ async def test_process_text_stream_forwards_incremental_text_per_choice(): "llo", "od", ] + assert [chunk["id"] for chunk in chunks] == ["public-request-id"] * 4 @pytest.mark.asyncio @@ -1811,3 +2301,77 @@ async def test_supported_sampling_reaches_engine(mode, n): assert [output["index"] for output in outputs] == list(range(n)) assert handler.engine.async_generate.await_args.kwargs["sampling_params"]["n"] == n assert all(output["finish_reason"] for output in outputs) + + +def test_prefill_dp_rank_kwargs_follows_engine_signature(): + from dynamo.sglang._compat import prefill_dp_rank_kwargs + + class Engine: + async def async_generate( + self, *, bootstrap_room=None, disagg_prefill_dp_rank=None + ): + pass + + class OlderEngine: + async def async_generate(self, *, bootstrap_room=None): + pass + + assert prefill_dp_rank_kwargs(Engine(), 3) == {"disagg_prefill_dp_rank": 3} + assert prefill_dp_rank_kwargs(Engine(), None) == {} + assert prefill_dp_rank_kwargs(OlderEngine(), 3) == {} + + +@pytest.mark.asyncio +async def test_disagg_decode_passes_prefill_dp_rank_to_engine(): + handler = _new_decode_handler() + handler.serving_mode = DisaggregationMode.DECODE + handler._enable_frontend_decoding = False + handler._mm_hashes_supported = False + handler._engine_supports_priority = False + handler._routed_experts_kwargs = {} + handler.enable_trace = False + handler._get_input_param = lambda request: {"input_ids": [1, 2]} + handler._resolve_lora = lambda request: None + chunk = { + "index": 0, + "output_ids": [42], + "meta_info": {"id": "sample-0", "finish_reason": {"type": "length"}}, + } + handler.engine = SimpleNamespace( + async_generate=AsyncMock(return_value=_stream([chunk])) + ) + context = SimpleNamespace( + id=lambda: "request-id", + trace_id="trace-id", + is_stopped=lambda: False, + notify_first_token=lambda: None, + ) + request = { + "sampling_options": {"n": 1}, + "stop_conditions": {"max_tokens": 1}, + "routing": {"dp_rank": 1, "prefill_dp_rank": 3}, + "bootstrap_info": { + "bootstrap_host": "prefill.invalid", + "bootstrap_port": 1234, + "bootstrap_room": 23, + }, + } + [output async for output in handler.generate(request, context)] + kwargs = handler.engine.async_generate.await_args.kwargs + assert kwargs["disagg_prefill_dp_rank"] == 3 + assert kwargs["data_parallel_rank"] == 1 + + +def test_build_native_generate_request_forwards_prefill_dp_rank(): + native = build_native_generate_request( + {"sampling_params": {"max_new_tokens": 1}}, + input_ids=[7, 8], + request_id="internal-request-id", + priority=None, + bootstrap_host="prefill.internal", + bootstrap_port=1234, + bootstrap_room=23, + routed_dp_rank=1, + prefill_dp_rank=3, + ) + assert native.disagg_prefill_dp_rank == 3 diff --git a/components/src/dynamo/sglang/tests/test_sglang_elastic_ep.py b/components/src/dynamo/sglang/tests/test_sglang_elastic_ep.py index dd107fa39f48..933fbd6cef50 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_elastic_ep.py +++ b/components/src/dynamo/sglang/tests/test_sglang_elastic_ep.py @@ -66,6 +66,10 @@ def _make_handler( tokenizer_manager = SimpleNamespace( server_args=SimpleNamespace(elastic_ep_backend=backend), + pause_generation=lambda: None, + continue_generation=lambda: None, + release_memory_occupation=lambda: None, + resume_memory_occupation=lambda: None, get_elastic_ep_state=MagicMock( return_value=state if state is not None diff --git a/components/src/dynamo/sglang/tests/test_sglang_frontend_decoding.py b/components/src/dynamo/sglang/tests/test_sglang_frontend_decoding.py index d8f2999e25af..c116e8a8719d 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_frontend_decoding.py +++ b/components/src/dynamo/sglang/tests/test_sglang_frontend_decoding.py @@ -450,7 +450,15 @@ def _new_decode_handler(*, enable_frontend_decoding: bool): @asynccontextmanager async def no_cancellation_monitor(*args, **kwargs): - yield None + async def wait_forever(): + await asyncio.Future() + + task = asyncio.create_task(wait_forever()) + try: + yield task + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) handler._cancellation_monitor = no_cancellation_monitor @@ -534,6 +542,7 @@ async def fake_async_generate(**kwargs): request = { "token_ids": [1, 2, 3], + "image_cache_scope": "session-42", "multi_modal_data": {"image_url": [{"Decoded": decoded_metadata}]}, } @@ -541,7 +550,7 @@ async def fake_async_generate(**kwargs): pass image_loader.load_image_batch.assert_awaited_once_with( - [{"Decoded": decoded_metadata}] + [{"Decoded": decoded_metadata}], cache_scope="session-42" ) assert captured["image_data"] == [pil_stub] @@ -676,7 +685,15 @@ def _new_prefill_handler() -> PrefillWorkerHandler: @asynccontextmanager async def no_cancellation_monitor(*args, **kwargs): - yield None + async def wait_forever(): + await asyncio.Future() + + task = asyncio.create_task(wait_forever()) + try: + yield task + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) handler._cancellation_monitor = no_cancellation_monitor @@ -687,6 +704,55 @@ async def no_cancellation_monitor(*args, **kwargs): return handler +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode", + [ + DisaggregationMode.AGGREGATED, + DisaggregationMode.PREFILL, + DisaggregationMode.DECODE, + ], +) +async def test_cache_salt_reaches_engine(mode, unused_tcp_port): + handler = ( + _new_prefill_handler() + if mode == DisaggregationMode.PREFILL + else _new_decode_handler(enable_frontend_decoding=False) + ) + handler.serving_mode = mode + recorder = _GenerateRecorder() + handler.engine = recorder + request = { + "token_ids": [1, 2, 3], + "routing": {"cache_salt": "tenant-a"}, + "extra_args": {"nvext": {"cache_salt": "body-salt"}}, + "bootstrap_info": { + "bootstrap_host": "prefill.invalid", + "bootstrap_port": unused_tcp_port, + "bootstrap_room": 7, + }, + } + routed_request = ( + {"request": request, "sampling_params": {}} + if mode == DisaggregationMode.PREFILL + else request + ) + + async for _ in handler.generate(routed_request, _Context()): + pass + + assert len(recorder.calls) == 1 + assert recorder.calls[0]["cache_salt"] == "tenant-a" + + del request["routing"] + del request["extra_args"] + async for _ in handler.generate(routed_request, _Context()): + pass + + assert len(recorder.calls) == 2 + assert "cache_salt" not in recorder.calls[1] + + @pytest.mark.asyncio async def test_aggregated_decode_omits_session_params_for_agent_context( session_agent_context: Dict[str, Any], diff --git a/components/src/dynamo/sglang/tests/test_sglang_gateway.py b/components/src/dynamo/sglang/tests/test_sglang_gateway.py new file mode 100644 index 000000000000..e118d67faf84 --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_sglang_gateway.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the multi-process SGLang gateway (no engine, no GPU).""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +pytest.importorskip("sglang", reason="sglang not installed in this container") + +from dynamo.sglang import gateway # noqa: E402 + +pytestmark = [ + pytest.mark.unit, + pytest.mark.sglang, + pytest.mark.gpu_0, + pytest.mark.pre_merge, +] + + +def _server_args(tokenizer_worker_num=1, node_rank=0, **kw): + return SimpleNamespace( + tokenizer_worker_num=tokenizer_worker_num, node_rank=node_rank, **kw + ) + + +def _dyn(gateway_workers=None, **kw): + return SimpleNamespace(gateway_workers=gateway_workers, **kw) + + +@pytest.fixture +def not_a_child(monkeypatch): + monkeypatch.delenv(gateway.ENV_PARENT_PID, raising=False) + monkeypatch.delenv(gateway.ENV_CHILD_INDEX, raising=False) + + +def test_effective_count_from_either_flag(not_a_child): + assert gateway.effective_gateway_workers(_server_args(1), _dyn()) == 1 + assert gateway.effective_gateway_workers(_server_args(8), _dyn()) == 8 + assert gateway.effective_gateway_workers(_server_args(1), _dyn(3)) == 3 + assert gateway.effective_gateway_workers(_server_args(3), _dyn(3)) == 3 + + +@pytest.mark.parametrize("tokenizer_workers,requested", [(8, 3), (2, 1)]) +def test_conflicting_explicit_counts_are_rejected( + not_a_child, tokenizer_workers, requested +): + with pytest.raises(ValueError, match="conflicts"): + gateway.effective_gateway_workers( + _server_args(tokenizer_workers), _dyn(requested) + ) + + +def test_gateway_worker_count_only_on_leader_and_not_in_children( + not_a_child, monkeypatch +): + assert gateway.gateway_worker_count(_server_args(8), _dyn()) == 8 + assert gateway.gateway_worker_count(_server_args(8, node_rank=1), _dyn()) == 1 + monkeypatch.setenv(gateway.ENV_PARENT_PID, "123") + assert gateway.is_gateway_child() + assert gateway.gateway_worker_count(_server_args(8), _dyn()) == 1 + + +def test_validate_rejects_unsupported_modes(not_a_child, monkeypatch): + monkeypatch.delenv("DYN_SNAPSHOT_CONTROL_DIR", raising=False) + gateway.validate_gateway_mode(_server_args(4), _dyn(), 4) + gateway.validate_gateway_mode(_server_args(1), _dyn(embedding_worker=True), 1) + with pytest.raises(ValueError, match="embedding-worker"): + gateway.validate_gateway_mode(_server_args(4), _dyn(embedding_worker=True), 4) + with pytest.raises(ValueError, match="enable-lora"): + gateway.validate_gateway_mode(_server_args(4, enable_lora=True), _dyn(), 4) + with pytest.raises(ValueError, match="forward-pass-metrics"): + gateway.validate_gateway_mode( + _server_args(4, enable_forward_pass_metrics=True), _dyn(), 4 + ) + monkeypatch.setenv("DYN_SNAPSHOT_CONTROL_DIR", "/snapshot-control") + with pytest.raises(ValueError, match="snapshot"): + gateway.validate_gateway_mode(_server_args(4), _dyn(), 4) + + +def test_validate_engine_routes_need_attach_api(monkeypatch): + monkeypatch.delenv("DYN_SNAPSHOT_CONTROL_DIR", raising=False) + monkeypatch.setattr( + gateway.sgl.Engine, "attach_tokenizer_worker", None, raising=False + ) + routes = ["flush_cache", "abort=abort_request:tm"] + gateway.validate_gateway_mode(_server_args(4), _dyn(engine_routes=routes[1:]), 4) + with pytest.raises(ValueError, match="attach_tokenizer_worker"): + gateway.validate_gateway_mode(_server_args(4), _dyn(engine_routes=routes), 4) + monkeypatch.setattr(gateway.sgl.Engine, "attach_tokenizer_worker", lambda pid: None) + gateway.validate_gateway_mode(_server_args(4), _dyn(engine_routes=routes), 4) + + +def test_child_index_decides_metrics_ownership_and_fanout(monkeypatch): + monkeypatch.delenv(gateway.ENV_PARENT_PID, raising=False) + monkeypatch.delenv(gateway.ENV_CHILD_INDEX, raising=False) + assert gateway.gateway_child_index() == 0 + assert gateway.owns_engine_metrics() + assert gateway.metrics_fanout_endpoint() is None + monkeypatch.setenv(gateway.ENV_PARENT_PID, "4242") + monkeypatch.setenv(gateway.ENV_CHILD_INDEX, "2") + assert not gateway.owns_engine_metrics() + assert gateway.metrics_fanout_endpoint().endswith("_4242") + + +def test_system_port_is_handed_to_children(not_a_child, monkeypatch): + monkeypatch.setenv(gateway.ENV_SYSTEM_PORT, "8081") + monkeypatch.delenv(gateway.ENV_SYSTEM_PORT_BASE, raising=False) + gateway.reserve_system_port_for_children() + import os + + assert os.environ[gateway.ENV_SYSTEM_PORT] == "-1" + env0, env2 = gateway.child_environment(0), gateway.child_environment(2) + assert env0[gateway.ENV_SYSTEM_PORT] == "8081" + assert env2[gateway.ENV_SYSTEM_PORT] == "0" + assert gateway.ENV_SYSTEM_PORT_BASE not in env2 + assert env2[gateway.ENV_CHILD_INDEX] == "2" + assert env2[gateway.ENV_PARENT_PID] == str(os.getpid()) + + +@pytest.mark.parametrize("value", [None, "-1", "0", "not-a-port"]) +def test_system_port_untouched_when_disabled(not_a_child, monkeypatch, value): + import os + + monkeypatch.delenv(gateway.ENV_SYSTEM_PORT_BASE, raising=False) + if value is None: + monkeypatch.delenv(gateway.ENV_SYSTEM_PORT, raising=False) + else: + monkeypatch.setenv(gateway.ENV_SYSTEM_PORT, value) + gateway.reserve_system_port_for_children() + assert os.environ.get(gateway.ENV_SYSTEM_PORT) == value + assert gateway.child_environment(1).get(gateway.ENV_SYSTEM_PORT) == value + assert gateway.ENV_SYSTEM_PORT_BASE not in os.environ + + +def test_parent_watchdog_fires_when_parent_is_gone(): + import threading + + died = threading.Event() + dead_pid = 2**22 - 7 # above pid_max on every default Linux, so never alive + thread = gateway.start_parent_watchdog( + dead_pid, on_parent_death=died.set, poll_seconds=0.01 + ) + assert died.wait(5) + thread.join(5) + assert not thread.is_alive() + + +def test_parent_watchdog_stays_quiet_while_parent_lives(): + import os + import threading + import time + + died = threading.Event() + gateway.start_parent_watchdog( + os.getppid(), on_parent_death=died.set, poll_seconds=0.01 + ) + time.sleep(0.1) + assert not died.is_set() + + +def test_gateway_engine_id_and_shutdown_budget(monkeypatch): + monkeypatch.delenv(gateway.ENV_PARENT_PID, raising=False) + assert gateway.gateway_engine_id() is None + monkeypatch.setenv(gateway.ENV_PARENT_PID, "777") + assert gateway.gateway_engine_id().endswith(":777") + monkeypatch.setenv("DYN_GRACEFUL_SHUTDOWN_GRACE_PERIOD_SECS", "12") + assert gateway.child_shutdown_timeout() == 12 + gateway.CHILD_DRAIN_AND_CLEANUP_SECS + + +def _child_env(monkeypatch, parent_pid="4321"): + monkeypatch.setenv(gateway.ENV_PARENT_PID, parent_pid) + monkeypatch.setenv(gateway.ENV_CHILD_INDEX, "1") + monkeypatch.setattr(gateway, "start_parent_watchdog", lambda pid: None) + + +def test_build_gateway_engine_prefers_attach_api(monkeypatch): + _child_env(monkeypatch) + seen = {} + + def attach(pid): + seen["pid"] = pid + return SimpleNamespace(port_args=SimpleNamespace(metrics_ipc_name="ipc:///p")) + + monkeypatch.setattr( + gateway.sgl.Engine, "attach_tokenizer_worker", attach, raising=False + ) + engine = gateway.build_gateway_engine() + assert seen["pid"] == 4321 + assert engine.port_args.metrics_ipc_name == "ipc:///p" + + +def test_build_gateway_engine_falls_back_to_facade(monkeypatch): + import sglang.srt.managers.multi_tokenizer_mixin as mixin + import sglang.srt.runtime_context as runtime_context + + _child_env(monkeypatch) + monkeypatch.setattr( + gateway.sgl.Engine, "attach_tokenizer_worker", None, raising=False + ) + port_args = SimpleNamespace( + metrics_ipc_name="ipc:///p", tokenizer_ipc_name="ipc:///t" + ) + server_args = SimpleNamespace(tokenizer_worker_num=2) + info = {"max_req_input_len": 64, "startup_time": {"t": 1}} + created = {} + + class FakeWorker: + def __init__(self, sa, pa): + created["port_args"] = pa + + def set_startup_time(self, t): + created["startup_time"] = t + + monkeypatch.setattr( + mixin, "read_from_shared_memory", lambda name: (port_args, server_args, info) + ) + monkeypatch.setattr(mixin, "get_tokenizer_worker_class", lambda sa: FakeWorker) + monkeypatch.setattr( + runtime_context, "publish", lambda sa, role: created.setdefault("role", role) + ) + + engine = gateway.build_gateway_engine() + assert isinstance(engine, gateway.GatewayEngine) + assert created["role"] == "tokenizer" + assert created["startup_time"] == {"t": 1} + assert created["port_args"].tokenizer_ipc_name != "ipc:///t" + assert engine.port_args.metrics_ipc_name == "ipc:///p" + assert engine.tokenizer_manager.max_req_input_len == 64 + + +def test_gateway_engine_facade_generates_through_tokenizer_manager(): + seen = {} + + class FakeTokenizerManager: + async def generate_request(self, obj, request): + seen["obj"] = obj + yield {"text": "ok"} + + facade = gateway.GatewayEngine( + FakeTokenizerManager(), + server_args=SimpleNamespace(), + port_args=SimpleNamespace(metrics_ipc_name="ipc:///tmp/x"), + scheduler_info={"max_req_input_len": 4}, + ) + assert facade._scheduler_init_result.scheduler_infos[0]["max_req_input_len"] == 4 + + async def run(): + gen = await facade.async_generate(input_ids=[1, 2, 3], stream=True) + return [chunk async for chunk in gen] + + assert asyncio.run(run()) == [{"text": "ok"}] + assert seen["obj"].input_ids == [1, 2, 3] + facade.shutdown() + + +class FakeProc: + def __init__(self, pid): + self.pid = pid + self.returncode = None + self.terminated = False + self.waited = False + + def poll(self): + return self.returncode + + def terminate(self): + self.terminated = True + self.returncode = -15 + + def wait(self, timeout=None): + self.waited = True + return self.returncode + + +class FakeShm: + unlinked = False + + def unlink(self): + self.unlinked = True + + +def _engine(): + return SimpleNamespace( + port_args=SimpleNamespace(), + server_args=SimpleNamespace(), + tokenizer_manager=SimpleNamespace(startup_time={"t": 0}), + _scheduler_init_result=SimpleNamespace( + scheduler_infos=[{"max_req_input_len": 8}] + ), + ) + + +def test_serve_via_gateway_children_spawns_and_fails_on_dead_child( + not_a_child, monkeypatch +): + spawned = [] + + def fake_popen(cmd, env): + p = FakeProc(4000 + len(spawned)) + p.cmd, p.env = cmd, env + spawned.append(p) + return p + + shm = FakeShm() + monkeypatch.setattr(gateway.subprocess, "Popen", fake_popen) + monkeypatch.setattr( + "sglang.srt.managers.multi_tokenizer_mixin.write_data_for_multi_tokenizer", + lambda port_args, server_args, info: shm, + ) + monkeypatch.setattr(gateway.sys, "argv", ["dynamo.sglang", "--model-path", "/m"]) + + async def fast_sleep(_): + spawned[1].returncode = 1 + + monkeypatch.setattr(gateway.asyncio, "sleep", fast_sleep) + + with pytest.raises(RuntimeError, match="exited rc=1"): + asyncio.run(gateway.serve_via_gateway_children(_engine(), 3, asyncio.Event())) + + assert len(spawned) == 3 + assert all(p.cmd[1:3] == ["-m", "dynamo.sglang"] for p in spawned) + assert all(p.cmd[3:] == ["--model-path", "/m"] for p in spawned) + assert [p.env[gateway.ENV_CHILD_INDEX] for p in spawned] == ["0", "1", "2"] + assert all(gateway.ENV_PARENT_PID in p.env for p in spawned) + assert all(p.terminated for p in spawned if p.pid != spawned[1].pid) + assert all(p.waited for p in spawned) + assert shm.unlinked + + +def test_serve_via_gateway_children_tolerates_child_exit_during_shutdown( + not_a_child, monkeypatch +): + spawned = [] + stop = asyncio.Event() + + def fake_popen(cmd, env): + p = FakeProc(4100 + len(spawned)) + spawned.append(p) + return p + + monkeypatch.setattr(gateway.subprocess, "Popen", fake_popen) + monkeypatch.setattr( + "sglang.srt.managers.multi_tokenizer_mixin.write_data_for_multi_tokenizer", + lambda port_args, server_args, info: FakeShm(), + ) + + async def fast_sleep(_): + stop.set() + spawned[0].returncode = 0 + + monkeypatch.setattr(gateway.asyncio, "sleep", fast_sleep) + asyncio.run(gateway.serve_via_gateway_children(_engine(), 2, stop)) + assert spawned[1].terminated and all(p.waited for p in spawned) + + +def test_serve_via_gateway_children_cleans_up_when_spawn_fails( + not_a_child, monkeypatch +): + spawned = [] + + def fake_popen(cmd, env): + if len(spawned) == 1: + raise OSError("no more pids") + p = FakeProc(5000 + len(spawned)) + spawned.append(p) + return p + + shm = FakeShm() + monkeypatch.setattr(gateway.subprocess, "Popen", fake_popen) + monkeypatch.setattr( + "sglang.srt.managers.multi_tokenizer_mixin.write_data_for_multi_tokenizer", + lambda port_args, server_args, info: shm, + ) + with pytest.raises(OSError, match="no more pids"): + asyncio.run(gateway.serve_via_gateway_children(_engine(), 3, asyncio.Event())) + assert len(spawned) == 1 and spawned[0].terminated and spawned[0].waited + assert shm.unlinked + + +def test_serve_via_gateway_children_reuses_engine_published_shm( + not_a_child, monkeypatch +): + """An SGLang Engine that already published its args (tokenizer_router set) owns + the shared memory; the gateway must not publish or unlink it.""" + calls = [] + monkeypatch.setattr( + "sglang.srt.managers.multi_tokenizer_mixin.write_data_for_multi_tokenizer", + lambda *a: calls.append(a), + ) + monkeypatch.setattr(gateway.subprocess, "Popen", lambda cmd, env: FakeProc(1)) + + class OwnedShm: + def unlink(self): + raise AssertionError("gateway must not unlink the engine's shm") + + engine = _engine() + engine._multi_tokenizer_shm = OwnedShm() + stop = asyncio.Event() + + async def run(): + task = asyncio.create_task(gateway.serve_via_gateway_children(engine, 1, stop)) + await asyncio.sleep(0) + stop.set() + await task + + asyncio.run(run()) + assert calls == [] + + +def test_child_zero_receives_the_leader_load_time(not_a_child, monkeypatch): + monkeypatch.delenv(gateway.ENV_LOAD_TIME, raising=False) + env0 = gateway.child_environment(0, load_time=12.5) + env1 = gateway.child_environment(1, load_time=12.5) + assert env0[gateway.ENV_LOAD_TIME] == "12.5" + assert gateway.ENV_LOAD_TIME not in env1 + assert gateway.ENV_LOAD_TIME not in gateway.child_environment(0) + assert gateway.attached_engine_load_time() is None + monkeypatch.setenv(gateway.ENV_LOAD_TIME, env0[gateway.ENV_LOAD_TIME]) + assert gateway.attached_engine_load_time() == 12.5 + + +def test_follow_pause_broadcasts_resyncs_after_each_broadcast(): + seen = [] + + class FakeWorker: + is_pause = False + + async def _apply_pause_continue_broadcast(self, obj): + self.is_pause = obj.is_pause + + worker = FakeWorker() + + async def on_change(): + seen.append(worker.is_pause) + + assert gateway.follow_pause_broadcasts(worker, on_change) + assert not gateway.follow_pause_broadcasts(SimpleNamespace(), on_change) + + async def run(): + await worker._apply_pause_continue_broadcast(SimpleNamespace(is_pause=True)) + await worker._apply_pause_continue_broadcast(SimpleNamespace(is_pause=False)) + + asyncio.run(run()) + assert seen == [True, False] diff --git a/components/src/dynamo/sglang/tests/test_sglang_image_diffusion_handler.py b/components/src/dynamo/sglang/tests/test_sglang_image_diffusion_handler.py index f8198b22368d..e8478c7e7cd7 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_image_diffusion_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_image_diffusion_handler.py @@ -440,32 +440,135 @@ async def test_generate_with_nvext(self, handler, mock_context): @pytest.mark.asyncio async def test_generate_i2i_passes_image_path( - self, handler, mock_context, tmp_path + self, handler, mock_context, tmp_path, monkeypatch ): - """Test that input_reference is passed as image_path to the generator.""" + """The generator gets the *resolved* path, never the raw reference. + + Resolution is what makes the confinement check meaningful: the string + handed to the generator has to be the one that was checked against + DYN_MM_LOCAL_PATH, not the one the client sent. + """ test_image = Image.new("RGB", (256, 256), color="green") handler.generator.generate = Mock( return_value=SimpleNamespace(frames=[test_image]) ) - input_ref = str(tmp_path / "test_input.png") + monkeypatch.setenv("DYN_MM_LOCAL_PATH", str(tmp_path)) + (tmp_path / "sub").mkdir() + input_ref = tmp_path / "test_input.png" + input_ref.write_bytes(b"x") + # A "sub/.." segment so the resolved form differs from what was sent; + # on Linux an already-normalized tmp_path resolves to itself, and the + # assertion below would hold even with no resolution at all. + sent = str(tmp_path / "sub" / ".." / "test_input.png") request = { "prompt": "Transform this image", "model": "test-model", "size": "256x256", "response_format": "b64_json", - "input_reference": input_ref, + "input_reference": sent, } results = [] async for result in handler.generate(request, mock_context): results.append(result) - # Verify image_path was passed to the generator - call_args = handler.generator.generate.call_args - sampling_params = call_args[1]["sampling_params_kwargs"] - assert sampling_params["image_path"] == input_ref + sampling_params = handler.generator.generate.call_args[1][ + "sampling_params_kwargs" + ] + assert sampling_params["image_path"] != sent + assert sampling_params["image_path"] == str(input_ref.resolve()) + + @pytest.mark.asyncio + async def test_generate_i2i_rejects_path_outside_allowed_dir( + self, handler, mock_context, tmp_path, monkeypatch + ): + """A rejected reference must be a 400, not a sanitized 500. + + UrlValidationError is a policy verdict on client input. It is also a + plain ValueError, so without the explicit mapping it reaches the runtime + as an unexpected failure and the client is told nothing. + """ + handler.generator.generate = Mock() + monkeypatch.setenv("DYN_MM_LOCAL_PATH", str(tmp_path)) + request = { + "prompt": "x", + "model": "test-model", + "size": "256x256", + "response_format": "b64_json", + "input_reference": "/etc/passwd", + } + + with pytest.raises(InvalidArgument, match="outside the allowed directory"): + async for _ in handler.generate(request, mock_context): + pass + + handler.generator.generate.assert_not_called() + + @pytest.mark.asyncio + async def test_generate_i2i_rejection_message_is_bounded( + self, handler, mock_context, tmp_path, monkeypatch + ): + """A rejected reference is echoed back bounded, not at full length. + + input_reference has no length limit, and the 400 body plus the error log + are two sinks per request: echoing it verbatim turns one request into an + arbitrarily large amplification. + """ + handler.generator.generate = Mock() + monkeypatch.setenv("DYN_MM_LOCAL_PATH", str(tmp_path)) + reference = "/nope/" + "A" * 200_000 + ".png" + request = { + "prompt": "x", + "model": "test-model", + "size": "256x256", + "response_format": "b64_json", + "input_reference": reference, + } + + with pytest.raises(InvalidArgument) as excinfo: + async for _ in handler.generate(request, mock_context): + pass + + assert len(str(excinfo.value)) < 500, "client-visible message is unbounded" + handler.generator.generate.assert_not_called() + + @pytest.mark.asyncio + async def test_generate_i2i_maps_a_4xx_fetch_to_invalid_argument( + self, handler, mock_context, tmp_path, monkeypatch + ): + """A 4xx from the origin the client chose is the client's problem. + + Everything that is not an InvalidArgument becomes a sanitized 500, so + without this the caller is told nothing about their own dead URL. The + URL is not echoed back -- it has no length limit and the video handler + puts str(exc) straight in its response body. + """ + from dynamo.common.http.base import HttpStatusError + + handler.generator.generate = Mock() + monkeypatch.setenv("DYN_MM_ALLOW_INTERNAL", "1") + url = "http://example.com/" + "u" * 5000 + + async def fake_fetch(u, timeout, *, policy=None, max_bytes=None): + raise HttpStatusError(404, "Not Found", u) + + monkeypatch.setattr("dynamo.common.http.fetch_bytes", fake_fetch) + request = { + "prompt": "x", + "model": "test-model", + "size": "256x256", + "response_format": "b64_json", + "input_reference": url, + } + + with pytest.raises(InvalidArgument, match="HTTP 404") as excinfo: + async for _ in handler.generate(request, mock_context): + pass + + assert "uuuu" not in str(excinfo.value) + handler.generator.generate.assert_not_called() @pytest.mark.asyncio async def test_generate_t2i_no_image_path(self, handler, mock_context): diff --git a/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py b/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py index 1174ee5f9f43..b6cb0e5103f0 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py +++ b/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py @@ -2,17 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio -import sys -import types from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest -from dynamo.sglang.request_handlers.handler_base import ( - BaseWorkerHandler, - SGLangEnginePauseController, -) +from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler pytestmark = [ pytest.mark.unit, @@ -22,24 +17,6 @@ ] -@pytest.fixture(autouse=True) -def _stub_sglang_io_struct(monkeypatch): - """Keep unit tests independent from CUDA-only sglang imports.""" - - io_struct = types.ModuleType("sglang.srt.managers.io_struct") - - class _Req: - def __init__(self, tags=None): - self.tags = tags - - io_struct.PauseGenerationReqInput = _Req - io_struct.ReleaseMemoryOccupationReqInput = _Req - io_struct.ResumeMemoryOccupationReqInput = _Req - io_struct.ContinueGenerationReqInput = _Req - - monkeypatch.setitem(sys.modules, "sglang.srt.managers.io_struct", io_struct) - - class _TestWorkerHandler(BaseWorkerHandler): async def generate(self, request, context): yield {} @@ -50,6 +27,7 @@ def handler(): handler = _TestWorkerHandler.__new__(_TestWorkerHandler) handler.engine = SimpleNamespace( tokenizer_manager=SimpleNamespace( + is_pause=False, pause_generation=AsyncMock(), release_memory_occupation=AsyncMock(), resume_memory_occupation=AsyncMock(), @@ -65,165 +43,103 @@ def handler(): unregister_endpoint_instance=AsyncMock(), register_endpoint_instance=AsyncMock(), ) - handler._pause_controller = SGLangEnginePauseController(handler.engine) - handler._pause_lock = asyncio.Lock() + handler._engine_route_lock = asyncio.Lock() return handler -@pytest.mark.asyncio -async def test_release_and_resume_are_idempotent(handler): - first_release = await handler.release_memory_occupation({}) - second_release = await handler.release_memory_occupation({}) - - first_resume = await handler.resume_memory_occupation({}) - second_resume = await handler.resume_memory_occupation({}) - - assert first_release["status"] == "ok" - assert second_release["status"] == "ok" - assert first_resume["status"] == "ok" - assert second_resume["status"] == "ok" - assert second_release["message"] == "Memory already released" - assert second_resume["message"] == "Memory already resumed" - - release_req = ( - handler.engine.tokenizer_manager.release_memory_occupation.await_args.args[0] - ) - resume_req = ( - handler.engine.tokenizer_manager.resume_memory_occupation.await_args.args[0] - ) - assert release_req.tags is None - assert resume_req.tags is None +def _registered_engine_routes(handler, configured_routes=None): + registered = {} - handler.engine.tokenizer_manager.pause_generation.assert_awaited_once() - handler.engine.tokenizer_manager.release_memory_occupation.assert_awaited_once() - handler.generate_endpoint.unregister_endpoint_instance.assert_awaited_once() + class Runtime: + def register_engine_route(self, path, route_handler): + registered[path] = route_handler - handler.engine.tokenizer_manager.resume_memory_occupation.assert_awaited_once() - handler.engine.tokenizer_manager.continue_generation.assert_awaited_once() - handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() + handler.config = SimpleNamespace( + dynamo_args=SimpleNamespace(engine_routes=configured_routes or []) + ) + handler.register_engine_routes(Runtime()) + return registered -@pytest.mark.asyncio -async def test_memory_occupation_handlers_forward_tags_exactly(handler): - await handler.release_memory_occupation({"tags": []}) - resume_result = await handler.resume_memory_occupation({"tags": []}) +def _make_native_manager_methods_routable(manager): + mocks = {} + for method_name in ( + "pause_generation", + "continue_generation", + "release_memory_occupation", + "resume_memory_occupation", + ): + method_mock = getattr(manager, method_name) + mocks[method_name] = method_mock - assert resume_result["status"] == "ok" - release_req = ( - handler.engine.tokenizer_manager.release_memory_occupation.await_args.args[0] - ) - resume_req = ( - handler.engine.tokenizer_manager.resume_memory_occupation.await_args.args[0] - ) - assert release_req.tags == [] - assert resume_req.tags == [] - - handler.engine.tokenizer_manager.pause_generation.reset_mock() - handler.engine.tokenizer_manager.release_memory_occupation.reset_mock() - handler.engine.tokenizer_manager.resume_memory_occupation.reset_mock() - handler.engine.tokenizer_manager.continue_generation.reset_mock() - handler.generate_endpoint.unregister_endpoint_instance.reset_mock() - handler.generate_endpoint.register_endpoint_instance.reset_mock() - - await handler.release_memory_occupation({"tags": ["weights"]}) - resume_result = await handler.resume_memory_occupation({}) - - assert resume_result["status"] == "ok" - release_req = ( - handler.engine.tokenizer_manager.release_memory_occupation.await_args.args[0] - ) - resume_req = ( - handler.engine.tokenizer_manager.resume_memory_occupation.await_args.args[0] - ) - assert release_req.tags == ["weights"] - assert resume_req.tags is None - handler.engine.tokenizer_manager.continue_generation.assert_awaited_once() - handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() + async def route_method(_method_mock=method_mock, **kwargs): + return await _method_mock(**kwargs) + + setattr(manager, method_name, route_method) + return mocks @pytest.mark.asyncio -async def test_resume_recovers_generation_pause_after_failed_release_rollback(handler): +async def test_native_memory_routes_follow_sglang_pause_state(handler): manager = handler.engine.tokenizer_manager - manager.release_memory_occupation = AsyncMock( - side_effect=RuntimeError("release failed") - ) - failed_continue = AsyncMock(side_effect=RuntimeError("continue failed")) - manager.continue_generation = failed_continue + method_mocks = _make_native_manager_methods_routable(manager) - release_result = await handler.release_memory_occupation({}) + async def pause_generation(**_kwargs): + manager.is_pause = True - assert release_result["status"] == "error" - assert handler._pause_controller.is_paused is False - assert handler._pause_controller.needs_resume_recovery is True - failed_continue.assert_awaited_once() - handler.generate_endpoint.unregister_endpoint_instance.assert_awaited_once() + async def continue_generation(**_kwargs): + manager.is_pause = False - manager.continue_generation = AsyncMock() - resume_result = await handler.resume_memory_occupation({}) + method_mocks["pause_generation"].side_effect = pause_generation + method_mocks["continue_generation"].side_effect = continue_generation + routes = _registered_engine_routes(handler) - assert resume_result["status"] == "ok" - manager.resume_memory_occupation.assert_not_awaited() - manager.continue_generation.assert_awaited_once() - handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() - assert handler._pause_controller.needs_resume_recovery is False + assert { + "pause_generation", + "continue_generation", + "release_memory_occupation", + "resume_memory_occupation", + }.issubset(routes) + await routes["pause_generation"]({}) + await routes["release_memory_occupation"]({"tags": ["weights", "kv_cache"]}) + await routes["resume_memory_occupation"]({"tags": ["weights"]}) -@pytest.mark.asyncio -async def test_release_reregisters_after_clean_pause_rollback(handler): - manager = handler.engine.tokenizer_manager - manager.release_memory_occupation = AsyncMock( - side_effect=RuntimeError("release failed") - ) + assert manager.is_pause is True + assert handler.generate_endpoint.unregister_endpoint_instance.await_count == 3 + handler.generate_endpoint.register_endpoint_instance.assert_not_awaited() - release_result = await handler.release_memory_occupation({}) + await routes["resume_memory_occupation"]({"tags": ["kv_cache"]}) - # Rollback continued generation cleanly, so the engine is serving-safe again. - assert release_result["status"] == "error" - assert handler._pause_controller.is_paused is False - assert handler._pause_controller.needs_resume_recovery is False - manager.continue_generation.assert_awaited_once() - handler.generate_endpoint.unregister_endpoint_instance.assert_awaited_once() - # The endpoint must rejoin the routing pool since resume will early-return. - handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() + assert handler.generate_endpoint.unregister_endpoint_instance.await_count == 4 + handler.generate_endpoint.register_endpoint_instance.assert_not_awaited() + await routes["continue_generation"]({}) -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("method_name", "endpoint_method"), - [ - ("release_memory_occupation", "unregister_endpoint_instance"), - ("resume_memory_occupation", "register_endpoint_instance"), - ], -) -async def test_memory_control_returns_error_without_pause_controller( - handler, - method_name, - endpoint_method, -): - handler.engine = None - handler._pause_controller = None + assert manager.is_pause is False + handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() - result = await getattr(handler, method_name)({}) - assert result == { - "status": "error", - "message": "memory control not supported on this worker", - } - getattr(handler.generate_endpoint, endpoint_method).assert_not_awaited() +@pytest.mark.asyncio +async def test_configured_engine_route_overrides_default_and_syncs_discovery(handler): + custom_pause_mock = AsyncMock(return_value={"custom": True}) + async def custom_pause(): + return await custom_pause_mock() -@pytest.mark.asyncio -async def test_resume_keeps_paused_state_when_register_fails(handler): - await handler.release_memory_occupation({}) - handler.generate_endpoint.register_endpoint_instance = AsyncMock( - side_effect=RuntimeError("discovery write timeout") + handler.engine.custom_pause = custom_pause + handler.engine.tokenizer_manager.is_pause = True + method_mocks = _make_native_manager_methods_routable( + handler.engine.tokenizer_manager ) + routes = _registered_engine_routes(handler, ["pause_generation=custom_pause"]) - result = await handler.resume_memory_occupation({}) + result = await routes["pause_generation"]({}) - assert result["status"] == "error" - assert handler._pause_controller is not None - assert handler._pause_controller.is_paused is True + assert result == {"custom": True} + custom_pause_mock.assert_awaited_once_with() + method_mocks["pause_generation"].assert_not_awaited() + handler.generate_endpoint.unregister_endpoint_instance.assert_awaited_once() @pytest.mark.asyncio diff --git a/components/src/dynamo/sglang/tests/test_sglang_multimodal_embedding_cache.py b/components/src/dynamo/sglang/tests/test_sglang_multimodal_embedding_cache.py index 4d53e0a7d1bc..e7eb141e5e1d 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_multimodal_embedding_cache.py +++ b/components/src/dynamo/sglang/tests/test_sglang_multimodal_embedding_cache.py @@ -18,6 +18,7 @@ HttpStatusError, HttpTimeoutError, ) +from dynamo.common.http.media_reference import DYN_MM_MAX_FILE_SIZE_MB from dynamo.common.http.url_validator import ( UrlValidationError, UrlValidationPolicy, @@ -181,6 +182,47 @@ async def test_encode_with_cache_partial_hit_and_reuse( assert torch.equal(full_embeddings2, full_embeddings) +@pytest.mark.asyncio +async def test_session_scope_partitions_image_embedding_keys( + cache_handler: MultimodalEncodeWorkerHandler, +) -> None: + """The same image URL in separate sessions must use separate entries.""" + cache_handler._session_scoped_cache = True + item = {"Url": "https://example.com/image.png"} + + _, first_keys, _ = await cache_handler._prepare_image_inputs( + [item], cache_scope="session-a" + ) + _, second_keys, _ = await cache_handler._prepare_image_inputs( + [item], cache_scope="session-b" + ) + _, first_keys_again, _ = await cache_handler._prepare_image_inputs( + [item], cache_scope="session-a" + ) + + assert first_keys[0] is not None + assert second_keys[0] is not None + assert first_keys[0] != second_keys[0] + assert first_keys_again == first_keys + + +@pytest.mark.asyncio +async def test_session_scoped_image_embedding_cache_bypasses_without_scope( + cache_handler: MultimodalEncodeWorkerHandler, +) -> None: + """Missing scope must not produce an image embedding-cache key.""" + cache_handler._session_scoped_cache = True + item = {"Url": "https://example.com/image.png"} + + _, missing_keys, _ = await cache_handler._prepare_image_inputs([item]) + _, blank_keys, _ = await cache_handler._prepare_image_inputs( + [item], cache_scope=" " + ) + + assert missing_keys == [None] + assert blank_keys == [None] + + def test_publish_cache_delta_delegates_to_publisher( cache_handler: MultimodalEncodeWorkerHandler, ) -> None: @@ -309,11 +351,15 @@ async def test_encode_with_cache_reencodes_only_unkeyed_items( @pytest.mark.asyncio async def test_video_requests_reuse_cached_embeddings( cache_handler: MultimodalEncodeWorkerHandler, + monkeypatch, ) -> None: """Second identical video request should reuse cached embeddings.""" video_url = "https://example.com/clip.mp4" video_token_id = cache_handler.video_token_id + video_bytes = b"video-bytes" + fetch = AsyncMock(return_value=video_bytes) + monkeypatch.setattr(f"{_HANDLER_MOD}.fetch_bytes", fetch) cache_handler.encoder.encode_mock.return_value = ( torch.tensor([2, 3, 4]), @@ -370,8 +416,9 @@ async def _responses(): outputs_second.append(item) cache_handler.encoder.encode_mock.assert_awaited_once_with( - [video_url], Modality.VIDEO + [video_bytes], Modality.VIDEO ) + fetch.assert_awaited_once() assert outputs == [{"token_ids": [7]}] assert outputs_second == [{"token_ids": [7]}] @@ -400,9 +447,13 @@ async def _responses(): @pytest.mark.asyncio async def test_video_request_skips_cache_key_when_cache_is_disabled( cache_handler: MultimodalEncodeWorkerHandler, + monkeypatch, ) -> None: video_url = "https://example.com/clip.mp4" video_token_id = cache_handler.video_token_id + video_bytes = b"video-bytes" + fetch = AsyncMock(return_value=video_bytes) + monkeypatch.setattr(f"{_HANDLER_MOD}.fetch_bytes", fetch) cache_handler._embedding_cache = None cache_handler._image_loader = None cache_handler._media_cache_key = Mock( @@ -452,8 +503,9 @@ async def _responses(): assert outputs == [{"token_ids": [7]}] cache_handler._media_cache_key.assert_not_called() cache_handler.encoder.encode_mock.assert_awaited_once_with( - [video_url], Modality.VIDEO + [video_bytes], Modality.VIDEO ) + fetch.assert_awaited_once() def test_aux_value_for_item_rejects_mismatched_batched_lists() -> None: @@ -474,10 +526,13 @@ def test_aux_value_for_item_rejects_mismatched_batched_lists() -> None: @pytest.mark.asyncio async def test_video_cache_key_includes_sampling_config( cache_handler: MultimodalEncodeWorkerHandler, + monkeypatch, ) -> None: """Changing video sampling config should force a cache miss for the same URL.""" video_url = "https://example.com/clip.mp4" + fetch = AsyncMock(return_value=b"video-bytes") + monkeypatch.setattr(f"{_HANDLER_MOD}.fetch_bytes", fetch) cache_handler.encoder.encode_mock.return_value = ( torch.tensor([[2, 3, 4]]), torch.arange(24, dtype=torch.float32).reshape(6, 4), @@ -502,6 +557,7 @@ async def test_video_cache_key_includes_sampling_config( assert first_key != second_key assert cache_handler.encoder.encode_mock.await_count == 2 + assert fetch.await_count == 2 assert cache_handler._embedding_cache.get(first_key) is not None assert cache_handler._embedding_cache.get(second_key) is not None @@ -555,8 +611,8 @@ async def test_disabled_nvdec_without_software_decoder_is_actionable( async def test_disabled_nvdec_with_software_decoder_passes_urls( nvdec_handler, monkeypatch ) -> None: - """NVDEC off but a software decoder exists: URLs pass through unchanged - (SGLang fetches and decodes them itself, as before).""" + """NVDEC off but a software decoder exists: non-HTTP URLs pass through + unchanged for SGLang to decode.""" monkeypatch.setenv("DYN_DISABLE_NVDEC", "1") monkeypatch.setattr(f"{_HANDLER_MOD}._software_video_decoder_imports", lambda: True) @@ -564,6 +620,28 @@ async def test_disabled_nvdec_with_software_decoder_passes_urls( assert out == [_INLINE_VIDEO] +@pytest.mark.asyncio +async def test_disabled_nvdec_fetches_remote_video_with_configured_limit( + nvdec_handler, monkeypatch +) -> None: + """Remote URLs must not escape Dynamo's limit through SGLang's downloader.""" + monkeypatch.setenv("DYN_DISABLE_NVDEC", "1") + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, "1") + url = "https://example.com/clip.webm" + fetch = AsyncMock(return_value=b"video-bytes") + monkeypatch.setattr(f"{_HANDLER_MOD}.fetch_bytes", fetch) + + out = await nvdec_handler._build_encode_inputs([url], "VIDEO") + + assert out == [b"video-bytes"] + fetch.assert_awaited_once_with( + url, + 30.0, + policy=nvdec_handler._url_policy, + max_bytes=1024 * 1024, + ) + + def test_nvdec_video_enabled_gating(nvdec_handler, monkeypatch) -> None: monkeypatch.setattr(f"{_HANDLER_MOD}.nvdec_available", lambda: True) diff --git a/components/src/dynamo/sglang/tests/test_sglang_multimodal_utils.py b/components/src/dynamo/sglang/tests/test_sglang_multimodal_utils.py index 2a0e0b92af11..290dce6e7bc1 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_multimodal_utils.py +++ b/components/src/dynamo/sglang/tests/test_sglang_multimodal_utils.py @@ -131,6 +131,26 @@ def test_raises_for_audio_url(self): def test_text_only_request_bypasses_guard(self): raise_if_unextracted_multimodal({"token_ids": [10, 20, 30]}) + def test_stripped_extra_args_still_uses_multi_modal_data(self): + data_url = "data:image/png;base64,AAAA" + request = { + "token_ids": [1, 2, 3], + "multi_modal_data": {"image_url": [{"Url": data_url}]}, + "extra_args": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": ""}}, + ], + } + ] + }, + } + raise_if_unextracted_multimodal(request) + assert build_disagg_mm_kwargs(request)["image_data"] == [data_url] + async def _stream(items): for item in items: diff --git a/components/src/dynamo/sglang/tests/test_sglang_multimodal_video.py b/components/src/dynamo/sglang/tests/test_sglang_multimodal_video.py index 3631a96335f5..0bf70476f948 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_multimodal_video.py +++ b/components/src/dynamo/sglang/tests/test_sglang_multimodal_video.py @@ -6,6 +6,7 @@ import json from contextlib import asynccontextmanager from types import SimpleNamespace +from unittest.mock import AsyncMock import numpy as np import pytest @@ -445,6 +446,7 @@ async def results(): async def test_multimodal_prefill_cleanup_awaits_consumers_before_engine_shutdown(): handler = MultimodalPrefillWorkerHandler.__new__(MultimodalPrefillWorkerHandler) handler.publisher = None + handler._abort_tasks = set() events = [] async def consumer(): @@ -659,7 +661,7 @@ async def test_vp9_without_software_decoder_is_actionable(monkeypatch): async def fake_validate(url, policy): return url - async def fake_fetch(url, timeout, policy=None): + async def fake_fetch(url, timeout, policy=None, max_bytes=None): return b"vp9-bytes" monkeypatch.setattr(ewh, "validate_media_url", fake_validate) @@ -687,7 +689,7 @@ async def test_vp9_with_software_decoder_passes_bytes_through(monkeypatch): async def fake_validate(url, policy): return url - async def fake_fetch(url, timeout, policy=None): + async def fake_fetch(url, timeout, policy=None, max_bytes=None): return b"vp9-bytes" monkeypatch.setattr(ewh, "validate_media_url", fake_validate) @@ -702,6 +704,24 @@ async def fake_fetch(url, timeout, policy=None): assert result == b"vp9-bytes" +@pytest.mark.asyncio +async def test_remote_video_fetch_honors_configured_media_limit(monkeypatch): + import dynamo.sglang.request_handlers.multimodal.encode_worker_handler as ewh + from dynamo.common.http.media_reference import DYN_MM_MAX_FILE_SIZE_MB + + monkeypatch.setenv(DYN_MM_MAX_FILE_SIZE_MB, "1") + fetch = AsyncMock(return_value=b"vp9-bytes") + monkeypatch.setattr(ewh, "fetch_bytes", fetch) + monkeypatch.setattr(ewh, "probe_video_codec", lambda b: "vp9") + monkeypatch.setattr(ewh, "should_use_nvdec", lambda c: False) + monkeypatch.setattr(ewh.importlib, "import_module", _selective_import({"decord"})) + + handler = _bare_handler() + await handler._maybe_nvdec_decoder("https://example.com/clip.webm") + + assert fetch.await_args.kwargs["max_bytes"] == 1024 * 1024 + + @pytest.mark.asyncio async def test_multimodal_decode_rejects_parallel_sampling_before_prefill_handoff(): """Reject multimodal n greater than one before contacting prefill.""" diff --git a/components/src/dynamo/sglang/tests/test_sglang_nixl_telemetry_ports.py b/components/src/dynamo/sglang/tests/test_sglang_nixl_telemetry_ports.py index a3e423055d44..3f2645558e4b 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_nixl_telemetry_ports.py +++ b/components/src/dynamo/sglang/tests/test_sglang_nixl_telemetry_ports.py @@ -292,6 +292,22 @@ def test_install_is_a_no_op_when_telemetry_is_disabled(self, telemetry_env): telemetry_env.setenv("NIXL_TELEMETRY_ENABLE", "n") install_per_rank_nixl_prometheus_ports() + @pytest.mark.parametrize( + ("env_name", "value"), + [ + ("NIXL_TELEMETRY_ENABLE", "maybe"), + (NIXL_TELEMETRY_PROMETHEUS_PORT_ENV, "not-a-port"), + ], + ) + def test_invalid_config_is_rejected_before_importing_sglang( + self, telemetry_env, env_name, value + ): + telemetry_env.setenv(env_name, value) + telemetry_env.setitem(sys.modules, "sglang", _LazyProxyModule()) + + with pytest.raises(ValueError, match=env_name): + install_per_rank_nixl_prometheus_ports() + def test_install_points_sglang_at_the_wrapper(self, telemetry_env): """The override has to land on the class, not on the ``sglang`` proxy.""" engine = type("Engine", (), {"run_scheduler_process_func": None}) diff --git a/components/src/dynamo/sglang/tests/test_sglang_prefill_handler.py b/components/src/dynamo/sglang/tests/test_sglang_prefill_handler.py new file mode 100644 index 000000000000..8da282f55f2e --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_sglang_prefill_handler.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from types import SimpleNamespace + +import pytest + +from dynamo.sglang.request_handlers import cancellation +from dynamo.sglang.request_handlers.llm.prefill_handler import PrefillWorkerHandler + +pytestmark = [ + pytest.mark.unit, + pytest.mark.sglang, + pytest.mark.fault_tolerance, + pytest.mark.gpu_0, + pytest.mark.profiled_vram_gib(0), + pytest.mark.pre_merge, +] + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("ordered_cancellation", [False, True]) +@pytest.mark.parametrize( + "phase", ["before_registration", "before_dispatch", "after_dispatch"] +) +async def test_prefill_cancellation_waits_for_dispatch_and_drains( + monkeypatch: pytest.MonkeyPatch, + phase: str, + native: bool, + ordered_cancellation: bool, +): + request_id = "native-request-id" if native else "request-id" + engine_request_id = "internal-request-id" + registry = {} + abort_calls = [] + started = asyncio.Event() + registered = asyncio.Event() + dispatched = asyncio.Event() + aborted = asyncio.Event() + drained = asyncio.Event() + cancelled = asyncio.Event() + polling = asyncio.Event() + allow_registration = asyncio.Event() + allow_dispatch = asyncio.Event() + if phase != "before_registration": + allow_registration.set() + if phase == "after_dispatch": + allow_dispatch.set() + + async def results(rid): + started.set() + try: + await allow_registration.wait() + state = SimpleNamespace( + time_stats=SimpleNamespace(api_server_dispatch_finish_time=0.0) + ) + registry[rid] = state + registered.set() + await allow_dispatch.wait() + state.time_stats.api_server_dispatch_finish_time = 1.0 + dispatched.set() + response = {"meta_info": {"id": rid}} + yield response + await aborted.wait() + finally: + registry.pop(rid, None) + drained.set() + + def abort_request(*, rid, abort_all): + abort_calls.append((rid, abort_all)) + if rid in registry and dispatched.is_set(): + aborted.set() + + def generate_request(request, _context): + return results(request.rid) + + class _Engine: + tokenizer_manager = SimpleNamespace( + rid_to_state=registry, + abort_request=abort_request, + generate_request=generate_request, + ) + + async def async_generate(self, **kwargs): + return results(kwargs["rid"]) + + handler = PrefillWorkerHandler.__new__(PrefillWorkerHandler) + handler.engine = _Engine() + handler.config = SimpleNamespace( + server_args=SimpleNamespace(pp_size=1, enable_dp_attention=False) + ) + handler.shutdown_event = None + handler.bootstrap_host = "127.0.0.1" + handler.bootstrap_port = 1234 + handler.enable_trace = False + handler._abort_tasks = set() + handler._consume_tasks = set() + handler._supports_ordered_cancellation = ordered_cancellation + handler._generate_bootstrap_room = lambda: 17 + handler._get_input_param = lambda request: {"input_ids": request["token_ids"]} + handler._resolve_lora = lambda request: None + handler._priority_kwargs = lambda priority: {} + monkeypatch.setattr( + "dynamo.sglang.request_handlers.llm.prefill_handler.require_reasoning_kwargs", + lambda engine, request: {}, + ) + monkeypatch.setattr( + cancellation, "resolved_server_args", lambda server_args: server_args + ) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.llm.prefill_handler.new_sglang_request_id", + lambda: engine_request_id, + ) + real_sleep = asyncio.sleep + + async def observed_sleep(delay): + polling.set() + await real_sleep(0) + + monkeypatch.setattr( + cancellation, + "asyncio", + SimpleNamespace(**{**vars(asyncio), "sleep": observed_sleep}), + ) + context = SimpleNamespace( + id=lambda: request_id, + trace_id=None, + trace_headers=lambda: {}, + async_killed_or_stopped=lambda: asyncio.create_task(cancelled.wait()), + ) + inner_request = {"token_ids": [1, 2, 3], "routing": {}} + if native: + inner_request["extra_args"] = {"sglang_tito": {"rid": request_id}} + stream = handler.generate( + { + "request": inner_request, + "sampling_params": {}, + }, + context, + ) + await anext(stream) + consumer = asyncio.create_task(_drain(stream)) + await asyncio.wait_for(started.wait(), timeout=1) + if phase != "before_registration": + await asyncio.wait_for(registered.wait(), timeout=1) + if phase == "after_dispatch": + await asyncio.wait_for(dispatched.wait(), timeout=1) + + cancelled.set() + if phase != "after_dispatch": + if ordered_cancellation: + await asyncio.wait_for(polling.wait(), timeout=1) + else: + await asyncio.sleep(0) + assert not abort_calls + allow_registration.set() + allow_dispatch.set() + + assert await asyncio.wait_for(consumer, timeout=1) == [] + assert abort_calls == [(engine_request_id, False)] + assert dispatched.is_set() + assert aborted.is_set() + assert drained.is_set() + assert not registry + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_native_batched_prefill_disables_ordered_cancellation(monkeypatch): + captured_submitted_id = None + captured_native_request = None + + async def native_results(): + if False: + yield None + + handler = PrefillWorkerHandler.__new__(PrefillWorkerHandler) + handler.engine = SimpleNamespace() + handler.bootstrap_host = "127.0.0.1" + handler.bootstrap_port = 1234 + handler.enable_trace = False + handler._consume_tasks = set() + handler._supports_ordered_cancellation = True + handler._generate_bootstrap_room = lambda: 17 + handler._get_input_param = lambda request: {"input_ids": request["token_ids"]} + handler._resolve_lora = lambda request: None + handler._priority_kwargs = lambda priority: {} + + async def capture_consume(_results, submitted_request_id, _context): + nonlocal captured_submitted_id + captured_submitted_id = submitted_request_id + + def capture_native_stream(_engine, request): + nonlocal captured_native_request + captured_native_request = request + return native_results() + + handler._consume_results = capture_consume + monkeypatch.setattr( + "dynamo.sglang.request_handlers.llm.prefill_handler.new_sglang_request_id", + lambda: "internal-request-id", + ) + monkeypatch.setattr( + "dynamo.sglang.request_handlers.llm.prefill_handler.native_generate_stream", + capture_native_stream, + ) + context = SimpleNamespace( + id=lambda: "request-id", trace_id=None, trace_headers=lambda: {} + ) + stream = handler.generate( + { + "request": { + "token_ids": [[1], [2]], + "routing": {}, + "extra_args": {"sglang_tito": {}}, + }, + "sampling_params": {}, + }, + context, + ) + + await anext(stream) + with pytest.raises(StopAsyncIteration): + await anext(stream) + assert captured_submitted_id is None + assert captured_native_request.input_ids == [[1], [2]] + + +async def _drain(stream): + return [item async for item in stream] diff --git a/components/src/dynamo/sglang/tests/test_sglang_publisher.py b/components/src/dynamo/sglang/tests/test_sglang_publisher.py index 9024ef5a5c30..40af00e49a9e 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_publisher.py +++ b/components/src/dynamo/sglang/tests/test_sglang_publisher.py @@ -12,7 +12,7 @@ from dynamo.sglang._disagg import SGLANG_WORKER_GROUP_ID_KEY, get_sglang_worker_group_id from dynamo.sglang.publisher import ( DynamoSglangPublisher, - _resolve_multinode_leader_worker_id, + _resolve_multinode_leader_worker_ids, get_local_dp_rank_range, handle_non_leader_node, set_forward_pass_metrics_worker_id, @@ -156,9 +156,9 @@ async def client(self): server_args = SimpleNamespace(nnodes=2, node_rank=1) - worker_id = await _resolve_multinode_leader_worker_id(FakeEndpoint(), server_args) + worker_ids = await _resolve_multinode_leader_worker_ids(FakeEndpoint(), server_args) - assert worker_id == 1234 + assert worker_ids == [1234] @pytest.mark.asyncio @@ -185,9 +185,9 @@ async def client(self): dist_timeout=5, ) - worker_id = await _resolve_multinode_leader_worker_id(FakeEndpoint(), server_args) + worker_ids = await _resolve_multinode_leader_worker_ids(FakeEndpoint(), server_args) - assert worker_id == 1234 + assert worker_ids == [1234] assert calls == [ ( SGLANG_WORKER_GROUP_ID_KEY, @@ -220,9 +220,9 @@ async def client(self): node_rank=1, ) - worker_id = await _resolve_multinode_leader_worker_id(FakeEndpoint(), server_args) + worker_ids = await _resolve_multinode_leader_worker_ids(FakeEndpoint(), server_args) - assert worker_id == 1234 + assert worker_ids == [1234] assert calls == [ ( SGLANG_WORKER_GROUP_ID_KEY, @@ -232,6 +232,54 @@ async def client(self): ] +@pytest.mark.asyncio +async def test_resolve_multinode_leader_worker_ids_waits_for_every_gateway(monkeypatch): + calls = [] + + class FakeClient: + async def wait_for_instances_by_runtime_data( + self, key, value, min_count, timeout_s=None + ): + calls.append((key, value, min_count, timeout_s)) + return [7, 5, 6] + + class FakeEndpoint: + async def client(self): + return FakeClient() + + monkeypatch.setattr( + publisher_mod, + "get_sglang_worker_group_id", + lambda server_args: "dist_init:tcp://10.0.0.1:2345", + ) + server_args = SimpleNamespace(nnodes=2, node_rank=1, dist_timeout=5) + + worker_ids = await _resolve_multinode_leader_worker_ids( + FakeEndpoint(), server_args, expected=3 + ) + + assert worker_ids == [7, 5, 6] + assert calls == [ + (SGLANG_WORKER_GROUP_ID_KEY, "dist_init:tcp://10.0.0.1:2345", 3, 5.0) + ] + + +@pytest.mark.asyncio +async def test_resolve_multinode_leader_worker_ids_needs_group_id_for_gateways(): + class FakeClient: + async def wait_for_instances(self): + return [1234] + + class FakeEndpoint: + async def client(self): + return FakeClient() + + with pytest.raises(RuntimeError, match="dist_init_addr"): + await _resolve_multinode_leader_worker_ids( + FakeEndpoint(), SimpleNamespace(nnodes=2, node_rank=1), expected=2 + ) + + @pytest.mark.asyncio async def test_resolve_multinode_leader_worker_id_ignores_ambiguous_instances(): class FakeClient: @@ -244,9 +292,9 @@ async def client(self): server_args = SimpleNamespace(nnodes=2, node_rank=1) - worker_id = await _resolve_multinode_leader_worker_id(FakeEndpoint(), server_args) + worker_ids = await _resolve_multinode_leader_worker_ids(FakeEndpoint(), server_args) - assert worker_id is None + assert worker_ids == [] @pytest.mark.asyncio @@ -310,11 +358,11 @@ def cleanup(self): @pytest.mark.asyncio async def test_handle_non_leader_node_skips_tp_only_kv_event_setup(monkeypatch): - resolve_leader = AsyncMock(return_value=1234) + resolve_leader = AsyncMock(return_value=[1234]) kv_event_publisher = Mock() monkeypatch.setattr( publisher_mod, - "_resolve_multinode_leader_worker_id", + "_resolve_multinode_leader_worker_ids", resolve_leader, ) monkeypatch.setattr(publisher_mod, "KvEventPublisher", kv_event_publisher) @@ -365,9 +413,9 @@ async def test_handle_non_leader_node_skips_kv_publish_without_resolved_worker( init_called = asyncio.Event() cleanup_called = asyncio.Event() - async def missing_resolution(generate_endpoint, server_args): + async def missing_resolution(generate_endpoint, server_args, expected=1): resolution_done.set() - return None + return [] class FakePublisher: generate_endpoint = object() @@ -383,7 +431,7 @@ def cleanup(self): monkeypatch.setattr( publisher_mod, - "_resolve_multinode_leader_worker_id", + "_resolve_multinode_leader_worker_ids", missing_resolution, ) metrics_task = asyncio.create_task(asyncio.Event().wait()) @@ -412,7 +460,7 @@ def cleanup(self): async def test_handle_non_leader_node_cleans_up_when_resolution_fails(monkeypatch): cleanup_called = asyncio.Event() - async def fail_resolution(generate_endpoint, server_args): + async def fail_resolution(generate_endpoint, server_args, expected=1): raise RuntimeError("resolution failed") class FakePublisher: @@ -425,7 +473,7 @@ def cleanup(self): monkeypatch.setattr( publisher_mod, - "_resolve_multinode_leader_worker_id", + "_resolve_multinode_leader_worker_ids", fail_resolution, ) metrics_task = asyncio.create_task(asyncio.Event().wait()) @@ -649,7 +697,10 @@ def shutdown(self): @pytest.mark.asyncio -async def test_setup_sgl_metrics_skips_chat_pipeline_for_embedding_worker(monkeypatch): +@pytest.mark.parametrize("role", ["embedding_worker", "rerank_worker"]) +async def test_setup_sgl_metrics_skips_chat_pipeline_for_pooling_worker( + monkeypatch, role +): """``setup_sgl_metrics`` short-circuits for embedding workers. Chat-shaped collectors (``sglang:*`` multiproc metrics, the Dynamo @@ -692,7 +743,7 @@ def _wrapped(*_a, **_kw): ) ) config = SimpleNamespace( - dynamo_args=SimpleNamespace(embedding_worker=True), + dynamo_args=SimpleNamespace(**{role: True}), server_args=engine.server_args, ) generate_endpoint = SimpleNamespace() @@ -808,3 +859,81 @@ async def _async_noop(): await task except asyncio.CancelledError: pass + + +def _publisher_with_sockets(sock, fanout, owner, stop_after_publish): + """A DynamoSglangPublisher with only what run() touches, no engine or runtime.""" + pub = publisher_mod.DynamoSglangPublisher.__new__( + publisher_mod.DynamoSglangPublisher + ) + pub._sock, pub._fanout = sock, fanout + pub._publishes_engine_gauges = owner + pub._running = True + pub.dp_rank = 0 + pub.server_args = SimpleNamespace(page_size=1) + published = [] + + def publish(dp_rank, kv_used_blocks): + published.append((dp_rank, kv_used_blocks)) + if stop_after_publish: + pub._running = False + + pub.metrics_publisher = SimpleNamespace(publish=publish) + pub.component_gauges = Mock() + return pub, published + + +@pytest.mark.asyncio +async def test_metrics_fanout_reaches_sibling_gateways(tmp_path, monkeypatch): + import zmq + import zmq.asyncio + + monkeypatch.setattr( + publisher_mod, "kv_metrics_block_values", lambda metrics, page_size: (7, 100) + ) + metrics_ep = f"ipc://{tmp_path}/metrics" + fanout_ep = f"ipc://{tmp_path}/fanout" + ctx = zmq.asyncio.Context() + sockets = [] + try: + owner_sock, fanout = publisher_mod._open_metrics_sockets( + ctx, metrics_ep, fanout_ep, owner=True + ) + sockets += [owner_sock, fanout] + assert fanout is not None + sibling_sock, sibling_fanout = publisher_mod._open_metrics_sockets( + ctx, metrics_ep, fanout_ep, owner=False + ) + sockets.append(sibling_sock) + assert sibling_fanout is None + await asyncio.sleep(0.3) # PUB/SUB slow joiner + + owner, owner_published = _publisher_with_sockets( + owner_sock, fanout, owner=True, stop_after_publish=True + ) + sibling, sibling_published = _publisher_with_sockets( + sibling_sock, None, owner=False, stop_after_publish=True + ) + + scheduler = ctx.socket(zmq.PUSH) + sockets.append(scheduler) + scheduler.connect(metrics_ep) + kv_metrics = SimpleNamespace(data_parallel_rank=3, gpu_cache_usage_perc=0.5) + await asyncio.wait_for(scheduler.send_pyobj(kv_metrics), 5) + + # The owner's run() consumes the scheduler message and relays it; the + # sibling's run() must see it through the fan-out with no other source. + await asyncio.wait_for(asyncio.gather(owner.run(), sibling.run()), 10) + + assert owner_published == [(3, 7)] + assert sibling_published == [(3, 7)] + owner.component_gauges.set_total_blocks.assert_called_once_with("3", 100) + sibling.component_gauges.set_total_blocks.assert_not_called() + + with pytest.raises(ValueError, match="fan-out"): + publisher_mod._open_metrics_sockets(ctx, metrics_ep, None, owner=False) + finally: + for s in sockets: + if s is not None: + s.close(linger=0) + ctx.term() diff --git a/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py new file mode 100644 index 000000000000..2d9c576ab28f --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for SGLang cross-encoder reranking.""" + +import argparse +import asyncio +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, Mock + +import pytest + +pytest.importorskip( + "sglang.srt.managers.io_struct", reason="sglang not installed in this container" +) + +pytest.importorskip( + "sglang.srt.speculative.spec_info", reason="full SGLang runtime is required" +) + +from dynamo.llm import ModelType # noqa: E402 +from dynamo.sglang import init_embedding as pooling_init # noqa: E402 +from dynamo.sglang.backend_args import ( # noqa: E402 + DynamoSGLangArgGroup, + DynamoSGLangConfig, +) +from dynamo.sglang.health_check import SglangRerankHealthCheckPayload # noqa: E402 +from dynamo.sglang.init_rerank import init_rerank # noqa: E402 +from dynamo.sglang.request_handlers.embedding import ( # noqa: E402 + EmbeddingWorkerHandler, +) +from dynamo.sglang.request_handlers.rerank import RerankWorkerHandler # noqa: E402 + +pytestmark = [ + pytest.mark.unit, + pytest.mark.sglang, + pytest.mark.core, + pytest.mark.gpu_0, + pytest.mark.profiled_vram_gib(0), + pytest.mark.pre_merge, +] + + +class _TokenizerManager: + def __init__( + self, + scores: list[float], + *, + chat_template: str = "", + model_path: str = "BAAI/bge-reranker-v2-m3", + ) -> None: + self.scores = scores + self.requests: list[tuple[Any, Any]] = [] + self.tokenizer = SimpleNamespace(chat_template=chat_template) + self.model_config = SimpleNamespace(model_path=model_path) + + async def generate_request(self, request: Any, raw_request: Any): + request.normalize_batch_and_arguments() + self.requests.append((request, raw_request)) + yield [ + { + "embedding": [score], + "meta_info": {"prompt_tokens": index + 1}, + } + for index, score in enumerate(self.scores) + ] + + +class _Engine: + def __init__(self, tokenizer_manager: _TokenizerManager) -> None: + self.tokenizer_manager = tokenizer_manager + + +class _Context: + trace_id = "rerank-trace" + + def trace_headers(self) -> dict[str, str]: + return {"traceparent": "00-test"} + + +def _handler( + scores: list[float], + *, + chat_template: str = "", + model_path: str = "BAAI/bge-reranker-v2-m3", +) -> tuple[RerankWorkerHandler, _TokenizerManager]: + manager = _TokenizerManager( + scores, chat_template=chat_template, model_path=model_path + ) + handler = RerankWorkerHandler.__new__(RerankWorkerHandler) + handler.engine = _Engine(manager) + handler.enable_trace = True + return handler, manager + + +@pytest.mark.asyncio +async def test_builds_pairs_sorts_scores_and_applies_top_n(): + handler, manager = _handler([0.2, 0.9, 0.5]) + outputs = [ + output + async for output in handler.generate( + { + "model": "reranker", + "query": "query", + "documents": ["zero", "one", "two"], + "top_n": 2, + "return_documents": True, + }, + _Context(), + ) + ] + + assert [item["index"] for item in outputs[0]] == [1, 2] + assert [item["document"] for item in outputs[0]] == ["one", "two"] + [(request, raw_request)] = manager.requests + assert raw_request is None + assert request.text == [["query", "zero"], ["query", "one"], ["query", "two"]] + assert request.is_cross_encoder_request is True + assert request.rid == ["rerank-trace-0", "rerank-trace-1", "rerank-trace-2"] + assert request.external_trace_header == {"traceparent": "00-test"} + + +@pytest.mark.asyncio +async def test_omits_documents_when_not_requested(): + handler, _ = _handler([0.1, 0.8]) + [output] = [ + output + async for output in handler.generate( + { + "model": "reranker", + "query": "query", + "documents": ["a", "b"], + "return_documents": False, + }, + _Context(), + ) + ] + assert [item["index"] for item in output] == [1, 0] + assert all("document" not in item for item in output) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "bad_request", + [ + {"model": "m", "query": " ", "documents": ["doc"]}, + {"model": "m", "query": "query", "documents": []}, + {"model": "m", "query": "query", "documents": [""]}, + {"model": "m", "query": "query", "documents": ["doc"], "top_n": 0}, + ], +) +async def test_rejects_invalid_requests_before_inference(bad_request): + handler, manager = _handler([0.5]) + with pytest.raises(ValueError): + _ = [output async for output in handler.generate(bad_request, _Context())] + assert manager.requests == [] + + +def test_rejects_decoder_only_qwen3_reranker_at_startup(): + manager = _TokenizerManager([0.5], model_path="Qwen/Qwen3-Reranker-0.6B") + with pytest.raises(ValueError, match="cross-encoder rerankers only"): + RerankWorkerHandler(_Engine(manager), SimpleNamespace()) + assert manager.requests == [] + + +@pytest.mark.asyncio +async def test_embedding_endpoint_does_not_dispatch_rerank(): + handler = EmbeddingWorkerHandler.__new__(EmbeddingWorkerHandler) + with pytest.raises(ValueError): + _ = [ + item + async for item in handler.generate( + {"model": "m", "query": "query", "documents": ["doc"]}, _Context() + ) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rerank", [False, True]) +@pytest.mark.parametrize("fail_startup", [False, True]) +async def test_dedicated_pooling_registration_health_and_cleanup( + monkeypatch, rerank, fail_startup +): + module = pooling_init + monkeypatch.delenv("DYN_HEALTH_CHECK_PAYLOAD", raising=False) + + engine = Mock() + engine.server_args = SimpleNamespace( + served_model_name="resolved-pooler", model_path="pooler" + ) + engine.tokenizer_manager.tokenizer.bos_token_id = 1 + endpoint = Mock() + endpoint.serve_endpoint = AsyncMock() + runtime = Mock() + runtime.endpoint.return_value = endpoint + config = SimpleNamespace( + server_args=SimpleNamespace(served_model_name="pooler", model_path="pooler"), + use_resolved_server_args=Mock(return_value=engine.server_args), + dynamo_args=SimpleNamespace( + namespace="test", + component="rerank" if rerank else "backend", + endpoint="generate", + use_sglang_tokenizer=False, + ), + ) + handler = Mock() + handlers = [Mock(return_value=handler), Mock(return_value=handler)] + monkeypatch.setattr(module, "EmbeddingWorkerHandler", handlers[0]) + monkeypatch.setattr(module, "RerankWorkerHandler", handlers[1]) + monkeypatch.setattr(module.sgl, "Engine", Mock(return_value=engine)) + metrics_task = asyncio.create_task(asyncio.Event().wait()) + monkeypatch.setattr( + module, "setup_sgl_metrics", AsyncMock(return_value=(None, metrics_task, [])) + ) + register = AsyncMock() + monkeypatch.setattr(module, "register_model_with_readiness_gate", register) + monkeypatch.setattr(module, "register_model_taint_route", Mock()) + monkeypatch.setattr(module, "register_engine_metrics_callback", Mock()) + monkeypatch.setattr(module, "init_embedding_metrics", Mock()) + deferred = AsyncMock() + shutdown_endpoints = [] + + init = init_rerank if rerank else module.init_embedding + if fail_startup: + handlers[int(rerank)].side_effect = ValueError("unsupported model") + with pytest.raises(ValueError, match="unsupported model"): + await init(runtime, config, asyncio.Event(), shutdown_endpoints, deferred) + register.assert_not_called() + endpoint.serve_endpoint.assert_not_called() + engine.shutdown.assert_called_once() + assert metrics_task.cancelled() + deferred.assert_awaited_once() + return + await init(runtime, config, asyncio.Event(), shutdown_endpoints, deferred) + + config.use_resolved_server_args.assert_called_once_with(engine.server_args) + assert register.call_args.args[2] is engine.server_args + assert register.call_args.kwargs["output_type"] == ( + ModelType.Rerank if rerank else ModelType.Embedding + ) + assert shutdown_endpoints == [endpoint] + handlers[int(rerank)].assert_called_once() + handlers[int(not rerank)].assert_not_called() + assert endpoint.serve_endpoint.call_args.args[0] == handler.generate + payload = endpoint.serve_endpoint.call_args.kwargs["health_check_payload"] + assert payload["model"] == "resolved-pooler" + if rerank: + assert payload["query"] and payload["documents"] + assert "input" not in payload + else: + assert payload["input"] == [1] + assert "query" not in payload + handler.cleanup.assert_called_once() + assert metrics_task.cancelled() + deferred.assert_awaited_once() + + +def test_rerank_flag_defaults_and_rejects_embedding_combination(monkeypatch): + monkeypatch.delenv("DYN_SGL_RERANK_WORKER", raising=False) + parser = argparse.ArgumentParser() + DynamoSGLangArgGroup().add_arguments(parser) + assert not parser.parse_args([]).rerank_worker + args = parser.parse_args(["--rerank-worker"]) + config = DynamoSGLangConfig.from_cli_args(args) + config.validate() + assert config.rerank_worker and not config.embedding_worker + args = parser.parse_args(["--rerank-worker", "--embedding-worker"]) + with pytest.raises(ValueError, match="cannot be combined"): + DynamoSGLangConfig.from_cli_args(args).validate() + + +@pytest.mark.asyncio +async def test_rerank_health_check_runs_cross_encoder(monkeypatch): + monkeypatch.delenv("DYN_HEALTH_CHECK_PAYLOAD", raising=False) + handler, manager = _handler([0.5]) + payload = SglangRerankHealthCheckPayload("pooler").to_dict() + outputs = [item async for item in handler.generate(payload, _Context())] + assert outputs[0][0]["score"] == 0.5 + assert manager.requests[0][0].is_cross_encoder_request diff --git a/components/src/dynamo/sglang/tests/test_sglang_unit.py b/components/src/dynamo/sglang/tests/test_sglang_unit.py index f45594a95ae7..bb81b198cb72 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_unit.py +++ b/components/src/dynamo/sglang/tests/test_sglang_unit.py @@ -21,6 +21,7 @@ from dynamo.common.constants import DisaggregationMode, EmbeddingTransferMode from dynamo.common.snapshot.constants import SNAPSHOT_CONTROL_DIR_ENV from dynamo.sglang._compat import ( + cache_salt_kwargs, ensure_sglang_tensor_image_size, filter_supported_async_generate_kwargs, get_sglang_model_config, @@ -165,6 +166,32 @@ def test_resolved_server_args_uses_declarative_view(monkeypatch): assert raw_server_args.page_size is None +def test_compat_detects_ordered_cancellation_support(monkeypatch): + request_stats = SimpleNamespace( + __dataclass_fields__={"api_server_dispatch_finish_time": object()} + ) + monkeypatch.setattr(sglang_compat, "APIServerReqTimeStats", request_stats) + engine = SimpleNamespace(tokenizer_manager=SimpleNamespace(rid_to_state={})) + + assert sglang_compat.supports_disagg_prefill_cancel_anytime(engine) + + +@pytest.mark.parametrize( + "engine", + [ + SimpleNamespace(), + SimpleNamespace(tokenizer_manager=SimpleNamespace(rid_to_state=[])), + ], +) +def test_compat_rejects_runtime_without_request_registry(monkeypatch, engine): + request_stats = SimpleNamespace( + __dataclass_fields__={"api_server_dispatch_finish_time": object()} + ) + monkeypatch.setattr(sglang_compat, "APIServerReqTimeStats", request_stats) + + assert not sglang_compat.supports_disagg_prefill_cancel_anytime(engine) + + def test_compat_uses_current_sglang_model_config_accessor(monkeypatch): expected = SimpleNamespace(is_multimodal=True) server_args = SimpleNamespace() @@ -315,7 +342,14 @@ def test_engine_generate_capability_registration_gate( def test_builtin_engine_routes_include_model_taint_update(monkeypatch): handler = object.__new__(DecodeWorkerHandler) - handler.engine = SimpleNamespace() + handler.engine = SimpleNamespace( + tokenizer_manager=SimpleNamespace( + pause_generation=lambda: None, + continue_generation=lambda: None, + release_memory_occupation=lambda: None, + resume_memory_occupation=lambda: None, + ) + ) handler.generate_endpoint = object() handler.config = SimpleNamespace(dynamo_args=SimpleNamespace(engine_routes=[])) @@ -340,6 +374,10 @@ def register_engine_route(self, path, route_handler): assert {path for path, _ in registered_routes} >= { "control/start_profile", "control/stop_profile", + "pause_generation", + "continue_generation", + "release_memory_occupation", + "resume_memory_occupation", } @@ -492,6 +530,179 @@ def resolve(parsed_args): assert config.server_args.incremental_streaming_output is True +def _dcp_server_args_stub(**overrides): + """Minimal resolved ServerArgs surface that parse_args reads.""" + stub = SimpleNamespace( + disaggregation_mode="null", + dllm_algorithm=None, + kv_events_config=None, + dcp_size=1, + attention_backend=None, + prefill_attention_backend=None, + decode_attention_backend=None, + use_mla_backend=lambda: False, + get_model_config=lambda: SimpleNamespace(is_multimodal=False), + ) + for name, value in overrides.items(): + setattr(stub, name, value) + return stub + + +@pytest.mark.asyncio +async def test_parse_args_rejects_dcp_on_backend_without_dcp_support( + monkeypatch, mock_sglang_cli, tmp_path +): + monkeypatch.setattr( + "dynamo.sglang.args.ServerArgs.from_cli_args", + lambda _: _dcp_server_args_stub(dcp_size=2, attention_backend="fa3"), + ) + mock_sglang_cli(model=str(tmp_path)) + + with pytest.raises(ValueError) as excinfo: + await parse_args(sys.argv[1:]) + + message = str(excinfo.value) + assert "'fa3'" in message + assert "--dcp-size 1" in message + assert "automatically" in message + + +@pytest.mark.asyncio +async def test_parse_args_rejects_dcp_on_decode_only_unsupported_backend( + monkeypatch, mock_sglang_cli, tmp_path +): + monkeypatch.setattr( + "dynamo.sglang.args.ServerArgs.from_cli_args", + lambda _: _dcp_server_args_stub( + dcp_size=2, + prefill_attention_backend="triton", + decode_attention_backend="fa3", + ), + ) + mock_sglang_cli(model=str(tmp_path)) + + with pytest.raises(ValueError) as excinfo: + await parse_args(sys.argv[1:]) + + message = str(excinfo.value) + assert "decode attention backend 'fa3'" in message + assert "prefill attention backend" not in message + + +@pytest.mark.asyncio +async def test_parse_args_does_not_call_a_phase_flag_automatic( + monkeypatch, mock_sglang_cli, tmp_path +): + monkeypatch.setattr( + "dynamo.sglang.args.ServerArgs.from_cli_args", + lambda _: _dcp_server_args_stub(dcp_size=2, decode_attention_backend="fa3"), + ) + mock_sglang_cli(model=str(tmp_path), decode_attention_backend="fa3") + + with pytest.raises(ValueError) as excinfo: + await parse_args(sys.argv[1:]) + + message = str(excinfo.value) + assert "decode attention backend 'fa3'" in message + assert "automatically" not in message + + +@pytest.mark.asyncio +async def test_resolved_server_args_rejects_an_automatic_unsupported_backend( + monkeypatch, mock_sglang_cli, tmp_path +): + """The backend SGLang chooses for itself is only readable from the engine. + + On the pinned SGLang release ``ServerArgs`` still holds what the caller + asked for, so a launch that passes no attention backend reaches parse_args + with all three fields unset. The engine resolves fa3 afterwards, and the + check on its configuration is the one that catches it. + """ + cli_server_args = _dcp_server_args_stub(dcp_size=2) + monkeypatch.setattr( + "dynamo.sglang.args.ServerArgs.from_cli_args", lambda _: cli_server_args + ) + mock_sglang_cli(model=str(tmp_path)) + + config = await parse_args(sys.argv[1:]) + assert config.server_args is cli_server_args + + engine_server_args = _dcp_server_args_stub(dcp_size=2, attention_backend="fa3") + with pytest.raises(ValueError) as excinfo: + config.use_resolved_server_args(engine_server_args) + + message = str(excinfo.value) + assert "'fa3'" in message + assert "automatically" in message + assert config.server_args is cli_server_args + + +@pytest.mark.asyncio +async def test_prepare_snapshot_engine_rejects_dcp_before_warmup(monkeypatch): + """Snapshot mode warms the engine with a real request of its own. + + That warmup runs before any init function reaches + ``use_resolved_server_args``, so the snapshot engine has to be checked + where it is built. + """ + from dynamo.sglang import snapshot as sglang_snapshot + + monkeypatch.setattr( + sglang_snapshot, + "SnapshotConfig", + SimpleNamespace(from_env=lambda: SimpleNamespace()), + ) + monkeypatch.setattr(sglang_snapshot, "configure_snapshot_capture_env", lambda: None) + monkeypatch.setattr(sglang_snapshot, "override_server_args", lambda *a, **kw: None) + + engine_server_args = _dcp_server_args_stub(dcp_size=2, attention_backend="fa3") + monkeypatch.setattr( + sglang_snapshot.sgl, + "Engine", + lambda server_args: SimpleNamespace(server_args=engine_server_args), + ) + + warmed = [] + + async def _warmup(engine, server_args): + warmed.append(engine) + + monkeypatch.setattr(sglang_snapshot, "warmup_engine", _warmup) + + config = sglang_args.Config(_dcp_server_args_stub(dcp_size=2), SimpleNamespace()) + + with pytest.raises(ValueError) as excinfo: + await sglang_snapshot.prepare_snapshot_engine(config) + + assert "'fa3'" in str(excinfo.value) + assert warmed == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides", + [ + {"dcp_size": 2, "attention_backend": "triton"}, + {"dcp_size": 1, "attention_backend": "fa3"}, + {"dcp_size": 2, "attention_backend": "fa3", "use_mla_backend": lambda: True}, + ], + ids=["dcp-capable-backend", "dcp-disabled", "mla-model"], +) +async def test_parse_args_accepts_supported_dcp_configurations( + monkeypatch, mock_sglang_cli, tmp_path, overrides +): + server_args = _dcp_server_args_stub(**overrides) + monkeypatch.setattr( + "dynamo.sglang.args.ServerArgs.from_cli_args", lambda _: server_args + ) + mock_sglang_cli(model=str(tmp_path)) + + config = await parse_args(sys.argv[1:]) + + assert config.server_args is server_args + assert config.use_resolved_server_args(server_args) is not None + + @pytest.mark.asyncio async def test_parse_args_applies_dynamo_defaults_before_resolution( monkeypatch, mock_sglang_cli @@ -603,6 +814,27 @@ async def async_generate(self, input_ids=None, sampling_params=None): } +@pytest.mark.parametrize("salt", [None, "", "tenant-a"]) +def test_cache_salt_kwargs_for_older_engines(salt): + class OldEngine: + async def async_generate(self, input_ids=None): + return None + + if salt: + with pytest.raises(ValueError, match="cache_salt is not supported"): + cache_salt_kwargs(OldEngine(), salt) + else: + assert cache_salt_kwargs(OldEngine(), salt) == {} + + +def test_cache_salt_kwargs_preserves_supported_salt(): + class SaltedEngine: + async def async_generate(self, cache_salt=None): + return None + + assert cache_salt_kwargs(SaltedEngine(), "tenant-a") == {"cache_salt": "tenant-a"} + + def test_compat_keeps_async_generate_kwargs_for_newer_engines(): class NewEngine: async def async_generate(self, return_routed_experts=False): @@ -1174,6 +1406,7 @@ async def test_invalid_fpm_trace_is_disabled_by_arg_parser( ("overrides", "role"), [ ({"embedding_worker": True}, "embedding"), + ({"rerank_worker": True}, "rerank"), ({"multimodal_encode_worker": True}, "dedicated multimodal"), ({"multimodal_worker": True}, "dedicated multimodal"), ({"image_diffusion_worker": True}, "image diffusion"), diff --git a/components/src/dynamo/sglang/tests/test_sglang_video_generation_handler.py b/components/src/dynamo/sglang/tests/test_sglang_video_generation_handler.py new file mode 100644 index 000000000000..437945a5d7d0 --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_sglang_video_generation_handler.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for VideoGenerationWorkerHandler's input_reference handling. + +The video (I2V) handler had the same unvalidated ``input_reference`` -> +``image_path`` passthrough as image diffusion and shares the fix, but had no +test of its own โ€” so the shared validation could be dropped from this handler +alone and every test would still pass. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, Mock + +import pytest + +from dynamo.common.http.url_validator import UrlValidationError +from dynamo.sglang.request_handlers.video_generation.video_generation_handler import ( + VideoGenerationWorkerHandler, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.sglang, + pytest.mark.gpu_0, # No GPU needed for unit tests + pytest.mark.pre_merge, + pytest.mark.parallel, +] + + +@pytest.fixture +def handler(): + config = MagicMock() + config.dynamo_args.media_output_fs_url = "file:///tmp/videos" + config.dynamo_args.media_output_http_url = "file:///tmp/videos" + handler = VideoGenerationWorkerHandler( + generator=MagicMock(), config=config, publisher=None, fs=MagicMock() + ) + handler._frames_to_video = AsyncMock(return_value=b"mp4-bytes") + return handler + + +def _generate_video(handler, **overrides): + kwargs = dict( + prompt="a cat", + width=256, + height=256, + num_frames=8, + fps=8, + num_inference_steps=10, + guidance_scale=5.0, + seed=42, + request_id="req-1", + ) + kwargs.update(overrides) + return handler._generate_video(**kwargs) + + +async def test_i2v_passes_the_resolved_path_to_the_generator( + handler, tmp_path, monkeypatch +) -> None: + """The generator gets the resolved path, not the client's reference.""" + monkeypatch.setenv("DYN_MM_LOCAL_PATH", str(tmp_path)) + (tmp_path / "sub").mkdir() + reference = tmp_path / "first_frame.png" + reference.write_bytes(b"x") + # Route through a "sub/.." segment so the resolved form differs from the + # string the client sent; otherwise this passes even with no resolution. + sent = str(tmp_path / "sub" / ".." / "first_frame.png") + handler.generator.generate = Mock(return_value=SimpleNamespace(frames=[object()])) + + await _generate_video(handler, input_reference=sent) + + args = handler.generator.generate.call_args[1]["sampling_params_kwargs"] + assert args["image_path"] != sent + assert args["image_path"] == str(reference.resolve()) + + +async def test_i2v_rejects_a_reference_outside_the_allowed_dir( + handler, tmp_path, monkeypatch +) -> None: + """Traversal is refused before the generator ever sees a path.""" + monkeypatch.setenv("DYN_MM_LOCAL_PATH", str(tmp_path)) + handler.generator.generate = Mock() + + with pytest.raises(UrlValidationError, match="outside the allowed directory"): + await _generate_video(handler, input_reference="/etc/passwd") + + handler.generator.generate.assert_not_called() + + +async def test_i2v_rejection_message_is_bounded(handler, tmp_path, monkeypatch) -> None: + """input_reference is unbounded and this message reaches the response body.""" + monkeypatch.setenv("DYN_MM_LOCAL_PATH", str(tmp_path)) + handler.generator.generate = Mock() + + with pytest.raises(UrlValidationError) as excinfo: + await _generate_video(handler, input_reference="/nope/" + "A" * 200_000) + + assert len(str(excinfo.value)) < 500 + + +async def test_t2v_passes_no_image_path(handler) -> None: + """Control: a text-only request must not gain an image_path.""" + handler.generator.generate = Mock(return_value=SimpleNamespace(frames=[object()])) + + await _generate_video(handler, input_reference=None) + + args = handler.generator.generate.call_args[1]["sampling_params_kwargs"] + assert "image_path" not in args diff --git a/components/src/dynamo/tokenspeed/args.py b/components/src/dynamo/tokenspeed/args.py index c499f45d1d50..5a9405d32c74 100644 --- a/components/src/dynamo/tokenspeed/args.py +++ b/components/src/dynamo/tokenspeed/args.py @@ -17,7 +17,12 @@ DynamoRuntimeArgGroup, DynamoRuntimeConfig, ) +from dynamo.common.constants import DisaggregationMode from dynamo.common.utils.runtime import parse_endpoint +from dynamo.tokenspeed.disagg import ( + resolve_disaggregation_mode, + validate_disagg_compatibility, +) DEFAULT_ENDPOINT_COMPONENT = "backend" DEFAULT_ENDPOINT_NAME = "generate" @@ -25,6 +30,7 @@ class Config(DynamoRuntimeConfig): component: str + disaggregation_mode: DisaggregationMode = DisaggregationMode.AGGREGATED use_kv_events: bool = False model: str @@ -33,7 +39,8 @@ class Config(DynamoRuntimeConfig): def validate(self) -> None: DynamoRuntimeConfig.validate(self) - self.use_kv_events = False + self.disaggregation_mode = resolve_disaggregation_mode(self.server_args) + validate_disagg_compatibility(self.disaggregation_mode, self.server_args) @register_encoder(Config) @@ -95,9 +102,14 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> Config: else: config.custom_jinja_template = None + component = ( + "prefill" + if config.disaggregation_mode == DisaggregationMode.PREFILL + else DEFAULT_ENDPOINT_COMPONENT + ) endpoint = ( config.endpoint - or f"dyn://{config.namespace}.{DEFAULT_ENDPOINT_COMPONENT}.{DEFAULT_ENDPOINT_NAME}" + or f"dyn://{config.namespace}.{component}.{DEFAULT_ENDPOINT_NAME}" ) parsed_namespace, parsed_component_name, parsed_endpoint_name = parse_endpoint( endpoint diff --git a/components/src/dynamo/tokenspeed/disagg.py b/components/src/dynamo/tokenspeed/disagg.py new file mode 100644 index 000000000000..b171e372bb0e --- /dev/null +++ b/components/src/dynamo/tokenspeed/disagg.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TokenSpeed's Mooncake bootstrap protocol for Dynamo P/D workers.""" + +from __future__ import annotations + +import ipaddress +import os +import secrets +import socket +from typing import Any + +from dynamo.common.constants import DisaggregationMode +from dynamo.llm.exceptions import InvalidArgument + +BOOTSTRAP_HOST_ENV = "DYN_TOKENSPEED_BOOTSTRAP_HOST" + + +def resolve_disaggregation_mode(server_args: Any) -> DisaggregationMode: + value = getattr(server_args, "disaggregation_mode", None) + if value in (None, "null", "aggregated"): + return DisaggregationMode.AGGREGATED + if value not in ("prefill", "decode"): + raise ValueError(f"Unsupported TokenSpeed disaggregation mode: {value!r}") + return DisaggregationMode(value) + + +def attention_dp_size(server_args: Any) -> int: + mapping = getattr(server_args, "mapping", None) + attention = getattr(mapping, "attn", None) + return int( + getattr(attention, "dp_size", None) + or getattr(server_args, "data_parallel_size", None) + or getattr(server_args, "dp_size", None) + or 1 + ) + + +def cache_block_size(server_args: Any) -> int | None: + # TokenSpeed renamed block_size to prefix_granularity when physical pages + # and reusable-prefix identity became separate concepts. Router blocks track + # the prefix identity carried by KV events, not physical allocation pages. + value = getattr(server_args, "prefix_granularity", None) + if value is None: + value = getattr(server_args, "block_size", None) + return int(value) if value is not None else None + + +def validate_disagg_compatibility(mode: DisaggregationMode, server_args: Any) -> None: + if mode == DisaggregationMode.AGGREGATED: + return + if ( + getattr(server_args, "disaggregation_transfer_backend", "mooncake") + != "mooncake" + ): + raise ValueError("Dynamo TokenSpeed disaggregation requires Mooncake") + if attention_dp_size(server_args) != 1: + # TokenSpeed assigns P/D requests by bootstrap_room modulo DP size. + # Dynamo's selected cache-owning DP rank must not be silently ignored. + raise ValueError( + "Dynamo TokenSpeed disaggregation requires attention DP=1; " + "use independent workers for multiple prefill/decode replicas" + ) + if (cache_block_size(server_args) or 0) <= 0: + raise ValueError( + "TokenSpeed disaggregation requires a positive --prefix-granularity (--block-size)" + ) + + +def runtime_disaggregated_endpoint(server_args: Any) -> tuple[str, int]: + """Advertise a peer-reachable host; the environment opts into single-host use.""" + port = getattr(server_args, "disaggregation_bootstrap_port", None) + if port is None or not 0 < int(port) < 65536: + raise ValueError("TokenSpeed prefill requires a valid bootstrap port") + explicit_host = os.environ.get(BOOTSTRAP_HOST_ENV) + host = explicit_host or getattr(server_args, "host", None) + try: + address = ipaddress.ip_address(host.strip("[]")) if host else None + except ValueError: + address = None # Hostnames are valid advertised addresses too. + if ( + not host + or (address is not None and address.is_unspecified) + or ( + not explicit_host + and ((address is not None and address.is_loopback) or host == "localhost") + ) + ): + # ServerArgs.host defaults to loopback. Mooncake's bootstrap server + # binds all interfaces, so derive the advertisement independently. + try: + host = socket.gethostbyname(socket.gethostname()) + except socket.gaierror as exc: + raise ValueError( + f"Cannot determine TokenSpeed bootstrap host; set {BOOTSTRAP_HOST_ENV} " + "to an address reachable by decode workers" + ) from exc + address = ipaddress.ip_address(host) + if address.is_loopback or address.is_unspecified: + raise ValueError( + "TokenSpeed bootstrap host resolves to a local-only address; " + f"set {BOOTSTRAP_HOST_ENV} to an address reachable by decode workers" + ) + return str(host), int(port) + + +def bootstrap_kwargs( + request: dict[str, Any], + mode: DisaggregationMode, + prefill_endpoint: tuple[str, int] | None = None, +) -> dict[str, Any]: + """Validate the handoff, creating a room for synchronous prefill if needed.""" + if mode == DisaggregationMode.AGGREGATED: + return {} + info = request.get("bootstrap_info") + if info is None and mode == DisaggregationMode.PREFILL and prefill_endpoint: + prefill_host, prefill_port = prefill_endpoint + info = { + "bootstrap_host": prefill_host, + "bootstrap_port": prefill_port, + # Match the router's signed-64-bit room range and attention DP=1. + "bootstrap_room": secrets.randbits(63), + } + if not isinstance(info, dict): + raise InvalidArgument( + f"TokenSpeed {mode.value} worker requires bootstrap_info from " + "Dynamo's prefill router" + ) + host, port, room = ( + info.get(f"bootstrap_{key}") for key in ("host", "port", "room") + ) + if ( + not isinstance(host, str) + or not host + or isinstance(port, bool) + or not isinstance(port, int) + or not 0 < port < 65536 + or isinstance(room, bool) + or not isinstance(room, int) + or not 0 <= room < 2**64 + ): + raise InvalidArgument("Invalid TokenSpeed bootstrap host, port, or room") + return {"bootstrap_host": host, "bootstrap_port": port, "bootstrap_room": room} diff --git a/components/src/dynamo/tokenspeed/kv_events.py b/components/src/dynamo/tokenspeed/kv_events.py new file mode 100644 index 000000000000..fead11d2314b --- /dev/null +++ b/components/src/dynamo/tokenspeed/kv_events.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Describe TokenSpeed's native KV event stream to the common Worker.""" + +from __future__ import annotations + +import json +from typing import Any + +from dynamo.common.backend.publisher import ZmqSource + + +def kv_events_config_dict(raw: Any) -> dict[str, Any]: + if raw is None or raw == "": + return {} + config = json.loads(raw) if isinstance(raw, str) else raw + if not isinstance(config, dict): + raise ValueError("TokenSpeed --kv-events-config must be a JSON object") + return dict(config) + + +def kv_events_enabled(config: dict[str, Any]) -> bool: + enabled = config.get("enable_kv_cache_events", False) + if not isinstance(enabled, bool): + raise ValueError("TokenSpeed enable_kv_cache_events must be a JSON boolean") + return enabled and config.get("publisher", "zmq") != "null" + + +def kv_event_source(config: dict[str, Any]) -> ZmqSource: + if config.get("replay_endpoint") is not None: + raise ValueError( + "Dynamo TokenSpeed KV events do not support replay_endpoint; " + "remove it from --kv-events-config" + ) + if (config.get("publisher") or "zmq") != "zmq": + raise ValueError("Dynamo TokenSpeed KV events require the zmq publisher") + endpoint = config["endpoint"] + if not isinstance(endpoint, str): + raise ValueError("TokenSpeed KV event endpoint must be a string") + # TokenSpeed binds wildcard TCP endpoints and IPC endpoints. Other TCP + # addresses make its publisher connect, which cannot pair with Worker's SUB. + if endpoint.startswith("tcp://*:"): + endpoint = endpoint.replace("tcp://*:", "tcp://127.0.0.1:", 1) + elif endpoint.startswith("tcp://[::]:"): + endpoint = endpoint.replace("tcp://[::]:", "tcp://[::1]:", 1) + elif not endpoint.startswith("ipc://"): + raise ValueError( + "TokenSpeed KV events require a binding endpoint: ipc://... or tcp://*:PORT" + ) + return ZmqSource(endpoint=endpoint, topic=config.get("topic", ""), dp_rank=0) diff --git a/components/src/dynamo/tokenspeed/llm_engine.py b/components/src/dynamo/tokenspeed/llm_engine.py index 2a9f74aa773d..837e0cc4d59f 100644 --- a/components/src/dynamo/tokenspeed/llm_engine.py +++ b/components/src/dynamo/tokenspeed/llm_engine.py @@ -9,6 +9,7 @@ import json import logging import re +import tempfile from collections.abc import AsyncGenerator from typing import Any @@ -20,12 +21,27 @@ LLMEngine, LlmRegistration, ) +from dynamo.common.backend.publisher import KvEventSource, ZmqSource from dynamo.common.backend.worker import WorkerConfig +from dynamo.common.constants import DisaggregationMode from dynamo.common.utils.engine_response import normalize_finish_reason from dynamo.common.utils.structural_tag import serialize_structural_tag from dynamo.llm import ModelInput from dynamo.llm.exceptions import InvalidArgument from dynamo.tokenspeed.args import parse_args +from dynamo.tokenspeed.disagg import ( + attention_dp_size, + bootstrap_kwargs, + cache_block_size, + resolve_disaggregation_mode, + runtime_disaggregated_endpoint, + validate_disagg_compatibility, +) +from dynamo.tokenspeed.kv_events import ( + kv_event_source, + kv_events_config_dict, + kv_events_enabled, +) logger = logging.getLogger(__name__) @@ -33,9 +49,14 @@ class TokenspeedLLMEngine(LLMEngine): def __init__(self, server_args: Any): self.server_args = server_args + self.disaggregation_mode = resolve_disaggregation_mode(server_args) + self._bootstrap_endpoint: tuple[str, int] | None = None + self._kv_source: ZmqSource | None = None + self._kv_event_dir: tempfile.TemporaryDirectory | None = None self.engine = None self._model_max_len: int | None = None self._active_rids_by_context: dict[str, list[str]] = {} + self._cancelled_contexts: set[str] = set() @classmethod async def from_args( @@ -52,7 +73,13 @@ async def from_args( return engine, worker_config async def start(self, worker_id: int) -> EngineConfig: - del worker_id # tokenspeed has no cluster-wide ID needs + del worker_id + validate_disagg_compatibility(self.disaggregation_mode, self.server_args) + bootstrap_host, bootstrap_port = None, None + if self.disaggregation_mode == DisaggregationMode.PREFILL: + self._bootstrap_endpoint = runtime_disaggregated_endpoint(self.server_args) + bootstrap_host, bootstrap_port = self._bootstrap_endpoint + self._configure_kv_events() # The Dynamo response layer expects per-chunk token deltas. self.server_args.stream_output = True self.engine = _tokenspeed_engine_cls()(server_args=self.server_args) @@ -66,7 +93,7 @@ async def start(self, worker_id: int) -> EngineConfig: or getattr(self.server_args, "max_model_len", None) ) - block_size = _optional_int(getattr(self.server_args, "block_size", None)) + block_size = cache_block_size(self.server_args) max_total_tokens = _optional_int( scheduler_info.get("max_total_num_tokens") or getattr(self.server_args, "max_total_tokens", None) @@ -88,6 +115,8 @@ async def start(self, worker_id: int) -> EngineConfig: served_model_name=self.server_args.served_model_name, llm=LlmRegistration( context_length=self._model_max_len, + bootstrap_host=bootstrap_host, + bootstrap_port=bootstrap_port, kv_cache_block_size=block_size, total_kv_blocks=total_kv_blocks, max_num_seqs=_optional_int( @@ -98,18 +127,62 @@ async def start(self, worker_id: int) -> EngineConfig: ), ) + def _configure_kv_events(self) -> None: + config = kv_events_config_dict( + getattr(self.server_args, "kv_events_config", None) + ) + if not kv_events_enabled(config): + return + if not getattr(self.server_args, "enable_prefix_caching", True): + # Native publisher construction is independent of prefix caching. + config["enable_kv_cache_events"] = False + config["publisher"] = "null" + self.server_args.kv_events_config = json.dumps(config) + logger.warning( + "TokenSpeed KV events were requested but enable_prefix_caching=False " + "(--disable-prefix-caching); no KV events will be published" + ) + return + if attention_dp_size(self.server_args) != 1: + raise ValueError( + "TokenSpeed KV-aware routing currently requires attention DP=1; " + "use independent worker replicas" + ) + if (cache_block_size(self.server_args) or 0) <= 0: + raise ValueError( + "TokenSpeed KV events require a positive --prefix-granularity (--block-size)" + ) + if not config.get("endpoint"): + self._kv_event_dir = tempfile.TemporaryDirectory( + prefix="dynamo-tokenspeed-" + ) + config["endpoint"] = f"ipc://{self._kv_event_dir.name}/kv-events" + self._kv_source = kv_event_source(config) + self.server_args.kv_events_config = json.dumps(config) + + async def kv_event_sources(self) -> list[KvEventSource]: + return [self._kv_source] if self._kv_source is not None else [] + async def generate( self, request: GenerateRequest, context: Context ) -> AsyncGenerator[GenerateChunk, None]: - assert self.engine is not None, "Engine not initialized" + if self.engine is None: + raise RuntimeError("Engine not initialized") + bootstrap = bootstrap_kwargs( + request, self.disaggregation_mode, self._bootstrap_endpoint + ) _validate_single_choice_sampling(request) sampling_params = build_sampling_params(request, self._model_max_len) + if self.disaggregation_mode == DisaggregationMode.PREFILL: + sampling_params["max_new_tokens"] = 1 + sampling_params.pop("min_new_tokens", None) token_ids = request.get("token_ids", []) obj = _generate_req_input_cls()( input_ids=token_ids, sampling_params=sampling_params, stream=True, + **bootstrap, ) request_id = context.id() @@ -119,6 +192,19 @@ async def generate( emitted_completion_tokens = 0 try: + if self.disaggregation_mode == DisaggregationMode.PREFILL: + # The router needs a handoff even when it supplied the bootstrap + # address. Preserve its trace identity; fallback rooms use the + # same validated parameters passed to the native prefill engine. + yield { + "token_ids": [], + "disaggregated_params": request.get("bootstrap_info") or bootstrap, + } + # Native abort_request ignores unknown RIDs. An abort while the + # handoff was yielded must therefore prevent native submission. + if request_id is not None and request_id in self._cancelled_contexts: + yield {"token_ids": [], "finish_reason": "cancelled"} + return async for out in self.engine.tokenizer_manager.generate_request(obj): delta_out, emitted_completion_tokens = _completion_delta_output( out, emitted_completion_tokens @@ -127,21 +213,32 @@ async def generate( finally: if request_id is not None: self._active_rids_by_context.pop(request_id, None) + self._cancelled_contexts.discard(request_id) async def abort(self, context: Context) -> None: request_id = context.id() if self.engine is None or request_id is None: return + if request_id in self._active_rids_by_context: + self._cancelled_contexts.add(request_id) rids = self._active_rids_by_context.get(request_id, [request_id]) for rid in rids: self.engine.tokenizer_manager.abort_request(rid) logger.debug("Aborted TokenSpeed request %s", rid) async def cleanup(self) -> None: - if self.engine is not None: - self.engine.shutdown() - logger.info("TokenSpeed engine shutdown") + engine, self.engine = self.engine, None + try: + if engine is not None: + engine.shutdown() + logger.info("TokenSpeed engine shutdown") + finally: + self._bootstrap_endpoint = None + self._kv_source = None + if self._kv_event_dir is not None: + self._kv_event_dir.cleanup() + self._kv_event_dir = None def build_sampling_params( @@ -305,7 +402,18 @@ def _finish_reason_type(finish_reason: Any) -> str: if hasattr(finish_reason, "to_json"): finish_reason = finish_reason.to_json() if isinstance(finish_reason, dict): - return str(finish_reason.get("type") or "unknown") + reason = str(finish_reason.get("type") or "unknown") + if reason == "abort": + message = finish_reason.get("message") or "Unknown backend error" + # ABORT_CODE.UnknownError (522) also represents client cancellation. + # Only TransferFailed (521), NumericalError (523), or an explicit + # transfer failure are errors. Current PD hooks use 522 with text. + transfer_failed = "transfer" in message.lower() and ( + "failed" in message.lower() or "timed out" in message.lower() + ) + if finish_reason.get("err_type") in (521, 523) or transfer_failed: + raise RuntimeError(f"TokenSpeed generation aborted: {message}") + return reason return str(finish_reason) diff --git a/components/src/dynamo/tokenspeed/tests/test_tokenspeed_disagg.py b/components/src/dynamo/tokenspeed/tests/test_tokenspeed_disagg.py new file mode 100644 index 000000000000..3e0670200d1f --- /dev/null +++ b/components/src/dynamo/tokenspeed/tests/test_tokenspeed_disagg.py @@ -0,0 +1,561 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Adapter contract tests; native TokenSpeed/GPU processes are replaced at the boundary.""" + +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from dynamo.common.constants import DisaggregationMode +from dynamo.llm.exceptions import InvalidArgument +from dynamo.tokenspeed import args, disagg, llm_engine +from dynamo.tokenspeed.disagg import BOOTSTRAP_HOST_ENV +from dynamo.tokenspeed.llm_engine import TokenspeedLLMEngine + +pytestmark = [pytest.mark.unit, pytest.mark.gpu_0, pytest.mark.pre_merge] + + +def server_args(**overrides): + values = dict( + model="test-model", + served_model_name="test-model", + disaggregation_mode="null", + disaggregation_transfer_backend="mooncake", + disaggregation_bootstrap_port=9000, + host="127.0.0.1", + mapping=SimpleNamespace(attn=SimpleNamespace(dp_size=1)), + prefix_granularity=64, + max_model_len=4096, + max_total_tokens=4096, + enable_prefix_caching=True, + kv_events_config=None, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +@pytest.fixture +def native_engine(monkeypatch): + monkeypatch.delenv(BOOTSTRAP_HOST_ENV, raising=False) + native = Mock() + native.scheduler_info = { + "max_model_len": 8192, + "max_total_num_tokens": 8192, + "max_num_seqs": 32, + } + constructor = Mock(return_value=native) + monkeypatch.setattr(llm_engine, "_tokenspeed_engine_cls", lambda: constructor) + monkeypatch.setattr(llm_engine, "_generate_req_input_cls", lambda: SimpleNamespace) + return native, constructor + + +@pytest.mark.parametrize( + "mode,component,expected", + [ + ("null", "backend", DisaggregationMode.AGGREGATED), + ("prefill", "prefill", DisaggregationMode.PREFILL), + ("decode", "backend", DisaggregationMode.DECODE), + ], +) +@pytest.mark.parametrize("explicit_endpoint", [False, True]) +async def test_cli_registers_worker_role( + monkeypatch, mode, component, expected, explicit_endpoint +): + class NativeArgs: + @staticmethod + def add_cli_args(parser): + parser.add_argument("--model") + parser.add_argument("--disaggregation-mode", default="null") + + @staticmethod + def from_cli_args(parsed): + return server_args(**vars(parsed)) + + monkeypatch.setattr(args, "_server_args_cls", lambda: NativeArgs) + cli = [ + "--model", + "test-model", + "--disaggregation-mode", + mode, + "--namespace", + "test", + ] + if explicit_endpoint: + cli += ["--endpoint", "dyn://custom.worker.generate"] + _, config = await TokenspeedLLMEngine.from_args(cli) + assert config.disaggregation_mode == expected + assert config.component == ("worker" if explicit_endpoint else component) + assert config.namespace == ("custom" if explicit_endpoint else "test") + assert config.endpoint == "generate" + + +@pytest.mark.parametrize("mode", ["prefill", "decode", "null"]) +@pytest.mark.parametrize("legacy_block_size", [False, True]) +async def test_registration_advertises_prefill_bootstrap_only( + monkeypatch, native_engine, mode, legacy_block_size +): + monkeypatch.setenv(BOOTSTRAP_HOST_ENV, "prefill.example") + native_args = server_args(disaggregation_mode=mode) + if legacy_block_size: + native_args.block_size = native_args.prefix_granularity + del native_args.prefix_granularity + engine = TokenspeedLLMEngine(native_args) + try: + registration = await engine.start(worker_id=1) + assert registration.llm.context_length == 8192 + assert registration.llm.total_kv_blocks == 128 + assert registration.llm.kv_cache_block_size == 64 + assert registration.llm.bootstrap_host == ( + "prefill.example" if mode == "prefill" else None + ) + assert registration.llm.bootstrap_port == (9000 if mode == "prefill" else None) + finally: + await engine.cleanup() + + +@pytest.mark.parametrize("mode", ["prefill", "decode"]) +async def test_generate_forwards_handoff_and_preserves_decode_budget( + native_engine, mode +): + native, _ = native_engine + inputs = [] + + async def generate(obj): + inputs.append(obj) + yield { + "output_ids": [10, 11, 20], + "meta_info": {"completion_tokens": 1, "finish_reason": {"type": "length"}}, + } + + native.tokenizer_manager.generate_request = generate + engine = TokenspeedLLMEngine( + server_args(disaggregation_mode=mode, host="prefill.example") + ) + await engine.start(worker_id=1) + request = { + "token_ids": [10, 11], + "stop_conditions": {"max_tokens": 30, "min_tokens": 10}, + "bootstrap_info": { + "bootstrap_host": "selected-prefill.example", + "bootstrap_port": 9010, + "bootstrap_room": 0, + "handoff_id": "1cfe6c6e-d8c8-4bea-a0a6-ccde2bd97c54", + }, + } + context = SimpleNamespace(id=lambda: "request-1") + try: + chunks = [chunk async for chunk in engine.generate(request, context)] + if mode == "prefill": + assert chunks[0] == { + "token_ids": [], + "disaggregated_params": request["bootstrap_info"], + } + assert chunks[1]["token_ids"] == [20] + else: + assert chunks[0]["token_ids"] == [20] + assert inputs[0].rid == "request-1" + assert inputs[0].bootstrap_host == "selected-prefill.example" + assert inputs[0].bootstrap_port == 9010 + assert inputs[0].bootstrap_room == 0 + assert not hasattr(inputs[0], "handoff_id") + if mode == "prefill": + assert inputs[0].sampling_params["max_new_tokens"] == 1 + assert "min_new_tokens" not in inputs[0].sampling_params + else: + assert inputs[0].sampling_params["max_new_tokens"] == 30 + assert inputs[0].sampling_params["min_new_tokens"] == 10 + assert engine._active_rids_by_context == {} + finally: + await engine.cleanup() + + +@pytest.mark.parametrize( + "bootstrap", + [ + None, + {}, + {"bootstrap_host": "", "bootstrap_port": 9000, "bootstrap_room": 1}, + {"bootstrap_host": "p", "bootstrap_port": 65536, "bootstrap_room": 1}, + {"bootstrap_host": "p", "bootstrap_port": 9000, "bootstrap_room": -1}, + {"bootstrap_host": "p", "bootstrap_port": 9000, "bootstrap_room": True}, + {"bootstrap_host": "p", "bootstrap_port": 9000, "bootstrap_room": 2**64}, + ], +) +async def test_bad_handoff_rejected_before_native_request(native_engine, bootstrap): + native, _ = native_engine + engine = TokenspeedLLMEngine(server_args(disaggregation_mode="decode")) + await engine.start(worker_id=1) + try: + with pytest.raises(InvalidArgument, match="bootstrap"): + async for _ in engine.generate( + {"token_ids": [1], "bootstrap_info": bootstrap}, + SimpleNamespace(id=lambda: "bad"), + ): + pass + native.tokenizer_manager.generate_request.assert_not_called() + finally: + await engine.cleanup() + + +@pytest.mark.parametrize( + "overrides,message", + [ + ( + {"mapping": SimpleNamespace(attn=SimpleNamespace(dp_size=2))}, + "attention DP=1", + ), + ({"disaggregation_transfer_backend": "fake"}, "Mooncake"), + ({"prefix_granularity": 0}, "positive --prefix-granularity"), + ({"prefix_granularity": -1}, "positive --prefix-granularity"), + ], +) +async def test_unsupported_disagg_fails_before_start(native_engine, overrides, message): + _, constructor = native_engine + engine = TokenspeedLLMEngine(server_args(disaggregation_mode="decode", **overrides)) + with pytest.raises(ValueError, match=message): + await engine.start(worker_id=1) + constructor.assert_not_called() + await engine.cleanup() + + +async def test_kv_sources_use_unique_ipc_endpoints_and_cleanup(native_engine): + engines = [ + TokenspeedLLMEngine( + server_args(kv_events_config='{"enable_kv_cache_events":true}') + ) + for _ in range(2) + ] + paths = [] + try: + for engine in engines: + await engine.start(worker_id=1) + [source] = await engine.kv_event_sources() + native_config = json.loads(engine.server_args.kv_events_config) + assert source.endpoint == native_config["endpoint"] + assert source.dp_rank == 0 + assert source.endpoint.startswith("ipc://") + paths.append(Path(source.endpoint.removeprefix("ipc://")).parent) + assert paths[-1].is_dir() + assert paths[0] != paths[1] + finally: + for engine in engines: + await engine.cleanup() + await engine.cleanup() + assert all(not path.exists() for path in paths) + + +@pytest.mark.parametrize( + "native_config,expected", + [ + ({"endpoint": "tcp://*:19000"}, "tcp://127.0.0.1:19000"), + ({"endpoint": "ipc:///test-events", "publisher": None}, "ipc:///test-events"), + ], +) +async def test_kv_source_matches_native_bind_and_topic( + native_engine, native_config, expected +): + native_config.update(enable_kv_cache_events=True, topic="kv") + engine = TokenspeedLLMEngine( + server_args(kv_events_config=json.dumps(native_config)) + ) + try: + await engine.start(worker_id=1) + [source] = await engine.kv_event_sources() + assert source.endpoint == expected + assert source.topic == "kv" + assert json.loads(engine.server_args.kv_events_config) == native_config + finally: + await engine.cleanup() + + +@pytest.mark.parametrize( + "overrides", + [ + {"kv_events_config": None}, + {"kv_events_config": '{"enable_kv_cache_events":false}'}, + {"kv_events_config": '{"enable_kv_cache_events":true,"publisher":"null"}'}, + { + "kv_events_config": '{"enable_kv_cache_events":true}', + "enable_prefix_caching": False, + }, + ], +) +async def test_disabled_kv_events_have_no_source(native_engine, overrides, caplog): + native, constructor = native_engine + native_configs = [] + + def construct(*, server_args): + native_configs.append(json.loads(server_args.kv_events_config or "{}")) + return native + + constructor.side_effect = construct + engine = TokenspeedLLMEngine(server_args(**overrides)) + try: + await engine.start(worker_id=1) + assert await engine.kv_event_sources() == [] + if overrides.get("enable_prefix_caching") is False: + assert native_configs[0]["enable_kv_cache_events"] is False + assert native_configs[0]["publisher"] == "null" + assert "enable_prefix_caching=False" in caplog.text + assert "KV events" in caplog.text + else: + assert caplog.text == "" + finally: + await engine.cleanup() + + +async def test_connecting_native_publisher_is_rejected(native_engine): + _, constructor = native_engine + engine = TokenspeedLLMEngine( + server_args( + kv_events_config=json.dumps( + { + "enable_kv_cache_events": True, + "endpoint": "tcp://127.0.0.1:19000", + } + ) + ) + ) + with pytest.raises(ValueError, match="binding endpoint"): + await engine.start(worker_id=1) + constructor.assert_not_called() + await engine.cleanup() + + +async def test_failed_start_cleans_up_event_endpoint(native_engine): + _, constructor = native_engine + constructor.side_effect = RuntimeError("native startup failed") + engine = TokenspeedLLMEngine( + server_args(kv_events_config='{"enable_kv_cache_events":true}') + ) + with pytest.raises(RuntimeError, match="native startup failed"): + await engine.start(worker_id=1) + path = Path( + json.loads(engine.server_args.kv_events_config)["endpoint"].removeprefix( + "ipc://" + ) + ).parent + await engine.cleanup() + assert not path.exists() + assert await engine.kv_event_sources() == [] + + +@pytest.mark.parametrize("enabled", ["false", "true", None, 1]) +async def test_non_boolean_event_flag_rejected_before_start(native_engine, enabled): + _, constructor = native_engine + engine = TokenspeedLLMEngine( + server_args( + kv_events_config=json.dumps( + { + "enable_kv_cache_events": enabled, + } + ) + ) + ) + with pytest.raises(ValueError, match="JSON boolean"): + await engine.start(worker_id=1) + constructor.assert_not_called() + await engine.cleanup() + + +@pytest.mark.parametrize("host", ["127.0.0.1", "0.0.0.0", "::", "::1", "localhost"]) +@pytest.mark.parametrize("address", ["127.0.1.1", "0.0.0.0"]) +async def test_prefill_rejects_local_only_derived_advertisement( + monkeypatch, native_engine, host, address +): + monkeypatch.delenv(BOOTSTRAP_HOST_ENV, raising=False) + monkeypatch.setattr(disagg.socket, "gethostbyname", lambda _: address) + _, constructor = native_engine + engine = TokenspeedLLMEngine(server_args(disaggregation_mode="prefill", host=host)) + with pytest.raises(ValueError, match=BOOTSTRAP_HOST_ENV): + await engine.start(worker_id=1) + constructor.assert_not_called() + await engine.cleanup() + + +@pytest.mark.parametrize("host", ["127.0.0.1", "0.0.0.0", "::", "::1", "localhost"]) +async def test_prefill_derives_non_loopback_advertisement( + monkeypatch, native_engine, host +): + monkeypatch.setattr(disagg.socket, "gethostbyname", lambda _: "192.0.2.10") + engine = TokenspeedLLMEngine(server_args(disaggregation_mode="prefill", host=host)) + try: + config = await engine.start(worker_id=1) + assert config.llm.bootstrap_host == "192.0.2.10" + finally: + await engine.cleanup() + + +async def test_explicit_loopback_prefill_allowed_for_single_host( + monkeypatch, native_engine +): + monkeypatch.setenv(BOOTSTRAP_HOST_ENV, "127.0.0.1") + engine = TokenspeedLLMEngine(server_args(disaggregation_mode="prefill")) + try: + config = await engine.start(worker_id=1) + assert config.llm.bootstrap_host == "127.0.0.1" + finally: + await engine.cleanup() + + +@pytest.mark.parametrize("explicit_null", [False, True]) +async def test_prefill_creates_handoff_when_router_has_no_bootstrap_endpoint( + monkeypatch, native_engine, explicit_null +): + """The synchronous path supplies the same fresh room to prefill and decode.""" + native, _ = native_engine + inputs = [] + rooms = Mock(side_effect=[17, 23]) + monkeypatch.setattr( + disagg, "secrets", SimpleNamespace(randbits=rooms), raising=False + ) + + async def generate(obj): + inputs.append(obj) + yield {"output_ids": [20], "meta_info": {"completion_tokens": 1}} + + native.tokenizer_manager.generate_request = generate + engine = TokenspeedLLMEngine( + server_args(disaggregation_mode="prefill", host="prefill.example") + ) + registration = await engine.start(worker_id=1) + try: + for index, room in enumerate([17, 23]): + request = {"token_ids": [10, 11], "stop_conditions": {"max_tokens": 30}} + if explicit_null: + request["bootstrap_info"] = None + stream = engine.generate(request, SimpleNamespace(id=lambda: "fallback")) + try: + first = await anext(stream) + assert first["token_ids"] == [] + handoff = first["disaggregated_params"] + assert handoff == { + "bootstrap_host": registration.llm.bootstrap_host, + "bootstrap_port": registration.llm.bootstrap_port, + "bootstrap_room": room, + } + assert len(inputs) == index + chunks = [chunk async for chunk in stream] + assert chunks[0]["token_ids"] == [20] + for key, value in handoff.items(): + assert getattr(inputs[index], key) == value + assert ( + disagg.bootstrap_kwargs( + {"bootstrap_info": handoff}, DisaggregationMode.DECODE + ) + == handoff + ) + assert engine._active_rids_by_context == {} + finally: + await stream.aclose() + assert rooms.call_count == 2 + rooms.assert_called_with(63) + finally: + await engine.cleanup() + + +@pytest.mark.parametrize("router_bootstrap", [False, True]) +@pytest.mark.parametrize("cancel", [False, True]) +async def test_prefill_cancellation_after_handoff_prevents_submission( + native_engine, router_bootstrap, cancel +): + """An abort before native admission survives resuming the handoff stream.""" + native, _ = native_engine + submitted = [] + + async def generate(obj): + submitted.append(obj.rid) + yield {"output_ids": [20], "meta_info": {"finish_reason": {"type": "length"}}} + + native.tokenizer_manager.generate_request = generate + engine = TokenspeedLLMEngine( + server_args(disaggregation_mode="prefill", host="prefill.example") + ) + registration = await engine.start(worker_id=1) + request = {"token_ids": [10, 11]} + if router_bootstrap: + request["bootstrap_info"] = { + "bootstrap_host": registration.llm.bootstrap_host, + "bootstrap_port": registration.llm.bootstrap_port, + "bootstrap_room": 1, + } + context = SimpleNamespace(id=lambda: "handoff-cancel") + stream = engine.generate(request, context) + try: + first = await anext(stream) + assert first["disaggregated_params"] + assert submitted == [] + if cancel: + await engine.abort(context) + await engine.abort(context) + chunks = [chunk async for chunk in stream] + assert submitted == ([] if cancel else ["handoff-cancel"]) + assert chunks[-1]["finish_reason"] == ("cancelled" if cancel else "length") + assert engine._active_rids_by_context == {} + # Reusing the context must not inherit cancellation from the old stream. + chunks = [chunk async for chunk in engine.generate(request, context)] + assert chunks[-1]["finish_reason"] == "length" + assert submitted[-1] == "handoff-cancel" + finally: + await stream.aclose() + await engine.cleanup() + + +async def test_cancelled_native_stream_preserves_cancelled_finish(native_engine): + native, _ = native_engine + aborted = set() + native.tokenizer_manager.abort_request.side_effect = aborted.add + + async def generate(obj): + yield {"output_ids": [20]} + assert obj.rid in aborted + yield { + "output_ids": [], + "meta_info": { + "finish_reason": { + "type": "abort", + "message": "AbortReq from client", + "err_type": 522, + } + }, + } + + native.tokenizer_manager.generate_request = generate + engine = TokenspeedLLMEngine(server_args()) + await engine.start(worker_id=1) + context = SimpleNamespace(id=lambda: "stream-cancel") + stream = engine.generate({"token_ids": [10]}, context) + try: + assert (await anext(stream))["token_ids"] == [20] + await engine.abort(context) + chunks = [chunk async for chunk in stream] + assert chunks[-1]["finish_reason"] == "cancelled" + assert engine._active_rids_by_context == {} + finally: + await stream.aclose() + await engine.cleanup() + + +async def test_kv_replay_rejected_before_native_start(native_engine, tmp_path): + _, constructor = native_engine + engine = TokenspeedLLMEngine( + server_args( + kv_events_config=json.dumps( + { + "enable_kv_cache_events": True, + "endpoint": f"ipc://{tmp_path}/events", + "replay_endpoint": f"ipc://{tmp_path}/replay", + } + ) + ) + ) + try: + with pytest.raises(ValueError, match="replay_endpoint"): + await engine.start(worker_id=1) + constructor.assert_not_called() + finally: + await engine.cleanup() diff --git a/components/src/dynamo/tokenspeed/tests/test_tokenspeed_llm_engine.py b/components/src/dynamo/tokenspeed/tests/test_tokenspeed_llm_engine.py index c7957fc647da..e915d378a57c 100644 --- a/components/src/dynamo/tokenspeed/tests/test_tokenspeed_llm_engine.py +++ b/components/src/dynamo/tokenspeed/tests/test_tokenspeed_llm_engine.py @@ -179,6 +179,41 @@ def test_convert_output_to_chunk_normalizes_abort_finish_reason(): assert chunk["finish_reason"] == "cancelled" +@pytest.mark.parametrize("as_object", [False, True]) +@pytest.mark.parametrize( + "message,err_type", + [ + ("PD/EPD remote transfer failed or timed out", "UnknownError"), + ("PD/EPD remote transfer failed or timed out", 522), + ("kv transfer failed", None), + ("remote peer disconnected", 521), + ("NaN in logits", 523), + ], +) +def test_convert_output_preserves_native_failure(as_object, message, err_type): + reason = { + "type": "abort", + "message": message, + "err_type": err_type, + } + if as_object: + + class FinishReason: + def to_json(self): + return reason + + finish_reason = FinishReason() + else: + finish_reason = reason + with pytest.raises(RuntimeError, match=message): + convert_output_to_chunk( + { + "output_ids": [], + "meta_info": {"finish_reason": finish_reason}, + } + ) + + def test_validate_single_choice_sampling_rejects_n_greater_than_one(): with pytest.raises(InvalidArgument, match="n=2"): _validate_single_choice_sampling( @@ -209,3 +244,32 @@ def test_completion_delta_output_preserves_later_token_delta(): assert emitted == 2 assert delta_out["output_ids"] == [100] + + +@pytest.mark.parametrize("as_object", [False, True]) +@pytest.mark.parametrize( + "message", + [ + "request was cancelled", + "generation stopped by content filter", + "Abort before prefill", + "AbortReq from client", + "Aborted by AbortReq.", + "Aborted by pause", + ], +) +@pytest.mark.parametrize("err_type", [None, 522]) +def test_convert_output_preserves_native_cancellation(as_object, message, err_type): + reason = {"type": "abort", "message": message, "err_type": err_type} + + class FinishReason: + def to_json(self): + return reason + + chunk = convert_output_to_chunk( + { + "output_ids": [], + "meta_info": {"finish_reason": FinishReason() if as_object else reason}, + } + ) + assert chunk["finish_reason"] == "cancelled" diff --git a/components/src/dynamo/triton/classification.py b/components/src/dynamo/triton/classification.py new file mode 100644 index 000000000000..c46e3fdfd57d --- /dev/null +++ b/components/src/dynamo/triton/classification.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KServe classification (top-K) for Triton model outputs. + +A requested output carrying the KServe ``classification`` parameter asks for the +output's top-K classes as ``":[: