diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000..03d52cc --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,5 @@ +# Global owners for the entire repo +* @hickeyma @vMaroon @elevran + +# vLLM extensions +/vllm/ @hickeyma diff --git a/.github/ISSUE_TEMPLATE/bug.yaml b/.github/ISSUE_TEMPLATE/bug.yaml new file mode 100644 index 0000000..787485b --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug.yaml @@ -0,0 +1,62 @@ +name: Bug Report +description: Create a report to help us improve +type: bug +body: + - type: markdown + attributes: + value: | + Thanks for taking the time to fill out this bug report! + + Please do **not** use this form to report a security vulnerability — + follow [SECURITY.md](https://github.com/llm-d/llm-d-api-extensions/blob/main/SECURITY.md) instead. + - type: dropdown + id: component + attributes: + label: Component + description: First, tell us which component we should focus on + options: + - I don't know + - vLLM / server introspection - config + - vLLM / server introspection - devices + - vLLM / server introspection - kv-cache + - Packaging and install + - Documentation + - CI + default: 0 + validations: + required: true + - type: textarea + id: describe + attributes: + label: Describe the bug + description: A clear and concise description of what the bug is + placeholder: Tell us what's wrong + validations: + required: true + - type: textarea + id: reproducer + attributes: + label: Steps to reproduce + description: Steps to reproduce the behavior + placeholder: I did (A), then (B) and then I saw (C) error + validations: + required: true + - type: textarea + id: environment + attributes: + label: Environment + description: How the server was launched and what it is running + placeholder: | + - Inference server and version (e.g. vLLM 0.11.0) + - llm-d-api-extensions-vllm version + - VLLM_PLUGINS value and any relevant `vllm serve` flags + - Python version, platform, accelerator + render: markdown + validations: + required: true + - type: textarea + id: extra + attributes: + label: Additional context or screenshots + description: Add any other context about the problem here + placeholder: Anything else you want to say to the report, attach screenshots, this is the place diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 0000000..64865cf --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,8 @@ +blank_issues_enabled: true +contact_links: + - name: Security vulnerability + url: https://github.com/llm-d/llm-d-api-extensions/blob/main/SECURITY.md + about: Please report vulnerabilities privately, not as a public issue. + - name: llm-d project discussion + url: https://github.com/llm-d/llm-d/issues + about: For questions about llm-d as a whole rather than these extensions. diff --git a/.github/ISSUE_TEMPLATE/feature.yaml b/.github/ISSUE_TEMPLATE/feature.yaml new file mode 100644 index 0000000..9727d98 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature.yaml @@ -0,0 +1,41 @@ +name: Feature Request +description: Suggest a new extension, endpoint or improvement +type: feature +body: + - type: markdown + attributes: + value: | + Thanks for taking the time to suggest an improvement! + - type: textarea + id: problem + attributes: + label: What problem does this solve? + description: Describe the use case. What does llm-d need that it cannot get today? + placeholder: Routing decisions need X, which is not exposed by any current endpoint + validations: + required: true + - type: textarea + id: proposal + attributes: + label: Proposed solution + description: | + If you are proposing an endpoint, sketch the route and an example response. + Note whether it needs an engine, and any server flags it would require. + validations: + required: true + - type: dropdown + id: server + attributes: + label: Inference server + options: + - vLLM + - Another server (please describe below) + - Not server specific + default: 0 + validations: + required: true + - type: textarea + id: alternatives + attributes: + label: Alternatives considered + description: Including whether this belongs upstream in the inference server instead diff --git a/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md b/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md new file mode 100644 index 0000000..b1523a2 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md @@ -0,0 +1,39 @@ +--- +name: Pull request +about: Create a pull request +--- + +## Description + +Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change. + +Fixes # (issue) + +## Type of change + +Please delete options that are not relevant. + +- [ ] Bug fix (non-breaking change which fixes an issue) +- [ ] New feature (non-breaking change which adds functionality) +- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) +- [ ] This change requires a documentation update + +## How Has This Been Tested? + +Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. + +### Test configuration + +- Inference server and version: +- Python version and platform: +- `VLLM_PLUGINS` value and relevant `vllm serve` flags: + +## Checklist + +- [ ] My changes follow the style guidelines of this project +- [ ] I have performed a self-review of my own changes +- [ ] I have signed off my commits (`git commit -s`, per the [DCO](https://developercertificate.org/)) +- [ ] `ruff check .` and `ruff format --check .` pass +- [ ] `pytest tests/ -v` passes +- [ ] Any new endpoint is opt-in, namespaced under `/plugins/llm-d-/...` and documented in `vllm/README.md` +- [ ] I have updated the documentation accordingly diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..978a717 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,26 @@ +# Dependabot configuration for llm-d-api-extensions +# Based on the canonical config from llm-d/llm-d-infra + +version: 2 +updates: + + # GitHub Actions dependencies + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + labels: + - "dependencies" + - "release-note-none" + commit-message: + prefix: "deps(actions)" + + # Python dependencies for the vLLM extensions + - package-ecosystem: "pip" + directory: "/vllm" + schedule: + interval: "weekly" + labels: + - "dependencies" + commit-message: + prefix: "deps(pip)" diff --git a/.github/workflows/ci-lint.yaml b/.github/workflows/ci-lint.yaml new file mode 100644 index 0000000..f5b6959 --- /dev/null +++ b/.github/workflows/ci-lint.yaml @@ -0,0 +1,37 @@ +name: CI - Lint and Format + +on: + pull_request: + branches: + - main + push: + branches: + - main + +permissions: + contents: read + +jobs: + ruff: + runs-on: ubuntu-latest + defaults: + run: + working-directory: vllm + steps: + - name: Checkout source + uses: actions/checkout@v7 + + # Checks only, never modifies files in CI. + - name: Run Ruff lint check + uses: astral-sh/ruff-action@v3 + with: + version: "0.15.11" + src: "./vllm" + args: "check --output-format=full" + + - name: Run Ruff format check + uses: astral-sh/ruff-action@v3 + with: + version: "0.15.11" + src: "./vllm" + args: "format --check --diff" diff --git a/.github/workflows/ci-test.yaml b/.github/workflows/ci-test.yaml new file mode 100644 index 0000000..3021783 --- /dev/null +++ b/.github/workflows/ci-test.yaml @@ -0,0 +1,47 @@ +name: CI - Unit Tests + +on: + pull_request: + branches: + - main + push: + branches: + - main + +permissions: + contents: read + +jobs: + unit-tests: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12"] + defaults: + run: + working-directory: vllm + steps: + - name: Checkout source + uses: actions/checkout@v7 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + + # `vllm` is a heavyweight GPU oriented dependency and the unit tests only + # need a handful of symbols from it (`vllm.logger`, `vllm.tasks`), so + # install it on its own rather than pulling the whole project dependency + # set through `pip install -e .`. + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install --extra-index-url https://download.pytorch.org/whl/cpu vllm + pip install fastapi 'pydantic>=2' pytest pytest-asyncio httpx + + # The opt-in e2e tests skip themselves unless RUN_VLLM_E2E is set. They + # download model weights and launch a real server which CI does not do. + - name: Run unit tests + run: pytest tests/ -v diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a7c4bf5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,41 @@ +# Python bytecode +__pycache__/ +*.py[cod] +*$py.class + +# Packaging / build artifacts +build/ +dist/ +*.egg-info/ +*.egg +.eggs/ +wheels/ + +# Virtual environments +.venv/ +venv/ +ENV/ + +# Test and tool caches +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +.coverage +.coverage.* +htmlcov/ +coverage.xml +.tox/ +.nox/ + +# Editors / IDEs +.idea/ +*.swp +*.swo +.vscode/ + +# OS +.DS_Store +Thumbs.db + +# Local model / server artifacts pulled down by the opt-in e2e tests +*.log diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 0000000..576122c --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,34 @@ +## Code of Conduct and Covenant + +## Our Pledge + +In the interest of fostering an open and welcoming environment, we as contributors and maintainers pledge to making participation in our project and our community a harassment-free experience for everyone, regardless of age, body size, disability, ethnicity, sex characteristics, gender identity and expression, level of experience, education, socio-economic status, nationality, personal appearance, race, religion, or sexual identity and orientation. + +## Our Standards + +Examples of behavior that contributes to creating a positive environment include: + +* Using welcoming and inclusive language +* Being respectful of differing viewpoints and experiences +* Gracefully accepting constructive criticism +* Focusing on what is best for the community +* Showing empathy towards other community members + +Examples of unacceptable behavior by participants include: + +* The use of sexualized language or imagery and unwelcome sexual attention or advances +* Trolling, insulting/derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or electronic address, without explicit permission +* Other conduct which could reasonably be considered inappropriate in a professional setting + +## Our Responsibilities + +Project maintainers have the right and responsibility to remove, edit, or reject comments, commits, code, wiki edits, issues, and other contributions that are not aligned to this Code of Conduct, or to ban temporarily or permanently any contributor for other behaviors that they deem inappropriate, threatening, offensive, or harmful. + +## Scope + +This Code of Conduct applies both within project spaces and in public spaces when an individual is representing the project or its community. Examples of representing a project or community include using an official project e-mail address, posting via an official social media account, or acting as an appointed representative at an online or offline event. Representation of a project may be further defined and clarified by project maintainers.Attribution +This Code of Conduct is adapted from the Contributor Covenant, version 1.4, available at + +For answers to common questions about this code of conduct, see diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..97dd7f5 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,74 @@ +# Contributing to llm-d-api-extensions + +Thank you for your interest in contributing. This repo hosts API extensions that llm-d loads into inference server. Contributions of all kinds — bug reports, docs, extensions for new servers — are welcome. + +Please also read the [Code of Conduct](CODE_OF_CONDUCT.md). Maintainers are listed in [MAINTAINERS.md](MAINTAINERS.md) and approvers/reviewers in [OWNERS](OWNERS). + +## Scope + +Extensions in this repo integrate with an inference server through that server's own documented extension points. A change that requires forking or patching the inference server does not belong here. Go upstream to the server instead. + +The top level of the repo is inference server agnostic. Server specific code lives in a directory named for that server (for example, [`vllm/`](vllm/)). + +## Reporting issues + +Before filing, search the [issue tracker](https://github.com/llm-d/llm-d-api-extensions/issues) for an existing report. When you do file one, include the inference server version, how the server was launched (the `VLLM_PLUGINS` value and relevant flags), what you expected, and what happened instead. + +Do **not** open a public issue for a security vulnerability. Follow [SECURITY.md](SECURITY.md). + +## Development + +The vLLM extensions are a standard Python package. From `vllm/`: + +```bash +pip install -e .[test] + +# fast unit tests: no engine, no model weights +pytest tests/ -v + +# lint and format, exactly what CI runs +ruff check . +ruff format --check . +``` + +Formatting and linting are enforced in CI by [ruff](https://docs.astral.sh/ruff/). The configuration lives in [`vllm/ruff.toml`](vllm/ruff.toml). Run `ruff format .` to fix formatting and `ruff check --fix .` for the auto-fixable lint rules before you push. + +The end-to-end tests launch a real vLLM server and download model weights, so they are opt-in and are not run by CI: + +```bash +RUN_VLLM_E2E=1 pytest tests/server_introspection/test_e2e.py -v +``` + +### Adding an endpoint + +New endpoints should: + +- Be **opt-in.** Register as a separate `vllm.endpoint_plugins` entry point so operators enable exactly what they want through `VLLM_PLUGINS`. +- **Namespace their routes** under `/plugins/llm-d-/...`. vLLM does not enforce route conflicts and a later route silently shadows an earlier one. +- Keep the entry point key identical to the plugin class's `name` attribute, since that key is what operators put in `VLLM_PLUGINS`. +- Set `required_tasks` honestly. Use `None` only if the endpoint genuinely needs no engine and return `503` rather than crashing when a dependency is absent. +- **Degrade, not crash.** If the endpoint depends on a vLLM API that is not in a released build yet, feature detect it and fall back. +- Ship a Pydantic response model and unit tests that need neither an engine nor model weights. +- Be documented in [`vllm/README.md`](vllm/README.md) with an example response. + +## Submitting a pull request + +1. **Discuss first for anything significant.** Opening an issue before a large change avoids rework and makes sure the direction fits llm-d's needs. +2. Fork the repo and work on a branch. +3. **Sign off your commits.** This project requires the [Developer Certificate of Origin](https://developercertificate.org/); use `git commit -s`. CI rejects unsigned commits. +4. Make sure `pytest`, `ruff check` and `ruff format --check` all pass locally. +5. Update the docs alongside the code. A new endpoint without a README entry is incomplete. +6. Open the PR with a clear description of what changed and why, linking the issue it addresses. + +## Code review + +Pull requests are reviewed by the maintainers listed in [OWNERS](OWNERS). Reviews are done on a best effort basis. Maintainers may prioritise their own work. Addressing review feedback promptly is the fastest path to a merge. + +## Community + +- **Slack:** the llm-d developer Slack at [llm-d.slack.com](https://llm-d.slack.com) +- **Project:** [llm-d/llm-d](https://github.com/llm-d/llm-d) + +## License + +By contributing, you agree that your contributions will be licensed under the [Apache 2.0 License](LICENSE). Source files carry an `# SPDX-License-Identifier: Apache-2.0` header. Please add that on new files. diff --git a/MAINTAINERS.md b/MAINTAINERS.md new file mode 100644 index 0000000..840cfe1 --- /dev/null +++ b/MAINTAINERS.md @@ -0,0 +1,19 @@ +# llm-d-api-extensions Maintainers + +Maintainers own the direction of this repo, review and approve pull requests and cut releases. See [CONTRIBUTING.md](CONTRIBUTING.md) for how to get a change reviewed. + +## Maintainers + +| Name | Employer | GitHub ID | +|----------------|----------|-------------| +| Martin Hickey | IBM | @hickeyma | +| Maroon Ayoub | IBM | @vMaroon | +| Etai Lev-Ran | IBM | @elevran | + +The machine readable approver and reviewer lists used by automation live in [OWNERS](OWNERS). + +## Escalation + +For project wide questions, see the [llm-d maintainers](https://github.com/llm-d/llm-d/blob/main/MAINTAINERS.md) and [SIGS.md](https://github.com/llm-d/llm-d/blob/main/SIGS.md). + +To report a security vulnerability, follow [SECURITY.md](SECURITY.md). Do NOT open a public issue. diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..b9d6972 --- /dev/null +++ b/Makefile @@ -0,0 +1,44 @@ +# llm-d-api-extensions +# +# Convenience targets that wrap the per server packages. Each package is also a +# standard Python project you can drive directly from its own directory. + +VLLM_DIR := vllm +RUFF_VERSION := 0.15.11 + +.DEFAULT_GOAL := help + +.PHONY: help +help: ## Show this help + @grep -hE '^[a-zA-Z0-9_-]+:.*?## ' $(MAKEFILE_LIST) \ + | awk 'BEGIN {FS = ":.*?## "}; {printf " \033[36m%-16s\033[0m %s\n", $$1, $$2}' + +.PHONY: install +install: ## Install the vLLM extensions with test extras (editable) + pip install -e '$(VLLM_DIR)[test]' + +.PHONY: lint +lint: ## Run ruff lint and format checks (what CI runs) + cd $(VLLM_DIR) && ruff check . && ruff format --check --diff . + +.PHONY: fix +fix: ## Apply ruff auto-fixes and formatting + cd $(VLLM_DIR) && ruff check --fix . && ruff format . + +.PHONY: test +test: ## Run the fast unit tests (no engine, no model weights) + cd $(VLLM_DIR) && pytest tests/ -v + +.PHONY: test-e2e +test-e2e: ## Run the opt-in e2e tests (downloads model weights, launches a server) + cd $(VLLM_DIR) && RUN_VLLM_E2E=1 pytest tests/server_introspection/test_e2e.py -v + +.PHONY: build +build: ## Build the vLLM extensions distribution + cd $(VLLM_DIR) && python -m build + +.PHONY: clean +clean: ## Remove build and test artifacts + rm -rf $(VLLM_DIR)/build $(VLLM_DIR)/dist $(VLLM_DIR)/*.egg-info + find . -name __pycache__ -type d -prune -exec rm -rf {} + + rm -rf $(VLLM_DIR)/.pytest_cache $(VLLM_DIR)/.ruff_cache diff --git a/OWNERS b/OWNERS new file mode 100644 index 0000000..3c65d5c --- /dev/null +++ b/OWNERS @@ -0,0 +1,9 @@ +# Project maintainers +approvers: +- hickeyma +- elevran + +# Project reviewers +reviewers: +- hickeyma +- elevran diff --git a/README.md b/README.md index 002c06f..f7de6c0 100644 --- a/README.md +++ b/README.md @@ -1 +1,43 @@ -# llm-d-api-extensions \ No newline at end of file +# llm-d API Extensions + +API extensions for inference servers, packaged for [llm-d](https://github.com/llm-d/llm-d). + +llm-d needs data out of the inference server. Information like how it was launched, what hardware it sits on, how big its KV cache turned out to be. This data is then used to make routing and scheduling decisions. Rather than adding those endpoints to an inference server's core, this repo hosts them as **extensions** that are installed alongside the server and loaded through the server's own extension points. + +For vLLM, that means using the [endpoint plugin framework](https://docs.vllm.ai/en/latest/design/endpoint_plugins/). The framework is an extension mechanism where a Python package is registered as a `vllm.endpoint_plugins` entry point and vLLM attaches its routes at startup. No fork, no patch and no core change required on the server. + +## Layout + +The repo is inference server agnostic at the top level. Each supported server gets its own directory holding an independently installable package: + +| Directory | Package | Server | +|---|---|---| +| [`vllm/`](vllm/) | `llm-d-api-extensions-vllm` | [vLLM](https://github.com/vllm-project/vllm) | + +Support for another server (for example SGLang) would land as a sibling directory once that server exposes a comparable extension point. + +## Extensions + +| Extension | Server | Endpoints | +|---|---|---| +| [Server introspection](vllm/README.md) | vLLM | `GET /plugins/llm-d-server-introspection/{config,devices,kv-cache}` | + +## Quick start + +```bash +pip install ./vllm + +VLLM_PLUGINS=llm_d_server_introspection_config vllm serve + +curl http://localhost:8000/plugins/llm-d-server-introspection/config +``` + +Extensions are opt-in. vLLM only loads the plugins named in `VLLM_PLUGINS`, so installing this package does not change the behaviour of a server that does not ask for it. See [`vllm/README.md`](vllm/README.md) for the full endpoint reference, the response schemas and the flags each extension needs. + +## Contributing + +Contributions are welcome. See [CONTRIBUTING.md](CONTRIBUTING.md) for the development workflow and [CODE_OF_CONDUCT.md](CODE_OF_CONDUCT.md) for community expectations. Maintainers are listed in [MAINTAINERS.md](MAINTAINERS.md). To report a vulnerability, follow [SECURITY.md](SECURITY.md) rather than opening an issue. + +## License + +Apache 2.0 — see [LICENSE](LICENSE). diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000..4a35bdc --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,3 @@ +# Security Reporting and Policy + +The llm-d project has [a common process and policy that can be found here](https://github.com/llm-d/llm-d/blob/main/SECURITY.md). \ No newline at end of file diff --git a/vllm/README.md b/vllm/README.md new file mode 100644 index 0000000..3923c54 --- /dev/null +++ b/vllm/README.md @@ -0,0 +1,174 @@ +# llm-d API extensions for vLLM + +`llm-d-api-extensions-vllm` packages llm-d's vLLM extensions. They integrate through vLLM's [endpoint plugin framework](https://docs.vllm.ai/en/latest/design/endpoint_plugins/) and `--worker-extension-cls`, so nothing here requires a change to vLLM core. + +## Server introspection + +Three read only endpoints describing the server llm-d is talking to. Each is a separate plugin so a deployment can enable only what it needs. + +### `GET /plugins/llm-d-server-introspection/config` + +Operator supplied, config time values of how the server was launched. Nothing profiled or derived from model internals. Requires no engine, so it also works on the CPU only render server. + +```jsonc +{ + "model": { + "name": "llama3", + "served_names": ["llama3"], + "dtype": "bfloat16", + "quantization": null, + "max_model_len": 8192 + }, + "kv_cache": { + "gpu_memory_utilization": 0.9, + "dtype": "bfloat16", + "enable_prefix_caching": true + }, + "scheduler": { + "max_num_seqs": 256, + "max_num_batched_tokens": 8192, + "enable_chunked_prefill": true, + "policy": "fcfs" + }, + "parallelism": { + "tensor_parallel_size": 2, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "data_parallel_rank": 0 + }, + "features": { + "speculative_decoding": false, + "lora": true, + "hma": true + }, + "kv_transfer": { + "kv_connector": "NixlConnector", + "kv_role": "kv_both", + "kv_connector_module_path": null, + "kv_buffer_device": "cuda", + "kv_buffer_size": 1000000000.0, + "kv_ip": "127.0.0.1", + "kv_port": 14579, + "kv_parallel_size": 1, + "kv_rank": null, + "engine_id": "engine-0", + "extra_config": {}, + "nixl_side_channel_host": "localhost", + "nixl_side_channel_port": 5600 + } +} +``` + +> [!NOTE] +> `kv_transfer` is `null` when no disaggregation/KV offloading connector is configured (`vllm_config.kv_transfer_config` is `None`). `nixl_side_channel_host`/`nixl_side_channel_port` are only populated when `kv_connector` is `"NixlConnector"`. They report the env derived *base* host/port (`VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). `NixlConnector` derives the actual per rank bound port as `base_port + rank_offset` inside the worker. + +> [!WARNING] +> `kv_connector_extra_config` is free form and operator controlled, so it may contain sensitive data. `extra_config` therefore reports its keys with every value replaced by `""`. Set `LLM_D_INTROSPECTION_EXPOSE_KV_EXTRA_CONFIG=1` in the server's environment to return the values verbatim. + +### `GET /plugins/llm-d-server-introspection/devices` + +Per rank hardware properties, gathered once at startup via `collective_rpc` and cached for the server's lifetime. Requires an engine (`503` on the CPU only render server) and the worker side `get_device_properties` method installed by `device_worker_ext.DeviceInfoWorkerExtension` via `--worker-extension-cls`. Without the worker extension, the server still starts and this endpoint returns `503`. + +```jsonc +{ + "devices": [ + { + "rank": 0, + "name": "A100-PCIE-40GB", + "total_memory_bytes": 42949672960, + "compute_capability": { "major": 8, "minor": 0 }, + "num_compute_units": 108 + } + ] +} +``` + +### `GET /plugins/llm-d-server-introspection/kv-cache` + +Post profiling KV cache capacity and attention group structure, gathered once at startup and cached for the server's lifetime. Requires an engine (`503` on the CPU only render server). `groups` is a discriminated union keyed on `kind`, correctly representing hybrid models with multiple attention/mamba groups. + +Against a vLLM build without `EngineClient.get_kv_cache_group_metadata()` ([vllm-project/vllm#48121](https://github.com/vllm-project/vllm/pull/48121)), this plugin falls back to capacity only fields read directly off `vllm_config.cache_config` and returns an empty `groups`. + +```jsonc +{ + "kv_cache_size_tokens": 393216, + "max_concurrency": 48.0, + "num_gpu_blocks": 24576, + "num_cpu_blocks": 0, + "groups": [ + { + "group_id": 0, + "kind": "full_attention", + "layer_count": 32, + "layer_names": ["model.layers.0.self_attn", "..."], + "block_size": 16, + "page_size_bytes": 131072, + "num_kv_heads": 8, + "head_size": 128, + "head_size_v": 128, + "dtype": "bfloat16", + "sliding_window": null, + "attention_chunk_size": null + } + ] +} +``` + +## Install + +```bash +pip install . + +# or for development +pip install -e . +``` + +## Enable + +Endpoint plugins load only when explicitly named in `VLLM_PLUGINS` (off by default), so installing this package alone changes nothing: + +```bash +# config only +VLLM_PLUGINS=llm_d_server_introspection_config vllm serve + +# devices only +VLLM_PLUGINS=llm_d_server_introspection_devices vllm serve \ + --worker-extension-cls llm_d_api_extensions_vllm.server_introspection.device_worker_ext.DeviceInfoWorkerExtension + +# kv cache only +VLLM_PLUGINS=llm_d_server_introspection_kv_cache vllm serve + +# all three +VLLM_PLUGINS=llm_d_server_introspection_config,llm_d_server_introspection_devices,llm_d_server_introspection_kv_cache \ + vllm serve \ + --worker-extension-cls llm_d_api_extensions_vllm.server_introspection.device_worker_ext.DeviceInfoWorkerExtension +``` + +```bash +curl http://localhost:8000/plugins/llm-d-server-introspection/config +curl http://localhost:8000/plugins/llm-d-server-introspection/devices +curl http://localhost:8000/plugins/llm-d-server-introspection/kv-cache +``` + +## Test + +```bash +pip install -e .[test] + +# fast unit tests (schema + FastAPI route, no engine, no real model) +pytest tests/ -v + +# real server e2e (downloads model weights, launches a subprocess server) +RUN_VLLM_E2E=1 pytest tests/server_introspection/test_e2e.py -v +``` + +Lint and format with [ruff](https://docs.astral.sh/ruff/), matching what CI runs: + +```bash +ruff check . +ruff format --check . +``` + +## Adding an extension + +Add a subpackage under `llm_d_api_extensions_vllm/`, register its plugin class as a `vllm.endpoint_plugins` entry point in [`pyproject.toml`](pyproject.toml) and namespace its routes under `/plugins/llm-d-/...`. Keep the entry point key identical to the plugin class's `name` attribute. The key is what operators put in `VLLM_PLUGINS`. diff --git a/vllm/llm_d_api_extensions_vllm/__init__.py b/vllm/llm_d_api_extensions_vllm/__init__.py new file mode 100644 index 0000000..db260c3 --- /dev/null +++ b/vllm/llm_d_api_extensions_vllm/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +"""llm-d API extensions for the vLLM inference server. + +Each subpackage is one extension family. Extensions integrate with vLLM +through documented extension points (`vllm.endpoint_plugins` entry points, +`--worker-extension-cls`) rather than by patching vLLM core. + +- `server_introspection`: read-only `GET /plugins/llm-d-server-introspection/*` + endpoints exposing how the server was launched, what hardware it is on and + what its KV cache looks like after profiling. +""" diff --git a/vllm/llm_d_api_extensions_vllm/server_introspection/__init__.py b/vllm/llm_d_api_extensions_vllm/server_introspection/__init__.py new file mode 100644 index 0000000..9881313 --- /dev/null +++ b/vllm/llm_d_api_extensions_vllm/server_introspection/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/vllm/llm_d_api_extensions_vllm/server_introspection/common.py b/vllm/llm_d_api_extensions_vllm/server_introspection/common.py new file mode 100644 index 0000000..5fd8f1c --- /dev/null +++ b/vllm/llm_d_api_extensions_vllm/server_introspection/common.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Helpers shared by the server introspection endpoint plugins.""" + +import logging + +from fastapi import HTTPException, Request +from vllm.logger import init_logger +from vllm.tasks import GENERATION_TASKS, POOLING_TASKS + +# Tasks for plugins that need an engine. Excludes the `render` frontend task +# because there is nothing to introspect on the CPU only render server. +ENGINE_TASKS: tuple[str, ...] = GENERATION_TASKS + POOLING_TASKS + +_UNSET = object() + + +def get_logger(name: str) -> logging.Logger: + # "vllm." prefix required. vLLM's default logging config only attaches a + # handler to the "vllm" logger tree (propagate=False), so a bare __name__ + # logger has no handler anywhere and silently drops every message. + return init_logger(f"vllm.{name}") + + +def no_engine_detail(endpoint: str) -> str: + return f"{endpoint} requires an engine, which this server does not have" + + +def cached_response(request: Request, state_attr: str, unavailable_detail: str): + """Return the response `init_state` cached on `app.state.`. + + Raises 500 when `init_state` never ran and 503 with `unavailable_detail` + when it ran but cached `None` (the endpoint can't serve on this server). + """ + response = getattr(request.app.state, state_attr, _UNSET) + if response is _UNSET: + raise HTTPException( + status_code=500, + detail=f"{state_attr.removesuffix('_response')} plugin state was never initialized", + ) + if response is None: + raise HTTPException(status_code=503, detail=unavailable_detail) + return response diff --git a/vllm/llm_d_api_extensions_vllm/server_introspection/config_plugin.py b/vllm/llm_d_api_extensions_vllm/server_introspection/config_plugin.py new file mode 100644 index 0000000..818f859 --- /dev/null +++ b/vllm/llm_d_api_extensions_vllm/server_introspection/config_plugin.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +"""`vllm.endpoint_plugins` entry point: `GET /plugins/llm-d-server-introspection/config`. + +Operator supplied, config time values of how the server was launched. +Nothing profiled or derived from model internals. + +`required_tasks` is `None`. This plugin needs no engine, so it is also +eligible on the CPU only render server (`init_state` receives +`engine_client=None` there which is fine since it is never touched). +""" + +import os +from argparse import Namespace +from typing import TYPE_CHECKING + +from fastapi import APIRouter, FastAPI, HTTPException, Request +from starlette.datastructures import State + +from .common import get_logger +from .schemas import ( + FeaturesInfo, + KVCacheInfo, + KVTransferInfo, + ModelInfo, + ParallelismInfo, + SchedulerInfo, + ServerConfigResponse, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.engine.protocol import EngineClient + +logger = get_logger(__name__) + + +# `kv_connector_extra_config` is free form and operator controlled, so it may +# hold credentials. Values are redacted (keys kept) unless this env var is "1". +EXPOSE_EXTRA_CONFIG_ENV = "LLM_D_INTROSPECTION_EXPOSE_KV_EXTRA_CONFIG" +REDACTED = "" + + +def _redact_extra_config(extra_config: dict | None) -> dict: + extra_config = extra_config or {} + if os.environ.get(EXPOSE_EXTRA_CONFIG_ENV) == "1": + return dict(extra_config) + return {key: REDACTED for key in extra_config} + + +def _dtype_str(dtype: object) -> str: + return str(dtype).removeprefix("torch.") + + +def _resolve_kv_cache_dtype(cache_dtype: str, model_dtype_str: str) -> str: + # "auto" must be resolved to the concrete model dtype, never left as "auto". + return model_dtype_str if cache_dtype == "auto" else cache_dtype + + +def _build_kv_transfer_info(vllm_config: "VllmConfig") -> KVTransferInfo | None: + kv_transfer_cfg = getattr(vllm_config, "kv_transfer_config", None) + if kv_transfer_cfg is None: + return None + + nixl_host = None + nixl_port = None + if kv_transfer_cfg.kv_connector == "NixlConnector": + from vllm import envs + + nixl_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST + nixl_port = envs.VLLM_NIXL_SIDE_CHANNEL_PORT + + return KVTransferInfo( + kv_connector=kv_transfer_cfg.kv_connector, + kv_role=kv_transfer_cfg.kv_role, + kv_connector_module_path=kv_transfer_cfg.kv_connector_module_path, + kv_buffer_device=kv_transfer_cfg.kv_buffer_device, + kv_buffer_size=kv_transfer_cfg.kv_buffer_size, + kv_ip=kv_transfer_cfg.kv_ip, + kv_port=kv_transfer_cfg.kv_port, + kv_parallel_size=kv_transfer_cfg.kv_parallel_size, + kv_rank=kv_transfer_cfg.kv_rank, + engine_id=kv_transfer_cfg.engine_id, + extra_config=_redact_extra_config(kv_transfer_cfg.kv_connector_extra_config), + nixl_side_channel_host=nixl_host, + nixl_side_channel_port=nixl_port, + ) + + +def _build_response(vllm_config: "VllmConfig", served_names: list[str]) -> ServerConfigResponse: + model_cfg = vllm_config.model_config + cache_cfg = vllm_config.cache_config + scheduler_cfg = vllm_config.scheduler_config + parallel_cfg = vllm_config.parallel_config + spec_cfg = vllm_config.speculative_config + lora_cfg = vllm_config.lora_config + + model_dtype_str = _dtype_str(model_cfg.dtype) + + return ServerConfigResponse( + model=ModelInfo( + name=served_names[0], + served_names=served_names, + dtype=model_dtype_str, + quantization=( + str(model_cfg.quantization) if model_cfg.quantization is not None else None + ), + max_model_len=model_cfg.max_model_len, + ), + kv_cache=KVCacheInfo( + gpu_memory_utilization=cache_cfg.gpu_memory_utilization, + dtype=_resolve_kv_cache_dtype(cache_cfg.cache_dtype, model_dtype_str), + enable_prefix_caching=cache_cfg.enable_prefix_caching, + ), + scheduler=SchedulerInfo( + max_num_seqs=scheduler_cfg.max_num_seqs, + max_num_batched_tokens=scheduler_cfg.max_num_batched_tokens, + enable_chunked_prefill=scheduler_cfg.enable_chunked_prefill, + policy=str(scheduler_cfg.policy), + ), + parallelism=ParallelismInfo( + tensor_parallel_size=parallel_cfg.tensor_parallel_size, + pipeline_parallel_size=parallel_cfg.pipeline_parallel_size, + data_parallel_size=parallel_cfg.data_parallel_size, + data_parallel_rank=parallel_cfg.data_parallel_rank, + ), + features=FeaturesInfo( + speculative_decoding=spec_cfg is not None, + lora=lora_cfg is not None, + hma=not bool(getattr(scheduler_cfg, "disable_hybrid_kv_cache_manager", False)), + ), + kv_transfer=_build_kv_transfer_info(vllm_config), + ) + + +class ServerConfigPlugin: + name = "llm_d_server_introspection_config" + required_tasks: tuple[str, ...] | None = None + + def attach_router(self, app: FastAPI) -> None: + router = APIRouter() + + @router.get( + "/plugins/llm-d-server-introspection/config", + response_model=ServerConfigResponse, + ) + async def get_server_config(raw_request: Request) -> ServerConfigResponse: + response = getattr(raw_request.app.state, "server_config_response", None) + if response is None: + raise HTTPException( + status_code=500, + detail="server_config plugin state was never initialized", + ) + return response + + app.include_router(router) + + async def init_state( + self, engine_client: "EngineClient | None", state: State, args: Namespace + ) -> None: + vllm_config: VllmConfig = state.vllm_config + if args.served_model_name: + served_names = list(args.served_model_name) + else: + served_names = [vllm_config.model_config.served_model_name] + response = _build_response(vllm_config, served_names) + state.server_config_response = response + logger.info( + "config plugin initialized: served_names=%s kv_transfer=%s", + served_names, + "configured" if response.kv_transfer is not None else "none", + ) diff --git a/vllm/llm_d_api_extensions_vllm/server_introspection/device_worker_ext.py b/vllm/llm_d_api_extensions_vllm/server_introspection/device_worker_ext.py new file mode 100644 index 0000000..df90938 --- /dev/null +++ b/vllm/llm_d_api_extensions_vllm/server_introspection/device_worker_ext.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Worker extension class integrated into vLLM workers via `--worker-extension-cls`. + +Installs `get_device_properties` and the worker side method the `devices` +endpoint plugin (`devices_plugin.py`) reaches over `engine_client.collective_rpc`. +This is vLLM's lever for adding a `collective_rpc` worker method from +external code with no core edit. `--worker-extension-cls` mixes this class +into the concrete `Worker` subclass's bases (see +`vllm.v1.worker.worker_base.WorkerWrapperBase`), so `self` below is the real +worker instance with `self.rank` / `self.local_rank` already set. + +Enable with: + --worker-extension-cls llm_d_api_extensions_vllm.server_introspection\ + .device_worker_ext.DeviceInfoWorkerExtension +""" + +from typing import Any + +from .common import get_logger + +logger = get_logger(__name__) + +# Worker method name the `devices` plugin calls over `collective_rpc`. Shared +# so a rename breaks at import time rather than at RPC dispatch. +GET_DEVICE_PROPERTIES = "get_device_properties" + + +def _safe(fn: Any, *args: Any) -> Any: + # Some platforms (e.g. CPU) don't implement every `current_platform` + # device introspection method, Set to `None` instead of failing the + # whole `collective_rpc` call for every rank. + try: + return fn(*args) + except NotImplementedError: + return None + + +class DeviceInfoWorkerExtension: + def get_device_properties(self) -> dict: + from vllm.platforms import current_platform + + device_id = self.local_rank + capability = _safe(current_platform.get_device_capability, device_id) + properties = { + "rank": self.rank, + "name": _safe(current_platform.get_device_name, device_id), + "total_memory_bytes": _safe(current_platform.get_device_total_memory, device_id), + "compute_capability": ( + {"major": capability.major, "minor": capability.minor} + if capability is not None + else None + ), + "num_compute_units": _safe(current_platform.num_compute_units, device_id), + } + logger.info("device_worker_ext: rank=%d name=%s", self.rank, properties["name"]) + return properties diff --git a/vllm/llm_d_api_extensions_vllm/server_introspection/devices_plugin.py b/vllm/llm_d_api_extensions_vllm/server_introspection/devices_plugin.py new file mode 100644 index 0000000..9adc39b --- /dev/null +++ b/vllm/llm_d_api_extensions_vllm/server_introspection/devices_plugin.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +"""`vllm.endpoint_plugins` entry point: `GET /plugins/llm-d-server-introspection/devices`. + +Per rank hardware properties (name, memory, compute capability, compute +units) gathered once at startup via `collective_rpc` and cached for the +server's lifetime (device properties are static). + +Requires the worker side `get_device_properties` method installed by +`device_worker_ext.DeviceInfoWorkerExtension` (see that module's docstring +for the `--worker-extension-cls` flag). Without it, `collective_rpc` raises. +The plugin catches that at `init_state` time, logs a warning and serves `503` +with the reason rather than crashing the server. + +`required_tasks` excludes the `render` frontend task. This plugin needs an +engine (there is nothing to introspect on the CPU only render server). +""" + +from argparse import Namespace +from typing import TYPE_CHECKING + +from fastapi import APIRouter, FastAPI, Request +from starlette.datastructures import State + +from .common import ENGINE_TASKS, cached_response, get_logger, no_engine_detail +from .device_worker_ext import GET_DEVICE_PROPERTIES +from .schemas import ComputeCapability, DeviceInfo, DevicesResponse + +if TYPE_CHECKING: + from vllm.engine.protocol import EngineClient + +logger = get_logger(__name__) + +_NO_ENGINE_DETAIL = no_engine_detail("devices") +_NO_WORKER_EXT_DETAIL = ( + "devices requires the DeviceInfoWorkerExtension worker extension " + "(--worker-extension-cls), which is not installed on this server" +) + + +def _build_response(raw_devices: list[dict]) -> DevicesResponse: + devices = [] + for d in raw_devices: + cap = d.get("compute_capability") + devices.append( + DeviceInfo( + rank=d["rank"], + name=d["name"], + total_memory_bytes=d["total_memory_bytes"], + compute_capability=( + ComputeCapability(major=cap["major"], minor=cap["minor"]) + if cap is not None + else None + ), + num_compute_units=d["num_compute_units"], + ) + ) + return DevicesResponse(devices=devices) + + +class ServerDevicesPlugin: + name = "llm_d_server_introspection_devices" + required_tasks: tuple[str, ...] | None = ENGINE_TASKS + + def attach_router(self, app: FastAPI) -> None: + router = APIRouter() + + @router.get( + "/plugins/llm-d-server-introspection/devices", + response_model=DevicesResponse, + ) + async def get_devices(raw_request: Request) -> DevicesResponse: + unavailable_detail = getattr( + raw_request.app.state, "server_devices_unavailable_detail", _NO_ENGINE_DETAIL + ) + return cached_response(raw_request, "server_devices_response", unavailable_detail) + + app.include_router(router) + + async def init_state( + self, engine_client: "EngineClient | None", state: State, args: Namespace + ) -> None: + if engine_client is None: + state.server_devices_response = None + state.server_devices_unavailable_detail = _NO_ENGINE_DETAIL + return + # Uses DeviceInfoWorkerExtension::get_device_properties() extension to call vLLM engine + # worker methods + try: + raw_devices = await engine_client.collective_rpc(GET_DEVICE_PROPERTIES) + except Exception: + logger.warning( + "devices plugin: collective_rpc(%r) failed " + "(is --worker-extension-cls set to DeviceInfoWorkerExtension?)", + GET_DEVICE_PROPERTIES, + exc_info=True, + ) + state.server_devices_response = None + state.server_devices_unavailable_detail = _NO_WORKER_EXT_DETAIL + return + response = _build_response(raw_devices) + state.server_devices_response = response + logger.info("devices plugin initialized: devices=%d", len(response.devices)) diff --git a/vllm/llm_d_api_extensions_vllm/server_introspection/kv_cache_plugin.py b/vllm/llm_d_api_extensions_vllm/server_introspection/kv_cache_plugin.py new file mode 100644 index 0000000..7da75ce --- /dev/null +++ b/vllm/llm_d_api_extensions_vllm/server_introspection/kv_cache_plugin.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +"""`vllm.endpoint_plugins` entry point: `GET /plugins/llm-d-server-introspection/kv-cache`. + +Post profiling KV cache capacity plus attention group structure. Capacity +fields (`kv_cache_size_tokens`, `max_concurrency`, `num_gpu_blocks`, +`num_cpu_blocks`) are read directly off `vllm_config.cache_config` & +`model_config` which are already post profiling values by the time +`init_state` runs (`EngineCore` mutates the shared config in place during +`_initialize_kv_caches`). + +`groups` is sourced separately from `engine_client.get_kv_cache_group_metadata()` +once at startup and cached for the server's lifetime (immutable once +profiling has run). That method is not yet part of a released `EngineClient`, +it is proposed upstream in vllm-project/vllm#48121 which this plugin's +group schema (`kind` values are the `KVCacheSpecKind` strings that PR +serializes groups under) and response shape are ported from. Feature +detected via `hasattr` so this package still installs against a vLLM build +that predates that method which means `groups` are empty for this. + +`required_tasks` excludes the `render` frontend task. This plugin needs an +engine (there is nothing to introspect on the CPU only render server). +""" + +from argparse import Namespace +from typing import TYPE_CHECKING + +from fastapi import APIRouter, FastAPI, Request +from starlette.datastructures import State + +from .common import ENGINE_TASKS, cached_response, get_logger, no_engine_detail +from .schemas import ( + ChunkedLocalAttentionGroupSpec, + CrossAttentionGroupSpec, + EncoderOnlyAttentionGroupSpec, + FullAttentionGroupSpec, + KVCacheResponse, + MambaGroupSpec, + MLAAttentionGroupSpec, + SinkFullAttentionGroupSpec, + SlidingWindowGroupSpec, + SlidingWindowMLAGroupSpec, + UnknownGroupSpec, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.engine.protocol import EngineClient + +logger = get_logger(__name__) + +_KIND_TO_MODEL: dict[str, type] = { + "full_attention": FullAttentionGroupSpec, + "mla_attention": MLAAttentionGroupSpec, + "sliding_window": SlidingWindowGroupSpec, + "sliding_window_mla": SlidingWindowMLAGroupSpec, + "chunked_local_attention": ChunkedLocalAttentionGroupSpec, + "mamba": MambaGroupSpec, + "cross_attention": CrossAttentionGroupSpec, + "encoder_only_attention": EncoderOnlyAttentionGroupSpec, + "sink_full_attention": SinkFullAttentionGroupSpec, + "unknown": UnknownGroupSpec, +} + + +def _build_group_spec(group: dict): + kind = group["kind"] + model_cls = _KIND_TO_MODEL.get(kind) + if model_cls is None: + raise ValueError(f"Unhandled KVCacheSpec kind: {kind!r}") + return model_cls(**group) + + +def _build_response(kv_cache_data: dict | None) -> KVCacheResponse: + if kv_cache_data is None: + return KVCacheResponse() + raw_groups = kv_cache_data.get("groups", []) + return KVCacheResponse( + kv_cache_size_tokens=kv_cache_data.get("kv_cache_size_tokens"), + max_concurrency=kv_cache_data.get("max_concurrency"), + num_gpu_blocks=kv_cache_data.get("num_gpu_blocks"), + num_cpu_blocks=kv_cache_data.get("num_cpu_blocks"), + groups=[_build_group_spec(g) for g in raw_groups], + ) + + +def _capacity_from_vllm_config(vllm_config: "VllmConfig") -> dict: + cache_cfg = vllm_config.cache_config + num_gpu_blocks = cache_cfg.num_gpu_blocks + block_size = cache_cfg.block_size + kv_cache_size_tokens = ( + num_gpu_blocks * block_size + if num_gpu_blocks is not None and block_size is not None + else None + ) + max_model_len = vllm_config.model_config.max_model_len + max_concurrency = ( + kv_cache_size_tokens / max_model_len + if kv_cache_size_tokens is not None and max_model_len + else None + ) + return { + "kv_cache_size_tokens": kv_cache_size_tokens, + "max_concurrency": max_concurrency, + "num_gpu_blocks": num_gpu_blocks, + "num_cpu_blocks": cache_cfg.num_cpu_blocks, + "groups": [], + } + + +class ServerKVCachePlugin: + name = "llm_d_server_introspection_kv_cache" + required_tasks: tuple[str, ...] | None = ENGINE_TASKS + + def attach_router(self, app: FastAPI) -> None: + router = APIRouter() + + @router.get( + "/plugins/llm-d-server-introspection/kv-cache", + response_model=KVCacheResponse, + ) + async def get_kv_cache(raw_request: Request) -> KVCacheResponse: + return cached_response( + raw_request, "server_kv_cache_response", no_engine_detail("kv-cache") + ) + + app.include_router(router) + + async def init_state( + self, engine_client: "EngineClient | None", state: State, args: Namespace + ) -> None: + if engine_client is None: + state.server_kv_cache_response = None + return + + kv_cache_data = _capacity_from_vllm_config(state.vllm_config) + if hasattr(engine_client, "get_kv_cache_group_metadata"): + kv_cache_data["groups"] = await engine_client.get_kv_cache_group_metadata() + + response = _build_response(kv_cache_data) + state.server_kv_cache_response = response + logger.info( + "kv_cache plugin initialized: num_gpu_blocks=%s groups=%d", + response.num_gpu_blocks, + len(response.groups), + ) diff --git a/vllm/llm_d_api_extensions_vllm/server_introspection/schemas.py b/vllm/llm_d_api_extensions_vllm/server_introspection/schemas.py new file mode 100644 index 0000000..1ff3a42 --- /dev/null +++ b/vllm/llm_d_api_extensions_vllm/server_introspection/schemas.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pydantic response models for the `/plugins/llm-d-server-introspection/*` introspection endpoints. + +Shared across all endpoint plugins in this package. +""" + +from typing import Annotated, Literal + +from pydantic import BaseModel, Field + + +class ModelInfo(BaseModel): + name: str + served_names: list[str] + dtype: str + quantization: str | None + max_model_len: int + + +class KVCacheInfo(BaseModel): + gpu_memory_utilization: float + dtype: str + enable_prefix_caching: bool + + +class SchedulerInfo(BaseModel): + max_num_seqs: int + max_num_batched_tokens: int + enable_chunked_prefill: bool + policy: str + + +class ParallelismInfo(BaseModel): + tensor_parallel_size: int + pipeline_parallel_size: int + data_parallel_size: int + data_parallel_rank: int + + +class FeaturesInfo(BaseModel): + speculative_decoding: bool + lora: bool + hma: bool + + +class KVTransferInfo(BaseModel): + kv_connector: str | None + kv_role: str | None + kv_connector_module_path: str | None = None + kv_buffer_device: str | None = None + kv_buffer_size: float | None = None + kv_ip: str | None = None + kv_port: int | None = None + kv_parallel_size: int | None = None + kv_rank: int | None = None + engine_id: str | None = None + extra_config: dict = Field(default_factory=dict) + # NIXL side channel base endpoint (env derived only when connector is + # NixlConnector). This is the base host/port. NixlConnector derives the + # actual per rank port as base_port + rank_offset inside the worker. + nixl_side_channel_host: str | None = None + nixl_side_channel_port: int | None = None + + +class ServerConfigResponse(BaseModel): + model: ModelInfo + kv_cache: KVCacheInfo + scheduler: SchedulerInfo + parallelism: ParallelismInfo + features: FeaturesInfo + kv_transfer: KVTransferInfo | None = None + + +class ComputeCapability(BaseModel): + major: int + minor: int + + +class DeviceInfo(BaseModel): + rank: int + name: str + total_memory_bytes: int + compute_capability: ComputeCapability | None = None + num_compute_units: int | None = None + + +class DevicesResponse(BaseModel): + devices: list[DeviceInfo] + + +# --------------------------------------------------------------------------- +# `/plugins/llm-d-server-introspection/kv-cache`: group specs are a +# discriminated union keyed on `kind`, matching the field names and +# `KVCacheSpecKind` string values `EngineClient.get_kv_cache_group_metadata()` +# serializes groups under (vllm-project/vllm#48121). Every group dict carries +# the full field superset (irrelevant fields sent as `None`). Pydantic's +# default `extra="ignore"` behavior drops whichever ones a given `kind` +# doesn't declare. +# --------------------------------------------------------------------------- + + +class _KVCacheGroupBase(BaseModel): + group_id: int + layer_count: int + layer_names: list[str] + block_size: int + page_size_bytes: int + # Present (non-None) only when the group fans out a UniformTypeKVCacheSpecs + # into its per layer specs. + layer_specs: list[dict] | None = None + + +class _FullAttentionBaseSpec(_KVCacheGroupBase): + num_kv_heads: int + head_size: int + head_size_v: int + dtype: str + sliding_window: int | None = None + attention_chunk_size: int | None = None + + +class FullAttentionGroupSpec(_FullAttentionBaseSpec): + kind: Literal["full_attention"] = "full_attention" + + +class MLAAttentionGroupSpec(_FullAttentionBaseSpec): + kind: Literal["mla_attention"] = "mla_attention" + cache_dtype_str: str | None = None + + +class SlidingWindowGroupSpec(_KVCacheGroupBase): + kind: Literal["sliding_window"] = "sliding_window" + num_kv_heads: int + head_size: int + dtype: str + sliding_window: int + + +class SlidingWindowMLAGroupSpec(_KVCacheGroupBase): + kind: Literal["sliding_window_mla"] = "sliding_window_mla" + num_kv_heads: int + head_size: int + head_size_v: int + dtype: str + sliding_window: int + cache_dtype_str: str | None = None + + +class ChunkedLocalAttentionGroupSpec(_KVCacheGroupBase): + kind: Literal["chunked_local_attention"] = "chunked_local_attention" + num_kv_heads: int + head_size: int + dtype: str + attention_chunk_size: int + + +class MambaGroupSpec(_KVCacheGroupBase): + kind: Literal["mamba"] = "mamba" + shapes: list[list[int]] + dtypes: list[str] + mamba_type: str + mamba_cache_mode: str + + +class CrossAttentionGroupSpec(_KVCacheGroupBase): + kind: Literal["cross_attention"] = "cross_attention" + num_kv_heads: int + head_size: int + dtype: str + + +class EncoderOnlyAttentionGroupSpec(_KVCacheGroupBase): + kind: Literal["encoder_only_attention"] = "encoder_only_attention" + num_kv_heads: int + head_size: int + dtype: str + + +class SinkFullAttentionGroupSpec(_FullAttentionBaseSpec): + kind: Literal["sink_full_attention"] = "sink_full_attention" + sink_len: int | None = None + + +class UnknownGroupSpec(_KVCacheGroupBase): + kind: Literal["unknown"] = "unknown" + + +KVCacheGroupSpec = Annotated[ + FullAttentionGroupSpec + | MLAAttentionGroupSpec + | SlidingWindowGroupSpec + | SlidingWindowMLAGroupSpec + | ChunkedLocalAttentionGroupSpec + | MambaGroupSpec + | CrossAttentionGroupSpec + | EncoderOnlyAttentionGroupSpec + | SinkFullAttentionGroupSpec + | UnknownGroupSpec, + Field(discriminator="kind"), +] + + +class KVCacheResponse(BaseModel): + kv_cache_size_tokens: int | None = None + max_concurrency: float | None = None + num_gpu_blocks: int | None = None + num_cpu_blocks: int | None = None + groups: list[KVCacheGroupSpec] = Field(default_factory=list) diff --git a/vllm/pyproject.toml b/vllm/pyproject.toml new file mode 100644 index 0000000..90709e5 --- /dev/null +++ b/vllm/pyproject.toml @@ -0,0 +1,38 @@ +[build-system] +requires = ["setuptools>=68.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "llm-d-api-extensions-vllm" +version = "0.1.0" +description = "llm-d API extensions for vLLM, integrated as endpoint plugins" +readme = "README.md" +requires-python = ">=3.10" +license = { text = "Apache-2.0" } +dependencies = [ + "vllm", + "fastapi", + "pydantic>=2", +] + +[project.urls] +Homepage = "https://github.com/llm-d/llm-d-api-extensions" +Source = "https://github.com/llm-d/llm-d-api-extensions/tree/main/vllm" +Issues = "https://github.com/llm-d/llm-d-api-extensions/issues" + +[project.optional-dependencies] +test = ["pytest", "pytest-asyncio", "httpx"] + +# Endpoint plugins are opt-in. vLLM only loads the ones named in `VLLM_PLUGINS`. +[project.entry-points."vllm.endpoint_plugins"] +llm_d_server_introspection_config = "llm_d_api_extensions_vllm.server_introspection.config_plugin:ServerConfigPlugin" +llm_d_server_introspection_devices = "llm_d_api_extensions_vllm.server_introspection.devices_plugin:ServerDevicesPlugin" +llm_d_server_introspection_kv_cache = "llm_d_api_extensions_vllm.server_introspection.kv_cache_plugin:ServerKVCachePlugin" + +[tool.setuptools.packages.find] +include = ["llm_d_api_extensions_vllm*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +# Lets `pytest tests/` run straight from a clone without an editable install first. +pythonpath = ["."] diff --git a/vllm/ruff.toml b/vllm/ruff.toml new file mode 100644 index 0000000..5538c0b --- /dev/null +++ b/vllm/ruff.toml @@ -0,0 +1,16 @@ +line-length = 100 +target-version = "py310" + +[lint] +select = [ + "E", # pycodestyle errors + "F", # pyflakes + "I", # isort + "UP", # pyupgrade + "B", # flake8-bugbear + "SIM", # flake8-simplify +] + +[lint.per-file-ignores] +# Tests deliberately poke at plugin internals (`_build_response`, ...). +"tests/**" = ["SIM117"] diff --git a/vllm/tests/server_introspection/kv_group_fixtures.py b/vllm/tests/server_introspection/kv_group_fixtures.py new file mode 100644 index 0000000..cf6b9a3 --- /dev/null +++ b/vllm/tests/server_introspection/kv_group_fixtures.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Serialized KV cache group dicts shared by the kv-cache plugin and schema tests. + +Mirrors what `get_kv_cache_group_metadata` produces per vllm-project/vllm#48121. +""" + +from llm_d_api_extensions_vllm.server_introspection.schemas import ( + ChunkedLocalAttentionGroupSpec, + CrossAttentionGroupSpec, + EncoderOnlyAttentionGroupSpec, + FullAttentionGroupSpec, + MambaGroupSpec, + MLAAttentionGroupSpec, + SinkFullAttentionGroupSpec, + SlidingWindowGroupSpec, + SlidingWindowMLAGroupSpec, + UnknownGroupSpec, +) + +BASE_GROUP = { + "group_id": 0, + "layer_count": 2, + "layer_names": ["model.layers.0.self_attn", "model.layers.1.self_attn"], + "block_size": 16, + "page_size_bytes": 131072, + "layer_specs": None, +} + +_ATTENTION = {"num_kv_heads": 8, "head_size": 128, "dtype": "bfloat16"} + +# Fields each kind adds on top of BASE_GROUP. +_KIND_FIELDS: dict[str, dict] = { + "full_attention": { + **_ATTENTION, + "head_size_v": 128, + "sliding_window": None, + "attention_chunk_size": None, + }, + "mla_attention": { + **_ATTENTION, + "head_size_v": 64, + "sliding_window": None, + "attention_chunk_size": None, + "cache_dtype_str": "float8_e4m3fn", + }, + "sliding_window": {**_ATTENTION, "dtype": "float16", "sliding_window": 4096}, + "sliding_window_mla": { + **_ATTENTION, + "head_size_v": 64, + "sliding_window": 4096, + "cache_dtype_str": "float8_e4m3fn", + }, + "chunked_local_attention": {**_ATTENTION, "dtype": "float16", "attention_chunk_size": 2048}, + "mamba": { + "block_size": 1, + "page_size_bytes": 4096, + "shapes": [[16, 128], [16, 64]], + "dtypes": ["float32", "float32"], + "mamba_type": "mamba2", + "mamba_cache_mode": "none", + }, + "cross_attention": dict(_ATTENTION), + "encoder_only_attention": dict(_ATTENTION), + "sink_full_attention": { + **_ATTENTION, + "head_size_v": 128, + "sliding_window": 2048, + "attention_chunk_size": None, + "sink_len": 4, + }, + "unknown": {}, +} + +SPEC_CLS_BY_KIND: dict[str, type] = { + "full_attention": FullAttentionGroupSpec, + "mla_attention": MLAAttentionGroupSpec, + "sliding_window": SlidingWindowGroupSpec, + "sliding_window_mla": SlidingWindowMLAGroupSpec, + "chunked_local_attention": ChunkedLocalAttentionGroupSpec, + "mamba": MambaGroupSpec, + "cross_attention": CrossAttentionGroupSpec, + "encoder_only_attention": EncoderOnlyAttentionGroupSpec, + "sink_full_attention": SinkFullAttentionGroupSpec, + "unknown": UnknownGroupSpec, +} + +ALL_KINDS = list(SPEC_CLS_BY_KIND) + + +def group_dict(kind: str, *, omit: tuple[str, ...] = (), **overrides) -> dict: + group = {**BASE_GROUP, "kind": kind, **_KIND_FIELDS[kind], **overrides} + for key in omit: + group.pop(key) + return group diff --git a/vllm/tests/server_introspection/test_config_plugin.py b/vllm/tests/server_introspection/test_config_plugin.py new file mode 100644 index 0000000..59b6775 --- /dev/null +++ b/vllm/tests/server_introspection/test_config_plugin.py @@ -0,0 +1,513 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for `GET /plugins/llm-d-server-introspection/config`. + +Fast, no engine and no real vLLM/model. `VllmConfig` and its sub configs are +`MagicMock`s so the tests exercise `_build_response` and the FastAPI route in +isolation. +""" + +import asyncio +from argparse import Namespace +from unittest.mock import MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from llm_d_api_extensions_vllm.server_introspection.config_plugin import ( + EXPOSE_EXTRA_CONFIG_ENV, + REDACTED, + ServerConfigPlugin, + _build_response, +) +from llm_d_api_extensions_vllm.server_introspection.schemas import ( + FeaturesInfo, + KVCacheInfo, + KVTransferInfo, + ModelInfo, + ParallelismInfo, + SchedulerInfo, + ServerConfigResponse, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _FakeTorchDtype: + """Stands in for `torch.dtype` without a torch dependency in tests. + + `str(torch.bfloat16) == "torch.bfloat16"`; `_build_response` does + `str(dtype).removeprefix("torch.")`, so this only needs to match that. + """ + + def __init__(self, name: str): + self._name = name + + def __str__(self) -> str: + return f"torch.{self._name}" + + +BFLOAT16 = _FakeTorchDtype("bfloat16") +FLOAT16 = _FakeTorchDtype("float16") +FLOAT32 = _FakeTorchDtype("float32") + + +def _make_vllm_config( + *, + model_name: str = "meta-llama/Llama-3.1-8B", + served_model_name: str = "meta-llama/Llama-3.1-8B", + model_dtype: _FakeTorchDtype = BFLOAT16, + quantization: str | None = None, + max_model_len: int = 32768, + gpu_memory_utilization: float = 0.9, + cache_dtype: str = "auto", + enable_prefix_caching: bool = True, + max_num_seqs: int = 128, + max_num_batched_tokens: int = 2048, + enable_chunked_prefill: bool = True, + policy: str = "fcfs", + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + data_parallel_size: int = 1, + data_parallel_rank: int = 0, + disable_hybrid_kv_cache_manager: bool = False, + speculative_config=None, + lora_config=None, + kv_transfer_config=None, +) -> MagicMock: + model_cfg = MagicMock() + model_cfg.model = model_name + model_cfg.served_model_name = served_model_name + model_cfg.dtype = model_dtype + model_cfg.quantization = quantization + model_cfg.max_model_len = max_model_len + + cache_cfg = MagicMock() + cache_cfg.gpu_memory_utilization = gpu_memory_utilization + cache_cfg.cache_dtype = cache_dtype + cache_cfg.enable_prefix_caching = enable_prefix_caching + + scheduler_cfg = MagicMock() + scheduler_cfg.max_num_seqs = max_num_seqs + scheduler_cfg.max_num_batched_tokens = max_num_batched_tokens + scheduler_cfg.enable_chunked_prefill = enable_chunked_prefill + scheduler_cfg.policy = policy + scheduler_cfg.disable_hybrid_kv_cache_manager = disable_hybrid_kv_cache_manager + + parallel_cfg = MagicMock() + parallel_cfg.tensor_parallel_size = tensor_parallel_size + parallel_cfg.pipeline_parallel_size = pipeline_parallel_size + parallel_cfg.data_parallel_size = data_parallel_size + parallel_cfg.data_parallel_rank = data_parallel_rank + + vllm_config = MagicMock() + vllm_config.model_config = model_cfg + vllm_config.cache_config = cache_cfg + vllm_config.scheduler_config = scheduler_cfg + vllm_config.parallel_config = parallel_cfg + vllm_config.speculative_config = speculative_config + vllm_config.lora_config = lora_config + vllm_config.kv_transfer_config = kv_transfer_config + + return vllm_config + + +def _make_kv_transfer_config( + *, + kv_connector: str | None = "NixlConnector", + kv_role: str | None = "kv_both", + kv_connector_module_path: str | None = None, + kv_buffer_device: str = "cuda", + kv_buffer_size: float = 1e9, + kv_ip: str = "127.0.0.1", + kv_port: int = 14579, + kv_parallel_size: int = 1, + kv_rank: int | None = None, + engine_id: str | None = "engine-0", + kv_connector_extra_config: dict | None = None, +) -> MagicMock: + cfg = MagicMock() + cfg.kv_connector = kv_connector + cfg.kv_role = kv_role + cfg.kv_connector_module_path = kv_connector_module_path + cfg.kv_buffer_device = kv_buffer_device + cfg.kv_buffer_size = kv_buffer_size + cfg.kv_ip = kv_ip + cfg.kv_port = kv_port + cfg.kv_parallel_size = kv_parallel_size + cfg.kv_rank = kv_rank + cfg.engine_id = engine_id + cfg.kv_connector_extra_config = kv_connector_extra_config or {} + return cfg + + +def _make_test_app( + vllm_config: MagicMock, + served_model_name: list[str] | None = None, +) -> FastAPI: + """FastAPI app with the plugin's route + state wired through both phases.""" + app = FastAPI() + plugin = ServerConfigPlugin() + plugin.attach_router(app) + + app.state.vllm_config = vllm_config + args = Namespace(served_model_name=served_model_name) + asyncio.run(plugin.init_state(None, app.state, args)) + return app + + +# --------------------------------------------------------------------------- +# _build_response unit tests +# --------------------------------------------------------------------------- + + +class TestBuildResponseModelSection: + def test_model_name_is_first_served_name(self): + cfg = _make_vllm_config(model_name="base-model") + resp = _build_response(cfg, ["alias-a", "alias-b"]) + assert resp.model.name == "alias-a" + + def test_model_served_names_matches_input(self): + cfg = _make_vllm_config() + resp = _build_response(cfg, ["m1", "m2", "m3"]) + assert resp.model.served_names == ["m1", "m2", "m3"] + + @pytest.mark.parametrize( + "dtype,expected", + [ + (BFLOAT16, "bfloat16"), + (FLOAT16, "float16"), + (FLOAT32, "float32"), + ], + ) + def test_model_dtype_format(self, dtype, expected): + cfg = _make_vllm_config(model_dtype=dtype) + resp = _build_response(cfg, ["m"]) + assert resp.model.dtype == expected + + def test_model_quantization_none(self): + cfg = _make_vllm_config(quantization=None) + resp = _build_response(cfg, ["m"]) + assert resp.model.quantization is None + + @pytest.mark.parametrize("quant", ["awq", "gptq", "fp8"]) + def test_model_quantization_string(self, quant): + cfg = _make_vllm_config(quantization=quant) + resp = _build_response(cfg, ["m"]) + assert resp.model.quantization == quant + + def test_model_max_model_len(self): + cfg = _make_vllm_config(max_model_len=131072) + resp = _build_response(cfg, ["m"]) + assert resp.model.max_model_len == 131072 + + +class TestBuildResponseKVCacheSection: + def test_gpu_memory_utilization(self): + cfg = _make_vllm_config(gpu_memory_utilization=0.85) + resp = _build_response(cfg, ["m"]) + assert resp.kv_cache.gpu_memory_utilization == pytest.approx(0.85) + + @pytest.mark.parametrize( + "model_dtype,expected", + [ + (BFLOAT16, "bfloat16"), + (FLOAT16, "float16"), + ], + ) + def test_cache_dtype_auto_resolves_to_model_dtype(self, model_dtype, expected): + cfg = _make_vllm_config(cache_dtype="auto", model_dtype=model_dtype) + resp = _build_response(cfg, ["m"]) + assert resp.kv_cache.dtype == expected + + def test_cache_dtype_explicit_passes_through(self): + cfg = _make_vllm_config(cache_dtype="fp8", model_dtype=BFLOAT16) + resp = _build_response(cfg, ["m"]) + assert resp.kv_cache.dtype == "fp8" + + def test_enable_prefix_caching_true(self): + cfg = _make_vllm_config(enable_prefix_caching=True) + resp = _build_response(cfg, ["m"]) + assert resp.kv_cache.enable_prefix_caching is True + + def test_enable_prefix_caching_false(self): + cfg = _make_vllm_config(enable_prefix_caching=False) + resp = _build_response(cfg, ["m"]) + assert resp.kv_cache.enable_prefix_caching is False + + +class TestBuildResponseSchedulerSection: + def test_max_num_seqs(self): + cfg = _make_vllm_config(max_num_seqs=256) + resp = _build_response(cfg, ["m"]) + assert resp.scheduler.max_num_seqs == 256 + + def test_max_num_batched_tokens(self): + cfg = _make_vllm_config(max_num_batched_tokens=4096) + resp = _build_response(cfg, ["m"]) + assert resp.scheduler.max_num_batched_tokens == 4096 + + def test_enable_chunked_prefill_true(self): + cfg = _make_vllm_config(enable_chunked_prefill=True) + resp = _build_response(cfg, ["m"]) + assert resp.scheduler.enable_chunked_prefill is True + + def test_enable_chunked_prefill_false(self): + cfg = _make_vllm_config(enable_chunked_prefill=False) + resp = _build_response(cfg, ["m"]) + assert resp.scheduler.enable_chunked_prefill is False + + @pytest.mark.parametrize("policy", ["fcfs", "priority"]) + def test_scheduler_policy(self, policy): + cfg = _make_vllm_config(policy=policy) + resp = _build_response(cfg, ["m"]) + assert resp.scheduler.policy == policy + + +class TestBuildResponseParallelismSection: + def test_defaults_single_device(self): + cfg = _make_vllm_config( + tensor_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + data_parallel_rank=0, + ) + resp = _build_response(cfg, ["m"]) + assert resp.parallelism.tensor_parallel_size == 1 + assert resp.parallelism.pipeline_parallel_size == 1 + assert resp.parallelism.data_parallel_size == 1 + assert resp.parallelism.data_parallel_rank == 0 + + def test_tensor_parallel(self): + cfg = _make_vllm_config(tensor_parallel_size=4) + resp = _build_response(cfg, ["m"]) + assert resp.parallelism.tensor_parallel_size == 4 + + def test_data_parallel_with_rank(self): + cfg = _make_vllm_config(data_parallel_size=4, data_parallel_rank=3) + resp = _build_response(cfg, ["m"]) + assert resp.parallelism.data_parallel_size == 4 + assert resp.parallelism.data_parallel_rank == 3 + + +class TestBuildResponseFeaturesSection: + def test_all_features_disabled_by_default(self): + cfg = _make_vllm_config( + speculative_config=None, + lora_config=None, + disable_hybrid_kv_cache_manager=False, + ) + resp = _build_response(cfg, ["m"]) + assert resp.features.speculative_decoding is False + assert resp.features.lora is False + assert resp.features.hma is True + + def test_speculative_decoding_enabled(self): + cfg = _make_vllm_config(speculative_config=MagicMock()) + resp = _build_response(cfg, ["m"]) + assert resp.features.speculative_decoding is True + + def test_lora_enabled(self): + cfg = _make_vllm_config(lora_config=MagicMock()) + resp = _build_response(cfg, ["m"]) + assert resp.features.lora is True + + def test_hma_disabled_when_manager_disabled(self): + cfg = _make_vllm_config(disable_hybrid_kv_cache_manager=True) + resp = _build_response(cfg, ["m"]) + assert resp.features.hma is False + + def test_hma_defaults_true_when_field_absent(self): + # Older vLLM without `disable_hybrid_kv_cache_manager` on + # SchedulerConfig: getattr(..., default=False) -> hma stays True. + cfg = _make_vllm_config() + del cfg.scheduler_config.disable_hybrid_kv_cache_manager + resp = _build_response(cfg, ["m"]) + assert resp.features.hma is True + + +class TestBuildResponseKVTransferSection: + def test_none_when_no_disaggregation_configured(self): + cfg = _make_vllm_config(kv_transfer_config=None) + resp = _build_response(cfg, ["m"]) + assert resp.kv_transfer is None + + def test_nixl_connector_populates_fields_and_env_derived_port(self, monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_NIXL_SIDE_CHANNEL_HOST", "10.0.0.5") + monkeypatch.setattr(envs, "VLLM_NIXL_SIDE_CHANNEL_PORT", 5600) + + kv_transfer_cfg = _make_kv_transfer_config( + kv_connector="NixlConnector", + kv_role="kv_producer", + kv_connector_extra_config={"backend": "UCX"}, + ) + cfg = _make_vllm_config(kv_transfer_config=kv_transfer_cfg) + resp = _build_response(cfg, ["m"]) + + assert resp.kv_transfer is not None + assert resp.kv_transfer.kv_connector == "NixlConnector" + assert resp.kv_transfer.kv_role == "kv_producer" + assert resp.kv_transfer.extra_config == {"backend": REDACTED} + assert resp.kv_transfer.nixl_side_channel_host == "10.0.0.5" + assert resp.kv_transfer.nixl_side_channel_port == 5600 + + def test_extra_config_values_redacted_by_default(self, monkeypatch): + monkeypatch.delenv(EXPOSE_EXTRA_CONFIG_ENV, raising=False) + kv_transfer_cfg = _make_kv_transfer_config( + kv_connector_extra_config={"api_key": "s3cret", "backend": "UCX"}, + ) + cfg = _make_vllm_config(kv_transfer_config=kv_transfer_cfg) + resp = _build_response(cfg, ["m"]) + assert resp.kv_transfer.extra_config == {"api_key": REDACTED, "backend": REDACTED} + + def test_extra_config_values_exposed_when_opted_in(self, monkeypatch): + monkeypatch.setenv(EXPOSE_EXTRA_CONFIG_ENV, "1") + kv_transfer_cfg = _make_kv_transfer_config( + kv_connector_extra_config={"api_key": "s3cret", "backend": "UCX"}, + ) + cfg = _make_vllm_config(kv_transfer_config=kv_transfer_cfg) + resp = _build_response(cfg, ["m"]) + assert resp.kv_transfer.extra_config == {"api_key": "s3cret", "backend": "UCX"} + + def test_empty_extra_config_stays_empty(self, monkeypatch): + monkeypatch.delenv(EXPOSE_EXTRA_CONFIG_ENV, raising=False) + cfg = _make_vllm_config(kv_transfer_config=_make_kv_transfer_config()) + resp = _build_response(cfg, ["m"]) + assert resp.kv_transfer.extra_config == {} + + def test_non_nixl_connector_leaves_nixl_fields_none(self): + kv_transfer_cfg = _make_kv_transfer_config(kv_connector="LMCacheConnectorV1") + cfg = _make_vllm_config(kv_transfer_config=kv_transfer_cfg) + resp = _build_response(cfg, ["m"]) + assert resp.kv_transfer is not None + assert resp.kv_transfer.kv_connector == "LMCacheConnectorV1" + assert resp.kv_transfer.nixl_side_channel_host is None + assert resp.kv_transfer.nixl_side_channel_port is None + + +class TestBuildResponseReturnType: + def test_returns_server_config_response(self): + cfg = _make_vllm_config() + resp = _build_response(cfg, ["m"]) + assert isinstance(resp, ServerConfigResponse) + assert isinstance(resp.model, ModelInfo) + assert isinstance(resp.kv_cache, KVCacheInfo) + assert isinstance(resp.scheduler, SchedulerInfo) + assert isinstance(resp.parallelism, ParallelismInfo) + assert isinstance(resp.features, FeaturesInfo) + assert resp.kv_transfer is None + + def test_returns_kv_transfer_info_when_configured(self): + cfg = _make_vllm_config(kv_transfer_config=_make_kv_transfer_config()) + resp = _build_response(cfg, ["m"]) + assert isinstance(resp.kv_transfer, KVTransferInfo) + + +# --------------------------------------------------------------------------- +# HTTP endpoint tests via TestClient +# --------------------------------------------------------------------------- + + +class TestGetServerConfigEndpoint: + def test_response_parses_as_server_config_response(self): + app = _make_test_app(_make_vllm_config(), served_model_name=["my-model"]) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/config") + parsed = ServerConfigResponse.model_validate(resp.json()) + assert parsed.model.name == "my-model" + + def test_served_names_from_args_multiple(self): + cfg = _make_vllm_config() + app = _make_test_app(cfg, served_model_name=["alias-a", "alias-b"]) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/config") + data = resp.json() + assert data["model"]["name"] == "alias-a" + assert data["model"]["served_names"] == ["alias-a", "alias-b"] + + def test_served_names_fallback_from_model_config(self): + cfg = _make_vllm_config(served_model_name="model-from-config") + app = _make_test_app(cfg, served_model_name=None) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/config") + data = resp.json() + assert data["model"]["name"] == "model-from-config" + assert data["model"]["served_names"] == ["model-from-config"] + + def test_content_type_is_json(self): + app = _make_test_app(_make_vllm_config(), served_model_name=["m"]) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/config") + assert "application/json" in resp.headers["content-type"] + + def test_top_level_keys_present(self): + app = _make_test_app(_make_vllm_config(), served_model_name=["m"]) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/config") + assert set(resp.json().keys()) == { + "model", + "kv_cache", + "scheduler", + "parallelism", + "features", + "kv_transfer", + } + + def test_wrong_method_returns_405(self): + app = _make_test_app(_make_vllm_config(), served_model_name=["m"]) + with TestClient(app) as client: + resp = client.post("/plugins/llm-d-server-introspection/config") + assert resp.status_code == 405 + + def test_missing_state_returns_500(self): + # attach_router without ever calling init_state. + # The handler must not crash on a missing attribute. + app = FastAPI() + ServerConfigPlugin().attach_router(app) + with TestClient(app, raise_server_exceptions=False) as client: + resp = client.get("/plugins/llm-d-server-introspection/config") + assert resp.status_code == 500 + + def test_response_cached_at_init_state_not_rebuilt_per_request(self): + cfg = _make_vllm_config(max_model_len=32768) + app = _make_test_app(cfg, served_model_name=["m"]) + cached = app.state.server_config_response + with TestClient(app) as client: + first = client.get("/plugins/llm-d-server-introspection/config") + # A rebuild per request would pick up this change. + cfg.model_config.max_model_len = 1024 + second = client.get("/plugins/llm-d-server-introspection/config") + assert first.json()["model"]["max_model_len"] == 32768 + assert second.json()["model"]["max_model_len"] == 32768 + assert app.state.server_config_response is cached + + def test_engine_client_none_on_render_server_is_fine(self): + # required_tasks=None -> eligible on the CPU only render server, + # which calls init_state with engine_client=None. This plugin never + # touches engine_client, so it must not raise. + app = FastAPI() + plugin = ServerConfigPlugin() + plugin.attach_router(app) + app.state.vllm_config = _make_vllm_config() + args = Namespace(served_model_name=["m"]) + asyncio.run(plugin.init_state(None, app.state, args)) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/config") + assert resp.status_code == 200 + + +# --------------------------------------------------------------------------- +# Plugin metadata +# --------------------------------------------------------------------------- + + +class TestPluginMetadata: + def test_name(self): + assert ServerConfigPlugin().name == "llm_d_server_introspection_config" + + def test_required_tasks_is_none(self): + assert ServerConfigPlugin().required_tasks is None diff --git a/vllm/tests/server_introspection/test_device_worker_ext.py b/vllm/tests/server_introspection/test_device_worker_ext.py new file mode 100644 index 0000000..3e367e0 --- /dev/null +++ b/vllm/tests/server_introspection/test_device_worker_ext.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for `DeviceInfoWorkerExtension.get_device_properties`. + +Fast with no GPU. `vllm.platforms.current_platform` is replaced by a stub so +both the populated path and the `_safe()` fallback (platforms that raise +`NotImplementedError` for device introspection) are exercised. +""" + +from types import SimpleNamespace + +import pytest +import vllm.platforms + +from llm_d_api_extensions_vllm.server_introspection.device_worker_ext import ( + DeviceInfoWorkerExtension, +) + + +class _GpuPlatform: + def __init__(self): + self.device_ids: list[int] = [] + + def get_device_capability(self, device_id): + self.device_ids.append(device_id) + return SimpleNamespace(major=9, minor=0) + + def get_device_name(self, device_id): + return "H100-SXM5-80GB" + + def get_device_total_memory(self, device_id): + return 85_899_345_920 + + def num_compute_units(self, device_id): + return 132 + + +class _UnsupportedPlatform: + def get_device_capability(self, device_id): + raise NotImplementedError + + def get_device_name(self, device_id): + raise NotImplementedError + + def get_device_total_memory(self, device_id): + raise NotImplementedError + + def num_compute_units(self, device_id): + raise NotImplementedError + + +class _BrokenNamePlatform(_GpuPlatform): + def get_device_name(self, device_id): + raise RuntimeError("boom") + + +class _Worker(DeviceInfoWorkerExtension): + """Stands in for the vLLM `Worker` the extension is mixed into.""" + + def __init__(self, rank: int, local_rank: int): + self.rank = rank + self.local_rank = local_rank + + +def _use_platform(monkeypatch, platform) -> None: + monkeypatch.setattr(vllm.platforms, "current_platform", platform) + + +class TestGetDeviceProperties: + def test_populated_platform(self, monkeypatch): + _use_platform(monkeypatch, _GpuPlatform()) + props = _Worker(rank=3, local_rank=1).get_device_properties() + assert props == { + "rank": 3, + "name": "H100-SXM5-80GB", + "total_memory_bytes": 85_899_345_920, + "compute_capability": {"major": 9, "minor": 0}, + "num_compute_units": 132, + } + + def test_queries_local_rank_device(self, monkeypatch): + platform = _GpuPlatform() + _use_platform(monkeypatch, platform) + _Worker(rank=5, local_rank=1).get_device_properties() + assert platform.device_ids == [1] + + def test_unsupported_platform_falls_back_to_none(self, monkeypatch): + _use_platform(monkeypatch, _UnsupportedPlatform()) + props = _Worker(rank=0, local_rank=0).get_device_properties() + assert props == { + "rank": 0, + "name": None, + "total_memory_bytes": None, + "compute_capability": None, + "num_compute_units": None, + } + + def test_other_exceptions_propagate(self, monkeypatch): + # `_safe()` only swallows NotImplementedError. Real failures must surface. + _use_platform(monkeypatch, _BrokenNamePlatform()) + with pytest.raises(RuntimeError, match="boom"): + _Worker(rank=0, local_rank=0).get_device_properties() diff --git a/vllm/tests/server_introspection/test_devices_plugin.py b/vllm/tests/server_introspection/test_devices_plugin.py new file mode 100644 index 0000000..3c6d42b --- /dev/null +++ b/vllm/tests/server_introspection/test_devices_plugin.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for `GET /plugins/llm-d-server-introspection/devices`. + +Fast with no engine and no real vLLM/model. `EngineClient` is a fake engine +exercising only `collective_rpc`, matching the pattern vLLM's own +`tests/plugins_tests/test_endpoint_plugins.py` uses for its `_FakeEngineClient`. +""" + +import asyncio +from argparse import Namespace +from typing import Any + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from llm_d_api_extensions_vllm.server_introspection.device_worker_ext import ( + GET_DEVICE_PROPERTIES, + DeviceInfoWorkerExtension, +) +from llm_d_api_extensions_vllm.server_introspection.devices_plugin import ( + ServerDevicesPlugin, + _build_response, +) +from llm_d_api_extensions_vllm.server_introspection.schemas import ( + ComputeCapability, + DeviceInfo, + DevicesResponse, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_A100_ENTRY = { + "rank": 0, + "name": "A100-PCIE-40GB", + "total_memory_bytes": 42_949_672_960, + "compute_capability": {"major": 8, "minor": 0}, + "num_compute_units": 108, +} + +_H100_ENTRY = { + "rank": 1, + "name": "H100-SXM5-80GB", + "total_memory_bytes": 85_899_345_920, + "compute_capability": {"major": 9, "minor": 0}, + "num_compute_units": 132, +} + +_NO_CAPABILITY_ENTRY = { + "rank": 0, + "name": "SomeNonCudaDevice", + "total_memory_bytes": 8_589_934_592, + "compute_capability": None, + "num_compute_units": None, +} + + +class _FakeEngineClient: + """Minimal stand in exercising `collective_rpc`. Not a real engine.""" + + def __init__(self, rpc_result: Any = None, rpc_error: Exception | None = None): + self.rpc_result = rpc_result + self.rpc_error = rpc_error + self.rpc_calls: list[tuple[str, tuple, dict]] = [] + + async def collective_rpc(self, method, timeout=None, args=(), kwargs=None): + self.rpc_calls.append((method, args, kwargs or {})) + if self.rpc_error is not None: + raise self.rpc_error + return self.rpc_result + + +def _make_test_app(engine_client: "_FakeEngineClient | None") -> FastAPI: + """FastAPI app with the plugin's route + state wired through both phases.""" + app = FastAPI() + plugin = ServerDevicesPlugin() + plugin.attach_router(app) + asyncio.run(plugin.init_state(engine_client, app.state, Namespace())) + return app + + +# --------------------------------------------------------------------------- +# _build_response unit tests +# --------------------------------------------------------------------------- + + +class TestBuildResponse: + def test_single_device(self): + result = _build_response([_A100_ENTRY]) + assert isinstance(result, DevicesResponse) + assert len(result.devices) == 1 + d = result.devices[0] + assert d.rank == 0 + assert d.name == "A100-PCIE-40GB" + assert d.total_memory_bytes == 42_949_672_960 + assert d.compute_capability == ComputeCapability(major=8, minor=0) + assert d.num_compute_units == 108 + + def test_multi_rank(self): + result = _build_response([_A100_ENTRY, _H100_ENTRY]) + assert len(result.devices) == 2 + assert result.devices[0].rank == 0 + assert result.devices[1].rank == 1 + assert result.devices[1].compute_capability == ComputeCapability(major=9, minor=0) + + def test_null_compute_capability_and_num_compute_units(self): + result = _build_response([_NO_CAPABILITY_ENTRY]) + d = result.devices[0] + assert d.compute_capability is None + assert d.num_compute_units is None + + def test_empty_list(self): + result = _build_response([]) + assert result.devices == [] + + def test_preserves_input_order(self): + entries = [ + {**_A100_ENTRY, "rank": 3}, + {**_A100_ENTRY, "rank": 1}, + {**_A100_ENTRY, "rank": 0}, + ] + result = _build_response(entries) + assert [d.rank for d in result.devices] == [3, 1, 0] + + def test_returns_device_info_instances(self): + result = _build_response([_A100_ENTRY]) + assert isinstance(result.devices[0], DeviceInfo) + + +# --------------------------------------------------------------------------- +# HTTP endpoint tests via TestClient +# --------------------------------------------------------------------------- + + +class TestGetDevicesEndpoint: + def test_returns_200_with_devices_from_engine_client(self): + fake_client = _FakeEngineClient(rpc_result=[_A100_ENTRY, _H100_ENTRY]) + app = _make_test_app(fake_client) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/devices") + assert resp.status_code == 200 + data = resp.json() + assert len(data["devices"]) == 2 + assert data["devices"][0]["name"] == "A100-PCIE-40GB" + + def test_calls_collective_rpc_with_get_device_properties(self): + fake_client = _FakeEngineClient(rpc_result=[_A100_ENTRY]) + app = _make_test_app(fake_client) + with TestClient(app) as client: + client.get("/plugins/llm-d-server-introspection/devices") + assert fake_client.rpc_calls == [("get_device_properties", (), {})] + + def test_response_cached_not_rebuilt_per_request(self): + fake_client = _FakeEngineClient(rpc_result=[_A100_ENTRY]) + app = _make_test_app(fake_client) + with TestClient(app) as client: + client.get("/plugins/llm-d-server-introspection/devices") + client.get("/plugins/llm-d-server-introspection/devices") + # collective_rpc only happens once, during init_state. + assert len(fake_client.rpc_calls) == 1 + + def test_engine_client_none_returns_503(self): + app = _make_test_app(None) + with TestClient(app, raise_server_exceptions=False) as client: + resp = client.get("/plugins/llm-d-server-introspection/devices") + assert resp.status_code == 503 + assert "requires an engine" in resp.json()["detail"] + + def test_missing_worker_extension_does_not_crash_init_and_returns_503(self): + # vLLM raises from collective_rpc when no worker has the method. + fake_client = _FakeEngineClient( + rpc_error=AttributeError("'Worker' object has no attribute 'get_device_properties'") + ) + app = _make_test_app(fake_client) + with TestClient(app, raise_server_exceptions=False) as client: + resp = client.get("/plugins/llm-d-server-introspection/devices") + assert resp.status_code == 503 + assert "--worker-extension-cls" in resp.json()["detail"] + + def test_missing_state_returns_500(self): + # attach_router without ever calling init_state. + # The handler must not crash on a missing attribute. + app = FastAPI() + ServerDevicesPlugin().attach_router(app) + with TestClient(app, raise_server_exceptions=False) as client: + resp = client.get("/plugins/llm-d-server-introspection/devices") + assert resp.status_code == 500 + + def test_empty_devices_list_returns_200(self): + fake_client = _FakeEngineClient(rpc_result=[]) + app = _make_test_app(fake_client) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/devices") + assert resp.status_code == 200 + assert resp.json() == {"devices": []} + + def test_content_type_is_json(self): + fake_client = _FakeEngineClient(rpc_result=[_A100_ENTRY]) + app = _make_test_app(fake_client) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/devices") + assert "application/json" in resp.headers["content-type"] + + def test_wrong_method_returns_405(self): + app = _make_test_app(_FakeEngineClient(rpc_result=[])) + with TestClient(app) as client: + resp = client.post("/plugins/llm-d-server-introspection/devices") + assert resp.status_code == 405 + + def test_null_capability_device_serializes_correctly(self): + fake_client = _FakeEngineClient(rpc_result=[_NO_CAPABILITY_ENTRY]) + app = _make_test_app(fake_client) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/devices") + data = resp.json() + assert data["devices"][0]["compute_capability"] is None + assert data["devices"][0]["num_compute_units"] is None + + +# --------------------------------------------------------------------------- +# Plugin metadata +# --------------------------------------------------------------------------- + + +class TestRpcMethodName: + def test_worker_extension_defines_rpc_method(self): + assert callable(getattr(DeviceInfoWorkerExtension, GET_DEVICE_PROPERTIES, None)) + + def test_plugin_calls_shared_rpc_method_name(self): + fake_client = _FakeEngineClient(rpc_result=[]) + _make_test_app(fake_client) + assert fake_client.rpc_calls[0][0] == GET_DEVICE_PROPERTIES + + +class TestPluginMetadata: + def test_name(self): + assert ServerDevicesPlugin().name == "llm_d_server_introspection_devices" + + def test_required_tasks_excludes_render(self): + assert "render" not in ServerDevicesPlugin().required_tasks + + def test_required_tasks_includes_generate(self): + assert "generate" in ServerDevicesPlugin().required_tasks diff --git a/vllm/tests/server_introspection/test_e2e.py b/vllm/tests/server_introspection/test_e2e.py new file mode 100644 index 0000000..ea5733f --- /dev/null +++ b/vllm/tests/server_introspection/test_e2e.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end tests for `GET /plugins/llm-d-server-introspection/{config,devices,kv-cache}` +against a real `vllm serve`. + +Mirrors `tests/plugins_tests/test_endpoint_plugins.py` in vllm-project/vllm: +`pip install -e` this package, launch a tiny model with `VLLM_PLUGINS` set, +assert a real HTTP 200 + schema. + +Unlike the unit tests in `test_config_plugin.py`/`test_devices_plugin.py`/ +`test_kv_cache_plugin.py`, this needs `vllm` importable and downloads model +weights over the network, so it is opt-in: + + pip install -e .[test] + RUN_VLLM_E2E=1 pytest tests/server_introspection/test_e2e.py -v + +It is skipped by default (no `vllm` install, no `RUN_VLLM_E2E`) so `pytest` +without extra setup only runs the fast unit tests. +""" + +import contextlib +import os +import socket +import subprocess +import sys +import time + +import httpx +import pytest + +pytest.importorskip("vllm") + +if not os.environ.get("RUN_VLLM_E2E"): + pytest.skip( + "set RUN_VLLM_E2E=1 to run the real server e2e test (downloads model " + "weights, launches a subprocess server)", + allow_module_level=True, + ) + +MODEL = "facebook/opt-125m" +STARTUP_TIMEOUT_S = 300 + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +@contextlib.contextmanager +def _running_server(*, plugins: str | None, worker_extension: bool = False): + port = _free_port() + env = dict(os.environ) + if plugins: + env["VLLM_PLUGINS"] = plugins + else: + env.pop("VLLM_PLUGINS", None) + args = [ + sys.executable, + "-m", + "vllm.entrypoints.openai.api_server", + "--model", + MODEL, + "--port", + str(port), + ] + if worker_extension: + args += [ + "--worker-extension-cls", + "llm_d_api_extensions_vllm.server_introspection.device_worker_ext.DeviceInfoWorkerExtension", + ] + proc = subprocess.Popen(args, env=env) + base_url = f"http://127.0.0.1:{port}" + try: + deadline = time.monotonic() + STARTUP_TIMEOUT_S + while time.monotonic() < deadline: + if proc.poll() is not None: + raise RuntimeError(f"server process exited early with {proc.returncode}") + try: + resp = httpx.get(f"{base_url}/health", timeout=5) + if resp.status_code == 200: + break + except httpx.HTTPError: + pass + time.sleep(1) + else: + raise TimeoutError("server did not become healthy in time") + yield base_url + finally: + proc.terminate() + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait(timeout=10) + + +def test_server_config_endpoint_returns_200_with_valid_schema(): + with _running_server(plugins="llm_d_server_introspection_config") as base_url: + resp = httpx.get(f"{base_url}/plugins/llm-d-server-introspection/config", timeout=10) + + assert resp.status_code == 200 + data = resp.json() + assert set(data.keys()) == { + "model", + "kv_cache", + "kv_transfer", + "scheduler", + "parallelism", + "features", + } + assert data["model"]["name"] + assert data["parallelism"]["tensor_parallel_size"] == 1 + + +def test_server_config_not_attached_without_allowlist(): + """No VLLM_PLUGINS set -> route must not exist (strict allowlist).""" + with _running_server(plugins=None) as base_url: + resp = httpx.get(f"{base_url}/plugins/llm-d-server-introspection/config", timeout=10) + + assert resp.status_code == 404 + + +def test_server_devices_endpoint_returns_200_with_valid_schema(): + with _running_server( + plugins="llm_d_server_introspection_devices", worker_extension=True + ) as base_url: + resp = httpx.get(f"{base_url}/plugins/llm-d-server-introspection/devices", timeout=10) + + assert resp.status_code == 200 + data = resp.json() + assert set(data.keys()) == {"devices"} + assert len(data["devices"]) == 1 + device = data["devices"][0] + assert device["rank"] == 0 + assert device["name"] + assert device["total_memory_bytes"] > 0 + + +def test_server_devices_not_attached_without_allowlist(): + """No VLLM_PLUGINS set -> route must not exist (strict allowlist).""" + with _running_server(plugins=None) as base_url: + resp = httpx.get(f"{base_url}/plugins/llm-d-server-introspection/devices", timeout=10) + + assert resp.status_code == 404 + + +def test_server_kv_cache_endpoint_returns_200_with_valid_schema(): + with _running_server(plugins="llm_d_server_introspection_kv_cache") as base_url: + resp = httpx.get(f"{base_url}/plugins/llm-d-server-introspection/kv-cache", timeout=10) + + assert resp.status_code == 200 + data = resp.json() + assert set(data.keys()) == { + "kv_cache_size_tokens", + "max_concurrency", + "num_gpu_blocks", + "num_cpu_blocks", + "groups", + } + # `EngineClient.get_kv_cache_group_metadata` (vllm-project/vllm#48121) + # isn't in a released vLLM yet, so this real server exercises the + # capacity only fallback path whioch means capacity fields populated, no group + # structure. + assert data["num_gpu_blocks"] > 0 + assert data["kv_cache_size_tokens"] > 0 + assert data["groups"] == [] + + +def test_server_kv_cache_not_attached_without_allowlist(): + """No VLLM_PLUGINS set -> route must not exist (strict allowlist).""" + with _running_server(plugins=None) as base_url: + resp = httpx.get(f"{base_url}/plugins/llm-d-server-introspection/kv-cache", timeout=10) + + assert resp.status_code == 404 diff --git a/vllm/tests/server_introspection/test_kv_cache_plugin.py b/vllm/tests/server_introspection/test_kv_cache_plugin.py new file mode 100644 index 0000000..62359a2 --- /dev/null +++ b/vllm/tests/server_introspection/test_kv_cache_plugin.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for `GET /plugins/llm-d-server-introspection/kv-cache`. + +Fast, no engine and no real vLLM/model. `_FakeEngineClient` optionally +exposes `get_kv_cache_group_metadata` so both the "new method present" and +"older vLLM falls back to capacity only" `init_state` paths get exercised. +""" + +import asyncio +from argparse import Namespace +from typing import Any +from unittest.mock import MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from kv_group_fixtures import ALL_KINDS, BASE_GROUP, SPEC_CLS_BY_KIND, group_dict + +from llm_d_api_extensions_vllm.server_introspection.kv_cache_plugin import ( + ServerKVCachePlugin, + _build_group_spec, + _build_response, + _capacity_from_vllm_config, +) +from llm_d_api_extensions_vllm.server_introspection.schemas import ( + FullAttentionGroupSpec, + KVCacheResponse, + MambaGroupSpec, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _full_attention_dict(**overrides) -> dict: + return group_dict("full_attention", **overrides) + + +def _mamba_dict(**overrides) -> dict: + return group_dict("mamba", **overrides) + + +def _capacity_data(**overrides) -> dict: + return { + "kv_cache_size_tokens": 16384, + "max_concurrency": 0.5, + "num_gpu_blocks": 1024, + "num_cpu_blocks": 256, + "groups": [], + **overrides, + } + + +def _kv_cache_data(**overrides) -> dict: + overrides.setdefault("groups", [_full_attention_dict()]) + return _capacity_data(**overrides) + + +class _FakeEngineClient: + """Minimal stand in exercising `get_kv_cache_group_metadata`. Not a real engine. + + `get_kv_cache_group_metadata` is only defined as an attribute when + `has_new_method=True`, so `hasattr` feature detection in the plugin can + be exercised both ways. + """ + + def __init__(self, groups: Any = None, has_new_method: bool = True): + self._groups = groups if groups is not None else [] + self.calls = 0 + if has_new_method: + self.get_kv_cache_group_metadata = self._get_kv_cache_group_metadata + + async def _get_kv_cache_group_metadata(self) -> Any: + self.calls += 1 + return self._groups + + +def _make_vllm_config( + *, + num_gpu_blocks: int | None = 1024, + block_size: int | None = 16, + num_cpu_blocks: int | None = 0, + max_model_len: int = 4096, +) -> MagicMock: + cache_cfg = MagicMock() + cache_cfg.num_gpu_blocks = num_gpu_blocks + cache_cfg.block_size = block_size + cache_cfg.num_cpu_blocks = num_cpu_blocks + + model_cfg = MagicMock() + model_cfg.max_model_len = max_model_len + + vllm_config = MagicMock() + vllm_config.cache_config = cache_cfg + vllm_config.model_config = model_cfg + return vllm_config + + +def _make_test_app( + engine_client: "_FakeEngineClient | None", vllm_config: MagicMock | None = None +) -> FastAPI: + """FastAPI app with the plugin's route + state wired through both phases.""" + app = FastAPI() + plugin = ServerKVCachePlugin() + plugin.attach_router(app) + app.state.vllm_config = vllm_config if vllm_config is not None else _make_vllm_config() + asyncio.run(plugin.init_state(engine_client, app.state, Namespace())) + return app + + +# --------------------------------------------------------------------------- +# _build_group_spec unit tests — discriminator dispatch +# --------------------------------------------------------------------------- + + +class TestBuildGroupSpecDispatch: + @pytest.mark.parametrize("kind", ALL_KINDS) + def test_dispatches_on_kind(self, kind): + result = _build_group_spec(group_dict(kind)) + assert type(result) is SPEC_CLS_BY_KIND[kind] + assert result.kind == kind + + def test_uniform_type_spec_resolves_to_inner_kind(self): + # A UniformTypeKVCacheSpecs group resolves to its inner kind (here + # full_attention) and carries a populated `layer_specs` list. + layer_specs = [ + {"head_size": 128, "dtype": "bfloat16"}, + {"head_size": 64, "dtype": "bfloat16"}, + ] + result = _build_group_spec(_full_attention_dict(layer_specs=layer_specs)) + assert isinstance(result, FullAttentionGroupSpec) + assert result.layer_specs == layer_specs + + def test_unhandled_kind_raises_value_error(self): + group = {**BASE_GROUP, "kind": "made_up_kind"} + with pytest.raises(ValueError, match="Unhandled KVCacheSpec kind"): + _build_group_spec(group) + + def test_preserves_base_fields(self): + result = _build_group_spec(_full_attention_dict(group_id=3)) + assert result.group_id == 3 + assert result.layer_count == 2 + assert result.layer_names == BASE_GROUP["layer_names"] + assert result.block_size == 16 + assert result.page_size_bytes == 131072 + + def test_ignores_irrelevant_always_present_fields(self): + # Real payloads send every field for every kind (irrelevant ones as + # None) - a full_attention group also carries mamba/sink fields. + group = _full_attention_dict(shapes=None, dtypes=None, mamba_type=None, sink_len=None) + result = _build_group_spec(group) + assert isinstance(result, FullAttentionGroupSpec) + + +# --------------------------------------------------------------------------- +# _build_response unit tests — pure function +# --------------------------------------------------------------------------- + + +class TestBuildResponseNullState: + def test_none_state(self): + resp = _build_response(None) + assert isinstance(resp, KVCacheResponse) + assert resp.kv_cache_size_tokens is None + assert resp.max_concurrency is None + assert resp.num_gpu_blocks is None + assert resp.num_cpu_blocks is None + assert resp.groups == [] + + +class TestBuildResponseCapacityFields: + def test_kv_cache_size_tokens(self): + resp = _build_response(_kv_cache_data(kv_cache_size_tokens=32768)) + assert resp.kv_cache_size_tokens == 32768 + + def test_max_concurrency(self): + resp = _build_response(_kv_cache_data(max_concurrency=1.25)) + assert resp.max_concurrency == pytest.approx(1.25) + + def test_num_gpu_blocks(self): + resp = _build_response(_kv_cache_data(num_gpu_blocks=2048)) + assert resp.num_gpu_blocks == 2048 + + def test_num_cpu_blocks_zero(self): + resp = _build_response(_kv_cache_data(num_cpu_blocks=0)) + assert resp.num_cpu_blocks == 0 + + def test_missing_capacity_fields_resolve_to_none(self): + resp = _build_response({"groups": []}) + assert resp.kv_cache_size_tokens is None + assert resp.num_gpu_blocks is None + + +class TestBuildResponseGroups: + def test_hybrid_model_two_groups(self): + groups = [_full_attention_dict(group_id=0), _mamba_dict(group_id=1)] + resp = _build_response(_kv_cache_data(groups=groups)) + assert len(resp.groups) == 2 + assert isinstance(resp.groups[0], FullAttentionGroupSpec) + assert isinstance(resp.groups[1], MambaGroupSpec) + + def test_empty_groups_list(self): + resp = _build_response(_kv_cache_data(groups=[])) + assert resp.groups == [] + + +# --------------------------------------------------------------------------- +# _capacity_from_vllm_config unit tests +# --------------------------------------------------------------------------- + + +class TestCapacityFromVllmConfig: + def test_computes_kv_cache_size_tokens(self): + cfg = _make_vllm_config(num_gpu_blocks=1024, block_size=16) + data = _capacity_from_vllm_config(cfg) + assert data["kv_cache_size_tokens"] == 16384 + assert data["num_gpu_blocks"] == 1024 + assert data["groups"] == [] + + def test_computes_max_concurrency(self): + cfg = _make_vllm_config(num_gpu_blocks=1024, block_size=16, max_model_len=4096) + data = _capacity_from_vllm_config(cfg) + assert data["max_concurrency"] == pytest.approx(16384 / 4096) + + def test_none_num_gpu_blocks_yields_none_fields(self): + cfg = _make_vllm_config(num_gpu_blocks=None) + data = _capacity_from_vllm_config(cfg) + assert data["kv_cache_size_tokens"] is None + assert data["max_concurrency"] is None + assert data["num_gpu_blocks"] is None + + def test_num_cpu_blocks_passthrough(self): + cfg = _make_vllm_config(num_cpu_blocks=512) + data = _capacity_from_vllm_config(cfg) + assert data["num_cpu_blocks"] == 512 + + +# --------------------------------------------------------------------------- +# HTTP endpoint tests via TestClient +# --------------------------------------------------------------------------- + + +class TestGetKVCacheEndpoint: + def test_returns_200_with_data_from_new_method(self): + fake_client = _FakeEngineClient([_full_attention_dict()], has_new_method=True) + vllm_config = _make_vllm_config(num_gpu_blocks=1024, block_size=16) + app = _make_test_app(fake_client, vllm_config=vllm_config) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/kv-cache") + assert resp.status_code == 200 + data = resp.json() + assert data["num_gpu_blocks"] == 1024 + assert len(data["groups"]) == 1 + assert fake_client.calls == 1 + + def test_calls_get_kv_cache_group_metadata_once_and_caches(self): + fake_client = _FakeEngineClient([_full_attention_dict()], has_new_method=True) + app = _make_test_app(fake_client) + with TestClient(app) as client: + client.get("/plugins/llm-d-server-introspection/kv-cache") + client.get("/plugins/llm-d-server-introspection/kv-cache") + assert fake_client.calls == 1 + + def test_degrades_to_capacity_only_without_new_method(self): + fake_client = _FakeEngineClient(has_new_method=False) + vllm_config = _make_vllm_config(num_gpu_blocks=2048, block_size=16) + app = _make_test_app(fake_client, vllm_config=vllm_config) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/kv-cache") + assert resp.status_code == 200 + data = resp.json() + assert data["num_gpu_blocks"] == 2048 + assert data["kv_cache_size_tokens"] == 32768 + assert data["groups"] == [] + + def test_engine_client_none_returns_503(self): + app = _make_test_app(None) + with TestClient(app, raise_server_exceptions=False) as client: + resp = client.get("/plugins/llm-d-server-introspection/kv-cache") + assert resp.status_code == 503 + + def test_missing_state_returns_500(self): + app = FastAPI() + ServerKVCachePlugin().attach_router(app) + with TestClient(app, raise_server_exceptions=False) as client: + resp = client.get("/plugins/llm-d-server-introspection/kv-cache") + assert resp.status_code == 500 + + def test_content_type_is_json(self): + fake_client = _FakeEngineClient([_full_attention_dict()]) + app = _make_test_app(fake_client) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/kv-cache") + assert "application/json" in resp.headers["content-type"] + + def test_wrong_method_returns_405(self): + app = _make_test_app(_FakeEngineClient([_full_attention_dict()])) + with TestClient(app) as client: + resp = client.post("/plugins/llm-d-server-introspection/kv-cache") + assert resp.status_code == 405 + + def test_groups_serialized_correctly(self): + groups = [_full_attention_dict(group_id=0), _mamba_dict(group_id=1)] + fake_client = _FakeEngineClient(groups) + app = _make_test_app(fake_client) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/kv-cache") + data = resp.json()["groups"] + assert data[0]["kind"] == "full_attention" + assert data[1]["kind"] == "mamba" + + def test_top_level_keys_present(self): + fake_client = _FakeEngineClient([_full_attention_dict()]) + app = _make_test_app(fake_client) + with TestClient(app) as client: + resp = client.get("/plugins/llm-d-server-introspection/kv-cache") + assert set(resp.json().keys()) == { + "kv_cache_size_tokens", + "max_concurrency", + "num_gpu_blocks", + "num_cpu_blocks", + "groups", + } + + +# --------------------------------------------------------------------------- +# Plugin metadata +# --------------------------------------------------------------------------- + + +class TestPluginMetadata: + def test_name(self): + assert ServerKVCachePlugin().name == "llm_d_server_introspection_kv_cache" + + def test_required_tasks_excludes_render(self): + assert "render" not in ServerKVCachePlugin().required_tasks + + def test_required_tasks_includes_generate(self): + assert "generate" in ServerKVCachePlugin().required_tasks diff --git a/vllm/tests/server_introspection/test_schemas.py b/vllm/tests/server_introspection/test_schemas.py new file mode 100644 index 0000000..6b64e8c --- /dev/null +++ b/vllm/tests/server_introspection/test_schemas.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pydantic model tests for the `/plugins/llm-d-server-introspection/config` response schema.""" + +import pytest +from kv_group_fixtures import ALL_KINDS, SPEC_CLS_BY_KIND, group_dict + +from llm_d_api_extensions_vllm.server_introspection.schemas import ( + ComputeCapability, + DeviceInfo, + DevicesResponse, + FeaturesInfo, + FullAttentionGroupSpec, + KVCacheInfo, + KVCacheResponse, + MambaGroupSpec, + ModelInfo, + ParallelismInfo, + SchedulerInfo, + ServerConfigResponse, + SinkFullAttentionGroupSpec, +) + + +def _server_config_response() -> ServerConfigResponse: + return ServerConfigResponse( + model=ModelInfo( + name="llama", + served_names=["llama"], + dtype="bfloat16", + quantization=None, + max_model_len=4096, + ), + kv_cache=KVCacheInfo( + gpu_memory_utilization=0.9, + dtype="bfloat16", + enable_prefix_caching=True, + ), + scheduler=SchedulerInfo( + max_num_seqs=64, + max_num_batched_tokens=1024, + enable_chunked_prefill=True, + policy="fcfs", + ), + parallelism=ParallelismInfo( + tensor_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + data_parallel_rank=0, + ), + features=FeaturesInfo( + speculative_decoding=False, + lora=False, + hma=True, + ), + ) + + +def test_server_config_response_roundtrip(): + original = _server_config_response() + restored = ServerConfigResponse.model_validate(original.model_dump()) + assert restored == original + + +def test_quantization_null_serialization(): + dumped = _server_config_response().model_dump() + assert dumped["model"]["quantization"] is None + + +def test_quantization_string_serialization(): + resp = _server_config_response() + resp.model.quantization = "gptq" + assert resp.model_dump()["model"]["quantization"] == "gptq" + + +def test_features_info_all_bool(): + feat = FeaturesInfo(speculative_decoding=True, lora=False, hma=True) + assert feat.model_dump() == { + "speculative_decoding": True, + "lora": False, + "hma": True, + } + + +def test_server_config_response_json_schema_has_required_sections(): + schema = ServerConfigResponse.model_json_schema() + assert set(schema["properties"]) == { + "model", + "kv_cache", + "scheduler", + "parallelism", + "features", + "kv_transfer", + } + + +def _device_info(**overrides) -> DeviceInfo: + fields = dict( + rank=0, + name="A100-PCIE-40GB", + total_memory_bytes=42_949_672_960, + compute_capability=ComputeCapability(major=8, minor=0), + num_compute_units=108, + ) + fields.update(overrides) + return DeviceInfo(**fields) + + +def test_devices_response_roundtrip(): + original = DevicesResponse(devices=[_device_info(), _device_info(rank=1)]) + restored = DevicesResponse.model_validate(original.model_dump()) + assert restored == original + + +def test_devices_response_empty(): + assert DevicesResponse(devices=[]).devices == [] + + +def test_device_info_compute_capability_defaults_to_none(): + device = DeviceInfo(rank=0, name="cpu", total_memory_bytes=1024) + assert device.compute_capability is None + assert device.num_compute_units is None + + +def test_device_info_null_capability_serialization(): + dumped = _device_info(compute_capability=None).model_dump() + assert dumped["compute_capability"] is None + + +def test_devices_response_json_schema_shape(): + schema = DevicesResponse.model_json_schema() + assert set(schema["properties"]) == {"devices"} + + +# --------------------------------------------------------------------------- +# kv-cache group spec + response tests +# --------------------------------------------------------------------------- + + +def _spec(kind: str, **kwargs): + return SPEC_CLS_BY_KIND[kind].model_validate(group_dict(kind, **kwargs)) + + +@pytest.mark.parametrize("kind", ALL_KINDS) +def test_group_spec_roundtrip(kind): + spec = _spec(kind) + restored = type(spec).model_validate(spec.model_dump()) + assert restored == spec + assert spec.kind == kind + + +def test_full_attention_optional_fields_default_none(): + spec = _spec("full_attention", omit=("sliding_window", "attention_chunk_size")) + assert spec.sliding_window is None + assert spec.attention_chunk_size is None + + +def test_sink_full_attention_sink_len_optional(): + spec = _spec("sink_full_attention", omit=("sink_len",)) + assert isinstance(spec, SinkFullAttentionGroupSpec) + assert spec.sink_len is None + + +def test_full_attention_layer_specs_populated_for_uniform_type_group(): + # A UniformTypeKVCacheSpecs group resolves to its inner kind but keeps + # layer_specs populated (None for a regular, non fanned out group). + spec = _spec("full_attention", layer_specs=[{"head_size": 128}, {"head_size": 64}]) + restored = FullAttentionGroupSpec.model_validate(spec.model_dump()) + assert restored == spec + assert spec.layer_specs == [{"head_size": 128}, {"head_size": 64}] + + +def test_kv_cache_response_defaults(): + resp = KVCacheResponse() + assert resp.kv_cache_size_tokens is None + assert resp.max_concurrency is None + assert resp.num_gpu_blocks is None + assert resp.num_cpu_blocks is None + assert resp.groups == [] + + +def test_kv_cache_response_roundtrip_with_groups(): + original = KVCacheResponse( + kv_cache_size_tokens=16384, + max_concurrency=0.5, + num_gpu_blocks=1024, + num_cpu_blocks=256, + groups=[_spec("full_attention")], + ) + restored = KVCacheResponse.model_validate(original.model_dump()) + assert restored == original + + +def test_kv_cache_response_json_schema_has_expected_properties(): + schema = KVCacheResponse.model_json_schema() + assert set(schema["properties"]) == { + "kv_cache_size_tokens", + "max_concurrency", + "num_gpu_blocks", + "num_cpu_blocks", + "groups", + } + + +def test_kv_cache_response_discriminates_group_spec_types_on_validate(): + # A discriminated union must reconstruct the correct concrete class from + # raw JSON, not just accept already typed model instances. + dumped = KVCacheResponse( + groups=[_spec("full_attention"), _spec("mamba", group_id=1)] + ).model_dump() + restored = KVCacheResponse.model_validate(dumped) + assert isinstance(restored.groups[0], FullAttentionGroupSpec) + assert isinstance(restored.groups[1], MambaGroupSpec)