From c3a0aecc90f6d6c4e0e6efde201fafc4bcbe638f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sun, 6 Sep 2026 11:39:18 -0700 Subject: [PATCH 1/3] Add the services: block for sidecars and standalone Mooncake stores One top-level list for every long-running process launched next to the job. Each entry has a type that selects a registered ServiceKind (@register_service, like @register_benchmark): generic launches exactly the argv written; mooncake-store runs a standalone Mooncake Store wired to the managed master. The kind supplies defaults (command, start phase, criticality) and the env it injects; ServiceStageMixin launches every kind the same way: placement.node -> physical nodes (head, infra, prefill, decode, agg, workers), optional clone/build of an immutable git source, one srun per node, optional TCP readiness gate, ManagedProcess into the shared registry. start: before_workers runs after the Mooncake master; after_frontend (default) runs once the frontend is healthy. Folds the two open sidecar PRs onto the 2.0 shape: the generic sidecar with source/build_command/inherit_discovery_env/critical from #374 (its --bash direct-runner half is gone with that path), and the standalone Mooncake Store services from #265 as a typed service instead of a backend.mooncake_kv_store.standalone sub-block with a per-role placements map. Per-role segment sizes are now two entries placed on prefill and decode; two services on one node with the same port are rejected before launch. Adds examples/features/services.yaml (an HTTP log browser gated on its port), docs/services.md, a config-reference section, dry-run output, the regenerated schema reference, and tests for the schema, kinds, and stage. --- CLAUDE.md | 18 + docs/README.md | 1 + docs/SUMMARY.md | 1 + docs/config-reference.md | 47 +++ docs/mooncake-kv-store.md | 24 ++ docs/schema-reference.md | 52 +++ docs/services.md | 284 ++++++++++++++++ examples/README.md | 1 + examples/features/services.yaml | 73 +++++ src/srtctl/cli/do_sweep.py | 12 + src/srtctl/cli/mixins/__init__.py | 3 + src/srtctl/cli/mixins/service_stage.py | 263 +++++++++++++++ src/srtctl/cli/submit.py | 28 ++ src/srtctl/core/schema.py | 23 ++ src/srtctl/services/__init__.py | 38 +++ src/srtctl/services/config.py | 226 +++++++++++++ src/srtctl/services/generic.py | 13 + src/srtctl/services/mooncake_store.py | 55 ++++ src/srtctl/services/registry.py | 100 ++++++ tests/test_dry_run.py | 85 ++++- tests/test_services.py | 438 +++++++++++++++++++++++++ 21 files changed, 1783 insertions(+), 2 deletions(-) create mode 100644 docs/services.md create mode 100644 examples/features/services.yaml create mode 100644 src/srtctl/cli/mixins/service_stage.py create mode 100644 src/srtctl/services/__init__.py create mode 100644 src/srtctl/services/config.py create mode 100644 src/srtctl/services/generic.py create mode 100644 src/srtctl/services/mooncake_store.py create mode 100644 src/srtctl/services/registry.py create mode 100644 tests/test_services.py diff --git a/CLAUDE.md b/CLAUDE.md index aaac48d1a..e83a57ffb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -183,6 +183,24 @@ backend: **Validation:** In disaggregated mode, srtslurm rejects configs that set `mooncake_kv_store` without `disaggregation-transfer-backend: mooncake` on `sglang_config.prefill` or `sglang_config.decode`. This catches the common misconfiguration where the master process gets launched but workers fall back to default transport. +### Services + +The top-level `services:` list declares long-running processes launched next to the job (see `docs/services.md`). Each entry has a `type` that selects a `ServiceKind` registered in `src/srtctl/services/` with `@register_service("")`; the kind supplies defaults (command, start phase, criticality) and the env it injects, and `ServiceStageMixin` (`src/srtctl/cli/mixins/service_stage.py`) launches every kind the same way: resolve `placement.node` to physical nodes, optional clone/build of `source`, one `srun` per node, optional TCP `readiness` gate, `ManagedProcess` into the shared registry. `start_services("before_workers")` runs after the Mooncake master; `start_services("after_frontend")` runs after the frontend is healthy. + +```yaml +services: + - name: store + type: mooncake-store # generic (default) | mooncake-store + placement: + node: workers # head | infra | prefill | decode | agg | workers + env: + MOONCAKE_GLOBAL_SEGMENT_SIZE: 100gb + readiness: + port: 8800 +``` + +Adding a kind: subclass `ServiceKind`, set `default_command` / `default_start` / `default_critical`, override `validate`, `container_fallback`, `default_environment`, `forced_environment` as needed, decorate, and import it from `src/srtctl/services/__init__.py`. `srtctl dry-run` prints every service; add a `tests/test_dry_run.py` case when a kind adds visible fields. + ### Host Setup `host_setup` runs commands on each node's **bare host, outside the container**, before any diff --git a/docs/README.md b/docs/README.md index d314ec34c..23939b342 100644 --- a/docs/README.md +++ b/docs/README.md @@ -57,3 +57,4 @@ Once allocated, workers launch inside containers, discover each other through ET - [Profiling](profiling.md) - Performance analysis with torch/nsys - [Analyzing Results](analyzing.md) - Dashboard and visualization - [SGLang Router](sglang-router.md) - Alternative to Dynamo for PD disaggregation +- [Services](services.md) - Sidecars and standalone stores launched next to the job diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md index d9c144664..9f0acf8f1 100644 --- a/docs/SUMMARY.md +++ b/docs/SUMMARY.md @@ -19,6 +19,7 @@ - [SGLang Router](sglang-router.md) - [vLLM Router](vllm-router.md) - [Mooncake KV Store](mooncake-kv-store.md) +- [Services](services.md) ## Benchmarking diff --git a/docs/config-reference.md b/docs/config-reference.md index d756d67a0..70c2e0ced 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -31,6 +31,7 @@ This page is the prose guide: what each block means, how the pieces interact, an - [srun_options](#srun_options) - [setup_script](#setup_script) - [host_setup](#host_setup) +- [services](#services) - [enable_config_dump](#enable_config_dump) - [Complete Examples](#complete-examples) @@ -1754,6 +1755,52 @@ host_setup: --- +## services + +Long-running processes srtctl launches and tracks next to the workers, frontend, and benchmark client. One list covers generic sidecars (an experimental router built from a PR) and typed services (a standalone Mooncake store per worker node). Full reference: [services.md](services.md). + +```yaml +services: + - name: my-sidecar + type: generic # generic (default) | mooncake-store + command: + - python3 + - -m + - my_package.my_sidecar + args: + - --port + - "9000" + container: my-image # alias or path; default: job container + env: + MY_FLAG: "1" + placement: + node: head # head | infra | prefill | decode | agg | workers + start: after_frontend # after_frontend | before_workers + readiness: + port: 9000 + timeout_seconds: 120 + inherit_discovery_env: true # ETCD_ENDPOINTS / NATS_SERVER + critical: false +``` + +| Field | Type | Default | Description | +| --- | --- | --- | --- | +| `name` | string | required | Unique; names `service_.out` and the tracked process | +| `type` | string | `generic` | Registered service kind; supplies defaults and injected env | +| `command` | list[string] | type default | Argv, not shell-interpreted; required for `generic` | +| `args` | list[string] | `[]` | Appended to `command` | +| `container` | string | type fallback, then job container | Image or `srtslurm.yaml` alias | +| `env` | dict | `{}` | Service environment; placeholders like `{node_ip}` are substituted | +| `placement.node` | string | `head` | One instance for `head`/`infra`; one per node for `prefill`/`decode`/`agg`/`workers` | +| `start` | string | type default | `after_frontend` (generic) or `before_workers` (mooncake-store) | +| `readiness` | object | none | `port` + `timeout_seconds`; the job waits for it on every service node | +| `inherit_discovery_env` | bool | `true` | Inject the Dynamo discovery env | +| `critical` | bool | type default | A crash fails the run when true | +| `source`, `build_command` | object, list[string] | none | Clone an immutable git rev and build once before launch; single-node placements only | +| `preamble`, `cpus_per_task`, `cpu_bind`, `srun_options` | | none | Pass-through launch knobs for this service | + +--- + ## enable_config_dump Enable dumping worker configuration to JSON for debugging. diff --git a/docs/mooncake-kv-store.md b/docs/mooncake-kv-store.md index 8d9f6353c..9efa2917b 100644 --- a/docs/mooncake-kv-store.md +++ b/docs/mooncake-kv-store.md @@ -9,6 +9,7 @@ First-class support for [Mooncake](https://github.com/kvcache-ai/Mooncake) as th - [Quick Start (vLLM)](#quick-start-vllm) - [What srtslurm Owns vs What You Set](#what-srtslurm-owns-vs-what-you-set) - [Configuration Reference](#configuration-reference) +- [Standalone Store Services](#standalone-store-services) - [Master Metrics Endpoint](#master-metrics-endpoint) - [Validation](#validation) - [Common Configurations](#common-configurations) @@ -173,6 +174,29 @@ backend: Older Mooncake versions do not recognize this option, so leave it out of those recipes. The existing `--eviction_high_watermark_ratio` controls memory eviction; the `--nof_...` option independently controls the NVMe-over-Fabrics SSD tier. +## Standalone Store Services + +Mooncake can run the Store as a standalone process per node, so workers use embedded clients with `MOONCAKE_GLOBAL_SEGMENT_SIZE=0` while dedicated stores own the DRAM segments. In srtslurm that is a `services:` entry with `type: mooncake-store`: it starts after the master is healthy and before workers, gets `MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER`, and `MOONCAKE_LOCAL_HOSTNAME` from the runtime, and defaults its container to `mooncake_kv_store.container`. + +```yaml +services: + - name: store + type: mooncake-store + placement: + node: workers # or prefill / decode for per-role segment sizes + args: + - --port + - "8800" + env: + MOONCAKE_PROTOCOL: rdma + MOONCAKE_DEVICE: "mlx5_0,mlx5_1" + MOONCAKE_GLOBAL_SEGMENT_SIZE: 100gb + readiness: + port: 8800 +``` + +The worker side stays in the backend's per-mode env (`prefill_environment` / `decode_environment`). See [Services](services.md#example-standalone-mooncake-stores) for the full shape, per-role entries, and the co-location rules. + ## Master Metrics Endpoint The `mooncake_master` admin HTTP server is always exposed on port `8702` on the infra node and starts before workers do (srtslurm waits for it). It serves: diff --git a/docs/schema-reference.md b/docs/schema-reference.md index 56b2044dd..cf9af703d 100644 --- a/docs/schema-reference.md +++ b/docs/schema-reference.md @@ -33,6 +33,7 @@ Top-level keys of a recipe YAML. | `enable_config_dump` | bool | `True` | | | `setup_script` | str \| None | `None` | Custom setup script (runs before dynamo install and worker startup) e.g. "custom-setup.sh" -> runs /configs/custom-setup.sh | | `host_setup` | [HostSetupConfig](#hostsetupconfig) | `HostSetupConfig()` | Commands run on each node's bare host, outside the container, before any worker starts. Cluster-wide default lives in srtslurm.yaml as default_host_setup; a recipe that sets this block replaces that default. | +| `services` | list[[ServiceConfig](#serviceconfig)] | `[]` | Long-running processes launched next to the job: generic sidecars (an experimental router built from a PR) and typed ones (a standalone Mooncake store per worker node). See docs/services.md. | | `identity` | [IdentityConfig](#identityconfig) | `IdentityConfig()` | Virtual identity — declares what *should* be running (verified against fingerprint) | | `reporting` | [ReportingConfig](#reportingconfig) \| None | `None` | Reporting configuration (status API, future: logs to S3, etc.) | @@ -261,6 +262,30 @@ Commands run on the bare host of each allocated node, outside the container. | `ignore_failure` | bool | `False` | When True, a failing node logs a warning instead of failing the job. | | `timeout_seconds` | int | `300` | Per-node wall-clock budget for commands and for teardown. | +### ServiceConfig + +One entry of the top-level ``services:`` list. + +| Key | Type | Default | Description | +|---|---|---|---| +| `name` | str | required | Unique label; names the log file (``service_.out``) and the tracked process. | +| `type` | str | `'generic'` | Service kind. ``generic`` (default) launches exactly what you wrote; ``mooncake-store`` runs a standalone Mooncake Store wired to the managed master. See ``docs/services.md`` for the kinds. | +| `command` | list[str] \| None | `None` | Argv to launch (not shell-interpreted). Required for ``generic``; typed kinds supply a default. | +| `args` | list[str] | `[]` | Extra argv appended to ``command``. | +| `container` | str \| None | `None` | Container image or ``srtslurm.yaml`` alias. Defaults to the kind's fallback (Mooncake's ``mooncake_kv_store.container``), then the job container. | +| `env` | dict[str, str] | `{}` | Environment for the service process, on top of what the kind injects. | +| `source` | [ServiceSourceConfig](#servicesourceconfig) \| None | `None` | Optional git source to clone before ``build_command`` and ``command`` run. Single-node placements only. | +| `build_command` | list[str] \| None | `None` | Argv run once inside the service container, from the clone, before ``command`` starts. Only meaningful with ``source``. | +| `placement` | [ServicePlacementConfig](#serviceplacementconfig) | `ServicePlacementConfig()` | Where the service runs. Default ``head``. | +| `start` | str \| None | `None` | ``after_frontend`` (default for ``generic``) or ``before_workers`` (default for ``mooncake-store``). | +| `readiness` | [ServiceReadinessConfig](#servicereadinessconfig) \| None | `None` | Optional TCP port gate; the job waits for it on every service node before continuing. | +| `inherit_discovery_env` | bool | `True` | Inject ``ETCD_ENDPOINTS`` / ``NATS_SERVER`` so the service can register with the job's Dynamo discovery plane. | +| `critical` | bool \| None | `None` | When true a crash fails the run, like a worker dying. Default false for ``generic`` (a dead sidecar costs its own log, not the run) and true for ``mooncake-store``. Set true for anything in the live request path. | +| `preamble` | str \| None | `None` | Shell run inside the container before ``command`` (``ulimit`` and friends). | +| `cpus_per_task` | int \| None | `None` | Optional ``srun --cpus-per-task``. | +| `cpu_bind` | str \| None | `None` | Optional ``srun --cpu-bind``. | +| `srun_options` | dict[str, str] | `{}` | Extra srun options for this service only. | + ### IdentityConfig Virtual identity for runtime verification and reproduction. @@ -326,6 +351,33 @@ Configuration for a metrics exporter deployed on worker nodes. | `port` | int | required | | | `command` | str \| None | `None` | | +### ServiceSourceConfig + +Git source to build a service from before launching it. + +| Key | Type | Default | Description | +|---|---|---|---| +| `git` | str | required | Repository URL to clone. | +| `rev` | str | required | Immutable ref to check out: a commit SHA, a tag, or ``refs/pull//head`` for an unmerged PR. Branch names are rejected because they move out from under a build. | +| `path` | str \| None | `None` | Optional subdirectory of the clone that ``build_command`` and ``command`` run from. Defaults to the repository root. | + +### ServicePlacementConfig + +Where a service runs. + +| Key | Type | Default | Description | +|---|---|---|---| +| `node` | str | `'head'` | ``head`` or ``infra`` (one instance), ``prefill`` / ``decode`` / ``agg`` (one instance per distinct physical node that role's workers use), or ``workers`` (one instance per worker node). | + +### ServiceReadinessConfig + +TCP readiness gate: the launch blocks until ``port`` accepts connections on every service node. + +| Key | Type | Default | Description | +|---|---|---|---| +| `port` | int | required | TCP port the service listens on. | +| `timeout_seconds` | int | `120` | How long to wait per node before failing the job. | + ### IdentityModelConfig Virtual model identity for runtime verification. diff --git a/docs/services.md b/docs/services.md new file mode 100644 index 000000000..8cc01f683 --- /dev/null +++ b/docs/services.md @@ -0,0 +1,284 @@ +# Services + +The top-level `services:` block declares long-running processes that srtctl launches and tracks next +to the inference workers, the frontend, and the benchmark client. One list, one shape, for anything +that is not a built-in component: an experimental router built from an unmerged PR, a standalone +Mooncake Store per worker node, a debugging HTTP server. Adding one is a recipe change, not a code +change. + +## Table of Contents + +- [Quick Start](#quick-start) +- [Configuration Reference](#configuration-reference) +- [Placement](#placement) +- [Start Order and Readiness](#start-order-and-readiness) +- [Environment](#environment) +- [Building From Source](#building-from-source) +- [Service Types](#service-types) +- [Example: a router from a PR](#example-a-router-from-a-pr) +- [Example: standalone Mooncake stores](#example-standalone-mooncake-stores) +- [Validation](#validation) +- [Limitations](#limitations) + +## Quick Start + +```yaml +services: + - name: my-sidecar + command: + - python3 + - -m + - my_package.my_sidecar + - --port + - "9000" + readiness: + port: 9000 +``` + +`name` and `command` are the only required fields for the default `generic` type. The service runs in +the job container on the head node, starts once workers and the frontend are healthy, and the job +waits until port 9000 answers before moving on. Its log is `service_my-sidecar.out` in the job's log +directory. `examples/features/services.yaml` is a runnable version of this. + +## Configuration Reference + +```yaml +services: + - name: my-sidecar # required, unique across the list + type: generic # generic (default) | mooncake-store + command: # argv, not shell-interpreted; required for generic + - python3 + - -m + - pkg + args: # appended to command + - --flag + container: my-image # image or srtslurm.yaml alias; default: job container + env: # environment for the service process + MY_FLAG: "1" + placement: + node: head # head | infra | prefill | decode | agg | workers + start: after_frontend # after_frontend | before_workers + readiness: # optional TCP gate, checked on every service node + port: 9000 + timeout_seconds: 120 + inherit_discovery_env: true # inject ETCD_ENDPOINTS / NATS_SERVER + critical: false # a crash fails the run when true + preamble: | # shell run before command, inside the container + ulimit -n 1048576 + cpus_per_task: 8 # srun --cpus-per-task + cpu_bind: none # srun --cpu-bind + srun_options: # extra srun options for this service only + exclusive: "" + source: # clone before build/launch; single-node placements only + git: https://github.com/org/repo + rev: + path: subdir + build_command: # run once from the clone, inside the container + - bash + - -lc + - pip install -e . +``` + +| Field | Default | Notes | +| --- | --- | --- | +| `name` | required | Unique. Names `service_.out` and the tracked process. | +| `type` | `generic` | Selects a [service type](#service-types) that supplies defaults and environment. | +| `command` | type default | Argv passed directly to the process. `generic` has no default, so it is required there. | +| `args` | `[]` | Appended to `command`. Handy with typed services that supply the command. | +| `container` | type fallback, then job container | Aliases resolve through `srtslurm.yaml` like every other container key. | +| `env` | `{}` | Merged over the type's defaults; see [Environment](#environment). | +| `placement.node` | `head` | See [Placement](#placement). | +| `start` | type default | `generic`: `after_frontend`. `mooncake-store`: `before_workers`. | +| `readiness` | none | TCP port gate per node. Timing out terminates what this stage started and fails the job. | +| `inherit_discovery_env` | `true` | Inject the same `ETCD_ENDPOINTS` / `NATS_SERVER` the Dynamo frontend gets. | +| `critical` | type default | `generic`: `false`. `mooncake-store`: `true`. | +| `preamble` | none | Shell run after the environment is exported and before `command`. | +| `cpus_per_task`, `cpu_bind`, `srun_options` | none | Pass-through srun knobs for this service's launches. | +| `source`, `build_command` | none | See [Building From Source](#building-from-source). | + +`command`, `args`, `env` values, and `preamble` may use these placeholders: `{node}`, `{node_ip}`, +`{node_id}` (position in the worker list), `{index}` (instance index within the service), `{role}` +(the `placement.node` value), `{head_node}`, `{head_ip}`, `{infra_node}`, `{infra_ip}`, +`{master_port}`, `{metadata_port}`. Only those names are substituted; other braces (JSON in an env +value) are left alone. + +## Placement + +`placement.node` picks the physical nodes. `head` and `infra` launch one instance. `prefill`, +`decode`, and `agg` launch one instance per distinct node that role's workers use, so two TP1 decode +workers on one node share one service. `workers` launches one instance per worker node. + +When a service launches on more than one node its processes and logs get a node suffix: +`service__`. Two services that declare the same `readiness.port` and land on the same +node are rejected before anything launches; give them disjoint placements or ports. + +## Start Order and Readiness + +Services launch in declaration order within a start phase: + +- `before_workers`: after etcd/NATS and the Mooncake master, before any worker. For things workers + connect to at startup. +- `after_frontend`: once workers and the frontend are healthy, before telemetry. For sidecars that + register into a running job. + +Within a phase, a service with `readiness` blocks until its port answers on each of its nodes; a +service without one is considered started when its `srun` is launched. Ongoing health is the shared +`ProcessRegistry` monitor, the same as every other process in the job. It tears the run down only for +`critical: true` services. Anything other components register under or route through should be +critical: a router that dies mid-run otherwise leaves the frontend silently talking to the raw +backend and the benchmark measuring something other than what it claims. + +## Environment + +The service process environment is built in layers, later ones winning: + +1. Discovery env, when `inherit_discovery_env` is true: `ETCD_ENDPOINTS=http://:2379`, + `NATS_SERVER=nats://:4222`. +2. The type's defaults (`mooncake-store` sets `MOONCAKE_LOCAL_HOSTNAME` to the node's IP). +3. The recipe's `env`, with placeholders substituted. +4. Values srtctl owns for the type (`mooncake-store`: `MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER`). + A recipe value for these is ignored. + +## Building From Source + +`source` plus `build_command` clone and build once before launch. The clone runs on the bare host of +the service node (git and network access are host concerns, and the job container may lack git) into +`/services//src`, which every container sees at `/logs/services//src`. +`build_command` and `command` then run inside the service container from that directory (or +`source.path` under it). + +`rev` must be immutable: a commit SHA, a tag, or `refs/pull//head` while iterating on an open +PR. `main`, `master`, and `HEAD` are rejected at load time. Because the build installs into one +container instance, `source` is only allowed with single-node placements (`head`, `infra`). + +## Service Types + +`type` selects a registered `ServiceKind` (`src/srtctl/services/`). A kind supplies defaults and the +environment its process needs; the launch path is shared by every kind. Register a new one with +`@register_service("")`, the same pattern as `@register_benchmark`. + +| Type | Default command | Start | Critical | Notes | +| --- | --- | --- | --- | --- | +| `generic` | none (required) | `after_frontend` | `false` | Launches exactly what you wrote. | +| `mooncake-store` | `python -m mooncake.mooncake_store_service` | `before_workers` | `true` | Requires `backend.mooncake_kv_store`. Container falls back to `mooncake_kv_store.container`. Injects the managed master's address. | + +## Example: a router from a PR + +```yaml +frontend: + type: dynamo + +services: + - name: thunderagent-router + source: + git: https://github.com/ai-dynamo/dynamo + rev: refs/pull/14000/head # switch to a commit SHA once merged + build_command: + - bash + - -lc + - "cd lib/bindings/python && maturin develop --uv && cd ../../.. && pip install -e ." + command: + - python3 + - -m + - dynamo.thunderagent_router + - --endpoint + - dyn://namespace.component.endpoint + - --model-name + - my-model + # The router is the thing under test: other components register under its + # endpoint, so running without it must fail the run, not degrade it. + critical: true +``` + +`inherit_discovery_env` defaults to true, so the router sees the same etcd/NATS as the Dynamo +frontend and workers with no extra configuration. + +## Example: standalone Mooncake stores + +Inference workers run embedded Mooncake clients with `MOONCAKE_GLOBAL_SEGMENT_SIZE=0` while dedicated +per-node stores own the DRAM segments. Decode nodes contribute host memory without an in-process +HiCache pool. One entry per role gives each role its own segment size: + +```yaml +backend: + type: sglang + prefill_environment: + MOONCAKE_PROTOCOL: rdma + MOONCAKE_DEVICE: "mlx5_0,mlx5_1" + MOONCAKE_GLOBAL_SEGMENT_SIZE: "0" + decode_environment: + MOONCAKE_PROTOCOL: rdma + MOONCAKE_DEVICE: "mlx5_0,mlx5_1" + MOONCAKE_GLOBAL_SEGMENT_SIZE: "0" + mooncake_kv_store: + container: mooncake # the master; also the stores' default container + sglang_config: + prefill: + disaggregation-transfer-backend: mooncake + decode: + disaggregation-transfer-backend: mooncake + +services: + - name: store-prefill + type: mooncake-store + placement: + node: prefill + args: + - --port + - "8800" + env: + MOONCAKE_PROTOCOL: rdma + MOONCAKE_DEVICE: "mlx5_0,mlx5_1" + MOONCAKE_GLOBAL_SEGMENT_SIZE: 100gb + preamble: | + ulimit -n 1048576 + ulimit -l unlimited + cpus_per_task: 8 + readiness: + port: 8800 + - name: store-decode + type: mooncake-store + placement: + node: decode + args: + - --port + - "8800" + env: + MOONCAKE_PROTOCOL: rdma + MOONCAKE_DEVICE: "mlx5_0,mlx5_1" + MOONCAKE_GLOBAL_SEGMENT_SIZE: 400gb + preamble: | + ulimit -n 1048576 + ulimit -l unlimited + cpus_per_task: 8 + readiness: + port: 8800 +``` + +Stores start after the master is healthy and before workers. Each store gets `MOONCAKE_MASTER`, +`MOONCAKE_TE_META_DATA_SERVER`, and `MOONCAKE_LOCAL_HOSTNAME` from the runtime. If prefill and decode +share a node, the two entries above collide on port 8800 and the job fails before launching; use one +entry with `placement.node: workers` and a single segment size instead. See +[Mooncake KV Store](mooncake-kv-store.md) for the worker side. + +## Validation + +Rejected at load time, so `srtctl dry-run` catches them: + +- Empty or duplicate `name`; unknown `type`. +- `generic` without `command`; a `command`/`args` entry that is blank; `build_command: []`. +- `placement.node` or `start` outside their vocabularies. +- `source` with a moving `rev`, or with a multi-node placement. +- `type: mooncake-store` without `backend.mooncake_kv_store`. + +Rejected at launch, before any service starts: two services listening on the same port on one node. + +`srtctl dry-run` prints every service's type, placement, start phase, criticality, command, +container, source, readiness, and env. + +## Limitations + +- Declared order is launch order, and `readiness` is the only wait. A service that needs another + service to be ready polls for it in its own `command`. +- `source` builds are single-node only. +- Services run on the sbatch/SLURM path only; there is no local dev mode in 2.0. diff --git a/examples/README.md b/examples/README.md index a7156d512..ce20e7acc 100644 --- a/examples/README.md +++ b/examples/README.md @@ -25,6 +25,7 @@ Aggregated examples run two TP1 workers; disaggregated examples run one TP1 pref | `features/sweep.yaml` | `sweep:` plus `{placeholder}` substitution; one job per combination | | `features/override.yaml` | `base` plus `override_*` and `zip_override_*` variants in one file | | `features/profiling.yaml` | `profiling:` torch capture on an aggregated worker | +| `features/services.yaml` | `services:` sidecar (an HTTP log browser on the head node) with a `readiness:` port gate | ## Cluster aliases diff --git a/examples/features/services.yaml b/examples/features/services.yaml new file mode 100644 index 000000000..2d9a71894 --- /dev/null +++ b/examples/features/services.yaml @@ -0,0 +1,73 @@ +# A `services:` sidecar next to a normal job. +# +# `services:` declares long-running processes srtctl launches and tracks alongside +# the workers, frontend, and benchmark client. This example runs a plain HTTP +# file server over the job's log directory on the head node, gated on its port, +# so you can watch logs from a browser while the benchmark runs. Swap the command +# for the real thing: an experimental router built from a PR (`source:` + +# `build_command:`), a metrics shim, a standalone Mooncake store +# (`type: mooncake-store`). See docs/services.md. +# +# Aliases (`qwen3-0.6b`, `sglang`) resolve through srtslurm.yaml; see examples/README.md. + +schema: 2 +name: "qwen3-0.6b-services" + +slurm: + time_limit: "00:30:00" + +model: + path: "qwen3-0.6b" + container: "sglang" + precision: "bf16" + +resources: + gpu_type: "h100" + gpus_per_node: 8 + +frontend: + type: sglang + enable_multiple_frontends: false + args: + policy: "cache_aware" + +backend: + type: sglang +roles: + agg: + nodes: 1 + workers: 2 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + args: + served-model-name: "Qwen/Qwen3-0.6B" + tensor-parallel-size: 1 + mem-fraction-static: 0.5 + context-length: 4096 + disable-piecewise-cuda-graph: true + +services: + - name: log-browser + # type: generic (default) launches exactly this argv inside the job container. + command: + - python3 + - -m + - http.server + - "9911" + - --directory + - /logs + placement: + node: head # head | infra | prefill | decode | agg | workers + # start: after_frontend (default). Use before_workers for something workers need. + readiness: + port: 9911 # the job waits until this port answers on the service node + timeout_seconds: 60 + inherit_discovery_env: false # no etcd/NATS here; the SGLang router runs without them + # critical: false (default). A dead sidecar costs its own log, not the run. + +benchmark: + type: "sa-bench" + isl: 128 + osl: 128 + concurrencies: "4x8" diff --git a/src/srtctl/cli/do_sweep.py b/src/srtctl/cli/do_sweep.py index 09f4f93fe..c3557ddc6 100644 --- a/src/srtctl/cli/do_sweep.py +++ b/src/srtctl/cli/do_sweep.py @@ -30,6 +30,7 @@ BenchmarkStageMixin, FrontendStageMixin, PostProcessStageMixin, + ServiceStageMixin, TelemetryStageMixin, WorkerStageMixin, ) @@ -83,6 +84,7 @@ class SweepOrchestrator( WorkerStageMixin, FrontendStageMixin, TelemetryStageMixin, + ServiceStageMixin, BenchmarkStageMixin, PostProcessStageMixin, ): @@ -778,6 +780,11 @@ def run(self) -> int: if mooncake_proc is not None: registry.add_process(mooncake_proc) + # Stage 1c: services that workers depend on (standalone Mooncake + # stores, anything with start: before_workers). See docs/services.md. + for proc in self.start_services("before_workers"): + registry.add_process(proc) + # Pre-worker: Ensure HF model is cached before starting workers. # 1. Clean stale lock files from previous crashed downloads # 2. Download model on a single node (blocks until complete) @@ -801,6 +808,11 @@ def run(self) -> int: for proc in frontend_procs: registry.add_process(proc) + # Stage 3b: sidecar services (start: after_frontend, the default), + # once workers and the frontend are healthy and before telemetry. + for proc in self.start_services("after_frontend"): + registry.add_process(proc) + if self.config.telemetry.enabled: if os.environ.get("EVAL_ONLY", "false").lower() == "true": # Eval-only runs skip the benchmark stage, so every expected diff --git a/src/srtctl/cli/mixins/__init__.py b/src/srtctl/cli/mixins/__init__.py index 2c9f084f0..c7550cb90 100644 --- a/src/srtctl/cli/mixins/__init__.py +++ b/src/srtctl/cli/mixins/__init__.py @@ -9,11 +9,13 @@ - FrontendStageMixin: Frontend/nginx orchestration - BenchmarkStageMixin: Benchmark execution - PostProcessStageMixin: Post-benchmark AI analysis +- ServiceStageMixin: The recipe's ``services:`` list (sidecars, standalone stores) """ from srtctl.cli.mixins.benchmark_stage import BenchmarkStageMixin from srtctl.cli.mixins.frontend_stage import FrontendStageMixin from srtctl.cli.mixins.postprocess_stage import PostProcessStageMixin +from srtctl.cli.mixins.service_stage import ServiceStageMixin from srtctl.cli.mixins.telemetry_stage import TelemetryStageMixin from srtctl.cli.mixins.worker_stage import WorkerStageMixin @@ -21,6 +23,7 @@ "BenchmarkStageMixin", "FrontendStageMixin", "PostProcessStageMixin", + "ServiceStageMixin", "TelemetryStageMixin", "WorkerStageMixin", ] diff --git a/src/srtctl/cli/mixins/service_stage.py b/src/srtctl/cli/mixins/service_stage.py new file mode 100644 index 000000000..2668162c4 --- /dev/null +++ b/src/srtctl/cli/mixins/service_stage.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Service stage mixin for ``SweepOrchestrator``: launches the recipe's ``services:`` list. + +Every service kind launches the same way: resolve the nodes its ``placement`` +selects, optionally clone and build a ``source`` once, then one ``srun`` per +node with the kind's environment merged around the recipe's ``env``, an +optional TCP readiness gate, and a ``ManagedProcess`` for the shared +``ProcessRegistry`` (which provides crash detection and teardown). The kind +(``srtctl.services.registry.ServiceKind``) never launches anything itself. + +``start_services("before_workers")`` runs after the Mooncake master and before +workers; ``start_services("after_frontend")`` runs once workers and the +frontend are healthy. See ``docs/services.md``. +""" + +from __future__ import annotations + +import logging +import shlex +from pathlib import Path +from typing import TYPE_CHECKING + +from srtctl.core.health import wait_for_port +from srtctl.core.processes import ManagedProcess +from srtctl.core.slurm import get_hostname_ip, start_srun_process +from srtctl.ports import ETCD_CLIENT_PORT, NATS_PORT +from srtctl.services.registry import ServiceLaunchContext, get_service_kind + +if TYPE_CHECKING: + from srtctl.core.runtime import RuntimeContext + from srtctl.core.schema import SrtConfig + from srtctl.core.topology import Endpoint + from srtctl.services.config import ServiceConfig + +logger = logging.getLogger(__name__) + + +def render_placeholders(value: str, replacements: dict[str, str]) -> str: + """Substitute known ``{placeholder}`` names only, leaving unrelated braces (JSON) untouched.""" + for key, replacement in replacements.items(): + value = value.replace(f"{{{key}}}", replacement) + return value + + +def _await_and_cd(work_dir: str) -> str: + """``cd work_dir`` tolerating a brief lag before the shared mount shows the checkout.""" + quoted = shlex.quote(work_dir) + return f"for _i in $(seq 1 20); do [ -d {quoted} ] && break; sleep 0.5; done; cd {quoted}" + + +class ServiceStageMixin: + """Launch the recipe's ``services:`` entries on the sbatch/SLURM path.""" + + config: SrtConfig + runtime: RuntimeContext + endpoints: list[Endpoint] + + # -- node resolution --------------------------------------------------------- + + def service_nodes(self, service: ServiceConfig) -> list[str]: + """Physical nodes a service's ``placement`` selects, in allocation order, deduplicated.""" + where = service.placement.node + if where == "head": + return [self.runtime.nodes.head] + if where == "infra": + return [self.runtime.nodes.infra] + if where == "workers": + return list(self.runtime.nodes.worker) + seen: dict[str, None] = {} + for endpoint in self.endpoints: + if endpoint.mode == where: + for node in endpoint.nodes: + seen.setdefault(node, None) + order = {node: i for i, node in enumerate(self.runtime.nodes.worker)} + return sorted(seen, key=lambda n: order.get(n, len(order))) + + def _check_port_collisions(self, services: list[ServiceConfig]) -> None: + """Two services that both listen on the same readiness port cannot share a node.""" + owners: dict[tuple[str, int], str] = {} + for service in services: + if service.readiness is None: + continue + for node in self.service_nodes(service): + key = (node, service.readiness.port) + other = owners.setdefault(key, service.name) + if other != service.name: + raise ValueError( + f"services[{service.name}] and services[{other}] both listen on port " + f"{service.readiness.port} on node {node}; give them disjoint placements or ports" + ) + + # -- source build ------------------------------------------------------------ + + def _service_container(self, service: ServiceConfig) -> str: + kind = get_service_kind(service.type) + return service.container or kind.container_fallback(self.config) or str(self.runtime.container_image) + + def _container_path(self, host_path: Path) -> str: + """Host path under ``log_dir`` as seen inside a container (``log_dir`` is mounted at ``/logs``).""" + return str(Path("/logs") / host_path.relative_to(self.runtime.log_dir)) + + def _clone_service_source(self, service: ServiceConfig, node: str) -> Path | None: + """Clone ``service.source`` once on the bare host of ``node``; returns the work dir (host path).""" + source = service.source + if source is None: + return None + checkout_root = self.runtime.log_dir / "services" / service.name / "src" + clone_log = self.runtime.log_dir / f"service_{service.name}.clone.out" + # HTTP/1.1 and no terminal prompt guard against the intermittent smart-HTTP stalls + # seen cloning github.com from compute nodes; 600s covers slow checkouts onto /logs. + git = "GIT_TERMINAL_PROMPT=0 timeout 600s git -c http.version=HTTP/1.1" + root = shlex.quote(str(checkout_root)) + clone_script = ( + f"set -e; mkdir -p {shlex.quote(str(checkout_root.parent))}; " + f"if [ ! -d {root} ]; then " + f"{git} clone --filter=blob:none {shlex.quote(source.git)} {root} && " + f"{git} -C {root} fetch origin {shlex.quote(source.rev)} && " + f"{git} -C {root} checkout FETCH_HEAD; " + "fi" + ) + logger.info("Cloning service %s source %s@%s on %s", service.name, source.git, source.rev, node) + proc = start_srun_process( + command=["bash", "-c", clone_script], + nodelist=[node], + output=str(clone_log), + container_image=None, # bare host: git and network access are host concerns + het_group=self.runtime.nodes.het_group_for(node), + ) + if proc.wait() != 0: + raise RuntimeError( + f"services[{service.name}] source clone failed (exit {proc.returncode}); see {clone_log}" + ) + return checkout_root / source.path if source.path else checkout_root + + def _build_service_source(self, service: ServiceConfig, node: str, work_dir: Path) -> None: + if not service.build_command: + return + build_log = self.runtime.log_dir / f"service_{service.name}.build.out" + logger.info("Building service %s: %s", service.name, shlex.join(service.build_command)) + proc = start_srun_process( + command=list(service.build_command), + nodelist=[node], + output=str(build_log), + container_image=self._service_container(service), + container_mounts=self.runtime.container_mounts, + srun_options=self.runtime.srun_options, + het_group=self.runtime.nodes.het_group_for(node), + bash_preamble=_await_and_cd(self._container_path(work_dir)), + ) + if proc.wait() != 0: + raise RuntimeError( + f"services[{service.name}] build_command failed (exit {proc.returncode}); see {build_log}" + ) + + # -- launch ------------------------------------------------------------------ + + def _service_environment(self, service: ServiceConfig, ctx: ServiceLaunchContext) -> dict[str, str]: + kind = get_service_kind(service.type) + template = ctx.template_vars() + env: dict[str, str] = {} + if service.inherit_discovery_env: + env["ETCD_ENDPOINTS"] = f"http://{self.runtime.nodes.infra}:{ETCD_CLIENT_PORT}" + env["NATS_SERVER"] = f"nats://{self.runtime.nodes.infra}:{NATS_PORT}" + env.update(kind.default_environment(service, ctx)) + env.update({k: render_placeholders(v, template) for k, v in service.env.items()}) + env.update(kind.forced_environment(service, ctx)) + return env + + def _launch_service_instance( + self, service: ServiceConfig, ctx: ServiceLaunchContext, work_dir: Path | None, instances: int + ) -> ManagedProcess: + template = ctx.template_vars() + command = [render_placeholders(part, template) for part in service.effective_command] + preamble_parts: list[str] = [] + if work_dir is not None: + preamble_parts.append(_await_and_cd(self._container_path(work_dir))) + if service.preamble: + preamble_parts.append(render_placeholders(service.preamble, template).rstrip()) + suffix = f"_{ctx.node}" if instances > 1 else "" + log_file = self.runtime.log_dir / f"service_{service.name}{suffix}.out" + + logger.info("Starting service %s (%s) on %s: %s", service.name, service.type, ctx.node, shlex.join(command)) + popen = start_srun_process( + command=command, + nodelist=[ctx.node], + output=str(log_file), + container_image=self._service_container(service), + container_mounts=self.runtime.container_mounts, + env_to_set=self._service_environment(service, ctx), + bash_preamble="; ".join(preamble_parts) or None, + cpus_per_task=service.cpus_per_task, + cpu_bind=service.cpu_bind, + srun_options={**self.runtime.srun_options, **service.srun_options}, + het_group=self.runtime.nodes.het_group_for(ctx.node), + ) + return ManagedProcess( + name=f"service_{service.name}{suffix}", + popen=popen, + log_file=log_file, + node=ctx.node, + critical=service.effective_critical, + ) + + def start_services(self, start: str) -> list[ManagedProcess]: + """Launch every service whose (effective) ``start`` matches, in declaration order. + + Readiness gates block per node; a gate that times out terminates every + process this call started and raises. Returned processes are for the + caller to register with the shared ``ProcessRegistry``. + """ + services = [s for s in self.config.services if s.effective_start == start] + if not services: + return [] + self._check_port_collisions(list(self.config.services)) + + worker_order = {node: i for i, node in enumerate(self.runtime.nodes.worker)} + started: list[ManagedProcess] = [] + try: + for service in services: + nodes = self.service_nodes(service) + if not nodes: + logger.warning( + "services[%s]: placement.node=%s selects no nodes in this allocation; skipping", + service.name, + service.placement.node, + ) + continue + work_dir = self._clone_service_source(service, nodes[0]) + if work_dir is not None: + self._build_service_source(service, nodes[0], work_dir) + + for index, node in enumerate(nodes): + ctx = ServiceLaunchContext( + runtime=self.runtime, + node=node, + node_ip=get_hostname_ip(node, self.runtime.network_interface), + node_id=worker_order.get(node, index), + index=index, + role=service.placement.node, + ) + started.append(self._launch_service_instance(service, ctx, work_dir, len(nodes))) + readiness = service.readiness + if readiness is not None: + logger.info( + "Waiting for service %s on %s (port %d, timeout %ds)", + service.name, + node, + readiness.port, + readiness.timeout_seconds, + ) + if not wait_for_port(node, readiness.port, timeout=readiness.timeout_seconds): + raise RuntimeError( + f"services[{service.name}] did not open port {readiness.port} on {node} within " + f"{readiness.timeout_seconds}s; see {started[-1].log_file}" + ) + logger.info("Service %s ready on %d node(s)", service.name, len(nodes)) + except Exception: + for proc in started: + proc.terminate() + raise + return started diff --git a/src/srtctl/cli/submit.py b/src/srtctl/cli/submit.py index b4e7070fe..fe79ac699 100755 --- a/src/srtctl/cli/submit.py +++ b/src/srtctl/cli/submit.py @@ -396,6 +396,34 @@ def show_config_details(config: SrtConfig) -> None: "outlives this allocation and is inherited by the next job on these nodes." ) + # --- services (see docs/services.md) --- + # Plain lines, not a Table: repo URLs and long argv overflow a narrow console and a + # Table would wrap or truncate them. crop=False keeps each value intact on one line. + if config.services: + console.print("[bold cyan]Services:[/]") + for service in config.services: + console.print( + f" [cyan]{service.name}[/] [dim]type={service.type} placement={service.placement.node} " + f"start={service.effective_start} critical={str(service.effective_critical).lower()}[/]" + ) + console.print(f" [yellow]command:[/] {shlex.join(service.effective_command)}", crop=False) + console.print(f" [yellow]container:[/] {service.container or ''}") + if service.source is not None: + console.print(f" [yellow]source:[/] {service.source.git} @ {service.source.rev}", crop=False) + if service.source.path: + console.print(f" [yellow]source.path:[/] {service.source.path}") + if service.build_command: + console.print(f" [yellow]build_command:[/] {shlex.join(service.build_command)}", crop=False) + if service.readiness is not None: + console.print( + f" [yellow]readiness:[/] tcp/{service.readiness.port}, timeout={service.readiness.timeout_seconds}s" + ) + if service.preamble: + console.print(f" [yellow]preamble:[/] {service.preamble.strip()}", crop=False) + console.print(f" [yellow]inherit_discovery_env:[/] {str(service.inherit_discovery_env).lower()}") + for var, val in sorted(service.env.items()): + console.print(f" [yellow]env.{var}:[/] {val}", crop=False) + # --- srun options --- if config.srun_options: opts = " ".join(f"--{k}={v}" if v else f"--{k}" for k, v in config.srun_options.items()) diff --git a/src/srtctl/core/schema.py b/src/srtctl/core/schema.py index 4d42a384f..6288fa34d 100755 --- a/src/srtctl/core/schema.py +++ b/src/srtctl/core/schema.py @@ -47,6 +47,7 @@ # Leaf module (stdlib-only imports), so this cannot cycle back into schema. from srtctl.core.power.contract import CONTAINER_LOG_DIR +from srtctl.services.config import ServiceConfig logger = logging.getLogger(__name__) @@ -1871,6 +1872,11 @@ class SrtConfig: # default_host_setup; a recipe that sets this block replaces that default. host_setup: HostSetupConfig = field(default_factory=HostSetupConfig) + # Long-running processes launched next to the job: generic sidecars (an + # experimental router built from a PR) and typed ones (a standalone Mooncake + # store per worker node). See docs/services.md. + services: list[ServiceConfig] = field(default_factory=list) + # Virtual identity — declares what *should* be running (verified against fingerprint) identity: IdentityConfig = field(default_factory=IdentityConfig) @@ -1893,8 +1899,25 @@ def __post_init__(self): self._validate_dynamo_sidecar() self._validate_host_setup() self._validate_benchmark_type() + self._validate_services() self._warn_dp_launch_mode() + def _validate_services(self) -> None: + """Whole-list checks for ``services:``: unique names, then each kind's recipe-level rules. + + Per-entry checks (empty command, moving-branch source rev, ...) live on + ``ServiceConfig.__post_init__``; a kind's ``validate`` sees the full + recipe (a ``mooncake-store`` needs ``backend.mooncake_kv_store``). + """ + from srtctl.services.registry import get_service_kind + + seen: set[str] = set() + for service in self.services: + if service.name in seen: + raise ValidationError(f"services[].name must be unique; duplicate: {service.name!r}") + seen.add(service.name) + get_service_kind(service.type).validate(service, self) + def _validate_benchmark_type(self) -> None: """Reject a benchmark.type that no runner is registered for. diff --git a/src/srtctl/services/__init__.py b/src/srtctl/services/__init__.py new file mode 100644 index 000000000..1ef269e3d --- /dev/null +++ b/src/srtctl/services/__init__.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The top-level ``services:`` block: user-declared long-running processes launched next to the job.""" + +# Import kinds to trigger registration. +from srtctl.services import generic, mooncake_store +from srtctl.services.config import ( + SERVICE_PLACEMENTS, + SERVICE_STARTS, + ServiceConfig, + ServicePlacementConfig, + ServiceReadinessConfig, + ServiceSourceConfig, +) +from srtctl.services.registry import ( + ServiceKind, + ServiceLaunchContext, + get_service_kind, + list_service_types, + register_service, +) + +__all__ = [ + "SERVICE_PLACEMENTS", + "SERVICE_STARTS", + "ServiceConfig", + "ServiceKind", + "ServiceLaunchContext", + "ServicePlacementConfig", + "ServiceReadinessConfig", + "ServiceSourceConfig", + "generic", + "get_service_kind", + "list_service_types", + "mooncake_store", + "register_service", +] diff --git a/src/srtctl/services/config.py b/src/srtctl/services/config.py new file mode 100644 index 000000000..4ea6e3605 --- /dev/null +++ b/src/srtctl/services/config.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Recipe dataclasses for the top-level ``services:`` block. + +A service is any long-running process srtctl launches next to the workers, +frontend, and benchmark client: an experimental router built from a PR, a +standalone Mooncake Store per worker node, a debugging HTTP server. One list, +one shape; ``type`` selects a :class:`~srtctl.services.registry.ServiceKind` +that supplies defaults and injects the environment that kind needs. See +``docs/services.md``. +""" + +import builtins +import logging +from dataclasses import field +from typing import ClassVar + +from marshmallow import Schema, ValidationError +from marshmallow_dataclass import dataclass + +logger = logging.getLogger(__name__) + +# Where a service runs. head / infra are one node; prefill / decode / agg are the +# distinct physical nodes the role's workers land on; workers is every worker node. +SERVICE_PLACEMENTS: tuple[str, ...] = ("head", "infra", "prefill", "decode", "agg", "workers") +SINGLE_NODE_PLACEMENTS: frozenset[str] = frozenset({"head", "infra"}) + +# When a service starts relative to the rest of the job. +SERVICE_STARTS: tuple[str, ...] = ("before_workers", "after_frontend") + +# Immutable-ref guard for services[].source.rev. +_MOVING_REFS: frozenset[str] = frozenset({"main", "master", "HEAD"}) + + +@dataclass(frozen=True) +class ServiceSourceConfig: + """Git source to build a service from before launching it. + + Attributes: + git: Repository URL to clone. + rev: Immutable ref to check out: a commit SHA, a tag, or + ``refs/pull//head`` for an unmerged PR. Branch names are + rejected because they move out from under a build. + path: Optional subdirectory of the clone that ``build_command`` and + ``command`` run from. Defaults to the repository root. + """ + + git: str + rev: str + path: str | None = None + + Schema: ClassVar[type[Schema]] = Schema + + def __post_init__(self) -> None: + if not self.git.strip(): + raise ValidationError("services[].source.git must be a non-empty repository URL") + if not self.rev.strip(): + raise ValidationError("services[].source.rev must be a non-empty immutable ref") + if self.rev.strip() in _MOVING_REFS: + raise ValidationError( + "services[].source.rev must be an immutable ref (commit SHA, tag, or refs/pull//head), " + f"not a moving branch name: {self.rev!r}" + ) + if self.path is not None and not self.path.strip(): + raise ValidationError("services[].source.path must not be blank when set") + + +@dataclass(frozen=True) +class ServicePlacementConfig: + """Where a service runs. + + Attributes: + node: ``head`` or ``infra`` (one instance), ``prefill`` / ``decode`` / + ``agg`` (one instance per distinct physical node that role's + workers use), or ``workers`` (one instance per worker node). + """ + + node: str = "head" + + Schema: ClassVar[type[Schema]] = Schema + + def __post_init__(self) -> None: + if self.node not in SERVICE_PLACEMENTS: + raise ValidationError( + f"services[].placement.node must be one of {', '.join(SERVICE_PLACEMENTS)}; got {self.node!r}" + ) + + +@dataclass(frozen=True) +class ServiceReadinessConfig: + """TCP readiness gate: the launch blocks until ``port`` accepts connections on every service node. + + Attributes: + port: TCP port the service listens on. + timeout_seconds: How long to wait per node before failing the job. + """ + + port: int + timeout_seconds: int = 120 + + Schema: ClassVar[type[Schema]] = Schema + + def __post_init__(self) -> None: + if not 1 <= self.port <= 65535: + raise ValidationError("services[].readiness.port must be between 1 and 65535") + if self.timeout_seconds <= 0: + raise ValidationError("services[].readiness.timeout_seconds must be positive") + + +@dataclass(frozen=True) +class ServiceConfig: + """One entry of the top-level ``services:`` list. + + Attributes: + name: Unique label; names the log file (``service_.out``) and the + tracked process. + type: Service kind. ``generic`` (default) launches exactly what you + wrote; ``mooncake-store`` runs a standalone Mooncake Store wired to + the managed master. See ``docs/services.md`` for the kinds. + command: Argv to launch (not shell-interpreted). Required for + ``generic``; typed kinds supply a default. + args: Extra argv appended to ``command``. + container: Container image or ``srtslurm.yaml`` alias. Defaults to the + kind's fallback (Mooncake's ``mooncake_kv_store.container``), then + the job container. + env: Environment for the service process, on top of what the kind injects. + source: Optional git source to clone before ``build_command`` and + ``command`` run. Single-node placements only. + build_command: Argv run once inside the service container, from the + clone, before ``command`` starts. Only meaningful with ``source``. + placement: Where the service runs. Default ``head``. + start: ``after_frontend`` (default for ``generic``) or + ``before_workers`` (default for ``mooncake-store``). + readiness: Optional TCP port gate; the job waits for it on every + service node before continuing. + inherit_discovery_env: Inject ``ETCD_ENDPOINTS`` / ``NATS_SERVER`` so + the service can register with the job's Dynamo discovery plane. + critical: When true a crash fails the run, like a worker dying. Default + false for ``generic`` (a dead sidecar costs its own log, not the + run) and true for ``mooncake-store``. Set true for anything in + the live request path. + preamble: Shell run inside the container before ``command`` + (``ulimit`` and friends). + cpus_per_task: Optional ``srun --cpus-per-task``. + cpu_bind: Optional ``srun --cpu-bind``. + srun_options: Extra srun options for this service only. + """ + + name: str + type: str = "generic" + command: list[str] | None = None + args: list[str] = field(default_factory=list) + container: str | None = None + env: dict[str, str] = field(default_factory=dict) + source: ServiceSourceConfig | None = None + build_command: list[str] | None = None + placement: ServicePlacementConfig = field(default_factory=ServicePlacementConfig) + start: str | None = None + readiness: ServiceReadinessConfig | None = None + inherit_discovery_env: bool = True + critical: bool | None = None + preamble: str | None = None + cpus_per_task: int | None = None + cpu_bind: str | None = None + srun_options: dict[str, str] = field(default_factory=dict) + + # builtins.type: the ``type`` field above shadows the builtin inside the class body. + Schema: ClassVar[builtins.type[Schema]] = Schema + + def __post_init__(self) -> None: + from srtctl.services.registry import get_service_kind, list_service_types + + if not self.name.strip(): + raise ValidationError("services[].name must be a non-empty string") + label = f"services[{self.name}]" + if self.type not in list_service_types(): + raise ValidationError( + f"{label}.type {self.type!r} is not a known service type (known: {', '.join(list_service_types())})" + ) + kind = get_service_kind(self.type) + if self.command is not None and not self.command: + raise ValidationError(f"{label}.command, if set, must be non-empty (omit it to use the type's default)") + if self.command is None and kind.default_command is None: + raise ValidationError(f"{label}.command is required for type {self.type!r}") + if any(not str(part).strip() for part in [*(self.command or []), *self.args]): + raise ValidationError(f"{label}.command/args must not contain empty arguments") + if self.build_command is not None and not self.build_command: + raise ValidationError(f"{label}.build_command, if set, must be non-empty (omit it entirely instead)") + if self.source is not None and self.placement.node not in SINGLE_NODE_PLACEMENTS: + raise ValidationError( + f"{label}.source requires a single-node placement (head or infra); got placement.node=" + f"{self.placement.node!r}" + ) + if self.source is not None and not self.build_command: + logger.warning( + "%s sets 'source' without 'build_command'; the source is cloned but nothing builds it " + "before 'command' runs. This is almost always a mistake.", + label, + ) + if self.start is not None and self.start not in SERVICE_STARTS: + raise ValidationError(f"{label}.start must be one of {', '.join(SERVICE_STARTS)}; got {self.start!r}") + if self.cpus_per_task is not None and self.cpus_per_task <= 0: + raise ValidationError(f"{label}.cpus_per_task must be positive") + + # -- effective values (type defaults applied) ------------------------------ + + @property + def effective_command(self) -> list[str]: + """``command`` plus ``args``, with the kind's default command when none is written.""" + from srtctl.services.registry import get_service_kind + + base = self.command if self.command is not None else list(get_service_kind(self.type).default_command or ()) + return [*base, *self.args] + + @property + def effective_start(self) -> str: + from srtctl.services.registry import get_service_kind + + return self.start if self.start is not None else get_service_kind(self.type).default_start + + @property + def effective_critical(self) -> bool: + from srtctl.services.registry import get_service_kind + + return self.critical if self.critical is not None else get_service_kind(self.type).default_critical diff --git a/src/srtctl/services/generic.py b/src/srtctl/services/generic.py new file mode 100644 index 000000000..090ee7ef2 --- /dev/null +++ b/src/srtctl/services/generic.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``type: generic``: launch exactly the command the recipe wrote.""" + +from __future__ import annotations + +from srtctl.services.registry import ServiceKind, register_service + + +@register_service("generic") +class GenericService(ServiceKind): + """A user-declared sidecar. No defaults beyond the shared ones: starts after the frontend, non-critical.""" diff --git a/src/srtctl/services/mooncake_store.py b/src/srtctl/services/mooncake_store.py new file mode 100644 index 000000000..61340380e --- /dev/null +++ b/src/srtctl/services/mooncake_store.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``type: mooncake-store``: a standalone Mooncake Store process wired to the managed master. + +Lets inference workers run embedded Mooncake clients with +``MOONCAKE_GLOBAL_SEGMENT_SIZE=0`` while dedicated per-node stores own the DRAM +segments (decode nodes contribute host memory without an in-process HiCache +pool). Requires ``backend.mooncake_kv_store``, which is what launches the master +the store registers with. Starts before workers and is critical by default. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from marshmallow import ValidationError + +from srtctl.ports import MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_MASTER_PORT +from srtctl.services.registry import ServiceKind, ServiceLaunchContext, register_service + +if TYPE_CHECKING: + from srtctl.core.schema import SrtConfig + from srtctl.services.config import ServiceConfig + + +@register_service("mooncake-store") +class MooncakeStoreService(ServiceKind): + """Standalone Mooncake Store; one instance per placed node, started before workers.""" + + default_command = ("python", "-m", "mooncake.mooncake_store_service") + default_start = "before_workers" + default_critical = True + + def validate(self, service: ServiceConfig, config: SrtConfig) -> None: + if getattr(config.backend, "mooncake_kv_store", None) is None: + raise ValidationError( + f"services[{service.name}] (type mooncake-store) requires backend.mooncake_kv_store, " + "which launches the master the store registers with" + ) + + def container_fallback(self, config: SrtConfig) -> str | None: + mooncake_cfg = getattr(config.backend, "mooncake_kv_store", None) + return getattr(mooncake_cfg, "container", None) + + def default_environment(self, service: ServiceConfig, ctx: ServiceLaunchContext) -> dict[str, str]: + # The recipe may pin a specific NIC identity by setting this itself. + return {"MOONCAKE_LOCAL_HOSTNAME": ctx.node_ip} + + def forced_environment(self, service: ServiceConfig, ctx: ServiceLaunchContext) -> dict[str, str]: + infra_ip = ctx.runtime.infra_node_ip + return { + "MOONCAKE_MASTER": f"{infra_ip}:{MOONCAKE_MASTER_PORT}", + "MOONCAKE_TE_META_DATA_SERVER": f"http://{infra_ip}:{MOONCAKE_HTTP_METADATA_PORT}/metadata", + } diff --git a/src/srtctl/services/registry.py b/src/srtctl/services/registry.py new file mode 100644 index 000000000..c54ab21d6 --- /dev/null +++ b/src/srtctl/services/registry.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Service kinds: what ``services[].type`` selects. + +A kind supplies defaults (command, start phase, criticality) and the environment +a service of that kind needs at launch. It never launches anything itself; the +``ServiceStageMixin`` does that uniformly for every kind. Register a new kind +with :func:`register_service`, the same pattern as ``@register_benchmark``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar + +if TYPE_CHECKING: + from srtctl.core.runtime import RuntimeContext + from srtctl.core.schema import SrtConfig + from srtctl.services.config import ServiceConfig + + +@dataclass(frozen=True) +class ServiceLaunchContext: + """Everything a kind may need to compute one service instance's environment and templates.""" + + runtime: RuntimeContext + node: str + node_ip: str + node_id: int # position of ``node`` in runtime.nodes.worker, or the instance index for head/infra + index: int # instance index within this service (0..n-1) + role: str # the service's placement.node value + + def template_vars(self) -> dict[str, str]: + """Placeholders substituted into command, args, env values, and preamble.""" + from srtctl.ports import MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_MASTER_PORT + + return { + "node": self.node, + "node_ip": self.node_ip, + "node_id": str(self.node_id), + "index": str(self.index), + "role": self.role, + "head_node": self.runtime.nodes.head, + "head_ip": self.runtime.head_node_ip, + "infra_node": self.runtime.nodes.infra, + "infra_ip": self.runtime.infra_node_ip, + "master_port": str(MOONCAKE_MASTER_PORT), + "metadata_port": str(MOONCAKE_HTTP_METADATA_PORT), + } + + +class ServiceKind: + """Base for a registered service type. Subclass, set the class attributes, override hooks as needed.""" + + type_name: ClassVar[str] = "" + # Argv used when the recipe omits ``command``. None means ``command`` is required. + default_command: ClassVar[tuple[str, ...] | None] = None + default_start: ClassVar[str] = "after_frontend" + default_critical: ClassVar[bool] = False + + def validate(self, service: ServiceConfig, config: SrtConfig) -> None: + """Whole-recipe checks for one service (raise ``marshmallow.ValidationError``).""" + + def container_fallback(self, config: SrtConfig) -> str | None: + """Image to use when the service sets no ``container``; None falls through to the job container.""" + return None + + def default_environment(self, service: ServiceConfig, ctx: ServiceLaunchContext) -> dict[str, str]: + """Environment the kind provides; the recipe's ``env`` overrides it.""" + return {} + + def forced_environment(self, service: ServiceConfig, ctx: ServiceLaunchContext) -> dict[str, str]: + """Environment srtctl owns for this kind; it overrides the recipe's ``env``.""" + return {} + + +_SERVICE_KINDS: dict[str, ServiceKind] = {} + + +def register_service(name: str): + """Class decorator registering a :class:`ServiceKind` under ``services[].type: ``.""" + + def decorator(cls: type[ServiceKind]) -> type[ServiceKind]: + cls.type_name = name + _SERVICE_KINDS[name] = cls() + return cls + + return decorator + + +def get_service_kind(name: str) -> ServiceKind: + try: + return _SERVICE_KINDS[name] + except KeyError: + raise ValueError(f"Unknown service type {name!r}. Known: {', '.join(list_service_types())}") from None + + +def list_service_types() -> list[str]: + return sorted(_SERVICE_KINDS) diff --git a/tests/test_dry_run.py b/tests/test_dry_run.py index 5d4b4cf68..dfcdda2fb 100644 --- a/tests/test_dry_run.py +++ b/tests/test_dry_run.py @@ -412,6 +412,86 @@ def test_vllm_mooncake_store_config_in_dry_run(self, capsys): assert "100GB" in output +class TestDryRunServices: + """services: command, env, source, and readiness must be visible before submitting.""" + + def test_service_command_env_and_readiness_shown(self, capsys): + config = _make_config( + { + "services": [ + { + "name": "thunderagent-router", + "command": ["python3", "-m", "dynamo.thunderagent_router", "--endpoint", "dyn://ns.comp.ep"], + "env": {"ROUTER_LOG_LEVEL": "debug"}, + "readiness": {"port": 9100}, + } + ] + } + ) + show_config_details(config) + output = capsys.readouterr().out + assert "Services:" in output + assert "thunderagent-router" in output + assert "dynamo.thunderagent_router" in output + assert "ROUTER_LOG_LEVEL" in output + assert "debug" in output + assert "tcp/9100" in output + assert "placement=head" in output + + def test_source_and_build_command_shown(self, capsys): + config = _make_config( + { + "services": [ + { + "name": "thunderagent-router", + "command": ["python3", "-m", "dynamo.thunderagent_router"], + "source": {"git": "https://github.com/ai-dynamo/dynamo", "rev": "refs/pull/14000/head"}, + "build_command": ["bash", "-lc", "maturin develop --uv && pip install -e ."], + } + ] + } + ) + show_config_details(config) + output = capsys.readouterr().out + assert "https://github.com/ai-dynamo/dynamo" in output + assert "refs/pull/14000/head" in output + assert "maturin develop" in output + + def test_mooncake_store_shows_type_defaults(self, capsys): + config = _make_config( + { + "backend": { + "type": "sglang", + "mooncake_kv_store": {"container": "mooncake.sqsh"}, + "sglang_config": { + "prefill": {"disaggregation-transfer-backend": "mooncake"}, + "decode": {"disaggregation-transfer-backend": "mooncake"}, + }, + }, + "services": [ + { + "name": "store", + "type": "mooncake-store", + "placement": {"node": "workers"}, + "env": {"MOONCAKE_GLOBAL_SEGMENT_SIZE": "100gb"}, + } + ], + } + ) + show_config_details(config) + output = capsys.readouterr().out + assert "mooncake.mooncake_store_service" in output + assert "type=mooncake-store" in output + assert "start=before_workers" in output + assert "critical=true" in output + assert "100gb" in output + + def test_no_services_omits_the_panel(self, capsys): + config = _make_config() + show_config_details(config) + assert "Services:" not in capsys.readouterr().out + + class TestDryRunHetJobs: """Het structure panel appears only when het is enabled.""" @@ -584,8 +664,9 @@ class TestInfmaxWorkspaceMount: --container-mounts against the failed arm's showed this single missing entry. """ - AGENTIC = {"benchmark": {"type": "custom", - "command": "bash /infmax-workspace/benchmarks/multi_node/agentic_srt.sh"}} + AGENTIC = { + "benchmark": {"type": "custom", "command": "bash /infmax-workspace/benchmarks/multi_node/agentic_srt.sh"} + } def test_mount_is_shown_when_the_variable_is_set(self, capsys): config = _make_config(self.AGENTIC) diff --git a/tests/test_services.py b/tests/test_services.py new file mode 100644 index 000000000..069790f2e --- /dev/null +++ b/tests/test_services.py @@ -0,0 +1,438 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the top-level ``services:`` block: schema, kinds, and the launch stage.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import yaml +from marshmallow import ValidationError + +from srtctl.cli.do_sweep import SweepOrchestrator +from srtctl.core.runtime import Nodes, RuntimeContext +from srtctl.core.schema import SrtConfig +from srtctl.core.topology import Endpoint +from srtctl.ports import MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_MASTER_PORT +from srtctl.services import ServiceConfig, ServiceSourceConfig, list_service_types + +SRUN = "srtctl.cli.mixins.service_stage.start_srun_process" +WAIT = "srtctl.cli.mixins.service_stage.wait_for_port" +HOST_IP = "srtctl.cli.mixins.service_stage.get_hostname_ip" + +DISAGG_HEAD = """ +name: services-test +model: + path: /model + container: /job.sqsh + precision: bf16 +resources: + gpu_type: b200 + gpus_per_node: 8 + prefill_nodes: 1 + decode_nodes: 2 + prefill_workers: 1 + decode_workers: 2 + gpus_per_prefill: 8 + gpus_per_decode: 8 +benchmark: + type: manual +""" + + +def _load(services_yaml: str, head: str = DISAGG_HEAD, backend: str = "backend:\n type: sglang\n") -> SrtConfig: + return SrtConfig.Schema().load(yaml.safe_load(head + backend + services_yaml)) + + +def _runtime(tmp_path: Path) -> RuntimeContext: + return RuntimeContext( + job_id="12345", + run_name="test-run", + nodes=Nodes(head="node0", bench="node0", infra="node0", worker=("node1", "node2", "node3")), + head_node_ip="10.0.0.10", + infra_node_ip="10.0.0.10", + log_dir=tmp_path, + model_path=Path("/model"), + container_image=Path("/job.sqsh"), + gpus_per_node=8, + network_interface="eth0", + container_mounts={}, + environment={}, + ) + + +def _proc(returncode: int = 0) -> MagicMock: + proc = MagicMock() + proc.wait.return_value = returncode + proc.returncode = returncode + proc.poll.return_value = None + return proc + + +# --- schema ------------------------------------------------------------------- + + +def test_registered_kinds() -> None: + assert list_service_types() == ["generic", "mooncake-store"] + + +def test_generic_service_loads_block_yaml_with_defaults() -> None: + config = _load( + """ +services: + - name: files + command: + - python3 + - -m + - http.server + args: + - "9911" + readiness: + port: 9911 + timeout_seconds: 30 +""" + ) + (svc,) = config.services + assert svc.type == "generic" + assert svc.effective_command == ["python3", "-m", "http.server", "9911"] + assert svc.placement.node == "head" + assert svc.effective_start == "after_frontend" + assert svc.effective_critical is False + assert svc.inherit_discovery_env is True + assert svc.readiness is not None and svc.readiness.port == 9911 + + +def test_generic_requires_command() -> None: + with pytest.raises(ValidationError, match="command is required"): + _load("services:\n - name: nothing\n") + + +def test_unknown_type_rejected() -> None: + with pytest.raises(ValidationError, match="not a known service type"): + _load("services:\n - name: x\n type: sidecar\n command: [/bin/true]\n") + + +def test_duplicate_names_rejected() -> None: + with pytest.raises(ValidationError, match="must be unique"): + _load("services:\n - name: a\n command: [/bin/true]\n - name: a\n command: [/bin/true]\n") + + +def test_invalid_placement_and_start_rejected() -> None: + with pytest.raises(ValidationError, match="placement.node must be one of"): + _load("services:\n - name: a\n command: [/bin/true]\n placement:\n node: everywhere\n") + with pytest.raises(ValidationError, match="start must be one of"): + _load("services:\n - name: a\n command: [/bin/true]\n start: eventually\n") + + +def test_source_rules() -> None: + with pytest.raises(ValidationError, match="immutable ref"): + ServiceSourceConfig(git="https://example.com/r", rev="main") + with pytest.raises(ValidationError, match="single-node placement"): + _load( + """ +services: + - name: router + command: [python3, -m, router] + placement: + node: workers + source: + git: https://example.com/r + rev: refs/pull/1/head +""" + ) + + +def test_mooncake_store_defaults_and_requires_master() -> None: + with pytest.raises(ValidationError, match="requires backend.mooncake_kv_store"): + _load("services:\n - name: store\n type: mooncake-store\n placement:\n node: workers\n") + + config = _load( + """ +services: + - name: store + type: mooncake-store + placement: + node: workers +""", + backend="backend:\n type: sglang\n mooncake_kv_store:\n container: /mooncake.sqsh\n" + " sglang_config:\n prefill:\n disaggregation-transfer-backend: mooncake\n" + " decode:\n disaggregation-transfer-backend: mooncake\n", + ) + (svc,) = config.services + assert svc.effective_command == ["python", "-m", "mooncake.mooncake_store_service"] + assert svc.effective_start == "before_workers" + assert svc.effective_critical is True + + +# --- stage --------------------------------------------------------------------- + + +def _orchestrator(config: SrtConfig, tmp_path: Path) -> SweepOrchestrator: + return SweepOrchestrator(config=config, runtime=_runtime(tmp_path)) + + +def test_no_matching_services_is_a_noop(tmp_path: Path) -> None: + orchestrator = _orchestrator(_load("services:\n - name: a\n command: [/bin/true]\n"), tmp_path) + with patch(SRUN) as srun: + assert orchestrator.start_services("before_workers") == [] + srun.assert_not_called() + + +def test_generic_launches_on_head_with_discovery_env(tmp_path: Path) -> None: + config = _load( + """ +services: + - name: router + command: [python3, -m, router, --node, "{node}", --infra, "{infra_ip}"] + env: + LOG_LEVEL: debug +""" + ) + orchestrator = _orchestrator(config, tmp_path) + with patch(SRUN, return_value=_proc()) as srun, patch(HOST_IP, return_value="10.0.0.10"): + procs = orchestrator.start_services("after_frontend") + + srun.assert_called_once() + kw = srun.call_args.kwargs + assert kw["nodelist"] == ["node0"] + assert kw["command"] == ["python3", "-m", "router", "--node", "node0", "--infra", "10.0.0.10"] + assert kw["container_image"] == "/job.sqsh" + assert kw["env_to_set"]["ETCD_ENDPOINTS"] == "http://node0:2379" + assert kw["env_to_set"]["NATS_SERVER"] == "nats://node0:4222" + assert kw["env_to_set"]["LOG_LEVEL"] == "debug" + assert kw["bash_preamble"] is None + (proc,) = procs + assert proc.name == "service_router" + assert proc.node == "node0" + assert proc.critical is False + assert proc.log_file == tmp_path / "service_router.out" + + +def test_container_alias_and_no_discovery_env(tmp_path: Path) -> None: + config = _load( + "services:\n - name: s\n command: [/bin/true]\n container: /mine.sqsh\n inherit_discovery_env: false\n" + " critical: true\n" + ) + with patch(SRUN, return_value=_proc()) as srun, patch(HOST_IP, return_value="10.0.0.10"): + (proc,) = _orchestrator(config, tmp_path).start_services("after_frontend") + assert srun.call_args.kwargs["container_image"] == "/mine.sqsh" + assert "ETCD_ENDPOINTS" not in srun.call_args.kwargs["env_to_set"] + assert proc.critical is True + + +def test_declared_order_is_launch_order(tmp_path: Path) -> None: + config = _load( + "services:\n - name: b\n command: [echo, b]\n - name: a\n command: [echo, a]\n" + " - name: c\n command: [echo, c]\n" + ) + with patch(SRUN, return_value=_proc()) as srun, patch(HOST_IP, return_value="10.0.0.10"): + _orchestrator(config, tmp_path).start_services("after_frontend") + assert [call.kwargs["command"][1] for call in srun.call_args_list] == ["b", "a", "c"] + + +def test_readiness_gate_blocks_and_failure_terminates_started(tmp_path: Path) -> None: + config = _load( + "services:\n - name: a\n command: [/bin/true]\n readiness:\n port: 9000\n timeout_seconds: 5\n" + ) + orchestrator = _orchestrator(config, tmp_path) + with ( + patch(SRUN, return_value=_proc()) as srun, + patch(HOST_IP, return_value="10.0.0.10"), + patch(WAIT, return_value=True) as wait, + ): + orchestrator.start_services("after_frontend") + wait.assert_called_once_with("node0", 9000, timeout=5) + + popen = _proc() + with ( + patch(SRUN, return_value=popen), + patch(HOST_IP, return_value="10.0.0.10"), + patch(WAIT, return_value=False), + pytest.raises(RuntimeError, match="did not open port 9000"), + ): + orchestrator.start_services("after_frontend") + popen.terminate.assert_called_once() + assert srun is not None + + +def test_source_is_cloned_on_bare_host_and_built_in_container(tmp_path: Path) -> None: + config = _load( + """ +services: + - name: router + command: [python3, -m, router] + source: + git: https://example.com/repo + rev: refs/pull/1/head + path: lib/router + build_command: [bash, -lc, "pip install -e ."] +""" + ) + with patch(SRUN, return_value=_proc()) as srun, patch(HOST_IP, return_value="10.0.0.10"): + _orchestrator(config, tmp_path).start_services("after_frontend") + + clone, build, launch = srun.call_args_list + assert clone.kwargs["container_image"] is None + assert "git -c http.version=HTTP/1.1 clone" in clone.kwargs["command"][-1] + assert "refs/pull/1/head" in clone.kwargs["command"][-1] + assert build.kwargs["container_image"] == "/job.sqsh" + assert build.kwargs["command"] == ["bash", "-lc", "pip install -e ."] + # Build and launch run inside the container, where log_dir is mounted at /logs. + assert "/logs/services/router/src/lib/router" in build.kwargs["bash_preamble"] + assert str(tmp_path) not in build.kwargs["bash_preamble"] + assert launch.kwargs["command"] == ["python3", "-m", "router"] + assert "/logs/services/router/src/lib/router" in launch.kwargs["bash_preamble"] + + +def test_clone_and_build_failures_raise(tmp_path: Path) -> None: + config = _load( + """ +services: + - name: router + command: [/bin/true] + source: + git: https://example.com/repo + rev: abc123 + build_command: [/bin/false] +""" + ) + orchestrator = _orchestrator(config, tmp_path) + with patch(SRUN, return_value=_proc(1)), pytest.raises(RuntimeError, match="source clone failed"): + orchestrator.start_services("after_frontend") + with ( + patch(SRUN, side_effect=[_proc(0), _proc(2)]), + pytest.raises(RuntimeError, match="build_command failed"), + ): + orchestrator.start_services("after_frontend") + + +MOONCAKE_BACKEND = """backend: + type: sglang + mooncake_kv_store: + container: /mooncake-master.sqsh + sglang_config: + prefill: + disaggregation-transfer-backend: mooncake + decode: + disaggregation-transfer-backend: mooncake +""" + +STORES = """ +services: + - name: store-prefill + type: mooncake-store + placement: + node: prefill + args: [--port, "8800", --label, "{role}-{node_id}"] + env: + MOONCAKE_PROTOCOL: rdma + MOONCAKE_MASTER: ignored:9999 + MOONCAKE_EXTRA_CONFIG: '{"prefetch_timeout_base": 4}' + MOONCAKE_GLOBAL_SEGMENT_SIZE: 100gb + preamble: | + ulimit -n 1048576 + echo starting-{role}-on-{node} + cpus_per_task: 8 + cpu_bind: none + srun_options: + exclusive: "" + readiness: + port: 8800 + timeout_seconds: 90 + - name: store-decode + type: mooncake-store + placement: + node: decode + container: /mooncake-store.sqsh + args: [--port, "8800"] + env: + MOONCAKE_GLOBAL_SEGMENT_SIZE: 400gb + readiness: + port: 8800 +""" + + +def test_mooncake_stores_launch_once_per_role_node_with_master_env(tmp_path: Path) -> None: + orchestrator = _orchestrator(_load(STORES, backend=MOONCAKE_BACKEND), tmp_path) + ips = {"node1": "10.0.0.11", "node2": "10.0.0.12", "node3": "10.0.0.13"} + with ( + patch(SRUN, side_effect=lambda **_: _proc()) as srun, + patch(HOST_IP, side_effect=lambda node, _iface: ips[node]), + patch(WAIT, return_value=True) as wait, + ): + procs = orchestrator.start_services("before_workers") + + # 1 prefill node + 2 decode nodes, no launches for the head. + assert [p.node for p in procs] == ["node1", "node2", "node3"] + assert [p.name for p in procs] == [ + "service_store-prefill", + "service_store-decode_node2", + "service_store-decode_node3", + ] + assert all(p.critical for p in procs) + assert wait.call_count == 3 + + prefill = srun.call_args_list[0].kwargs + assert prefill["container_image"] == "/mooncake-master.sqsh" # falls back to mooncake_kv_store.container + assert prefill["command"] == [ + "python", + "-m", + "mooncake.mooncake_store_service", + "--port", + "8800", + "--label", + "prefill-0", + ] + env = prefill["env_to_set"] + assert env["MOONCAKE_LOCAL_HOSTNAME"] == "10.0.0.11" + assert env["MOONCAKE_GLOBAL_SEGMENT_SIZE"] == "100gb" + assert env["MOONCAKE_EXTRA_CONFIG"] == '{"prefetch_timeout_base": 4}' + assert env["MOONCAKE_MASTER"] == f"10.0.0.10:{MOONCAKE_MASTER_PORT}" # srtctl always wins + assert env["MOONCAKE_TE_META_DATA_SERVER"] == f"http://10.0.0.10:{MOONCAKE_HTTP_METADATA_PORT}/metadata" + assert prefill["bash_preamble"] == "ulimit -n 1048576\necho starting-prefill-on-node1" + assert prefill["cpus_per_task"] == 8 + assert prefill["cpu_bind"] == "none" + assert prefill["srun_options"] == {"exclusive": ""} + + decode = srun.call_args_list[1].kwargs + assert decode["container_image"] == "/mooncake-store.sqsh" + assert decode["env_to_set"]["MOONCAKE_GLOBAL_SEGMENT_SIZE"] == "400gb" + assert decode["env_to_set"]["MOONCAKE_LOCAL_HOSTNAME"] == "10.0.0.12" + + +def test_colocated_roles_with_same_port_rejected_before_launch(tmp_path: Path) -> None: + orchestrator = _orchestrator(_load(STORES, backend=MOONCAKE_BACKEND), tmp_path) + orchestrator.__dict__["endpoints"] = [ + Endpoint(mode="prefill", index=0, nodes=("node1",)), + Endpoint(mode="decode", index=0, nodes=("node1",)), + ] + with patch(SRUN) as srun, pytest.raises(ValueError, match="both listen on port 8800 on node node1"): + orchestrator.start_services("before_workers") + srun.assert_not_called() + + +def test_workers_placement_deduplicates_shared_nodes(tmp_path: Path) -> None: + config = _load( + "services:\n - name: store\n type: mooncake-store\n placement:\n node: workers\n" + " readiness:\n port: 8800\n", + backend=MOONCAKE_BACKEND, + ) + orchestrator = _orchestrator(config, tmp_path) + with ( + patch(SRUN, side_effect=lambda **_: _proc()) as srun, + patch(HOST_IP, return_value="10.0.0.11"), + patch(WAIT, return_value=True), + ): + procs = orchestrator.start_services("before_workers") + assert srun.call_count == 3 + assert {p.node for p in procs} == {"node1", "node2", "node3"} + + +def test_service_config_direct_construction() -> None: + svc = ServiceConfig(name="x", command=["true"], start="before_workers") + assert svc.effective_start == "before_workers" + with pytest.raises(ValidationError, match="must not contain empty arguments"): + ServiceConfig(name="x", command=["python", ""]) From b1e73b067f87fdd7263f9c3c05268695fbdd88f5 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sun, 6 Sep 2026 11:39:18 -0700 Subject: [PATCH 2/3] fix(sglang): kv_events_config: true also covers aggregated workers The global-bool shortcut only matched prefill and decode, so an aggregated topology never got --kv-events-config and the sidecar's cache-overlap score stayed at zero. Carried over from #374 with its regression test. --- src/srtctl/backends/sglang.py | 7 +++++-- tests/test_configs.py | 7 +++++-- tests/test_sidecar_backends.py | 21 +++++++++++++++++++++ 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/src/srtctl/backends/sglang.py b/src/srtctl/backends/sglang.py index ef62c7e1f..702812879 100644 --- a/src/srtctl/backends/sglang.py +++ b/src/srtctl/backends/sglang.py @@ -222,9 +222,12 @@ def get_kv_events_config_for_mode(self, mode: WorkerMode) -> dict[str, str] | No if not self.kv_events_config: return None - # Global bool: enable for prefill+decode with defaults + # Global bool: enable for every worker mode with defaults. Aggregated + # workers publish too; without this, `kv_events_config: true` on an agg + # topology silently dropped --kv-events-config and the router's cache + # overlap stayed at zero. if self.kv_events_config is True: - if mode in ("prefill", "decode"): + if mode in ("prefill", "decode", "agg"): return {"publisher": "zmq", "topic": "kv-events"} return None diff --git a/tests/test_configs.py b/tests/test_configs.py index 9f4394753..34a1d7b82 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -538,7 +538,7 @@ def test_get_environment_for_mode(self): assert config.get_environment_for_mode("agg") == {} def test_kv_events_config_global_bool(self): - """Test kv_events_config=True enables prefill+decode with defaults.""" + """Test kv_events_config=True enables prefill+decode+aggregated with defaults.""" config = SGLangProtocol(kv_events_config=True) assert config.get_kv_events_config_for_mode("prefill") == { @@ -549,7 +549,10 @@ def test_kv_events_config_global_bool(self): "publisher": "zmq", "topic": "kv-events", } - assert config.get_kv_events_config_for_mode("agg") is None + assert config.get_kv_events_config_for_mode("agg") == { + "publisher": "zmq", + "topic": "kv-events", + } def test_kv_events_config_per_mode(self): """Test kv_events_config per-mode control.""" diff --git a/tests/test_sidecar_backends.py b/tests/test_sidecar_backends.py index 55328aba5..be21735fd 100644 --- a/tests/test_sidecar_backends.py +++ b/tests/test_sidecar_backends.py @@ -73,6 +73,27 @@ def test_sglang_sidecar_owns_leader_and_couples_lifecycle() -> None: assert "dynamo.sglang.sidecar" not in follower_command +def test_sglang_sidecar_kv_events_config_true_covers_aggregated_mode() -> None: + # Regression: the kv_events_config=True shortcut only matched prefill/decode, so an + # aggregated topology never got --kv-events-config and the sidecar's + # kv_event_sources stayed at 0 (every routed request scored 0.00 cache overlap). + process = _process(mode="agg", kv_events_port=5557) + backend = SGLangProtocol( + kv_events_config=True, + sglang_config=SGLangServerConfig(aggregated={"tensor-parallel-size": 8}), + ) + + with patch("srtctl.core.slurm.get_hostname_ip", return_value="10.0.0.1"): + command = backend.build_worker_command(process, [process], _runtime()) + + leader_script = command[2] + assert "--kv-events-config" in leader_script + after_flag = leader_script.split("--kv-events-config ", 1)[1] + kv_config = json.loads(after_flag.split("'", 2)[1]) + assert kv_config["endpoint"] == "tcp://*:5557" + assert kv_config["publisher"] == "zmq" + + def test_vllm_sidecar_exposes_one_complete_multi_node_dp_group() -> None: backend = VLLMProtocol( connector=None, From 2c868fa45ffe8290c6cbd1684906fc0b29827f67 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Sun, 6 Sep 2026 11:59:17 -0700 Subject: [PATCH 3/3] services: register every srun at launch, bound clone/build, fail fast on death Close the windows where a service process could run untracked: - start_services takes the ProcessRegistry and registers each service process the moment its srun exists, not after the whole stage returns. A SIGTERM during a readiness wait now finds every launched process in the registry; the stage also terminates what it started on any BaseException before re-raising. - The one-shot clone and build sruns are ManagedProcesses in the registry too (non-critical), and wait under a wall-clock bound: the clone by the sum of its per-command git timeouts, the build by a new build_timeout_seconds (default 1800). A step that overruns is killed and the job fails pointing at its log instead of holding the allocation. - The readiness gate polls in short slices and checks the process between them, so a service that dies before opening its port fails at once with its exit code rather than after the full timeout. Documented under a Cleanup section in docs/services.md. --- docs/config-reference.md | 1 + docs/schema-reference.md | 1 + docs/services.md | 20 ++++ src/srtctl/cli/do_sweep.py | 9 +- src/srtctl/cli/mixins/service_stage.py | 143 ++++++++++++++++++------- src/srtctl/services/config.py | 6 ++ tests/test_services.py | 73 ++++++++++++- 7 files changed, 208 insertions(+), 45 deletions(-) diff --git a/docs/config-reference.md b/docs/config-reference.md index 70c2e0ced..11f7dc570 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -1797,6 +1797,7 @@ services: | `inherit_discovery_env` | bool | `true` | Inject the Dynamo discovery env | | `critical` | bool | type default | A crash fails the run when true | | `source`, `build_command` | object, list[string] | none | Clone an immutable git rev and build once before launch; single-node placements only | +| `build_timeout_seconds` | int | `1800` | `build_command` is killed when this runs out so a hung build cannot hold the allocation | | `preamble`, `cpus_per_task`, `cpu_bind`, `srun_options` | | none | Pass-through launch knobs for this service | --- diff --git a/docs/schema-reference.md b/docs/schema-reference.md index cf9af703d..8b17d33f8 100644 --- a/docs/schema-reference.md +++ b/docs/schema-reference.md @@ -285,6 +285,7 @@ One entry of the top-level ``services:`` list. | `cpus_per_task` | int \| None | `None` | Optional ``srun --cpus-per-task``. | | `cpu_bind` | str \| None | `None` | Optional ``srun --cpu-bind``. | | `srun_options` | dict[str, str] | `{}` | Extra srun options for this service only. | +| `build_timeout_seconds` | int | `1800` | Kill ``build_command`` after this many seconds. | ### IdentityConfig diff --git a/docs/services.md b/docs/services.md index 8cc01f683..ac9558888 100644 --- a/docs/services.md +++ b/docs/services.md @@ -95,6 +95,7 @@ services: | `preamble` | none | Shell run after the environment is exported and before `command`. | | `cpus_per_task`, `cpu_bind`, `srun_options` | none | Pass-through srun knobs for this service's launches. | | `source`, `build_command` | none | See [Building From Source](#building-from-source). | +| `build_timeout_seconds` | `1800` | `build_command` is killed when this runs out, so a hung build cannot hold the allocation. | `command`, `args`, `env` values, and `preamble` may use these placeholders: `{node}`, `{node_ip}`, `{node_id}` (position in the worker list), `{index}` (instance index within the service), `{role}` @@ -276,6 +277,25 @@ Rejected at launch, before any service starts: two services listening on the sam `srtctl dry-run` prints every service's type, placement, start phase, criticality, command, container, source, readiness, and env. +## Cleanup + +Nothing a service launches outlives the job: + +- Every `srun` the stage starts, including the one-shot clone and build steps, is registered with the + job's `ProcessRegistry` the moment it exists, not when the stage returns. The registry's cleanup + runs on normal completion, on any failed stage, from the SIGTERM handler (`scancel`), and from the + crash monitor when a critical process dies, and it terminates then kills each tracked `srun`. Slurm + cancels the step, which kills the whole step cgroup inside the container, so forked or daemonized + children of the service go with it. +- A readiness gate that fails, or a signal that arrives during one, terminates everything the stage + already launched before the error propagates. +- A service whose process exits before its readiness port answers fails immediately with its exit + code instead of waiting out the readiness timeout. +- The clone step is bounded by the `timeout` on each git command (600s each) and the build step by + `build_timeout_seconds`; a step that overruns is killed and the job fails with a pointer to its log. +- When the batch script exits, Slurm releases the allocation and reaps any remaining step, so even a + cleanup path srtctl never reaches cannot leave a service running on a compute node. + ## Limitations - Declared order is launch order, and `readiness` is the only wait. A service that needs another diff --git a/src/srtctl/cli/do_sweep.py b/src/srtctl/cli/do_sweep.py index c3557ddc6..a2a0ba472 100644 --- a/src/srtctl/cli/do_sweep.py +++ b/src/srtctl/cli/do_sweep.py @@ -781,9 +781,9 @@ def run(self) -> int: registry.add_process(mooncake_proc) # Stage 1c: services that workers depend on (standalone Mooncake - # stores, anything with start: before_workers). See docs/services.md. - for proc in self.start_services("before_workers"): - registry.add_process(proc) + # stores, anything with start: before_workers). The stage registers + # each process as it launches. See docs/services.md. + self.start_services("before_workers", registry) # Pre-worker: Ensure HF model is cached before starting workers. # 1. Clean stale lock files from previous crashed downloads @@ -810,8 +810,7 @@ def run(self) -> int: # Stage 3b: sidecar services (start: after_frontend, the default), # once workers and the frontend are healthy and before telemetry. - for proc in self.start_services("after_frontend"): - registry.add_process(proc) + self.start_services("after_frontend", registry) if self.config.telemetry.enabled: if os.environ.get("EVAL_ONLY", "false").lower() == "true": diff --git a/src/srtctl/cli/mixins/service_stage.py b/src/srtctl/cli/mixins/service_stage.py index 2668162c4..fbe5fc4f8 100644 --- a/src/srtctl/cli/mixins/service_stage.py +++ b/src/srtctl/cli/mixins/service_stage.py @@ -10,6 +10,13 @@ ``ProcessRegistry`` (which provides crash detection and teardown). The kind (``srtctl.services.registry.ServiceKind``) never launches anything itself. +Nothing a service launches may outlive the job. Every srun this stage starts, +including the one-shot clone and build steps, is registered with the +``ProcessRegistry`` the moment it exists, so ``registry.cleanup()`` (normal +exit, a failed stage, the SIGTERM handler, the crash monitor) reaches it +without depending on this stage returning. The clone and build steps also +run under a wall-clock timeout so a hung build cannot hold the allocation. + ``start_services("before_workers")`` runs after the Mooncake master and before workers; ``start_services("after_frontend")`` runs once workers and the frontend are healthy. See ``docs/services.md``. @@ -19,11 +26,12 @@ import logging import shlex +import subprocess from pathlib import Path from typing import TYPE_CHECKING from srtctl.core.health import wait_for_port -from srtctl.core.processes import ManagedProcess +from srtctl.core.processes import ManagedProcess, ProcessRegistry, terminate_and_reap from srtctl.core.slurm import get_hostname_ip, start_srun_process from srtctl.ports import ETCD_CLIENT_PORT, NATS_PORT from srtctl.services.registry import ServiceLaunchContext, get_service_kind @@ -32,10 +40,16 @@ from srtctl.core.runtime import RuntimeContext from srtctl.core.schema import SrtConfig from srtctl.core.topology import Endpoint - from srtctl.services.config import ServiceConfig + from srtctl.services.config import ServiceConfig, ServiceReadinessConfig logger = logging.getLogger(__name__) +# The clone script runs three git commands, each under its own `timeout 600s`. +CLONE_TIMEOUT_SECONDS = 3 * 600 + 60 +# Readiness polling slice: how long one wait_for_port call may block before we +# re-check that the service process is still alive. +_READINESS_SLICE_SECONDS = 5 + def render_placeholders(value: str, replacements: dict[str, str]) -> str: """Substitute known ``{placeholder}`` names only, leaving unrelated braces (JSON) untouched.""" @@ -91,7 +105,27 @@ def _check_port_collisions(self, services: list[ServiceConfig]) -> None: f"{service.readiness.port} on node {node}; give them disjoint placements or ports" ) - # -- source build ------------------------------------------------------------ + # -- one-shot steps (clone, build) ---------------------------------------------- + + @staticmethod + def _run_step( + step: ManagedProcess, *, timeout: float, registry: ProcessRegistry | None, what: str, log: Path + ) -> None: + """Wait for a one-shot srun, tracked and bounded. + + Registered before waiting so a signal or a crash elsewhere tears it down + with everything else; killed on timeout so a hung step cannot hold the + allocation until walltime. + """ + if registry is not None: + registry.add_process(step) + try: + returncode = step.popen.wait(timeout=timeout) + except subprocess.TimeoutExpired: + terminate_and_reap(step.popen) + raise RuntimeError(f"{what} timed out after {int(timeout)}s and was killed; see {log}") from None + if returncode != 0: + raise RuntimeError(f"{what} failed (exit {returncode}); see {log}") def _service_container(self, service: ServiceConfig) -> str: kind = get_service_kind(service.type) @@ -101,7 +135,7 @@ def _container_path(self, host_path: Path) -> str: """Host path under ``log_dir`` as seen inside a container (``log_dir`` is mounted at ``/logs``).""" return str(Path("/logs") / host_path.relative_to(self.runtime.log_dir)) - def _clone_service_source(self, service: ServiceConfig, node: str) -> Path | None: + def _clone_service_source(self, service: ServiceConfig, node: str, registry: ProcessRegistry | None) -> Path | None: """Clone ``service.source`` once on the bare host of ``node``; returns the work dir (host path).""" source = service.source if source is None: @@ -121,25 +155,33 @@ def _clone_service_source(self, service: ServiceConfig, node: str) -> Path | Non "fi" ) logger.info("Cloning service %s source %s@%s on %s", service.name, source.git, source.rev, node) - proc = start_srun_process( + popen = start_srun_process( command=["bash", "-c", clone_script], nodelist=[node], output=str(clone_log), container_image=None, # bare host: git and network access are host concerns het_group=self.runtime.nodes.het_group_for(node), ) - if proc.wait() != 0: - raise RuntimeError( - f"services[{service.name}] source clone failed (exit {proc.returncode}); see {clone_log}" - ) + step = ManagedProcess( + name=f"service_{service.name}.clone", popen=popen, log_file=clone_log, node=node, critical=False + ) + self._run_step( + step, + timeout=CLONE_TIMEOUT_SECONDS, + registry=registry, + what=f"services[{service.name}] source clone", + log=clone_log, + ) return checkout_root / source.path if source.path else checkout_root - def _build_service_source(self, service: ServiceConfig, node: str, work_dir: Path) -> None: + def _build_service_source( + self, service: ServiceConfig, node: str, work_dir: Path, registry: ProcessRegistry | None + ) -> None: if not service.build_command: return build_log = self.runtime.log_dir / f"service_{service.name}.build.out" logger.info("Building service %s: %s", service.name, shlex.join(service.build_command)) - proc = start_srun_process( + popen = start_srun_process( command=list(service.build_command), nodelist=[node], output=str(build_log), @@ -149,10 +191,16 @@ def _build_service_source(self, service: ServiceConfig, node: str, work_dir: Pat het_group=self.runtime.nodes.het_group_for(node), bash_preamble=_await_and_cd(self._container_path(work_dir)), ) - if proc.wait() != 0: - raise RuntimeError( - f"services[{service.name}] build_command failed (exit {proc.returncode}); see {build_log}" - ) + step = ManagedProcess( + name=f"service_{service.name}.build", popen=popen, log_file=build_log, node=node, critical=False + ) + self._run_step( + step, + timeout=service.build_timeout_seconds, + registry=registry, + what=f"services[{service.name}] build_command", + log=build_log, + ) # -- launch ------------------------------------------------------------------ @@ -203,12 +251,40 @@ def _launch_service_instance( critical=service.effective_critical, ) - def start_services(self, start: str) -> list[ManagedProcess]: + @staticmethod + def _wait_ready(proc: ManagedProcess, service: ServiceConfig, readiness: ServiceReadinessConfig) -> None: + """Block until the service's port answers, failing fast if the process dies first.""" + assert proc.node is not None + logger.info( + "Waiting for service %s on %s (port %d, timeout %ds)", + service.name, + proc.node, + readiness.port, + readiness.timeout_seconds, + ) + waited = 0 + while waited < readiness.timeout_seconds: + step = min(_READINESS_SLICE_SECONDS, readiness.timeout_seconds - waited) + if wait_for_port(proc.node, readiness.port, timeout=step): + return + waited += step + if not proc.is_running: + raise RuntimeError( + f"services[{service.name}] exited with code {proc.exit_code} on {proc.node} before opening " + f"port {readiness.port}; see {proc.log_file}" + ) + raise RuntimeError( + f"services[{service.name}] did not open port {readiness.port} on {proc.node} within " + f"{readiness.timeout_seconds}s; see {proc.log_file}" + ) + + def start_services(self, start: str, registry: ProcessRegistry | None = None) -> list[ManagedProcess]: """Launch every service whose (effective) ``start`` matches, in declaration order. - Readiness gates block per node; a gate that times out terminates every - process this call started and raises. Returned processes are for the - caller to register with the shared ``ProcessRegistry``. + Each process is added to ``registry`` as soon as its srun exists, so a + readiness wait interrupted by a signal still leaves nothing untracked. + A readiness gate that fails terminates every process this call started + and raises. The started processes are also returned. """ services = [s for s in self.config.services if s.effective_start == start] if not services: @@ -227,9 +303,9 @@ def start_services(self, start: str) -> list[ManagedProcess]: service.placement.node, ) continue - work_dir = self._clone_service_source(service, nodes[0]) + work_dir = self._clone_service_source(service, nodes[0], registry) if work_dir is not None: - self._build_service_source(service, nodes[0], work_dir) + self._build_service_source(service, nodes[0], work_dir, registry) for index, node in enumerate(nodes): ctx = ServiceLaunchContext( @@ -240,23 +316,16 @@ def start_services(self, start: str) -> list[ManagedProcess]: index=index, role=service.placement.node, ) - started.append(self._launch_service_instance(service, ctx, work_dir, len(nodes))) - readiness = service.readiness - if readiness is not None: - logger.info( - "Waiting for service %s on %s (port %d, timeout %ds)", - service.name, - node, - readiness.port, - readiness.timeout_seconds, - ) - if not wait_for_port(node, readiness.port, timeout=readiness.timeout_seconds): - raise RuntimeError( - f"services[{service.name}] did not open port {readiness.port} on {node} within " - f"{readiness.timeout_seconds}s; see {started[-1].log_file}" - ) + proc = self._launch_service_instance(service, ctx, work_dir, len(nodes)) + started.append(proc) + if registry is not None: + registry.add_process(proc) + if service.readiness is not None: + self._wait_ready(proc, service, service.readiness) logger.info("Service %s ready on %d node(s)", service.name, len(nodes)) - except Exception: + except BaseException: + # Belt and braces: the registry already tracks these, but terminate + # here too so a failure inside this stage never depends on the caller. for proc in started: proc.terminate() raise diff --git a/src/srtctl/services/config.py b/src/srtctl/services/config.py index 4ea6e3605..a758b7899 100644 --- a/src/srtctl/services/config.py +++ b/src/srtctl/services/config.py @@ -145,6 +145,7 @@ class ServiceConfig: cpus_per_task: Optional ``srun --cpus-per-task``. cpu_bind: Optional ``srun --cpu-bind``. srun_options: Extra srun options for this service only. + build_timeout_seconds: Kill ``build_command`` after this many seconds. """ name: str @@ -164,6 +165,9 @@ class ServiceConfig: cpus_per_task: int | None = None cpu_bind: str | None = None srun_options: dict[str, str] = field(default_factory=dict) + # Wall-clock budget for build_command; the build srun is killed when it runs out + # so a hung build cannot hold the allocation until walltime. + build_timeout_seconds: int = 1800 # builtins.type: the ``type`` field above shadows the builtin inside the class body. Schema: ClassVar[builtins.type[Schema]] = Schema @@ -202,6 +206,8 @@ def __post_init__(self) -> None: raise ValidationError(f"{label}.start must be one of {', '.join(SERVICE_STARTS)}; got {self.start!r}") if self.cpus_per_task is not None and self.cpus_per_task <= 0: raise ValidationError(f"{label}.cpus_per_task must be positive") + if self.build_timeout_seconds <= 0: + raise ValidationError(f"{label}.build_timeout_seconds must be positive") # -- effective values (type defaults applied) ------------------------------ diff --git a/tests/test_services.py b/tests/test_services.py index 069790f2e..84457dc63 100644 --- a/tests/test_services.py +++ b/tests/test_services.py @@ -5,6 +5,7 @@ from __future__ import annotations +import subprocess from pathlib import Path from unittest.mock import MagicMock, patch @@ -239,7 +240,7 @@ def test_readiness_gate_blocks_and_failure_terminates_started(tmp_path: Path) -> ) orchestrator = _orchestrator(config, tmp_path) with ( - patch(SRUN, return_value=_proc()) as srun, + patch(SRUN, return_value=_proc()), patch(HOST_IP, return_value="10.0.0.10"), patch(WAIT, return_value=True) as wait, ): @@ -247,15 +248,50 @@ def test_readiness_gate_blocks_and_failure_terminates_started(tmp_path: Path) -> wait.assert_called_once_with("node0", 9000, timeout=5) popen = _proc() + registry = MagicMock() with ( patch(SRUN, return_value=popen), patch(HOST_IP, return_value="10.0.0.10"), patch(WAIT, return_value=False), pytest.raises(RuntimeError, match="did not open port 9000"), ): - orchestrator.start_services("after_frontend") + orchestrator.start_services("after_frontend", registry) + popen.terminate.assert_called_once() + # Registered before the readiness wait, so a signal during the wait still finds it. + (registered,) = [call.args[0] for call in registry.add_process.call_args_list] + assert registered.popen is popen + + +def test_readiness_fails_fast_when_the_process_dies(tmp_path: Path) -> None: + config = _load( + "services:\n - name: a\n command: [/bin/true]\n readiness:\n port: 9000\n timeout_seconds: 600\n" + ) + dead = _proc() + dead.poll.return_value = 127 + with ( + patch(SRUN, return_value=dead), + patch(HOST_IP, return_value="10.0.0.10"), + patch(WAIT, return_value=False) as wait, + pytest.raises(RuntimeError, match="exited with code 127 .* before opening port 9000"), + ): + _orchestrator(config, tmp_path).start_services("after_frontend") + # One 5s slice, not the full 600s budget. + wait.assert_called_once_with("node0", 9000, timeout=5) + + +def test_signal_during_readiness_wait_terminates_started(tmp_path: Path) -> None: + # The SIGTERM handler raises SystemExit inside whatever the orchestrator is doing; + # the stage must still tear down what it launched. + config = _load("services:\n - name: a\n command: [/bin/true]\n readiness:\n port: 9000\n") + popen = _proc() + with ( + patch(SRUN, return_value=popen), + patch(HOST_IP, return_value="10.0.0.10"), + patch(WAIT, side_effect=SystemExit(1)), + pytest.raises(SystemExit), + ): + _orchestrator(config, tmp_path).start_services("after_frontend") popen.terminate.assert_called_once() - assert srun is not None def test_source_is_cloned_on_bare_host_and_built_in_container(tmp_path: Path) -> None: @@ -309,6 +345,37 @@ def test_clone_and_build_failures_raise(tmp_path: Path) -> None: orchestrator.start_services("after_frontend") +def test_clone_and_build_steps_are_registered_and_bounded(tmp_path: Path) -> None: + config = _load( + """ +services: + - name: router + command: [/bin/true] + source: + git: https://example.com/repo + rev: abc123 + build_command: [make] + build_timeout_seconds: 7 +""" + ) + registry = MagicMock() + hung = _proc() + hung.wait.side_effect = subprocess.TimeoutExpired(cmd="make", timeout=7) + hung.poll.return_value = None + with ( + patch(SRUN, side_effect=[_proc(0), hung]), + patch("srtctl.cli.mixins.service_stage.terminate_and_reap") as reap, + pytest.raises(RuntimeError, match="build_command timed out after 7s"), + ): + _orchestrator(config, tmp_path).start_services("after_frontend", registry) + + hung.wait.assert_called_once_with(timeout=7) + reap.assert_called_once_with(hung) + names = [call.args[0].name for call in registry.add_process.call_args_list] + assert names == ["service_router.clone", "service_router.build"] + assert all(not call.args[0].critical for call in registry.add_process.call_args_list) + + MOONCAKE_BACKEND = """backend: type: sglang mooncake_kv_store: