diff --git a/CLAUDE.md b/CLAUDE.md index e3ef4b4ca..a22a8e20c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -289,11 +289,28 @@ with patch.dict(os.environ, H100Rack.slurm_env()): - `build_worker_command(process, runtime)` - Command construction 3. Export from `backends/__init__.py` 4. Add polymorphic deserialization in `BackendConfigField` in `schema.py` +5. Register the per-role argument key in `core/roles.py` and backend class in + `core/schema_docs.py`, then regenerate schema docs. Recipes use `engine:` and + `roles..args` / `.env`. + +Optional role-image and finite preparation hooks are described in +[TileRT's extension-point reference](docs/tilert.md#backend-extension-points). +Keep engine-specific commands in the backend; use the shared worker launcher +and process registry for lifecycle handling. **Current backends:** - **SGLang**: Per-process srun launching, supports prefill/decode/aggregated modes - **TRTLLM**: MPI-style launching (one srun per endpoint with all nodes), prefill/decode only +### Adding a New Frontend + +Implement `FrontendProtocol` (or extend `StaticRouterFrontend` for a router with +explicit worker URLs), export it from `frontends/__init__.py`, and add its factory +entry in `frontends/base.py`. Validate backend pairing in `SrtConfig`. Supply the +frontend health parser and direct backend health URLs when router health alone +is insufficient. Use `start_srun_process` and `ManagedProcess` with matching +`step_name` values so normal signal cleanup reaches the router. + ### Adding a New Benchmark 1. Create `benchmarks/mybench.py` inheriting from `BenchmarkRunner` diff --git a/configs/tilert_setup.sh b/configs/tilert_setup.sh new file mode 100644 index 000000000..63776d640 --- /dev/null +++ b/configs/tilert_setup.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +readonly TILERT_VERSION="${TILERT_VERSION:-0.1.5.post3}" +readonly NIXL_VERSION="${TILERT_NIXL_VERSION:-1.3.1}" +readonly FASTAPI_VERSION="${TILERT_FASTAPI_VERSION:-0.141.1}" +readonly UVICORN_VERSION="${TILERT_UVICORN_VERSION:-0.52.4}" +readonly TRANSFORMERS_VERSION="${TILERT_TRANSFORMERS_VERSION:-5.16.1}" + +if [[ -d /opt/conda/envs/tilert/bin ]]; then + export PATH="/opt/conda/envs/tilert/bin:${PATH}" +fi + +PYTHON="$(command -v python || command -v python3)" +readonly PYTHON + +installed_version() { + "${PYTHON}" - "$1" <<'PY' +import sys +from importlib.metadata import PackageNotFoundError, version + +try: + print(version(sys.argv[1])) +except PackageNotFoundError: + pass +PY +} + +install_exact() { + local package="$1" + local version="$2" + shift 2 + if [[ "$(installed_version "${package}")" == "${version}" ]]; then + return + fi + "${PYTHON}" -m pip install --no-cache-dir "$@" "${package}==${version}" +} + +case "${TILERT_ROLE:-}" in + decode) + install_exact tilert "${TILERT_VERSION}" + install_exact nixl "${NIXL_VERSION}" + # TileRT's wheel does not declare the HTTP service dependencies used + # by its native decode server. Install them explicitly before the + # import preflight so the image is made runnable deterministically. + install_exact fastapi "${FASTAPI_VERSION}" + install_exact uvicorn "${UVICORN_VERSION}" + "${PYTHON}" -c 'import tilert.pd_vllm.decode_server' + ;; + prefill) + "${PYTHON}" -c 'import vllm' + install_exact tilert "${TILERT_VERSION}" --no-deps + install_exact nixl "${NIXL_VERSION}" + "${PYTHON}" -c 'import tilert.pd_vllm.prefill_connector' + ;; + router) + install_exact tilert "${TILERT_VERSION}" + # GLM-5.1 tokenizer metadata names the Transformers v5 + # TokenizersBackend class. TileRT's broad >=4.46.3 dependency leaves + # its native router on the image's incompatible 4.x release unless we + # make the router-side tokenizer contract explicit. + install_exact transformers "${TRANSFORMERS_VERSION}" + # The native router uses the same undeclared FastAPI/Uvicorn stack as + # the decode server and must be self-contained on the frontend host. + install_exact fastapi "${FASTAPI_VERSION}" + install_exact uvicorn "${UVICORN_VERSION}" + "${PYTHON}" -c 'from transformers import TokenizersBackend' + "${PYTHON}" -c 'import tilert.pd_vllm.pd_router' + ;; + *) + echo "TILERT_ROLE must be prefill, decode, or router (got ${TILERT_ROLE:-unset})" >&2 + exit 1 + ;; +esac diff --git a/docs/legacy-v1.md b/docs/legacy-v1.md index e2669996e..fa44033b6 100644 --- a/docs/legacy-v1.md +++ b/docs/legacy-v1.md @@ -40,6 +40,7 @@ Recipes without `schema: 2` (or with `schema: 1`) use the pre-2.0 layout: worker | `backend.trtllm_config` | `roles..args` (one mapping per role; the `prefill` / `decode` / `aggregated` keys) | | `backend.vllm_config` | `roles..args` (one mapping per role; the `prefill` / `decode` / `aggregated` keys) | | `backend.mocker_config` | `roles..args` (one mapping per role; the `prefill` / `decode` / `aggregated` keys) | +| `backend.tilert_config` | `roles..args` (one mapping per role; the `prefill` / `decode` / `aggregated` keys) | ## v1 values that changed meaning @@ -144,6 +145,16 @@ These keys exist in both layouts, but the value means something else in 2.0. A s | `aggregated_environment` | dict[str, str] | `{}` | | | `mocker_config` | [MockerServerConfig](#mockerserverconfig) \| None | `None` | Per-mode CLI overrides | +### TileRTProtocol + +`backend.type: tilert` + +| Key | Type | Default | Description | +|---|---|---|---| +| `prefill_environment` | dict[str, str] | `{}` | | +| `decode_environment` | dict[str, str] | `{}` | | +| `tilert_config` | [TileRTServerConfig](#tilertserverconfig) \| None | `None` | | + ### SGLangServerConfig SGLang server CLI configuration per mode (prefill/decode/aggregated). @@ -205,6 +216,15 @@ Mocker CLI configuration per mode (prefill/decode/aggregated). | `decode` | dict[str, Any] \| None | `None` | | | `aggregated` | dict[str, Any] \| None | `None` | | +### TileRTServerConfig + +Role-specific CLI arguments not owned by the topology adapter. + +| Key | Type | Default | Description | +|---|---|---|---| +| `prefill` | dict[str, Any] \| None | `None` | | +| `decode` | dict[str, Any] \| None | `None` | | + ## infra ### InfraConfig diff --git a/docs/schema-reference.md b/docs/schema-reference.md index 23a38efa3..62122f7f1 100644 --- a/docs/schema-reference.md +++ b/docs/schema-reference.md @@ -13,7 +13,7 @@ Top-level keys of a recipe YAML. | `name` | str | required | | | `model` | [ModelConfig](#modelconfig) | required | | | `resources` | [ResourceConfig](#resourceconfig) | required | | -| `engine` | str \| mapping | required | The engine type (`sglang`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types). | +| `engine` | str \| mapping | required | The engine type (`sglang`, `trtllm`, `vllm`, `mocker`, `tilert`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types). | | `roles` | mapping of role -> [Role](#roles) | required | One block per worker role (`prefill`, `decode`, `agg`): topology, env, and engine args. | | `schema` | int | `2` | Recipe schema version. Write `schema: 2` for this layout. | | `slurm` | [SlurmConfig](#slurmconfig) | `SlurmConfig()` | | @@ -44,7 +44,7 @@ Three vocabularies are specific to the 2.0 layout. They are normalized into the ### engine -`engine: ` or `engine: {type: , ...}`. `type` is one of `sglang`, `trtllm`, `vllm`, `mocker`; the remaining keys are that engine's knobs, listed under [Engine types](#engine-types). +`engine: ` or `engine: {type: , ...}`. `type` is one of `sglang`, `trtllm`, `vllm`, `mocker`, `tilert`; the remaining keys are that engine's knobs, listed under [Engine types](#engine-types). ### roles @@ -112,7 +112,7 @@ Frontend/router configuration. | Key | Type | Default | Description | |---|---|---|---| -| `type` | str | `'dynamo'` | Frontend type - "dynamo" (default); "sglang-router" (SGLang Model Gateway) and "vllm-router" (static routers); "sglang", "vllm", and "trtllm_serve" (direct: the single aggregate worker binds the public port, no router process). In schema 1 recipes "sglang" still means the router and loads as "sglang-router". | +| `type` | str | `'dynamo'` | Frontend type - "dynamo" (default); "sglang-router" (SGLang Model Gateway), "vllm-router", and "tilert-router" (static routers); "sglang", "vllm", and "trtllm_serve" (direct: the single aggregate worker binds the public port, no router process). In schema 1 recipes "sglang" still means the router and loads as "sglang-router". | | `enable_multiple_frontends` | bool | `True` | Scale with nginx + multiple routers. When ``True`` (default), srtctl stands up nginx and fans out to ``num_additional_frontends + 1`` router replicas. When ``False``, there is NO nginx proxy — the benchmark must target the single master router (or a worker) directly at ``http://localhost:``. ``benchmark.command`` has no placeholder substitution, so write the URL out literally. | | `num_additional_frontends` | int | `9` | Additional routers beyond master (default: 9) | | `nginx_container` | str | `'nginx:1.27.4'` | Custom nginx container image (default: nginx:1.27.4) | @@ -595,6 +595,27 @@ Dynamo Mocker protocol - implements BackendProtocol. | `enable_chunked_prefill` | bool | `True` | | | `preemption_mode` | str \| None | `None` | | +### TileRTProtocol + +`engine.type: tilert` + +Run vLLM prefill workers and TileRT decode workers without Dynamo. + +| Key | Type | Default | Description | +|---|---|---|---| +| `type` | one of `'tilert'` | `'tilert'` | | +| `served_model_name` | str \| None | `None` | OpenAI model name; defaults to the final component of model.path. | +| `prefill_container` | str \| None | `None` | vLLM image or cluster container alias for prefill workers only. | +| `model_profile` | str | `'glm5'` | TileRT P/D model profile (for example glm5). | +| `weight_model_type` | str | `'glm-5'` | Model type passed to the TileRT weight converter. | +| `weights_dir` | str | `'/tilert_weights/converted'` | Dedicated shared converted-weight directory, mounted at this path in all decode containers. | +| `max_seq_len` | int | `202752` | Maximum sequence length shared by prefill and decode. | +| `kv_cache_dtype` | str | `'fp8'` | TileRT decode KV cache dtype. | +| `prefill_kv_cache_dtype` | str | `'fp8_ds_mla'` | vLLM KV cache dtype; must be wire-compatible with the decode dtype. | +| `transport` | one of `'nixl'` | `'nixl'` | KV transfer transport (the adapter supports NIXL). | +| `with_mtp` | bool | `True` | Enable MTP on both sides of the P/D transfer. | +| `speculative_tokens` | int | `1` | Number of speculative tokens advertised by vLLM prefill. | + ## Cluster config Top-level keys of `srtslurm.yaml`. Recipes inherit these defaults and resolve aliases through them. diff --git a/docs/tilert.md b/docs/tilert.md new file mode 100644 index 000000000..efb4c50e1 --- /dev/null +++ b/docs/tilert.md @@ -0,0 +1,95 @@ +# TileRT with vLLM prefill + +`engine.type: tilert` pairs a native vLLM prefill server with native TileRT decode +servers. `frontend.type: tilert-router` runs TileRT's own P/D router. It does not +start Dynamo, etcd, or NATS. The adapter implements `BackendProtocol`; the router +extends `StaticRouterFrontend` and uses its worker discovery, Slurm launch, and +process cleanup. + +See [`examples/tilert/glm5-disagg.yaml`](https://github.com/SemiAnalysisAI/srt-slurm/blob/main/examples/tilert/glm5-disagg.yaml) +for a schema 2 recipe. This is a launch example, not a claim of a hardware-validated +benchmark result. The native interfaces follow +[TileRT's P/D implementation](https://github.com/tile-ai/TileRT/tree/main/tilert/pd_vllm). + +## Images and setup + +- `model.container` selects the TileRT decode image. +- `engine.prefill_container` selects the vLLM prefill image. It accepts the same + cluster `containers:` aliases as `model.container`. +- `frontend.container_image` optionally selects a router image, otherwise it + uses the decode image. + +The example uses `ghcr.io/tile-ai/tilert:0.1.5` and +`vllm/vllm-openai:v0.26.0`. `setup_script: tilert_setup.sh` installs +TileRT `0.1.5.post3` and NIXL `1.3.1`; it adds the HTTP dependencies to decode and +router containers and Transformers 5 to the router for GLM-5.1 tokenizers. +The prefill install uses `--no-deps` for TileRT to preserve vLLM's dependencies. +The script accepts `TILERT_VERSION`, `TILERT_NIXL_VERSION`, +`TILERT_FASTAPI_VERSION`, `TILERT_UVICORN_VERSION`, and +`TILERT_TRANSFORMERS_VERSION` overrides. Omit setup when using prepared images. +The adapter supplies `TILERT_ROLE` to each worker and to the router setup. + +Use `roles.prefill.args` / `roles.decode.args` and `roles..env` for normal +engine arguments and environment. srtctl owns ports, model paths, prefill TP, +KV transfer configuration, and MTP flags. Attempts to override those arguments +are rejected. `engine.served_model_name` overrides the OpenAI model name; +otherwise the final component of `model.path` is used. + +## Weight conversion and mounts + +TileRT requires converted weights. Before workers launch, the backend requests +one GPU on the first decode worker's node and runs TileRT's weight converter in +the decode image. The step is tracked in the normal process registry, including +signal cleanup, failure propagation, and a six-hour timeout. Its log is +`tilert_weight_conversion.out`. + +Set `engine.weights_dir` to a dedicated directory under a shared writable mount, +visible at the same container path to decode workers **and the router**. For +example, add `/shared/tilert:/tilert_weights` to the cluster's mounts. Use a +separate directory for every model, revision, and converter version; the adapter +does not infer cache identity from model contents. Do not point it at source +weights or a directory with unrelated files. + +Conversion uses a filesystem lock shared by concurrent jobs, writes to a sibling +temporary directory, copies tokenizer metadata and the chat template, and +publishes only after checking the index and every referenced shard. Complete +caches are reused. A cache with missing shards or tokenizer metadata is rebuilt. +Failed conversions leave the previous target untouched. The router loads its +tokenizer from this prepared directory, including for local model paths. +Hugging Face sources are completed through `snapshot_download` before conversion; +local sources use the worker's mounted model path. + +## Topology, readiness, and metrics + +The native router supports exactly one prefill worker and one or more decode +workers. Every worker must fit on one node. The GLM-5 decode kernels require +exactly eight GPUs per decode worker. Use one frontend +(`enable_multiple_frontends: false`); aggregate mode, Dynamo sidecars, and Torch +profiling are rejected. Nsys profiling remains available. + +The backend configures the TileRT NIXL connector on vLLM, derives TP from the +allocated prefill GPUs, and supplies matching context length and transport to +both roles. `engine.with_mtp` controls MTP on both sides. Choose compatible +`kv_cache_dtype` and `prefill_kv_cache_dtype` values for the model profile. + +The router is started after every worker `/health` endpoint is ready. Before +benchmark traffic, srtctl checks the router's `/health` plus all worker endpoints; +a healthy router alone does not prove the workers are ready. AIPerf and Tachometer receive +only vLLM prefill Prometheus URLs because TileRT decode does not expose `/metrics`. + +## Backend extension points + +TileRT uses two optional backend hooks alongside the required `BackendProtocol` +methods: + +- `get_container_image_for_mode(mode, default)` selects a role image for worker + and preparation launches. Backends without it keep `model.container`. +- `get_preparation(runtime, processes)` returns a `BackendPreparation` describing + a finite pre-worker command, selected node/mode, log, timeout, and optional GPU + request. The orchestrator runs it after HF download/model staging and before + workers, using that role's image, setup, mounts, and environment. + +`get_metrics_port(process)` optionally limits AIPerf discovery to real metrics +endpoints for both AIPerf and Tachometer. Returning `None` omits that process from metrics, without changing +its serving or profiling endpoint. Backends without it retain existing discovery. +A frontend may set `has_metrics = False` to omit its own Tachometer target. diff --git a/examples/README.md b/examples/README.md index 25ebbcbe1..37ffc3290 100644 --- a/examples/README.md +++ b/examples/README.md @@ -32,6 +32,12 @@ Every example is written in the 2.0 layout: `engine:` names the engine (a string | `features/infra-services.yaml` | etcd and NATS as declared services on a dedicated node with a NATS payload limit; the implied exporters overridden or switched off | | `features/dynamo-source.yaml` | `dynamo.source:` building Dynamo from a git tag (or a PR head via `--set dynamo.source.rev=refs/pull//head`), pinned to a commit at submit | +## TileRT + +`tilert/glm5-disagg.yaml` runs vLLM prefill and TileRT decode with a native +TileRT router on two B200 nodes. See [the TileRT guide](../docs/tilert.md) for +image and shared converted-weight mount requirements. + ## Cluster aliases The examples reference two kinds of aliases that `srtslurm.yaml` resolves: diff --git a/examples/tilert/glm5-disagg.yaml b/examples/tilert/glm5-disagg.yaml new file mode 100644 index 000000000..fe2679de6 --- /dev/null +++ b/examples/tilert/glm5-disagg.yaml @@ -0,0 +1,65 @@ +# Native vLLM prefill + TileRT decode, one 8-GPU B200 node per role. +# Mount a shared writable directory at /tilert_weights on every node. +# See docs/tilert.md for images, setup, and cache requirements. +schema: 2 +name: glm5-tilert-disagg + +model: + path: hf:zai-org/GLM-5.1-FP8 + container: ghcr.io/tile-ai/tilert:0.1.5 + precision: fp8 + +identity: + model: + repo: zai-org/GLM-5.1-FP8 + revision: f396cf805182f4ca10fa675e1a99815b3ca384db + container: + image: ghcr.io/tile-ai/tilert:0.1.5 + frameworks: + tilert: 0.1.5.post3 + vllm: 0.26.0 + nixl: 1.3.1 + +resources: + gpu_type: b200 + gpus_per_node: 8 + +engine: + type: tilert + prefill_container: vllm/vllm-openai:v0.26.0 + model_profile: glm5 + weight_model_type: glm-5 + weights_dir: /tilert_weights/glm5.1-fp8 + max_seq_len: 202752 + with_mtp: true + +roles: + prefill: + nodes: 1 + workers: 1 + gpus: 8 + args: + enforce-eager: true + trust-remote-code: true + gpu-memory-utilization: 0.75 + decode: + nodes: 1 + workers: 1 + gpus: 8 + +frontend: + type: tilert-router + enable_multiple_frontends: false + args: + parser: none + queue-timeout: 1800 + +setup_script: tilert_setup.sh +srun_options: + container-writable: "" +health_check: + max_attempts: 720 + interval_seconds: 5 + +benchmark: + type: manual diff --git a/mkdocs.yml b/mkdocs.yml index 64f0ffd58..be7475c62 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -59,6 +59,7 @@ nav: - Architecture: architecture.md - SGLang Router: sglang-router.md - vLLM Router: vllm-router.md + - TileRT: tilert.md - Mooncake KV Store: mooncake-kv-store.md - Services: services.md - Benchmarking: diff --git a/src/srtctl/backends/__init__.py b/src/srtctl/backends/__init__.py index 99618d70a..10b5cd689 100644 --- a/src/srtctl/backends/__init__.py +++ b/src/srtctl/backends/__init__.py @@ -12,11 +12,12 @@ from .base import BackendProtocol, BackendType, SrunConfig from .mocker import MockerProtocol, MockerServerConfig from .sglang import SGLangProtocol, SGLangServerConfig +from .tilert import TileRTProtocol, TileRTServerConfig from .trtllm import TRTLLMProtocol, TRTLLMServerConfig from .vllm import VLLMProtocol, VLLMServerConfig # Union type for all backend configs -BackendConfig = SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol +BackendConfig = SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol | TileRTProtocol __all__ = [ "BackendConfig", @@ -33,6 +34,9 @@ # TRTLLM "TRTLLMProtocol", "TRTLLMServerConfig", + # TileRT + "TileRTProtocol", + "TileRTServerConfig", # vLLM "VLLMProtocol", "VLLMServerConfig", diff --git a/src/srtctl/backends/base.py b/src/srtctl/backends/base.py index d71d3b63c..add7fa899 100644 --- a/src/srtctl/backends/base.py +++ b/src/srtctl/backends/base.py @@ -27,6 +27,7 @@ class BackendType(str, Enum): TRTLLM = "trtllm" VLLM = "vllm" MOCKER = "mocker" + TILERT = "tilert" @dataclass @@ -47,6 +48,18 @@ class SrunConfig: cpu_bind: str | None = None +@dataclass(frozen=True) +class BackendPreparation: + """One synchronous backend-owned preparation step before worker launch.""" + + command: list[str] + node: str + mode: str + log_name: str + timeout_seconds: int = 21600 + gpus_per_task: int | None = None + + class BackendProtocol(Protocol): """Protocol that all backend configurations must implement. @@ -55,6 +68,11 @@ class BackendProtocol(Protocol): 1. Allocating logical endpoints (serving units) 2. Converting endpoints to physical processes 3. Building commands to start those processes + + Optional hooks supported by the worker/telemetry stages (backends without + them retain the defaults): get_container_image_for_mode(mode, default), + get_preparation(runtime, processes) -> BackendPreparation | None, and + get_metrics_port(process) -> int | None. See docs/tilert.md. """ @property diff --git a/src/srtctl/backends/tilert.py b/src/srtctl/backends/tilert.py new file mode 100644 index 000000000..b662638bf --- /dev/null +++ b/src/srtctl/backends/tilert.py @@ -0,0 +1,418 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TileRT decode with vLLM prefill for native P/D disaggregation.""" + +from __future__ import annotations + +import builtins +import json +from collections.abc import Sequence +from dataclasses import field +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar, Literal + +from marshmallow import Schema, ValidationError +from marshmallow_dataclass import dataclass + +from srtctl.ports import DYN_SYSTEM_PORT_BASE + +if TYPE_CHECKING: + from srtctl.backends.base import BackendPreparation, SrunConfig + from srtctl.core.runtime import RuntimeContext + from srtctl.core.schema import ProfilingConfig + from srtctl.core.topology import Endpoint, NodePortAllocator, Process + +WorkerMode = Literal["prefill", "decode", "agg"] + +_PREFILL_MANAGED_ARGS = { + "host", + "port", + "served-model-name", + "tensor-parallel-size", + "max-model-len", + "return-tokens-as-token-ids", + "kv-cache-dtype", + "kv-transfer-config", + "speculative-config", +} + +_DECODE_MANAGED_ARGS = { + "engine", + "model", + "model-weights-dir", + "max-seq-len", + "kv-cache-dtype", + "transport", + "ctrl-port", + "http-port", + "with-mtp", +} + + +@dataclass(frozen=True) +class TileRTServerConfig: + """Role-specific CLI arguments not owned by the topology adapter.""" + + prefill: dict[str, Any] | None = None + decode: dict[str, Any] | None = None + + Schema: ClassVar[type[Schema]] = Schema + + +@dataclass(frozen=True) +class TileRTProtocol: + """Run vLLM prefill workers and TileRT decode workers without Dynamo. + + TileRT's public P/D implementation intentionally uses different runtime + images for the two roles. ``model.container`` remains the decode image; + ``prefill_container`` selects the vLLM image used only for prefill sruns. + """ + + type: Literal["tilert"] = "tilert" + # OpenAI model name; defaults to the final component of model.path. + served_model_name: str | None = None + # vLLM image or cluster container alias for prefill workers only. + prefill_container: str | None = None + prefill_environment: dict[str, str] = field(default_factory=dict) + decode_environment: dict[str, str] = field(default_factory=dict) + tilert_config: TileRTServerConfig | None = None + # TileRT P/D model profile (for example glm5). + model_profile: str = "glm5" + # Model type passed to the TileRT weight converter. + weight_model_type: str = "glm-5" + # Dedicated shared converted-weight directory, mounted at this path in all decode containers. + weights_dir: str = "/tilert_weights/converted" + # Maximum sequence length shared by prefill and decode. + max_seq_len: int = 202752 + # TileRT decode KV cache dtype. + kv_cache_dtype: str = "fp8" + # vLLM KV cache dtype; must be wire-compatible with the decode dtype. + prefill_kv_cache_dtype: str = "fp8_ds_mla" + # KV transfer transport (the adapter supports NIXL). + transport: Literal["nixl"] = "nixl" + # Enable MTP on both sides of the P/D transfer. + with_mtp: bool = True + # Number of speculative tokens advertised by vLLM prefill. + speculative_tokens: int = 1 + + Schema: ClassVar[builtins.type[Schema]] = Schema + + def validate_recipe(self, config: Any) -> None: + """Reject unsupported topology and adapter overrides before allocating Slurm nodes.""" + resources = config.resources + if config.frontend.type != "tilert-router": + raise ValidationError("engine.type: tilert requires frontend.type: tilert-router") + if config.frontend.enable_multiple_frontends: + raise ValidationError("TileRT requires frontend.enable_multiple_frontends: false") + if resources.num_agg or resources.num_prefill != 1 or resources.num_decode < 1: + raise ValidationError("TileRT requires exactly one prefill worker and at least one decode worker") + if max(resources.gpus_per_prefill, resources.gpus_per_decode) > resources.gpus_per_node: + raise ValidationError("Each TileRT prefill and decode worker must fit on one node") + if self.model_profile in {"glm5", "glm_5", "glm-5"} and resources.gpus_per_decode != 8: + raise ValidationError("TileRT GLM-5 decode kernels require exactly 8 GPUs per worker") + if not self.prefill_container: + raise ValidationError("engine.prefill_container must select the vLLM prefill image") + if not Path(self.weights_dir).is_absolute() or self.weights_dir == "/": + raise ValidationError("engine.weights_dir must be an absolute, dedicated converted-weight directory") + if self.max_seq_len < 1 or self.speculative_tokens < 1: + raise ValidationError("TileRT max_seq_len and speculative_tokens must be positive") + if config.profiling.is_torch: + raise ValidationError("TileRT does not support torch profiling; use nsys") + # Command builders own these flags even when supplied through roles.*.args. + for mode, reserved in (("prefill", _PREFILL_MANAGED_ARGS), ("decode", _DECODE_MANAGED_ARGS)): + try: + _reject_managed_args(self.get_config_for_mode(mode), reserved) + except ValueError as error: + raise ValidationError(str(error)) from error + + def get_srun_config(self) -> SrunConfig: + from srtctl.backends.base import SrunConfig + + return SrunConfig(mpi=None, oversubscribe=False, launch_per_endpoint=False) + + def get_config_for_mode(self, mode: WorkerMode) -> dict[str, Any]: + if self.tilert_config is None: + return {} + values = {"prefill": self.tilert_config.prefill, "decode": self.tilert_config.decode} + return dict(values.get(mode) or {}) + + def get_environment_for_mode(self, mode: WorkerMode) -> dict[str, str]: + values = {"prefill": self.prefill_environment, "decode": self.decode_environment} + return dict(values.get(mode) or {}) + + def get_process_environment(self, process: Process) -> dict[str, str]: + return {"TILERT_ROLE": process.endpoint_mode} + + def get_container_image_for_mode(self, mode: WorkerMode, default: str) -> str: + if mode == "prefill" and self.prefill_container: + return self.prefill_container + return default + + def get_served_model_name(self, default: str) -> str: + return self.served_model_name or default + + def get_metrics_port(self, process: Process) -> int | None: + """Return only real Prometheus endpoints. + + The vLLM prefill API exports ``/metrics``. TileRT's decode API exports + health and token-transfer routes but no Prometheus handler. + """ + return process.http_port if process.endpoint_mode == "prefill" else None + + def get_preparation(self, runtime: RuntimeContext, processes: list[Process]) -> BackendPreparation: + """Materialize converted TileRT weights once under a shared lock.""" + from srtctl.backends.base import BackendPreparation + + decode_leaders = [process for process in processes if process.endpoint_mode == "decode" and process.is_leader] + if not decode_leaders: + raise ValueError("TileRT model preparation requires a decode worker") + source = str(runtime.model_path) if runtime.is_hf_model else runtime.worker_model_arg + if not runtime.is_hf_model and Path(source) == Path(self.weights_dir): + raise ValueError("TileRT weights_dir must differ from the source model directory") + script = _weight_preparation_script( + source=source, + is_hf_model=runtime.is_hf_model, + weights_dir=self.weights_dir, + weight_model_type=self.weight_model_type, + ) + return BackendPreparation( + command=["bash", "-ceu", script], + node=decode_leaders[0].node, + mode="decode", + log_name="tilert_weight_conversion.out", + gpus_per_task=1, + ) + + def allocate_endpoints( + self, + num_prefill: int, + num_decode: int, + num_agg: int, + gpus_per_prefill: int, + gpus_per_decode: int, + gpus_per_agg: int, + gpus_per_node: int, + available_nodes: Sequence[str], + spread_workers: bool = False, + ) -> list[Endpoint]: + if num_agg: + raise ValueError("TileRT supports prefill/decode topology only") + from srtctl.core.topology import allocate_endpoints + + return allocate_endpoints( + num_prefill=num_prefill, + num_decode=num_decode, + num_agg=num_agg, + gpus_per_prefill=gpus_per_prefill, + gpus_per_decode=gpus_per_decode, + gpus_per_agg=gpus_per_agg, + gpus_per_node=gpus_per_node, + available_nodes=available_nodes, + spread_workers=spread_workers, + ) + + def endpoints_to_processes( + self, + endpoints: list[Endpoint], + base_sys_port: int = DYN_SYSTEM_PORT_BASE, + port_allocator: NodePortAllocator | None = None, + frontend_type: str = "tilert-router", + dynamo_sidecar: bool = False, + ) -> list[Process]: + if dynamo_sidecar: + raise ValueError("TileRT does not support Dynamo sidecars") + if frontend_type != "tilert-router": + raise ValueError(f"backend.type: tilert requires frontend.type: tilert-router (got {frontend_type!r})") + from srtctl.core.topology import endpoints_to_processes + + return endpoints_to_processes(endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator) + + def build_worker_command( + self, + process: Process, + endpoint_processes: list[Process], + runtime: RuntimeContext, + frontend_type: str = "tilert-router", + nsys_prefix: list[str] | None = None, + dump_config_path: Path | None = None, + profiling: ProfilingConfig | None = None, + ) -> list[str]: + del dump_config_path, profiling + if frontend_type != "tilert-router": + raise ValueError(f"backend.type: tilert requires frontend.type: tilert-router (got {frontend_type!r})") + if len({item.node for item in endpoint_processes}) != 1: + raise ValueError("TileRT currently requires every prefill and decode endpoint to fit on one node") + if process.endpoint_mode not in {"prefill", "decode"}: + raise ValueError("TileRT supports prefill/decode workers only") + + config = self.get_config_for_mode(process.endpoint_mode) + if process.endpoint_mode == "decode": + _reject_managed_args(config, _DECODE_MANAGED_ARGS) + command = [*(nsys_prefix or []), "python", "-m", "tilert.pd_vllm.decode_server"] + command.extend( + [ + "--engine", + "tilert", + "--model", + self.model_profile, + "--model-weights-dir", + self.weights_dir, + "--max-seq-len", + str(self.max_seq_len), + "--kv-cache-dtype", + self.kv_cache_dtype, + "--transport", + self.transport, + "--ctrl-port", + str(process.sys_port), + "--http-port", + str(process.http_port), + ] + ) + if self.with_mtp: + command.append("--with-mtp") + command.extend(_config_to_cli_args(config)) + return command + + _reject_managed_args(config, _PREFILL_MANAGED_ARGS) + transfer = { + "kv_connector": "TileRTConnector", + "kv_connector_module_path": "tilert.pd_vllm.prefill_connector", + "kv_role": "kv_producer", + "kv_connector_extra_config": { + "tilert_ctrl_port": process.sys_port, + "tilert_model": self.model_profile, + "tilert_max_seq_len": self.max_seq_len, + "tilert_transport": self.transport, + }, + } + speculative = {"method": "mtp", "num_speculative_tokens": self.speculative_tokens} + command = [*(nsys_prefix or []), "vllm", "serve", runtime.worker_model_arg] + command.extend( + [ + "--served-model-name", + self.get_served_model_name(Path(str(runtime.model_path)).name), + "--host", + "0.0.0.0", + "--port", + str(process.http_port), + "--tensor-parallel-size", + str(len(process.gpu_indices)), + "--max-model-len", + str(self.max_seq_len), + "--return-tokens-as-token-ids", + "--kv-cache-dtype", + self.prefill_kv_cache_dtype, + "--kv-transfer-config", + json.dumps(transfer, separators=(",", ":")), + ] + ) + if self.with_mtp: + command.extend(["--speculative-config", json.dumps(speculative, separators=(",", ":"))]) + command.extend(_config_to_cli_args(config)) + return command + + +def _config_to_cli_args(config: dict[str, Any]) -> list[str]: + args: list[str] = [] + for key, value in config.items(): + flag = key if key.startswith("-") else f"--{key.replace('_', '-')}" + if value is True: + args.append(flag) + elif value is False or value is None: + continue + elif isinstance(value, list): + for item in value: + args.extend([flag, str(item)]) + else: + args.extend([flag, str(value)]) + return args + + +def _reject_managed_args(config: dict[str, Any], reserved: set[str]) -> None: + overlap = {key.lstrip("-").replace("_", "-") for key in config}.intersection(reserved) + if overlap: + raise ValueError(f"TileRT config cannot override srtctl-managed argument(s): {sorted(overlap)}") + + +def _weight_preparation_script(*, source: str, is_hf_model: bool, weights_dir: str, weight_model_type: str) -> str: + """Build a lock-safe, atomic TileRT conversion command.""" + # Backend preparation is the final completeness boundary before conversion. + # The best-effort orchestrator pre-download may time out, and + # ``local_files_only=True`` can still resolve a partially populated snapshot. + # Let snapshot_download resume missing files here under its normal HF cache + # locks; a complete cache remains a metadata-only fast path. + resolve_source = ( + f"from huggingface_hub import snapshot_download; source = snapshot_download({source!r})" + if is_hf_model + else f"source = {source!r}" + ) + python = f""" +import fcntl +import json +import os +import shutil +import subprocess +from pathlib import Path + +{resolve_source} +source_path = Path(source) +target = Path({weights_dir!r}) +target.parent.mkdir(parents=True, exist_ok=True) +lock_path = target.parent / f".{{target.name}}.convert.lock" +with lock_path.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + def complete(path): + index = path / "model.safetensors.index.json" + if not index.is_file() or not (path / "chat_template.jinja").is_file(): + return False + if not ((path / "tokenizer.json").is_file() or (path / "tokenizer_config.json").is_file()): + return False + try: + shards = set(json.loads(index.read_text())["weight_map"].values()) + return bool(shards) and all((path / shard).is_file() for shard in shards) + except (ValueError, KeyError, TypeError): + return False + + if complete(target): + print(f"TileRT weight cache hit: {{target}}", flush=True) + raise SystemExit(0) + tmp = target.parent / f".{{target.name}}.tmp.{{os.environ.get('SLURM_JOB_ID', os.getpid())}}" + shutil.rmtree(tmp, ignore_errors=True) + try: + subprocess.run( + [ + "python", + "-m", + "tilert.models.preprocess.weight_converter", + "--model_type", + {weight_model_type!r}, + "--model_dir", + str(source_path), + "--save_dir", + str(tmp), + ], + check=True, + ) + for item in source_path.iterdir(): + if not item.is_file() or item.suffix == ".safetensors" or item.name == "model.safetensors.index.json": + continue + shutil.copy2(item, tmp / item.name) + converted_index = tmp / "model.safetensors.index.json" + if not converted_index.is_file(): + raise RuntimeError(f"TileRT conversion produced no index: {{converted_index}}") + if not (tmp / "chat_template.jinja").is_file(): + raise RuntimeError(f"TileRT conversion produced no chat template: {{tmp}}") + if not ((tmp / "tokenizer.json").is_file() or (tmp / "tokenizer_config.json").is_file()): + raise RuntimeError(f"TileRT conversion produced no tokenizer metadata: {{tmp}}") + if not complete(tmp): + raise RuntimeError(f"TileRT conversion produced incomplete shards: {{tmp}}") + if target.exists(): + shutil.rmtree(target) + os.replace(tmp, target) + print(f"TileRT weight cache ready: {{target}}", flush=True) + finally: + shutil.rmtree(tmp, ignore_errors=True) +""" + return f"python - <<'PY'\n{python}PY\n" diff --git a/src/srtctl/cli/do_sweep.py b/src/srtctl/cli/do_sweep.py index 6a7193227..4ea1fefbc 100644 --- a/src/srtctl/cli/do_sweep.py +++ b/src/srtctl/cli/do_sweep.py @@ -669,6 +669,8 @@ def run(self) -> int: if self.runtime.staged_model_path is not None: self._stage_model() + self.prepare_backend(registry) + # Stage 2: Workers reporter.report(JobStatus.WORKERS, JobStage.WORKERS, "Starting workers") worker_procs = self.start_all_workers() diff --git a/src/srtctl/cli/mixins/benchmark_stage.py b/src/srtctl/cli/mixins/benchmark_stage.py index d03ee7c27..9a4ec2a70 100644 --- a/src/srtctl/cli/mixins/benchmark_stage.py +++ b/src/srtctl/cli/mixins/benchmark_stage.py @@ -643,7 +643,14 @@ def _get_aiperf_server_metrics_env( # OpenAI port (GET /metrics there is JSON iteration stats, not # exposition text); every other frontend serves it at /metrics. metrics_path = "/prometheus/metrics" if self.config.frontend.type == "trtllm_serve" else "/metrics" - if logical_workers_only: + metrics_port = getattr(self.config.backend, "get_metrics_port", None) + if metrics_port is not None: + for process in self.backend_processes: + port = metrics_port(process) + if port is not None and port > 0: + host = get_hostname_ip(process.node, self.runtime.network_interface) + urls.append(f"http://{host}:{port}{metrics_path}") + elif logical_workers_only: if logical_endpoints is None: logical_endpoints = self._logical_worker_endpoints() # Sidecars use native worker commands, so publish_metrics does not diff --git a/src/srtctl/cli/mixins/telemetry_stage.py b/src/srtctl/cli/mixins/telemetry_stage.py index 5c0ebcaad..31d16e00c 100644 --- a/src/srtctl/cli/mixins/telemetry_stage.py +++ b/src/srtctl/cli/mixins/telemetry_stage.py @@ -25,6 +25,7 @@ from srtctl.core.schema import TelemetryExporterConfig from srtctl.core.slurm import start_srun_process from srtctl.core.telemetry import TACHOMETER_STORAGE_PARENT, generate_tachometer_config +from srtctl.frontends import get_frontend if TYPE_CHECKING: from srtctl.core.runtime import RuntimeContext @@ -560,6 +561,8 @@ def start_tachometer(self) -> list[ManagedProcess]: dcgm_exporter=dcgm_exporter, frontend_type=self.config.frontend.type, frontend_metrics_port=self._frontend_metrics_port(), + worker_metrics_port=getattr(self.config.backend, "get_metrics_port", None), + frontend_metrics_enabled=getattr(get_frontend(self.config.frontend.type), "has_metrics", True), ) ) diff --git a/src/srtctl/cli/mixins/worker_stage.py b/src/srtctl/cli/mixins/worker_stage.py index 7aa365fa8..34bba749e 100644 --- a/src/srtctl/cli/mixins/worker_stage.py +++ b/src/srtctl/cli/mixins/worker_stage.py @@ -9,6 +9,7 @@ import logging import shlex +import subprocess from collections import defaultdict from concurrent.futures import ThreadPoolExecutor, as_completed from typing import TYPE_CHECKING, Any @@ -22,6 +23,7 @@ from srtctl.services.implicit import discovery_env if TYPE_CHECKING: + from srtctl.core.processes import ProcessRegistry from srtctl.core.runtime import RuntimeContext from srtctl.core.schema import SrtConfig from srtctl.core.topology import Endpoint, Process @@ -155,6 +157,60 @@ def _get_worker_environment_for_mode(self, mode: str) -> dict[str, str]: environment.setdefault("SGLANG_PYSPY_DUMP_BEFORE_CRASH", "0") return environment + def _worker_container(self, mode: str) -> str: + """Resolve an optional backend-owned role image; default to model.container.""" + select = getattr(self.backend, "get_container_image_for_mode", None) + default = str(self.runtime.container_image) + return select(mode, default) if select is not None else default + + def prepare_backend(self, registry: "ProcessRegistry") -> None: + """Run an optional backend preparation step through the normal Slurm lifecycle.""" + prepare = getattr(self.backend, "get_preparation", None) + preparation = prepare(self.runtime, self.backend_processes) if prepare is not None else None + if preparation is None: + return + process = next( + item + for item in self.backend_processes + if item.node == preparation.node and item.endpoint_mode == preparation.mode + ) + environment = self._get_worker_environment_for_mode(preparation.mode) + environment.update(self.runtime.environment) + environment.update(self.backend.get_process_environment(process)) + options = dict(self.runtime.srun_options) + if preparation.gpus_per_task is not None: + options["gpus-per-task"] = str(preparation.gpus_per_task) + log_path = self.runtime.log_dir / preparation.log_name + step_name = "backend_preparation" + proc = start_srun_process( + command=preparation.command, + nodelist=[preparation.node], + output=str(log_path), + container_image=self._worker_container(preparation.mode), + container_mounts=self.runtime.container_mounts, + env_to_set=environment, + bash_preamble=self._build_worker_preamble(), + srun_options=options, + het_group=process.het_group, + step_name=step_name, + ) + managed = ManagedProcess( + name=step_name, + popen=proc, + log_file=log_path, + node=preparation.node, + critical=True, + step_name=step_name, + ) + registry.add_process(managed) + try: + return_code = proc.wait(timeout=preparation.timeout_seconds) + except subprocess.TimeoutExpired as error: + managed.terminate() + raise RuntimeError(f"Backend preparation timed out; see {log_path}") from error + if return_code != 0: + raise RuntimeError(f"Backend preparation failed with exit code {return_code}; see {log_path}") + def start_worker(self, process: "Process", endpoint_processes: list["Process"]) -> ManagedProcess: """Start a single worker process (one srun per node, used by SGLang).""" mode = process.endpoint_mode @@ -277,7 +333,7 @@ def __missing__(self, key: str) -> str: command=cmd, nodelist=[process.node], output=str(worker_log), - container_image=str(self.runtime.container_image), + container_image=self._worker_container(mode), container_mounts=self.runtime.container_mounts, env_to_set=env_to_set, env_to_unset=env_to_unset, @@ -438,7 +494,7 @@ def start_endpoint_worker(self, endpoint_processes: list["Process"]) -> ManagedP ntasks=total_gpus, nodelist=endpoint_nodes, output=str(worker_log), - container_image=str(self.runtime.container_image), + container_image=self._worker_container(mode), container_mounts=self.runtime.container_mounts, env_to_set=env_to_set, bash_preamble=bash_preamble, diff --git a/src/srtctl/cli/submit.py b/src/srtctl/cli/submit.py index d2ec44bd0..241cc0f77 100755 --- a/src/srtctl/cli/submit.py +++ b/src/srtctl/cli/submit.py @@ -262,6 +262,16 @@ def show_config_details(config: SrtConfig) -> None: "to avoid confusion (the configured value is ignored)." ) + select_container = getattr(config.backend, "get_container_image_for_mode", None) + if select_container is not None: + images = Table(title="Worker Containers") + images.add_column("Role") + images.add_column("Image") + for mode in ("prefill", "decode", "agg"): + if getattr(config.resources, f"num_{mode}"): + images.add_row(mode, select_container(mode, str(config.model.container))) + console.print(images) + # --- Container Mounts --- mounts_table = Table(title="Container Mounts", show_lines=False, pad_edge=False) mounts_table.add_column("Source", style="dim", width=14) diff --git a/src/srtctl/core/config.py b/src/srtctl/core/config.py index 4a08f8e5d..b9bb5b38c 100755 --- a/src/srtctl/core/config.py +++ b/src/srtctl/core/config.py @@ -81,7 +81,9 @@ def load_cluster_config() -> dict[str, Any] | None: # Keys whose string values name a container image. Any such leaf anywhere in a # recipe resolves against the cluster `containers:` alias map. -CONTAINER_ALIAS_KEYS: frozenset[str] = frozenset({"container", "container_image", "image", "nginx_container"}) +CONTAINER_ALIAS_KEYS: frozenset[str] = frozenset( + {"container", "container_image", "image", "nginx_container", "prefill_container"} +) # Sub-trees the alias walker never enters: `identity` declares the pullable image a # run *should* be using (verification only, never an alias); the rest are free-form @@ -107,6 +109,7 @@ def load_cluster_config() -> dict[str, Any] | None: "vllm_config", "trtllm_config", "mocker_config", + "tilert_config", "store_config", } ) diff --git a/src/srtctl/core/fingerprint.py b/src/srtctl/core/fingerprint.py index 15a50fa4a..7d3bc1c93 100644 --- a/src/srtctl/core/fingerprint.py +++ b/src/srtctl/core/fingerprint.py @@ -46,6 +46,8 @@ # Used in both native Python probes and the bash capture script. FRAMEWORK_PACKAGES: dict[str, str] = { "vllm": "vllm", + "tilert": "tilert", + "nixl": "nixl", "sglang": "sglang", "tensorrt_llm": "tensorrt-llm", "dynamo": "ai-dynamo", diff --git a/src/srtctl/core/roles.py b/src/srtctl/core/roles.py index d877ecdb3..5e554f0ff 100644 --- a/src/srtctl/core/roles.py +++ b/src/srtctl/core/roles.py @@ -57,6 +57,7 @@ "vllm": "vllm_config", "trtllm": "trtllm_config", "mocker": "mocker_config", + "tilert": "tilert_config", } _ALL_ENGINE_CONFIG_KEYS = frozenset(ENGINE_CONFIG_KEY.values()) diff --git a/src/srtctl/core/schema.py b/src/srtctl/core/schema.py index 7f2d743c0..2dabb4828 100755 --- a/src/srtctl/core/schema.py +++ b/src/srtctl/core/schema.py @@ -38,6 +38,7 @@ BackendConfig, MockerProtocol, SGLangProtocol, + TileRTProtocol, TRTLLMProtocol, VLLMProtocol, ) @@ -394,7 +395,7 @@ def _deserialize( # Default to SGLang return SGLangProtocol() - if isinstance(value, SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol): + if isinstance(value, SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol | TileRTProtocol): return value if not isinstance(value, dict): @@ -412,12 +413,14 @@ def _deserialize( elif backend_type == "vllm": schema = VLLMProtocol.Schema() return schema.load(value) + elif backend_type == "tilert": + return TileRTProtocol.Schema().load(value) elif backend_type == "mocker": schema = MockerProtocol.Schema() return schema.load(value) else: raise ValidationError( - f"Unknown backend type: {backend_type!r}. Supported types: sglang, trtllm, vllm, mocker" + f"Unknown backend type: {backend_type!r}. Supported types: sglang, trtllm, vllm, mocker, tilert" ) def _serialize(self, value: Any | None, attr: str | None, obj: Any, **kwargs) -> Any: @@ -426,6 +429,8 @@ def _serialize(self, value: Any | None, attr: str | None, obj: Any, **kwargs) -> return None if isinstance(value, SGLangProtocol): return SGLangProtocol.Schema().dump(value) + if isinstance(value, TileRTProtocol): + return TileRTProtocol.Schema().dump(value) if isinstance(value, TRTLLMProtocol): return TRTLLMProtocol.Schema().dump(value) if isinstance(value, VLLMProtocol): @@ -1931,7 +1936,7 @@ class FrontendConfig: Attributes: type: Frontend type - "dynamo" (default); "sglang-router" (SGLang Model - Gateway) and "vllm-router" (static routers); "sglang", "vllm", and + Gateway), "vllm-router", and "tilert-router" (static routers); "sglang", "vllm", and "trtllm_serve" (direct: the single aggregate worker binds the public port, no router process). In schema 1 recipes "sglang" still means the router and loads as "sglang-router". @@ -2133,6 +2138,8 @@ def __post_init__(self): self._validate_vllm_frontend() self._validate_sglang_direct_frontend() self._validate_static_router_frontend() + if isinstance(self.backend, TileRTProtocol): + self.backend.validate_recipe(self) self._validate_dynamo_sidecar() self._validate_host_setup() self._validate_benchmark_type() @@ -2331,7 +2338,9 @@ def _validate_sglang_direct_frontend(self): def _validate_static_router_frontend(self): """Validate static-router/backend pairings and vLLM DP ownership.""" - required_backend = {"sglang-router": "sglang", "vllm-router": "vllm"}.get(self.frontend.type) + required_backend = {"sglang-router": "sglang", "vllm-router": "vllm", "tilert-router": "tilert"}.get( + self.frontend.type + ) if required_backend is None: return if self.backend_type != required_backend: @@ -2340,6 +2349,13 @@ def _validate_static_router_frontend(self): f"got {self.backend_type!r}" ) + if self.frontend.type == "tilert-router": + from srtctl.frontends.tilert_router import TileRTRouterFrontend + + try: + TileRTRouterFrontend().get_managed_frontend_args(self, self.backend, []) + except ValueError as error: + raise ValidationError(str(error)) from error if self.frontend.type != "vllm-router": return if not isinstance(self.backend, VLLMProtocol): diff --git a/src/srtctl/core/schema_docs.py b/src/srtctl/core/schema_docs.py index 5cf567b3e..611c44a46 100644 --- a/src/srtctl/core/schema_docs.py +++ b/src/srtctl/core/schema_docs.py @@ -47,6 +47,8 @@ MockerServerConfig, SGLangProtocol, SGLangServerConfig, + TileRTProtocol, + TileRTServerConfig, TRTLLMProtocol, TRTLLMServerConfig, VLLMProtocol, @@ -83,6 +85,7 @@ ("trtllm", TRTLLMProtocol), ("vllm", VLLMProtocol), ("mocker", MockerProtocol), + ("tilert", TileRTProtocol), ) _BACKEND_CLASSES: set[type] = {cls for _, cls in BACKEND_TYPES} @@ -103,7 +106,7 @@ class FieldDoc: type_label="str \\| mapping", default="required", description=( - "The engine type (`sglang`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` " + "The engine type (`sglang`, `trtllm`, `vllm`, `mocker`, `tilert`) as a string, or a mapping with `type` " "plus the engine-wide knobs listed under [Engine types](#engine-types)." ), ), @@ -398,6 +401,7 @@ def _present(cls: type, mapping: dict[str, str]) -> dict[str, str]: SGLangProtocol: _present(SGLangProtocol, _backend_legacy_fields(ENGINE_CONFIG_KEY["sglang"])), TRTLLMProtocol: _present(TRTLLMProtocol, _backend_legacy_fields(ENGINE_CONFIG_KEY["trtllm"])), VLLMProtocol: _present(VLLMProtocol, _backend_legacy_fields(ENGINE_CONFIG_KEY["vllm"])), + TileRTProtocol: _present(TileRTProtocol, _backend_legacy_fields(ENGINE_CONFIG_KEY["tilert"])), MockerProtocol: _present(MockerProtocol, _backend_legacy_fields(ENGINE_CONFIG_KEY["mocker"])), } @@ -415,6 +419,7 @@ def _present(cls: type, mapping: dict[str, str]) -> dict[str, str]: TRTLLMServerConfig, VLLMServerConfig, MockerServerConfig, + TileRTServerConfig, MooncakeKVStoreConfig, VLLMMooncakeKVStoreConfig, } diff --git a/src/srtctl/core/telemetry.py b/src/srtctl/core/telemetry.py index 309865ce6..0b7a8b4d8 100644 --- a/src/srtctl/core/telemetry.py +++ b/src/srtctl/core/telemetry.py @@ -6,6 +6,7 @@ from __future__ import annotations import json +from collections.abc import Callable from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -48,6 +49,8 @@ def generate_tachometer_config( dcgm_exporter: TelemetryExporterConfig | None = None, frontend_type: str = "dynamo", frontend_metrics_port: int | None = None, + worker_metrics_port: Callable[[Process], int | None] | None = None, + frontend_metrics_enabled: bool = True, ) -> str: """Generate Tachometer TOML from backend and frontend topology. @@ -138,7 +141,12 @@ def generate_tachometer_config( # the HTTP server that carries /metrics. continue node_ip = get_hostname_ip(process.node, runtime.network_interface) - if frontend_type in ("vllm", "sglang") and process.endpoint_mode == "agg": + if worker_metrics_port is not None: + selected_port = worker_metrics_port(process) + if selected_port is None: + continue + port = selected_port + elif frontend_type in ("vllm", "sglang") and process.endpoint_mode == "agg": # Direct modes: the aggregate leader binds the public port itself. port = FRONTEND_PUBLIC_PORT elif frontend_type in ("vllm-router", "trtllm_serve", "sglang-router"): @@ -163,7 +171,7 @@ def generate_tachometer_config( ) ) - frontend_nodes = frontend_topology.frontend_nodes + frontend_nodes = frontend_topology.frontend_nodes if frontend_metrics_enabled else [] if frontend_type in ("vllm", "sglang"): # Direct vLLM / SGLang have no separate frontend process. The public endpoint is # the aggregate leader, which may differ from the Slurm/orchestrator diff --git a/src/srtctl/frontends/__init__.py b/src/srtctl/frontends/__init__.py index d4ab7f432..bfe12d2d5 100644 --- a/src/srtctl/frontends/__init__.py +++ b/src/srtctl/frontends/__init__.py @@ -20,6 +20,7 @@ from srtctl.frontends.dynamo import DynamoFrontend from srtctl.frontends.sglang import SGLangRouterFrontend from srtctl.frontends.sglang_direct import SGLangFrontend +from srtctl.frontends.tilert_router import TileRTRouterFrontend from srtctl.frontends.trtllm_serve import TRTLLMServeFrontend from srtctl.frontends.vllm import VLLMFrontend from srtctl.frontends.vllm_router import VLLMRouterFrontend @@ -31,6 +32,7 @@ "SGLangFrontend", "SGLangRouterFrontend", "TRTLLMServeFrontend", + "TileRTRouterFrontend", "VLLMFrontend", "VLLMRouterFrontend", "get_frontend", diff --git a/src/srtctl/frontends/base.py b/src/srtctl/frontends/base.py index ba49b1771..efb88949c 100644 --- a/src/srtctl/frontends/base.py +++ b/src/srtctl/frontends/base.py @@ -20,7 +20,7 @@ from srtctl.core.topology import Process # Supported frontend types - extensible by adding new literals -FrontendType = Literal["dynamo", "sglang", "sglang-router", "trtllm_serve", "vllm", "vllm-router"] +FrontendType = Literal["dynamo", "sglang", "sglang-router", "trtllm_serve", "vllm", "vllm-router", "tilert-router"] class FrontendProtocol(Protocol): @@ -106,6 +106,7 @@ def get_frontend(frontend_type: str) -> FrontendProtocol: from srtctl.frontends.dynamo import DynamoFrontend from srtctl.frontends.sglang import SGLangRouterFrontend from srtctl.frontends.sglang_direct import SGLangFrontend + from srtctl.frontends.tilert_router import TileRTRouterFrontend from srtctl.frontends.trtllm_serve import TRTLLMServeFrontend from srtctl.frontends.vllm import VLLMFrontend from srtctl.frontends.vllm_router import VLLMRouterFrontend @@ -120,9 +121,11 @@ def get_frontend(frontend_type: str) -> FrontendProtocol: return TRTLLMServeFrontend() elif frontend_type == "vllm": return VLLMFrontend() + elif frontend_type == "tilert-router": + return TileRTRouterFrontend() elif frontend_type == "vllm-router": return VLLMRouterFrontend() else: raise ValueError( - f"Unknown frontend type: {frontend_type!r}. Supported: dynamo, sglang, sglang-router, trtllm_serve, vllm, vllm-router" + f"Unknown frontend type: {frontend_type!r}. Supported: dynamo, sglang, sglang-router, trtllm_serve, vllm, vllm-router, tilert-router" ) diff --git a/src/srtctl/frontends/static_router.py b/src/srtctl/frontends/static_router.py index e6868f3e0..35f0f25c6 100644 --- a/src/srtctl/frontends/static_router.py +++ b/src/srtctl/frontends/static_router.py @@ -227,3 +227,20 @@ def start_frontends( ) ) return processes + + +def setup_script_preamble(config: Any) -> str | None: + """Run a recipe setup script inside a static router container.""" + setup_script = getattr(config, "setup_script", None) + if not setup_script: + return None + script_name = shlex.quote(setup_script) + return ( + f"setup_script={script_name} && " + 'script_path="/configs/${setup_script}" && ' + 'patch_script_path="/configs/patches/${setup_script}" && ' + 'echo "Running setup script: ${script_path} (fallback ${patch_script_path})" && ' + 'if [ -f "${script_path}" ]; then bash "${script_path}"; ' + 'elif [ -f "${patch_script_path}" ]; then bash "${patch_script_path}"; ' + 'else echo "WARNING: ${script_path} or ${patch_script_path} not found"; fi' + ) diff --git a/src/srtctl/frontends/tilert_router.py b/src/srtctl/frontends/tilert_router.py new file mode 100644 index 000000000..d008a0837 --- /dev/null +++ b/src/srtctl/frontends/tilert_router.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TileRT's native vLLM-prefill/TileRT-decode router frontend.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, ClassVar +from urllib.parse import urlsplit + +from srtctl.core.health import WorkerHealthResult, wait_for_http_endpoints +from srtctl.frontends.static_router import RouterWorker, StaticRouterFrontend + +if TYPE_CHECKING: + import threading + + from srtctl.core.processes import ManagedProcess + from srtctl.core.runtime import RuntimeContext + from srtctl.core.topology import Process + + +class TileRTRouterFrontend(StaticRouterFrontend): + """Launch ``tilert.pd_vllm.pd_router`` from an explicit worker topology.""" + + type: ClassVar[str] = "tilert-router" + has_metrics: ClassVar[bool] = False + backend_type: ClassVar[str] = "tilert" + executable: ClassVar[tuple[str, ...]] = ("python", "-m", "tilert.pd_vllm.pd_router") + pd_flag: ClassVar[str] = "" + process_name: ClassVar[str] = "tilert_router" + + def build_bash_preamble(self, config: Any) -> str | None: + """Apply the same recipe setup in the router image with its role selected.""" + from srtctl.frontends.static_router import setup_script_preamble + + setup = setup_script_preamble(config) + return f"export TILERT_ROLE=router && {setup}" if setup else None + + def start_frontends( + self, + topology: Any, + runtime: RuntimeContext, + config: Any, + backend: Any, + backend_processes: list[Process], + stop_event: threading.Event | None = None, + ) -> list[ManagedProcess]: + """TileRT's router connects during startup, so workers must be healthy first.""" + ready = wait_for_http_endpoints( + self.get_backend_health_urls(backend, backend_processes, runtime.network_interface), + poll_interval=config.health_check.interval_seconds, + timeout=config.health_check.max_attempts * config.health_check.interval_seconds, + stop_event=stop_event, + ) + if not ready: + raise RuntimeError("TileRT workers did not become ready before router startup") + return super().start_frontends(topology, runtime, config, backend, backend_processes, stop_event) + + @property + def health_endpoint(self) -> str: + return "/health" + + def parse_health( + self, + response_json: dict, + expected_prefill: int, + expected_decode: int, + ) -> WorkerHealthResult: + ready = response_json.get("status") == "ok" + return WorkerHealthResult( + ready=ready, + message="TileRT P/D router healthy" if ready else "TileRT P/D router not ready", + prefill_ready=expected_prefill if ready else 0, + prefill_expected=expected_prefill, + decode_ready=expected_decode if ready else 0, + decode_expected=expected_decode, + ) + + def worker_bootstrap_port(self, backend: Any, process: Process) -> int | None: + del backend + return process.sys_port if process.endpoint_mode == "decode" else None + + def get_backend_health_urls( + self, + backend: Any, + backend_processes: list[Process], + network_interface: str | None = None, + ) -> list[str]: + return [ + f"{worker.url.rstrip('/')}/health" + for worker in self.collect_workers(backend, backend_processes, network_interface) + ] + + def build_router_command( + self, + workers: list[RouterWorker], + host: str, + port: int, + ) -> list[str]: + prefills = [worker for worker in workers if worker.mode == "prefill"] + decodes = [worker for worker in workers if worker.mode == "decode"] + aggregate = [worker for worker in workers if worker.mode == "agg"] + if aggregate or len(prefills) != 1 or not decodes: + raise ValueError("TileRT router requires exactly one prefill worker and at least one decode worker") + + command = [*self.executable, "--vllm-url", prefills[0].url] + for worker in decodes: + parsed = urlsplit(worker.url) + if parsed.hostname is None or parsed.port is None or worker.bootstrap_port is None: + raise ValueError(f"Incomplete TileRT decode endpoint: {worker}") + command.extend(["--decode", f"{parsed.hostname}:{worker.bootstrap_port}:{parsed.port}"]) + command.extend(["--host", host, "--port", str(port)]) + return command + + def get_managed_frontend_args( + self, + config: Any, + backend: Any, + backend_processes: list[Process], + ) -> list[str]: + del backend_processes + normalized = {str(key).lstrip("-").replace("_", "-") for key in (config.frontend.args or {})} + reserved = {"vllm-url", "decode", "host", "port", "model-path"} + overlap = normalized.intersection(reserved) + if overlap: + raise ValueError(f"frontend.args cannot override srtctl-managed TileRT argument(s): {sorted(overlap)}") + # Conversion copied tokenizer metadata into the shared decode-weight mount. + return ["--model-path", backend.weights_dir] diff --git a/src/srtctl/frontends/vllm_router.py b/src/srtctl/frontends/vllm_router.py index a0dc8ecf4..63b0bfa95 100644 --- a/src/srtctl/frontends/vllm_router.py +++ b/src/srtctl/frontends/vllm_router.py @@ -5,10 +5,9 @@ from __future__ import annotations -import shlex from typing import TYPE_CHECKING, Any, ClassVar -from srtctl.frontends.static_router import StaticRouterFrontend +from srtctl.frontends.static_router import StaticRouterFrontend, setup_script_preamble if TYPE_CHECKING: from srtctl.core.topology import Process @@ -57,19 +56,7 @@ class VLLMRouterFrontend(StaticRouterFrontend): def build_bash_preamble(self, config: Any) -> str | None: """Run the recipe setup script in the vLLM Router container.""" - setup_script = getattr(config, "setup_script", None) - if not setup_script: - return None - script_name = shlex.quote(setup_script) - return ( - f"setup_script={script_name} && " - 'script_path="/configs/${setup_script}" && ' - 'patch_script_path="/configs/patches/${setup_script}" && ' - 'echo "Running setup script: ${script_path} (fallback ${patch_script_path})" && ' - 'if [ -f "${script_path}" ]; then bash "${script_path}"; ' - 'elif [ -f "${patch_script_path}" ]; then bash "${patch_script_path}"; ' - 'else echo "WARNING: ${script_path} or ${patch_script_path} not found"; fi' - ) + return setup_script_preamble(config) def get_backend_health_urls( self, diff --git a/tests/test_container_aliases.py b/tests/test_container_aliases.py index 08bcfdbc4..b108f2386 100644 --- a/tests/test_container_aliases.py +++ b/tests/test_container_aliases.py @@ -5,7 +5,7 @@ from __future__ import annotations -from srtctl.core.config import CONTAINER_ALIAS_KEYS, resolve_config_with_defaults, resolve_container_aliases +from srtctl.core.config import resolve_config_with_defaults, resolve_container_aliases CONTAINERS = { "sglang": "/sqsh/sglang.sqsh", @@ -111,7 +111,3 @@ def test_no_cluster_containers_map_leaves_aliases_as_written() -> None: resolved = resolve_config_with_defaults(_recipe(), {"default_account": "acct"}) assert resolved["model"]["container"] == "sglang" assert resolve_config_with_defaults(_recipe(), None)["frontend"]["nginx_container"] == "nginx" - - -def test_alias_key_set_is_the_documented_one() -> None: - assert set(CONTAINER_ALIAS_KEYS) == {"container", "container_image", "image", "nginx_container"} diff --git a/tests/test_power_collector.py b/tests/test_power_collector.py index ca016fc32..db0f7fd1e 100644 --- a/tests/test_power_collector.py +++ b/tests/test_power_collector.py @@ -17,6 +17,7 @@ from srtctl.cli.do_sweep import SweepOrchestrator from srtctl.cli.mixins.benchmark_stage import BenchmarkStageMixin from srtctl.cli.mixins.telemetry_stage import TelemetryStageMixin +from srtctl.backends import SGLangProtocol from srtctl.core.power.contract import MANIFEST_FILENAME, SAMPLES_FILENAME, WINDOWS_DIRNAME, Reason from srtctl.core.power.manifest import ExpectedWindow from srtctl.core.power.samples import read_samples @@ -799,6 +800,7 @@ class TestRequiredReadinessGate: def _orchestrator(self, tmp_path, *, required, ready): config = MagicMock() + config.backend = SGLangProtocol() config.telemetry.enabled = True config.telemetry.required = required config.telemetry.cpu_power_exporter = None diff --git a/tests/test_tilert.py b/tests/test_tilert.py new file mode 100644 index 000000000..fc9745eeb --- /dev/null +++ b/tests/test_tilert.py @@ -0,0 +1,349 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exercise the TileRT adapter using real normalization, launch planning, and conversion.""" + +import json +import os +import re +import subprocess +import sys +import threading +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from marshmallow import ValidationError + +from srtctl.backends import TileRTProtocol +from srtctl.cli.mixins.benchmark_stage import BenchmarkStageMixin +from srtctl.cli.mixins.worker_stage import WorkerStageMixin +from srtctl.cli.submit import show_config_details +from srtctl.core.config import resolve_container_aliases +from srtctl.core.processes import ProcessRegistry +from srtctl.core.roles import expand_roles +from srtctl.core.schema import SrtConfig, TachometerConfig +from srtctl.core.telemetry import generate_tachometer_config +from srtctl.core.topology import Process +from srtctl.frontends import get_frontend +from srtctl.frontends.static_router import RouterWorker + + +def recipe(): + return { + "schema": 2, + "name": "tile-pd", + "model": {"path": "hf:org/model", "container": "decode", "precision": "fp8"}, + "resources": {"gpu_type": "b200", "gpus_per_node": 8}, + "engine": {"type": "tilert", "prefill_container": "prefill", "weights_dir": "/weights/model"}, + "roles": { + "prefill": {"nodes": 1, "workers": 1, "gpus": 8, "args": {"enforce-eager": True}}, + "decode": {"nodes": 1, "workers": 1, "gpus": 8}, + }, + "frontend": {"type": "tilert-router", "enable_multiple_frontends": False}, + "benchmark": {"type": "manual"}, + } + + +def load(data): + resolve_container_aliases(data, {"prefill": "prefill.sqsh", "decode": "decode.sqsh"}) + return SrtConfig.Schema().load(expand_roles(data)) + + +def workers(backend): + endpoints = backend.allocate_endpoints(1, 1, 0, 8, 8, 0, 8, ["p0", "d0"]) + return backend.endpoints_to_processes(endpoints, base_sys_port=7200) + + +@pytest.mark.parametrize("mtp", [True, False]) +def test_roles_render_prefill_transfer_and_decode_commands(mtp): + data = recipe() + data["engine"].update(with_mtp=mtp, max_seq_len=8192, served_model_name="served") + backend = load(data).backend + prefill, decode = workers(backend) + runtime = SimpleNamespace(worker_model_arg="org/model", model_path="org/model") + prefill_cmd = backend.build_worker_command(prefill, [prefill], runtime) + decode_cmd = backend.build_worker_command(decode, [decode], runtime) + assert prefill_cmd[:3] == ["vllm", "serve", "org/model"] + assert prefill_cmd[prefill_cmd.index("--served-model-name") + 1] == "served" + assert prefill_cmd[prefill_cmd.index("--tensor-parallel-size") + 1] == "8" + assert "--enforce-eager" in prefill_cmd + assert json.loads(prefill_cmd[prefill_cmd.index("--kv-transfer-config") + 1]) == { + "kv_connector": "TileRTConnector", + "kv_connector_module_path": "tilert.pd_vllm.prefill_connector", + "kv_role": "kv_producer", + "kv_connector_extra_config": { + "tilert_ctrl_port": 7200, + "tilert_model": "glm5", + "tilert_max_seq_len": 8192, + "tilert_transport": "nixl", + }, + } + assert decode_cmd == [ + "python", + "-m", + "tilert.pd_vllm.decode_server", + "--engine", + "tilert", + "--model", + "glm5", + "--model-weights-dir", + "/weights/model", + "--max-seq-len", + "8192", + "--kv-cache-dtype", + "fp8", + "--transport", + "nixl", + "--ctrl-port", + "7201", + "--http-port", + "6100", + ] + (["--with-mtp"] if mtp else []) + assert ("--speculative-config" in prefill_cmd) == mtp + + +@pytest.mark.parametrize( + ("section", "key", "value", "message"), + [ + ("frontend", "type", "dynamo", "requires frontend.type"), + ("frontend", "enable_multiple_frontends", True, "multiple_frontends"), + ("frontend", "args", {"decode": "wrong:1:2"}, "managed"), + ("engine", "prefill_container", None, "prefill image"), + ("engine", "weights_dir", "/", "dedicated"), + ("engine", "max_seq_len", 0, "positive"), + ], +) +def test_invalid_settings_rejected_before_submission(section, key, value, message): + data = recipe() + data[section][key] = value + with pytest.raises(ValidationError, match=message): + load(data) + + +@pytest.mark.parametrize("mutation", ["two_prefills", "cross_node", "decode_gpu_count", "managed_arg"]) +def test_invalid_role_topology_and_managed_overrides(mutation): + data = recipe() + if mutation == "two_prefills": + data["roles"]["prefill"].update(nodes=2, workers=2) + elif mutation == "cross_node": + data["roles"]["decode"].update(nodes=2, gpus=16) + elif mutation == "decode_gpu_count": + data["roles"]["decode"]["gpus"] = 4 + else: + data["roles"]["prefill"]["args"]["--tensor_parallel_size"] = 2 + with pytest.raises(ValidationError, match="exactly one|fit on one|managed|exactly 8"): + load(data) + + +def test_router_routes_multiple_decodes_and_uses_prepared_tokenizer(): + frontend = get_frontend("tilert-router") + config = load(recipe()) + command = frontend.build_router_command( + [ + RouterWorker("prefill", "http://10.0.0.1:6100"), + RouterWorker("decode", "http://10.0.0.2:6100", 7201), + RouterWorker("decode", "http://10.0.0.3:6110", 7202), + ], + "0.0.0.0", + 8000, + ) + assert command == [ + "python", + "-m", + "tilert.pd_vllm.pd_router", + "--vllm-url", + "http://10.0.0.1:6100", + "--decode", + "10.0.0.2:7201:6100", + "--decode", + "10.0.0.3:7202:6110", + "--host", + "0.0.0.0", + "--port", + "8000", + ] + assert frontend.get_managed_frontend_args(config, config.backend, []) == ["--model-path", "/weights/model"] + assert frontend.parse_health({"status": "ok", "decode_free": 0}, 1, 2).ready + assert not frontend.parse_health({"status": "starting"}, 1, 2).ready + with pytest.raises(ValueError, match="exactly one"): + frontend.build_router_command([RouterWorker("decode", "http://d:6100", 7201)], "0.0.0.0", 8000) + + +def test_router_waits_for_workers_and_aborts_before_launch(): + config = load(recipe()) + frontend = get_frontend("tilert-router") + stopped = threading.Event() + stopped.set() + with ( + patch("srtctl.frontends.static_router.get_hostname_ip", side_effect=lambda node, _: node), + pytest.raises(RuntimeError, match="did not become ready"), + ): + frontend.start_frontends( + None, SimpleNamespace(network_interface="eth0"), config, config.backend, workers(config.backend), stopped + ) + + +class Harness(WorkerStageMixin, BenchmarkStageMixin): + @property + def backend_processes(self): + return workers(self.backend) + + +def test_runtime_selects_role_images_and_only_real_metrics(capsys): + config = load(recipe()) + runner = Harness() + runner.config = config + runner.runtime = SimpleNamespace( + container_image="decode.sqsh", network_interface="eth0", job_id="1", run_name="tile", log_dir=Path("/logs") + ) + assert runner._worker_container("prefill") == "prefill.sqsh" + assert runner._worker_container("decode") == "decode.sqsh" + with patch("srtctl.cli.mixins.benchmark_stage.get_hostname_ip", side_effect=lambda node, _: node): + assert runner._get_aiperf_server_metrics_env(logical_workers_only=True) == { + "AIPERF_SERVER_METRICS_URLS": "http://p0:6100/metrics" + } + with patch("srtctl.core.telemetry.get_hostname_ip", side_effect=lambda node, _: node): + toml = generate_tachometer_config( + processes=runner.backend_processes, + frontend_topology=SimpleNamespace(frontend_nodes=["p0"], frontend_port=8000), + runtime=runner.runtime, + tachometer=TachometerConfig(default_exporters=False), + frontend_type="tilert-router", + worker_metrics_port=config.backend.get_metrics_port, + frontend_metrics_enabled=get_frontend("tilert-router").has_metrics, + ) + assert re.findall(r'url = "([^"\n]+)"', toml) == ["http://p0:6100/metrics"] + show_config_details(config) + output = capsys.readouterr().out + assert "prefill.sqsh" in output and "decode.sqsh" in output + + +def test_preparation_uses_decode_gpu_step_and_propagates_failure(tmp_path, monkeypatch): + runner = Harness() + runner.config = load(recipe()) + runner.runtime = SimpleNamespace( + container_image="decode.sqsh", + environment={}, + srun_options={}, + container_mounts={}, + model_path="org/model", + worker_model_arg="org/model", + is_hf_model=True, + log_dir=tmp_path, + ) + launched = [] + + def srun(**kwargs): + launched.append(kwargs) + return subprocess.Popen([sys.executable, "-c", "raise SystemExit(7)"]) + + monkeypatch.setattr("srtctl.cli.mixins.worker_stage.start_srun_process", srun) + registry = ProcessRegistry("1") + with pytest.raises(RuntimeError, match="exit code 7"): + runner.prepare_backend(registry) + step = launched[0] + assert (step["nodelist"], step["container_image"], step["srun_options"], step["env_to_set"]) == ( + ["d0"], + "decode.sqsh", + {"gpus-per-task": "1"}, + {"TILERT_ROLE": "decode"}, + ) + assert registry.get_process("backend_preparation").exit_code == 7 + + +@pytest.mark.parametrize("broken", [False, True]) +def test_weight_conversion_executes_and_reuses_only_complete_cache(tmp_path, broken): + """Stub only TileRT's external converter; execute the generated preparation verbatim.""" + source = tmp_path / "source" + source.mkdir() + (source / "chat_template.jinja").write_text("{{ messages }}") + (source / "tokenizer.json").write_text('{"version": "1.0"}') + module = tmp_path / "packages" / "tilert" / "models" / "preprocess" + module.mkdir(parents=True) + (module / "weight_converter.py").write_text(""" +import argparse, json, os +from pathlib import Path +p = argparse.ArgumentParser() +p.add_argument('--model_type'); p.add_argument('--model_dir'); p.add_argument('--save_dir') +a = p.parse_args() +target = Path(a.save_dir); target.mkdir(parents=True) +with open(os.environ['CALLS'], 'a') as f: f.write(a.model_type + '\\n') +(target / 'model.safetensors.index.json').write_text(json.dumps({'weight_map': {'weight': 'shard.safetensors'}})) +if os.environ['BROKEN'] == '0': (target / 'shard.safetensors').write_bytes(b'converted') +""") + target = tmp_path / "cache" + backend = TileRTProtocol(weights_dir=str(target)) + runtime = SimpleNamespace(is_hf_model=False, worker_model_arg=str(source)) + preparation = backend.get_preparation(runtime, [Process("d0", frozenset(range(8)), 7201, 6100, "decode", 0)]) + calls = tmp_path / "calls" + env = dict(os.environ, PYTHONPATH=str(tmp_path / "packages"), CALLS=str(calls), BROKEN=str(int(broken))) + result = subprocess.run(preparation.command, env=env, capture_output=True, text=True, check=False) + if broken: + assert result.returncode != 0 + assert "incomplete shards" in result.stderr + assert not target.exists() + assert list(tmp_path.glob(".cache.tmp.*")) == [] + else: + assert result.returncode == 0, result.stderr + assert (target / "shard.safetensors").read_bytes() == b"converted" + assert (target / "chat_template.jinja").read_text() == "{{ messages }}" + assert (target / "tokenizer.json").read_text() == '{"version": "1.0"}' + subprocess.run(preparation.command, env=env, check=True, capture_output=True) + assert calls.read_text() == "glm-5\n" + (target / "shard.safetensors").unlink() + subprocess.run(preparation.command, env=env, check=True, capture_output=True) + assert calls.read_text() == "glm-5\nglm-5\n" + + +def test_native_frontend_launches_after_health_checks(tmp_path, monkeypatch): + config = load(recipe()) + frontend = get_frontend("tilert-router") + runtime = SimpleNamespace( + network_interface="eth0", + log_dir=tmp_path, + container_image="decode.sqsh", + container_mounts={}, + environment={}, + nodes=SimpleNamespace(het_group_for=lambda node: None), + ) + events = [] + + def request(url, **kwargs): + events.append(url) + return SimpleNamespace(status_code=200) + + def launch(**kwargs): + events.append(kwargs) + return subprocess.Popen([sys.executable, "-c", "pass"]) + + monkeypatch.setattr("srtctl.core.health.requests.get", request) + monkeypatch.setattr("srtctl.frontends.static_router.get_hostname_ip", lambda node, _: node) + monkeypatch.setattr("srtctl.frontends.static_router.start_srun_process", launch) + processes = frontend.start_frontends( + SimpleNamespace(frontend_nodes=["p0"], frontend_port=8000), + runtime, + config, + config.backend, + workers(config.backend), + ) + processes[0].popen.wait(timeout=5) + assert events[:2] == ["http://p0:6100/health", "http://d0:6100/health"] + assert events[2]["command"] == [ + "python", + "-m", + "tilert.pd_vllm.pd_router", + "--vllm-url", + "http://p0:6100", + "--decode", + "d0:7201:6100", + "--host", + "0.0.0.0", + "--port", + "8000", + "--model-path", + "/weights/model", + ] + assert events[2]["container_image"] == "decode.sqsh" + assert processes[0].step_name == "tilert_router_0"