From 334b52a30e13f12d3cfe0bf0b7b821d4267efb61 Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Thu, 13 Aug 2026 19:49:37 +0000 Subject: [PATCH 01/11] feat(lm): declare an SGLang endpoint per graph via a top-level `lm:` block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `stargraph run` now resolves an LM endpoint before the first node executes: it probes `/v1/models` and attaches when a server already serves the requested model (leaving it running — it is not ours), otherwise it spawns `python -m sglang.launch_server`, waits for the endpoint to answer, and terminates it at run end. It never imports sglang and never kills a server it did not start. The endpoint is expressible two ways: `--sglang-model/--sglang-port/ --sglang-arg` flags, or a top-level `lm:` block in the authoring format (lowered to the `SGLangServer` IR model). The block is not part of the structural graph hash — like `--lm-url`, an endpoint is an environment binding, not graph topology. A graph-declared block is partly untrusted: `_check_graph_declared` refuses YAML-supplied `args` (argv into a subprocess) and a non-loopback `host` (it would receive `--lm-key` and every prompt) unless the operator re-states them as flags. This is the only path by which a graph reaches `subprocess` — every process-spawning std tool sits behind the default-deny capability gate. Adds `LMServerError` for the attach/spawn failure modes: a port serving a different model, a launch that exits before answering, and a startup that exceeds `startup_timeout_s`. Signed-off-by: Sean Mauk --- docs/how-to/authoring-format.md | 41 +++ docs/reference/cli.md | 30 ++ docs/reference/ir-schema.md | 1 + docs/reference/openapi.json | 54 ++++ src/stargraph/authoring.py | 30 +- src/stargraph/cli/run.py | 189 +++++++++++- src/stargraph/errors/__init__.py | 2 + src/stargraph/errors/_hierarchy.py | 12 + src/stargraph/ir/__init__.py | 2 + src/stargraph/ir/_models.py | 46 +++ src/stargraph/lm/__init__.py | 12 + src/stargraph/lm/sglang.py | 232 +++++++++++++++ src/stargraph/schemas/ir-v1-draft7.json | 54 ++++ src/stargraph/schemas/ir-v1.json | 54 ++++ tests/integration/test_lm_sglang_spawn.py | 128 ++++++++ tests/unit/cli/test_run_sglang.py | 343 ++++++++++++++++++++++ tests/unit/test_authoring_compiler.py | 37 +++ tests/unit/test_errors_walker.py | 3 + tests/unit/test_lm_sglang.py | 128 ++++++++ 19 files changed, 1391 insertions(+), 7 deletions(-) create mode 100644 src/stargraph/lm/__init__.py create mode 100644 src/stargraph/lm/sglang.py create mode 100644 tests/integration/test_lm_sglang_spawn.py create mode 100644 tests/unit/cli/test_run_sglang.py create mode 100644 tests/unit/test_lm_sglang.py diff --git a/docs/how-to/authoring-format.md b/docs/how-to/authoring-format.md index 5ec485af..a0adbc6e 100644 --- a/docs/how-to/authoring-format.md +++ b/docs/how-to/authoring-format.md @@ -51,6 +51,44 @@ works: the [prebuilt kinds](../reference/nodes/prebuilt.md), `tool`, `NodeBase`. Bare tool ids get `@1` appended (`std.web_search` → `std.web_search@1`). +### `lm` (optional) + +Pin the model the graph's LLM nodes run against, so the graph carries its +own endpoint instead of depending on the caller's flags: + +```yaml +lm: + provider: sglang # the only provider today + model: microsoft/phi-4 # --model-path, and the id the server must report + port: 41002 # default 30000 + args: [--attention-backend, triton] # passed to sglang.launch_server verbatim + startup_timeout_s: 600 # weights take minutes on big models +``` + +`stargraph run` resolves it before the first node: it attaches to a server +already serving that model on the port (and leaves it running), otherwise +it launches `python -m sglang.launch_server`, waits for the endpoint to +answer, and terminates it at the end of the run. The derived base URL + +model configure the DSPy LM, so `--lm-url`/`--lm-model` become unnecessary +(and are rejected alongside it). `--sglang-*` flags override this block +field by field — `--sglang-port 41010` re-points it without editing the +graph. + +Two fields are **operator-only**, because a graph file can be less trusted +than the person running it — and this block is the only way a graph reaches +a subprocess at all (every process-spawning std tool sits behind the +default-deny capability gate): + +- `args` — passthrough argv into `sglang.launch_server`. A graph-declared + value is refused; re-state it as `--sglang-arg` to allow it. +- a non-loopback `host` — the derived endpoint receives `--lm-key` and every + prompt. Refused unless the operator passes the same `--sglang-host`. + +`model`, `port` and `startup_timeout_s` stay graph-declarable. + +The block is not part of the graph hash: an endpoint is an environment +binding, not topology. + ### `routes` Declaration order is the default flow: with no rule firing, execution @@ -111,6 +149,9 @@ overwritten. - Value routes branch on `verdict` only — standardize on it (both `classify` and `judge` already emit it). - One graph per file; `state` fields are flat primitives/containers. +- `lm:` is honoured by `stargraph run` only — `stargraph serve` binds its + LM from its own `--lm-*` flags at boot, one endpoint for every graph it + serves. - For anything the sugar can't say (custom fact templates, multi-field `when` conditions, verifiers), write IR — see [Build a graph](build-graph.md). diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 732a37b0..b68c7d76 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -77,10 +77,35 @@ on `failed`. | `--lm-model NAME` | str | _(none)_ | LLM model identifier (e.g. `gpt-oss:20b`). | | `--lm-key KEY` | str | `placeholder` | API key for the LLM endpoint (`placeholder` works for ollama). | | `--lm-timeout SEC` | int | `60` | LLM call timeout in seconds. | +| `--sglang-model NAME`| str | _(none)_ | Serve this model with SGLang for the run; sets `--lm-url`/`--lm-model` from it. | +| `--sglang-host HOST` | str | `127.0.0.1` | SGLang bind/probe host. | +| `--sglang-port PORT` | int | `30000` | SGLang port. | +| `--sglang-arg ARG` | str (repeatable) | _(empty)_ | Extra argv passed through to `sglang.launch_server` verbatim. | +| `--sglang-timeout SEC` | int | `600` | Seconds to wait for a launched SGLang server to answer. | `--quiet` and `--verbose` are mutually exclusive. `--lm-url` and `--lm-model` must be supplied together (or neither). +The `--sglang-*` flags bind the run to a local +[SGLang](https://docs.sglang.ai/) server, and derive `--lm-url` / +`--lm-model` from it — so they conflict with those two flags. Before the +first node runs, `stargraph run` probes `http://host:port/v1/models`: + +- a server already serving that model is **attached to** and left running; +- a server serving a *different* model is a loud error (pick another port); +- nothing listening means one is launched + (`python -m sglang.launch_server --model-path ...`, plus every + `--sglang-arg` verbatim), waited on until it answers, and terminated — + process group included — when the run ends. + +The same binding can be declared in the graph itself as an `lm:` block +(see [Author a graph in simple YAML](../how-to/authoring-format.md)); the +flags override it field by field. A graph-declared block may set +`model`/`port`/`startup_timeout_s` only — passthrough `args` and a +non-loopback `host` are operator-only and must be re-stated as +`--sglang-arg` / `--sglang-host`. Neither is part of the graph hash: like +`--lm-url`, an endpoint is an environment binding, not topology. + **Examples** ```bash @@ -92,6 +117,11 @@ stargraph run graphs/triage.yaml --inspect # Bind a local LLM for dspy nodes stargraph run graphs/triage.yaml --lm-url http://localhost:11434 --lm-model gpt-oss:20b + +# Boot SGLang for the run (attaches instead if :41002 already serves it) +stargraph run graphs/triage.yaml \ + --sglang-model microsoft/phi-4 --sglang-port 41002 \ + --sglang-arg=--attention-backend --sglang-arg=triton ``` See also: [Concepts: IR](../concepts/ir.md), diff --git a/docs/reference/ir-schema.md b/docs/reference/ir-schema.md index 5ba220b8..11527c04 100644 --- a/docs/reference/ir-schema.md +++ b/docs/reference/ir-schema.md @@ -61,6 +61,7 @@ and `nodes`; every other section defaults to an empty list / dict. | `parallel` | `list[ParallelBlock]` | no | `[]` | Top-level parallel/join declarations. | | `governance` | `list[PackMount]` | no | `[]` | Mounted Bosun rule packs. | | `migrate` | `list[MigrateBlock]` | no | `[]` | Hash-to-hash migration descriptors for resume. | +| `lm` | `SGLangServer \| None` | no | `None` | Local LM endpoint bound to the run (attach-or-launch SGLang). Environment binding, not topology — excluded from the structural hash. | ```yaml ir_version: "1.0.0" diff --git a/docs/reference/openapi.json b/docs/reference/openapi.json index b362b5fe..7d680370 100644 --- a/docs/reference/openapi.json +++ b/docs/reference/openapi.json @@ -1942,6 +1942,49 @@ "title": "RuleSpec", "type": "object" }, + "SGLangServer": { + "additionalProperties": false, + "description": "Declared LM endpoint: an SGLang server bound to this graph for a run.\n\n``stargraph run`` resolves this block before the first node executes: it\nprobes ``http://{host}:{port}/v1/models`` and **attaches** when a server\nalready serves ``model`` (left running afterwards -- it is not ours),\notherwise it spawns ``python -m sglang.launch_server``, waits for the\nendpoint to answer, and terminates it when the run ends. The derived\nbase URL + ``model`` configure the DSPy LM, so declaring this block is\nequivalent to passing ``--lm-url``/``--lm-model`` at a live endpoint.\n\nNot part of the structural graph hash: like ``--lm-url``, the endpoint is\nan environment binding, not graph topology.\n\nAttributes:\n provider: Only ``\"sglang\"`` today; the discriminator for future\n providers.\n model: ``--model-path`` value; must equal the id the server reports\n in ``/v1/models`` when attaching to an already-running server.\n host: Bind address / probe host.\n port: Bind port (SGLang's own default is 30000).\n args: Extra argv passed through to ``sglang.launch_server``\n verbatim (e.g. ``[\"--attention-backend\", \"triton\"]``).\n Operator-only: a graph-declared value is refused unless the\n operator re-states it as ``--sglang-arg`` (argv into a\n subprocess is a code-execution surface, and a graph file can\n be less trusted than the operator running it). The same holds\n for a non-loopback ``host``, which would receive the API key\n and every prompt.\n startup_timeout_s: How long to wait for a spawned server to answer\n before failing the run. Big models take minutes to load.", + "properties": { + "args": { + "items": { + "type": "string" + }, + "title": "Args", + "type": "array" + }, + "host": { + "default": "127.0.0.1", + "title": "Host", + "type": "string" + }, + "model": { + "title": "Model", + "type": "string" + }, + "port": { + "default": 30000, + "title": "Port", + "type": "integer" + }, + "provider": { + "const": "sglang", + "default": "sglang", + "title": "Provider", + "type": "string" + }, + "startup_timeout_s": { + "default": 600, + "title": "Startup Timeout S", + "type": "integer" + } + }, + "required": [ + "model" + ], + "title": "SGLangServer", + "type": "object" + }, "SkillRef": { "additionalProperties": false, "description": "Reference to a skill by namespaced id.", @@ -2046,6 +2089,17 @@ "title": "Ir Version", "type": "string" }, + "lm": { + "anyOf": [ + { + "$ref": "#/$defs/SGLangServer" + }, + { + "type": "null" + } + ], + "default": null + }, "migrate": { "items": { "$ref": "#/$defs/MigrateBlock" diff --git a/src/stargraph/authoring.py b/src/stargraph/authoring.py index 211ff123..cdee3100 100644 --- a/src/stargraph/authoring.py +++ b/src/stargraph/authoring.py @@ -46,12 +46,13 @@ import fathom from pydantic import BaseModel, create_model +from pydantic import ValidationError as PydanticValidationError from stargraph.checkpoint.sqlite import SQLiteCheckpointer from stargraph.errors import CheckpointError, IRValidationError from stargraph.fathom import FathomAdapter from stargraph.graph import Graph, GraphRun -from stargraph.ir import IRDocument, NodeSpec, RuleSpec +from stargraph.ir import IRDocument, NodeSpec, RuleSpec, SGLangServer from stargraph.ir._mirror import Mirror from stargraph.ir._models import GotoAction, HaltAction from stargraph.ir._when import compile_when as _compile_when @@ -502,7 +503,7 @@ async def _consume() -> None: "dict": (dict, {}), } -_AUTHORING_KEYS = frozenset({"id", "state", "nodes", "routes"}) +_AUTHORING_KEYS = frozenset({"id", "state", "nodes", "routes", "lm"}) _NAME_RE = _re.compile(r"^[a-z0-9][a-z0-9_.\-]*$") @@ -601,6 +602,28 @@ def _compile_node(name: str, raw: Any) -> NodeSpec: ) +def _compile_lm(raw: Any) -> SGLangServer | None: + """Lower the optional top-level ``lm:`` block to an IR endpoint spec.""" + if raw is None: + return None + if not isinstance(raw, dict): + raise _authoring_error( + f"`lm` must be a mapping, got {type(raw).__name__}", + hint="keys: provider (sglang), model, host, port, args, startup_timeout_s", + ) + try: + return SGLangServer.model_validate(raw) + except PydanticValidationError as exc: + problems = "; ".join( + f"{'.'.join(str(loc) for loc in err['loc']) or 'lm'}: {err['msg']}" + for err in exc.errors() + ) + raise _authoring_error( + f"`lm` block is invalid -- {problems}", + hint="keys: provider (sglang), model, host, port, args, startup_timeout_s", + ) from exc + + def _compile_routes( routes: dict[str, Any], node_names: list[str], @@ -680,6 +703,8 @@ def compile_authoring(doc: dict[str, Any], *, default_id: str = "authored") -> I if not isinstance(routes_raw, dict): raise _authoring_error("`routes` must be a mapping of node -> target") + lm = _compile_lm(doc.get("lm")) + node_names = [str(name) for name in nodes_map] nodes = [_compile_node(str(name), raw) for name, raw in nodes_map.items()] rules = _compile_routes(cast("dict[str, Any]", routes_raw), node_names, routed_fields) @@ -690,6 +715,7 @@ def compile_authoring(doc: dict[str, Any], *, default_id: str = "authored") -> I state_class=state_ref, nodes=nodes, rules=rules, + lm=lm, ) diff --git a/src/stargraph/cli/run.py b/src/stargraph/cli/run.py index 083ab6c4..09044864 100644 --- a/src/stargraph/cli/run.py +++ b/src/stargraph/cli/run.py @@ -30,6 +30,7 @@ import asyncio import contextlib +import ipaddress from pathlib import Path from typing import TYPE_CHECKING, Annotated, Any @@ -54,7 +55,7 @@ from stargraph.errors import StargraphError from stargraph.fathom import build_ir_routing from stargraph.graph import Graph, GraphRun -from stargraph.ir import IRDocument +from stargraph.ir import IRDocument, SGLangServer from stargraph.ir._ids import new_run_id # Node-kind resolution moved to stargraph.nodes.registry (the CLI is one @@ -69,6 +70,8 @@ ) if TYPE_CHECKING: + from collections.abc import Callable, Generator + from stargraph.checkpoint.protocol import RunSummary __all__ = ["cmd", "node_kinds"] @@ -76,6 +79,12 @@ _build_node_registry = build_node_registry +def _check_lm_pairing(lm_url: str | None, lm_model: str | None) -> None: + """Fail loud when exactly one of --lm-url / --lm-model is set.""" + if (lm_url is None) != (lm_model is None): + raise typer.BadParameter("--lm-url and --lm-model must be specified together (or neither)") + + def _configure_lm( lm_url: str | None, lm_model: str | None, @@ -88,8 +97,7 @@ def _configure_lm( call entirely when both are None lets graphs without DSPy nodes run without dragging in dspy at all. """ - if (lm_url is None) != (lm_model is None): - raise typer.BadParameter("--lm-url and --lm-model must be specified together (or neither)") + _check_lm_pairing(lm_url, lm_model) if lm_url is None: return import dspy # pyright: ignore[reportMissingTypeStubs] @@ -104,6 +112,117 @@ def _configure_lm( ) +def _is_loopback(host: str) -> bool: + """True for hosts that cannot leave the box.""" + if host == "localhost": + return True + try: + return ipaddress.ip_address(host).is_loopback + except ValueError: + return False + + +def _check_graph_declared( + declared: SGLangServer, *, args_from_flags: bool, host_from_flags: str | None +) -> None: + """Refuse the parts of a graph-declared ``lm:`` block that escape the box. + + A graph file can be less trusted than the operator running it, and this + is the only path by which a graph reaches ``subprocess`` at all (every + process-spawning std tool is behind the default-deny capability gate). + Two fields are therefore operator-only, and must be re-stated as flags + (which are trusted) before they take effect: + + * ``args`` -- passthrough argv into ``sglang.launch_server``; a graph + that could set it could turn a run into arbitrary code execution + (``--trust-remote-code`` executes the model repo's own Python). + * a non-loopback ``host`` -- the derived base URL receives ``--lm-key`` + and every prompt, so a graph-chosen host is an exfiltration channel + (and binding a launched server off-loopback publishes it). + + ``model``/``port``/``startup_timeout_s`` stay graph-declarable: they + bound the run to a specific local endpoint without handing the graph + argv or a network destination. + """ + if declared.args and not args_from_flags: + raise typer.BadParameter( + f"graph declares lm.args {declared.args} -- passthrough argv is " + "operator-only; re-state them as --sglang-arg to allow them" + ) + if host_from_flags is None and not _is_loopback(declared.host): + raise typer.BadParameter( + f"graph declares a non-loopback lm.host {declared.host!r} -- the " + "endpoint receives --lm-key and every prompt; pass " + f"--sglang-host {declared.host} to allow it" + ) + + +def _resolve_sglang( + ir: IRDocument, + *, + model: str | None, + host: str | None, + port: int | None, + args: list[str] | None, + timeout: int | None, +) -> SGLangServer | None: + """Merge the graph's ``lm:`` block with the ``--sglang-*`` flags. + + The flags override field-by-field, so ``--sglang-port`` alone re-points a + graph-declared server while ``--sglang-model`` alone declares one for a + graph that has no ``lm:`` block. Returns ``None`` when neither source + asks for a server. + """ + overrides: dict[str, object] = {} + if model is not None: + overrides["model"] = model + if host is not None: + overrides["host"] = host + if port is not None: + overrides["port"] = port + if args: + overrides["args"] = list(args) + if timeout is not None: + overrides["startup_timeout_s"] = timeout + + if ir.lm is not None: + _check_graph_declared(ir.lm, args_from_flags=args is not None, host_from_flags=host) + return ir.lm.model_copy(update=overrides) + if "model" in overrides: + return SGLangServer.model_validate(overrides) + if overrides: + raise typer.BadParameter( + "--sglang-* flags need --sglang-model (or an `lm:` block in the graph)" + ) + return None + + +@contextlib.contextmanager +def _lm_endpoint( + spec: SGLangServer | None, + lm_url: str | None, + lm_model: str | None, + lm_key: str, + lm_timeout: int, + echo: Callable[[str], None], +) -> Generator[None]: + """Hold the run's LM endpoint open: boot/attach sglang, then configure dspy. + + A declared ``spec`` supplies both halves of the DSPy binding (base URL + + model), so ``--lm-url``/``--lm-model`` are rejected alongside it by the + caller. With no spec this is exactly the pre-existing ``_configure_lm`` + behaviour. + """ + with contextlib.ExitStack() as stack: + if spec is not None: + from stargraph.lm.sglang import sglang_server + + lm_url = stack.enter_context(sglang_server(spec, echo=echo)) + lm_model = spec.model + _configure_lm(lm_url, lm_model, lm_key, lm_timeout) + yield + + def _build_audit_sink(log_file: Path) -> AuditSink: """Choose chained vs legacy sink for ``log_file`` (chain-write, dual-read). @@ -273,6 +392,43 @@ def cmd( help="LLM call timeout in seconds.", ), ] = 60, + sglang_model: Annotated[ + str | None, + typer.Option( + "--sglang-model", + help=( + "Serve this model with SGLang for the run: attaches to an " + "already-running server on the port, else launches one and " + "shuts it down at the end. Sets --lm-url/--lm-model for you." + ), + ), + ] = None, + sglang_host: Annotated[ + str | None, + typer.Option("--sglang-host", help="SGLang bind/probe host (default: 127.0.0.1)."), + ] = None, + sglang_port: Annotated[ + int | None, + typer.Option("--sglang-port", help="SGLang port (default: 30000)."), + ] = None, + sglang_arg: Annotated[ + list[str] | None, + typer.Option( + "--sglang-arg", + help=( + "Extra argument passed through to sglang.launch_server " + "(repeatable, e.g. --sglang-arg=--attention-backend " + "--sglang-arg=triton)." + ), + ), + ] = None, + sglang_timeout: Annotated[ + int | None, + typer.Option( + "--sglang-timeout", + help="Seconds to wait for a launched SGLang server to answer (default: 600).", + ), + ] = None, ) -> None: """Run a Stargraph graph end-to-end (FR-8 POC). @@ -287,7 +443,9 @@ def cmd( if quiet and verbose: raise typer.BadParameter("--quiet and --verbose are mutually exclusive") - _configure_lm(lm_url, lm_model, lm_key, lm_timeout) + # The LM is configured later (inside ``_lm_endpoint``, once a declared + # sglang server is up); flag pairing is still checked before any work. + _check_lm_pairing(lm_url, lm_model) ir_dict = yaml.safe_load(graph.read_text(encoding="utf-8")) if is_authoring_format(ir_dict): @@ -297,6 +455,20 @@ def cmd( else: ir = IRDocument.model_validate(ir_dict) + sglang_spec = _resolve_sglang( + ir, + model=sglang_model, + host=sglang_host, + port=sglang_port, + args=sglang_arg, + timeout=sglang_timeout, + ) + if sglang_spec is not None and (lm_url is not None or lm_model is not None): + raise typer.BadParameter( + "--lm-url/--lm-model conflict with an sglang endpoint -- the URL and " + "model are derived from it" + ) + # Builtin-seeded ToolRegistry on the Graph: the ``kind: tool`` # ToolCallNode resolves its tool id here at execute time. from stargraph.registry.tools import ToolRegistry @@ -373,8 +545,15 @@ async def _bootstrap_and_drive() -> RunSummary: if audit_sink is not None: await audit_sink.close() + def _echo(message: str) -> None: + if not quiet: + console.print(message, style="dim", highlight=False, markup=False) + try: - summary = asyncio.run(_bootstrap_and_drive()) + # The endpoint outlives the whole run: a spawned sglang server is torn + # down only once the loop is done (or has raised). + with _lm_endpoint(sglang_spec, lm_url, lm_model, lm_key, lm_timeout, _echo): + summary = asyncio.run(_bootstrap_and_drive()) except KeyboardInterrupt: console.print("[yellow]cancelled[/yellow]") raise typer.Exit(code=130) from None diff --git a/src/stargraph/errors/__init__.py b/src/stargraph/errors/__init__.py index 58904d90..452b4f2c 100644 --- a/src/stargraph/errors/__init__.py +++ b/src/stargraph/errors/__init__.py @@ -17,6 +17,7 @@ IncompatibleSchemaError, IncompatibleSklearnVersion, IRValidationError, + LMServerError, MemoryScopeError, MigrationNotSupported, MLNodeError, @@ -51,6 +52,7 @@ "IncompatibleModelHashError", "IncompatibleSchemaError", "IncompatibleSklearnVersion", + "LMServerError", "MLNodeError", "MemoryScopeError", "MigrationNotSupported", diff --git a/src/stargraph/errors/_hierarchy.py b/src/stargraph/errors/_hierarchy.py index 1ee17031..6b53ff65 100644 --- a/src/stargraph/errors/_hierarchy.py +++ b/src/stargraph/errors/_hierarchy.py @@ -168,6 +168,18 @@ class RLNodeError(StargraphRuntimeError): """ +class LMServerError(StargraphRuntimeError): + """Raised when a declared LM endpoint cannot be attached to or spawned. + + Covers the ``stargraph.lm.sglang`` launcher: an already-running server + on the requested port serving a *different* model (never silently used), + a launch subprocess that exits before answering, and a startup that + exceeds ``startup_timeout_s``. ``context`` carries ``base_url`` plus, on + a failed spawn, ``exit_code`` and ``log`` (path to the captured server + output). + """ + + class SimulationError(StargraphRuntimeError): """Raised by :meth:`stargraph.graph.Graph.simulate` for fixture-input violations. diff --git a/src/stargraph/ir/__init__.py b/src/stargraph/ir/__init__.py index ea395958..59b7d1b1 100644 --- a/src/stargraph/ir/__init__.py +++ b/src/stargraph/ir/__init__.py @@ -25,6 +25,7 @@ RetractAction, RetryAction, RuleSpec, + SGLangServer, SkillRef, SkillSpec, SlotDef, @@ -57,6 +58,7 @@ "RetractAction", "RetryAction", "RuleSpec", + "SGLangServer", "SkillRef", "SkillSpec", "SlotDef", diff --git a/src/stargraph/ir/_models.py b/src/stargraph/ir/_models.py index 0a91c968..ab411bef 100644 --- a/src/stargraph/ir/_models.py +++ b/src/stargraph/ir/_models.py @@ -44,6 +44,7 @@ "RetractAction", "RetryAction", "RuleSpec", + "SGLangServer", "SkillRef", "SkillSpec", "SlotDef", @@ -337,6 +338,47 @@ class CheckpointBlock(IRBase): store: str = "sqlite:./.stargraph/run.sqlite" +class SGLangServer(IRBase): + """Declared LM endpoint: an SGLang server bound to this graph for a run. + + ``stargraph run`` resolves this block before the first node executes: it + probes ``http://{host}:{port}/v1/models`` and **attaches** when a server + already serves ``model`` (left running afterwards -- it is not ours), + otherwise it spawns ``python -m sglang.launch_server``, waits for the + endpoint to answer, and terminates it when the run ends. The derived + base URL + ``model`` configure the DSPy LM, so declaring this block is + equivalent to passing ``--lm-url``/``--lm-model`` at a live endpoint. + + Not part of the structural graph hash: like ``--lm-url``, the endpoint is + an environment binding, not graph topology. + + Attributes: + provider: Only ``"sglang"`` today; the discriminator for future + providers. + model: ``--model-path`` value; must equal the id the server reports + in ``/v1/models`` when attaching to an already-running server. + host: Bind address / probe host. + port: Bind port (SGLang's own default is 30000). + args: Extra argv passed through to ``sglang.launch_server`` + verbatim (e.g. ``["--attention-backend", "triton"]``). + Operator-only: a graph-declared value is refused unless the + operator re-states it as ``--sglang-arg`` (argv into a + subprocess is a code-execution surface, and a graph file can + be less trusted than the operator running it). The same holds + for a non-loopback ``host``, which would receive the API key + and every prompt. + startup_timeout_s: How long to wait for a spawned server to answer + before failing the run. Big models take minutes to load. + """ + + provider: Literal["sglang"] = "sglang" + model: str + host: str = "127.0.0.1" + port: int = 30000 + args: list[str] = Field(default_factory=list[str]) + startup_timeout_s: int = 600 + + # --------------------------------------------------------------------------- # IRDocument -- top-level IR shell. Required: ir_version, id, nodes. # --------------------------------------------------------------------------- @@ -369,6 +411,10 @@ class IRDocument(IRBase): governance: list[PackMount] = Field(default_factory=list[PackMount]) migrate: list[MigrateBlock] = Field(default_factory=list[MigrateBlock]) checkpoints: CheckpointBlock | None = None + # Optional LM endpoint binding (env, not topology -- excluded from the + # structural hash by construction: `structural_hash` reads topology, + # node signatures, state schema and rule packs only). + lm: SGLangServer | None = None # --------------------------------------------------------------------------- diff --git a/src/stargraph/lm/__init__.py b/src/stargraph/lm/__init__.py new file mode 100644 index 00000000..d874658a --- /dev/null +++ b/src/stargraph/lm/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LM endpoint lifecycle: bind a graph run to a local inference server. + +Today one provider: :mod:`stargraph.lm.sglang`, which attaches to (or boots +and tears down) an SGLang OpenAI-compatible server for the length of a run. +""" + +from __future__ import annotations + +from stargraph.lm.sglang import base_url, served_models, sglang_server + +__all__ = ["base_url", "served_models", "sglang_server"] diff --git a/src/stargraph/lm/sglang.py b/src/stargraph/lm/sglang.py new file mode 100644 index 00000000..75c7317d --- /dev/null +++ b/src/stargraph/lm/sglang.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Attach to -- or boot and tear down -- an SGLang server for one graph run. + +``stargraph run`` resolves an :class:`~stargraph.ir.SGLangServer` spec (from +the graph's ``lm:`` block or the ``--sglang-*`` flags) into a live +OpenAI-compatible base URL via :func:`sglang_server`: + +1. **Probe.** ``GET {base_url}/models``. If something answers, the endpoint is + *not ours*: attach when it already serves the requested model, and leave it + running when the run ends. If it answers with a different model, fail loud + (:class:`~stargraph.errors.LMServerError`) -- silently running a graph + against the wrong weights is worse than not running it. +2. **Spawn.** Nothing answering means we own it: launch + ``python -m sglang.launch_server --model-path ... --host ... --port ...`` + plus the spec's passthrough ``args``, in its own process group, with stdout + and stderr captured to a log file. +3. **Wait.** Poll the endpoint until it serves the model, the subprocess dies + (error carries the exit code + a tail of the log), or + ``startup_timeout_s`` elapses. Weight loading takes minutes for big models. +4. **Teardown.** ``SIGTERM`` the whole process group, then ``SIGKILL`` after a + grace period -- SGLang forks scheduler/detokenizer children that a bare + ``proc.terminate()`` would orphan on the GPU. + +SGLang itself is never imported here (it is a heavy, GPU-only dependency); it +is invoked as a subprocess of the *current* interpreter, so a missing install +surfaces as a launch failure with the pip hint attached. +""" + +from __future__ import annotations + +import contextlib +import os +import signal +import subprocess +import sys +import tempfile +import time +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +import httpx + +from stargraph.errors import LMServerError + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from stargraph.ir import SGLangServer + +__all__ = ["base_url", "served_models", "sglang_server"] + +_PROBE_TIMEOUT_S = 2.0 +_POLL_INTERVAL_S = 2.0 +_TERM_GRACE_S = 30.0 +_LOG_TAIL_LINES = 20 + + +def base_url(spec: SGLangServer) -> str: + """OpenAI-compatible base URL implied by ``spec`` (no trailing slash).""" + return f"http://{spec.host}:{spec.port}/v1" + + +def served_models(url: str, *, timeout: float = _PROBE_TIMEOUT_S) -> list[str] | None: + """Model ids served at ``url``, or ``None`` when nothing answers there. + + ``None`` (connection refused / timeout / non-2xx) means "no server": + the caller may spawn one. An empty list means a server answered but + serves nothing -- a real, distinct condition the caller must not + confuse with an unused port. + """ + try: + resp = httpx.get(f"{url}/models", timeout=timeout) + resp.raise_for_status() + payload = cast("dict[str, Any]", resp.json()) + except (httpx.HTTPError, ValueError): + return None + data = cast("list[dict[str, Any]]", payload.get("data", [])) + return [str(entry.get("id", "")) for entry in data] + + +def _launch_argv(spec: SGLangServer) -> list[str]: + """Argv for the launch subprocess (monkeypatched in tests).""" + return [ + sys.executable, + "-m", + "sglang.launch_server", + "--model-path", + spec.model, + "--host", + spec.host, + "--port", + str(spec.port), + *spec.args, + ] + + +def _log_tail(log_path: Path) -> str: + """Last :data:`_LOG_TAIL_LINES` lines of the captured server output.""" + try: + lines = log_path.read_text(encoding="utf-8", errors="replace").splitlines() + except OSError: + return "" + return "\n".join(lines[-_LOG_TAIL_LINES:]) + + +def _spawn(spec: SGLangServer, log_path: Path) -> subprocess.Popen[bytes]: + argv = _launch_argv(spec) + handle = log_path.open("wb") + try: + return subprocess.Popen( + argv, + stdout=handle, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + except OSError as exc: + raise LMServerError( + f"failed to launch sglang: {exc}", + hint=f"is sglang installed for {sys.executable}? `pip install sglang`", + argv=" ".join(argv), + ) from exc + finally: + # The child holds its own dup of the fd; ours is dead weight either way. + with contextlib.suppress(OSError): + handle.close() + + +def _await_ready( + proc: subprocess.Popen[bytes], + spec: SGLangServer, + url: str, + log_path: Path, + echo: Callable[[str], None] | None, +) -> None: + """Block until ``url`` serves ``spec.model``; raise on death or timeout.""" + deadline = time.monotonic() + spec.startup_timeout_s + while True: + code = proc.poll() + if code is not None: + raise LMServerError( + f"sglang exited with code {code} before serving {spec.model!r}", + hint=f"server output: {log_path}", + base_url=url, + exit_code=code, + log=str(log_path), + tail=_log_tail(log_path), + ) + models = served_models(url) + if models is not None: + if spec.model not in models: + raise LMServerError( + f"sglang on {url} serves {models} but the graph asked for {spec.model!r}", + hint="pass --served-model-name via args, or fix the model id", + base_url=url, + ) + if echo is not None: + echo(f"sglang ready on {url} ({spec.model})") + return + if time.monotonic() >= deadline: + raise LMServerError( + f"sglang did not answer on {url} within {spec.startup_timeout_s}s", + hint=(f"raise startup_timeout_s for a big model; server output: {log_path}"), + base_url=url, + log=str(log_path), + tail=_log_tail(log_path), + ) + time.sleep(_POLL_INTERVAL_S) + + +def _terminate(proc: subprocess.Popen[bytes]) -> None: + """SIGTERM the process group, SIGKILL what survives the grace period.""" + if proc.poll() is not None: + return + _signal_group(proc, signal.SIGTERM) + try: + proc.wait(timeout=_TERM_GRACE_S) + except subprocess.TimeoutExpired: + _signal_group(proc, signal.SIGKILL) + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=_TERM_GRACE_S) + + +def _signal_group(proc: subprocess.Popen[bytes], sig: int) -> None: + """Signal the child's process group, falling back to the child alone.""" + with contextlib.suppress(ProcessLookupError, PermissionError, OSError): + os.killpg(os.getpgid(proc.pid), sig) + return + with contextlib.suppress(ProcessLookupError, OSError): + proc.send_signal(sig) + + +@contextlib.contextmanager +def sglang_server( + spec: SGLangServer, + *, + log_path: Path | None = None, + echo: Callable[[str], None] | None = None, +) -> Generator[str]: + """Yield the base URL of a server for ``spec``, booting one if needed. + + Attaches to an already-running server that serves ``spec.model`` and + leaves it running on exit; otherwise spawns one and terminates it (and + its process group) when the block ends. ``log_path`` receives the + spawned server's stdout+stderr (default: a temp file); ``echo`` gets + one-line progress messages. + """ + url = base_url(spec) + existing = served_models(url) + if existing is not None: + if spec.model not in existing: + raise LMServerError( + f"a server already listening on {url} serves {existing}, not {spec.model!r}", + hint="pick a free --sglang-port, or stop that server", + base_url=url, + ) + if echo is not None: + echo(f"attached to running sglang on {url} ({spec.model})") + yield url + return + + if log_path is None: + fd, name = tempfile.mkstemp(prefix=f"sglang-{spec.port}-", suffix=".log") + os.close(fd) + log_path = Path(name) + if echo is not None: + echo(f"starting sglang ({spec.model}) on {url}; output -> {log_path}") + proc = _spawn(spec, log_path) + try: + _await_ready(proc, spec, url, log_path, echo) + yield url + finally: + _terminate(proc) diff --git a/src/stargraph/schemas/ir-v1-draft7.json b/src/stargraph/schemas/ir-v1-draft7.json index 94a76a3e..888e382e 100644 --- a/src/stargraph/schemas/ir-v1-draft7.json +++ b/src/stargraph/schemas/ir-v1-draft7.json @@ -471,6 +471,49 @@ "title": "RuleSpec", "type": "object" }, + "SGLangServer": { + "additionalProperties": false, + "description": "Declared LM endpoint: an SGLang server bound to this graph for a run.\n\n``stargraph run`` resolves this block before the first node executes: it\nprobes ``http://{host}:{port}/v1/models`` and **attaches** when a server\nalready serves ``model`` (left running afterwards -- it is not ours),\notherwise it spawns ``python -m sglang.launch_server``, waits for the\nendpoint to answer, and terminates it when the run ends. The derived\nbase URL + ``model`` configure the DSPy LM, so declaring this block is\nequivalent to passing ``--lm-url``/``--lm-model`` at a live endpoint.\n\nNot part of the structural graph hash: like ``--lm-url``, the endpoint is\nan environment binding, not graph topology.\n\nAttributes:\n provider: Only ``\"sglang\"`` today; the discriminator for future\n providers.\n model: ``--model-path`` value; must equal the id the server reports\n in ``/v1/models`` when attaching to an already-running server.\n host: Bind address / probe host.\n port: Bind port (SGLang's own default is 30000).\n args: Extra argv passed through to ``sglang.launch_server``\n verbatim (e.g. ``[\"--attention-backend\", \"triton\"]``).\n Operator-only: a graph-declared value is refused unless the\n operator re-states it as ``--sglang-arg`` (argv into a\n subprocess is a code-execution surface, and a graph file can\n be less trusted than the operator running it). The same holds\n for a non-loopback ``host``, which would receive the API key\n and every prompt.\n startup_timeout_s: How long to wait for a spawned server to answer\n before failing the run. Big models take minutes to load.", + "properties": { + "provider": { + "const": "sglang", + "default": "sglang", + "title": "Provider", + "type": "string" + }, + "model": { + "title": "Model", + "type": "string" + }, + "host": { + "default": "127.0.0.1", + "title": "Host", + "type": "string" + }, + "port": { + "default": 30000, + "title": "Port", + "type": "integer" + }, + "args": { + "items": { + "type": "string" + }, + "title": "Args", + "type": "array" + }, + "startup_timeout_s": { + "default": 600, + "title": "Startup Timeout S", + "type": "integer" + } + }, + "required": [ + "model" + ], + "title": "SGLangServer", + "type": "object" + }, "SkillRef": { "additionalProperties": false, "description": "Reference to a skill by namespaced id.", @@ -642,6 +685,17 @@ } ], "default": null + }, + "lm": { + "anyOf": [ + { + "$ref": "#/definitions/SGLangServer" + }, + { + "type": "null" + } + ], + "default": null } }, "required": [ diff --git a/src/stargraph/schemas/ir-v1.json b/src/stargraph/schemas/ir-v1.json index bf66a746..55440bc4 100644 --- a/src/stargraph/schemas/ir-v1.json +++ b/src/stargraph/schemas/ir-v1.json @@ -471,6 +471,49 @@ "title": "RuleSpec", "type": "object" }, + "SGLangServer": { + "additionalProperties": false, + "description": "Declared LM endpoint: an SGLang server bound to this graph for a run.\n\n``stargraph run`` resolves this block before the first node executes: it\nprobes ``http://{host}:{port}/v1/models`` and **attaches** when a server\nalready serves ``model`` (left running afterwards -- it is not ours),\notherwise it spawns ``python -m sglang.launch_server``, waits for the\nendpoint to answer, and terminates it when the run ends. The derived\nbase URL + ``model`` configure the DSPy LM, so declaring this block is\nequivalent to passing ``--lm-url``/``--lm-model`` at a live endpoint.\n\nNot part of the structural graph hash: like ``--lm-url``, the endpoint is\nan environment binding, not graph topology.\n\nAttributes:\n provider: Only ``\"sglang\"`` today; the discriminator for future\n providers.\n model: ``--model-path`` value; must equal the id the server reports\n in ``/v1/models`` when attaching to an already-running server.\n host: Bind address / probe host.\n port: Bind port (SGLang's own default is 30000).\n args: Extra argv passed through to ``sglang.launch_server``\n verbatim (e.g. ``[\"--attention-backend\", \"triton\"]``).\n Operator-only: a graph-declared value is refused unless the\n operator re-states it as ``--sglang-arg`` (argv into a\n subprocess is a code-execution surface, and a graph file can\n be less trusted than the operator running it). The same holds\n for a non-loopback ``host``, which would receive the API key\n and every prompt.\n startup_timeout_s: How long to wait for a spawned server to answer\n before failing the run. Big models take minutes to load.", + "properties": { + "provider": { + "const": "sglang", + "default": "sglang", + "title": "Provider", + "type": "string" + }, + "model": { + "title": "Model", + "type": "string" + }, + "host": { + "default": "127.0.0.1", + "title": "Host", + "type": "string" + }, + "port": { + "default": 30000, + "title": "Port", + "type": "integer" + }, + "args": { + "items": { + "type": "string" + }, + "title": "Args", + "type": "array" + }, + "startup_timeout_s": { + "default": 600, + "title": "Startup Timeout S", + "type": "integer" + } + }, + "required": [ + "model" + ], + "title": "SGLangServer", + "type": "object" + }, "SkillRef": { "additionalProperties": false, "description": "Reference to a skill by namespaced id.", @@ -642,6 +685,17 @@ } ], "default": null + }, + "lm": { + "anyOf": [ + { + "$ref": "#/$defs/SGLangServer" + }, + { + "type": "null" + } + ], + "default": null } }, "required": [ diff --git a/tests/integration/test_lm_sglang_spawn.py b/tests/integration/test_lm_sglang_spawn.py new file mode 100644 index 00000000..94f86ce6 --- /dev/null +++ b/tests/integration/test_lm_sglang_spawn.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Real spawn -> ready -> teardown for :func:`stargraph.lm.sglang.sglang_server`. + +SGLang is a GPU-only dependency, so the launch argv is swapped for a stub +OpenAI-compatible server (stdlib ``http.server``) that reports one model id. +Everything else is the production path: a real subprocess in its own process +group, real HTTP readiness polling, real SIGTERM teardown. +""" + +from __future__ import annotations + +import socket +import textwrap +from typing import TYPE_CHECKING + +import pytest + +from stargraph.errors import LMServerError +from stargraph.ir import SGLangServer +from stargraph.lm import sglang as sg + +if TYPE_CHECKING: + from pathlib import Path + +pytestmark = pytest.mark.integration + +_STUB = textwrap.dedent( + ''' + """Minimal OpenAI-compatible stub: GET /v1/models -> one model id.""" + import json + import sys + from http.server import BaseHTTPRequestHandler, HTTPServer + + MODEL = sys.argv[1] + PORT = int(sys.argv[2]) + + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): # noqa: N802 + body = json.dumps({"data": [{"id": MODEL}]}).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *_args): # keep the captured log quiet + return + + + HTTPServer(("127.0.0.1", PORT), Handler).serve_forever() + ''' +) + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +@pytest.fixture +def stub(tmp_path: Path) -> Path: + path = tmp_path / "stub_server.py" + path.write_text(_STUB, encoding="utf-8") + return path + + +def test_spawns_waits_for_ready_then_tears_down( + stub: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import sys + + port = _free_port() + spec = SGLangServer(model="stub/model", port=port, startup_timeout_s=30) + + def _stub_argv(spec: SGLangServer) -> list[str]: + return [sys.executable, str(stub), spec.model, str(spec.port)] + + monkeypatch.setattr(sg, "_launch_argv", _stub_argv) + log = tmp_path / "sglang.log" + + with sg.sglang_server(spec, log_path=log) as url: + assert url == f"http://127.0.0.1:{port}/v1" + assert sg.served_models(url) == ["stub/model"] + + # Teardown is real: nothing answers on the port once the block exits. + assert sg.served_models(f"http://127.0.0.1:{port}/v1", timeout=1.0) is None + + +def test_launch_that_dies_reports_exit_code_and_log( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import sys + + spec = SGLangServer(model="stub/model", port=_free_port(), startup_timeout_s=30) + + def _dying_argv(_spec: SGLangServer) -> list[str]: + return [sys.executable, "-c", 'import sys; print("CUDA go boom"); sys.exit(3)'] + + monkeypatch.setattr(sg, "_launch_argv", _dying_argv) + log = tmp_path / "sglang.log" + + with ( + pytest.raises(LMServerError, match="exited with code 3") as excinfo, + sg.sglang_server(spec, log_path=log), + ): + pytest.fail("must not yield when the launch dies") + + assert excinfo.value.context["exit_code"] == 3 + assert "CUDA go boom" in excinfo.value.context["tail"] + + +def test_startup_timeout_is_loud(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + import sys + + spec = SGLangServer(model="stub/model", port=_free_port(), startup_timeout_s=1) + + def _hanging_argv(_spec: SGLangServer) -> list[str]: + return [sys.executable, "-c", "import time; time.sleep(60)"] + + monkeypatch.setattr(sg, "_launch_argv", _hanging_argv) + + with ( + pytest.raises(LMServerError, match="did not answer"), + sg.sglang_server(spec, log_path=tmp_path / "sglang.log"), + ): + pytest.fail("must not yield before the endpoint answers") diff --git a/tests/unit/cli/test_run_sglang.py b/tests/unit/cli/test_run_sglang.py new file mode 100644 index 00000000..8f4887bb --- /dev/null +++ b/tests/unit/cli/test_run_sglang.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``stargraph run --sglang-*`` / graph ``lm:`` block wiring. + +Covers spec resolution (flags over the graph's block), the conflict with +``--lm-url``/``--lm-model``, and the end-to-end contract: the endpoint is +held open around the whole run and its URL + model configure the DSPy LM. +The launcher itself is stubbed here -- see ``tests/unit/test_lm_sglang.py`` +and ``tests/integration/test_lm_sglang_spawn.py`` for that. +""" + +from __future__ import annotations + +import contextlib +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import pytest +import typer +from typer.testing import CliRunner + +from stargraph.cli.run import _resolve_sglang, cmd # pyright: ignore[reportPrivateUsage] +from stargraph.ir import IRDocument, SGLangServer + +if TYPE_CHECKING: + from collections.abc import Generator + +pytestmark = pytest.mark.unit + +REPO_ROOT = Path(__file__).resolve().parents[3] +SAMPLE_GRAPH = REPO_ROOT / "tests" / "fixtures" / "sample-graph.yaml" + +_GRAPH_BODY = """\ +id: sglang-demo +state: + message: str +nodes: + work: + kind: echo +routes: + work: done +""" + +_GRAPH_WITH_LM = ( + _GRAPH_BODY + + """\ +lm: + provider: sglang + model: Qwen/Qwen3-8B + port: 41002 +""" +) + +_GRAPH_WITH_LM_ARGS = _GRAPH_WITH_LM + ' args: ["--trust-remote-code"]\n' + +_GRAPH_WITH_REMOTE_LM = ( + _GRAPH_BODY + + """\ +lm: + provider: sglang + model: Qwen/Qwen3-8B + host: evil.example.com +""" +) + + +def _make_app() -> typer.Typer: + app = typer.Typer() + app.command()(cmd) + return app + + +def _ir(lm: SGLangServer | None = None) -> IRDocument: + return IRDocument(ir_version="1.0.0", id="run:t", nodes=[], lm=lm) + + +def _resolve(ir: IRDocument, **overrides: Any) -> SGLangServer | None: + kwargs: dict[str, Any] = { + "model": None, + "host": None, + "port": None, + "args": None, + "timeout": None, + } + kwargs.update(overrides) + return _resolve_sglang(ir, **kwargs) + + +# -------------------------------------------------------------------------- +# spec resolution +# -------------------------------------------------------------------------- + + +def test_no_block_and_no_flags_means_no_server() -> None: + assert _resolve(_ir()) is None + + +def test_flags_alone_declare_a_server() -> None: + spec = _resolve(_ir(), model="Qwen/Qwen3-8B", port=41002, args=["--tp", "2"]) + + assert spec is not None + assert (spec.model, spec.port, spec.args) == ("Qwen/Qwen3-8B", 41002, ["--tp", "2"]) + assert spec.host == "127.0.0.1" + + +def test_graph_block_alone_declares_a_server() -> None: + spec = _resolve(_ir(SGLangServer(model="phi-4", port=41002))) + + assert spec is not None + assert (spec.model, spec.port) == ("phi-4", 41002) + + +def test_flags_override_the_graph_block_field_by_field() -> None: + declared = SGLangServer(model="phi-4", port=41002, startup_timeout_s=900) + + spec = _resolve(_ir(declared), port=41010) + + assert spec is not None + assert spec.port == 41010 + assert (spec.model, spec.startup_timeout_s) == ("phi-4", 900) # untouched + + +def test_sglang_flag_without_a_model_is_loud() -> None: + with pytest.raises(typer.BadParameter, match="--sglang-model"): + _resolve(_ir(), port=41002) + + +# -------------------------------------------------------------------------- +# CLI +# -------------------------------------------------------------------------- + + +def _stub_launcher(monkeypatch: pytest.MonkeyPatch, calls: list[Any]) -> None: + import stargraph.lm.sglang as sglang_mod + + @contextlib.contextmanager + def _fake(spec: SGLangServer, **_kwargs: Any) -> Generator[str]: + calls.append(("enter", spec)) + try: + yield "http://stub:41002/v1" + finally: + calls.append(("exit", spec)) + + monkeypatch.setattr(sglang_mod, "sglang_server", _fake) + + +def _stub_dspy(monkeypatch: pytest.MonkeyPatch, captured: dict[str, Any]) -> None: + import dspy # pyright: ignore[reportMissingTypeStubs] + + class FakeLM: + def __init__(self, model: str, **kwargs: Any) -> None: + captured["model"] = model + captured["lm_kwargs"] = kwargs + + def _configure(**kwargs: Any) -> None: + captured["configured"] = kwargs + + monkeypatch.setattr(dspy, "configure", _configure) + monkeypatch.setattr(dspy, "LM", FakeLM) + + +def test_sglang_model_flag_binds_the_derived_endpoint( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls: list[Any] = [] + captured: dict[str, Any] = {} + _stub_launcher(monkeypatch, calls) + _stub_dspy(monkeypatch, captured) + + result = CliRunner().invoke( + _make_app(), + [ + str(SAMPLE_GRAPH), + "--checkpoint", + str(tmp_path / "ck.sqlite"), + "--sglang-model", + "Qwen/Qwen3-8B", + "--sglang-port", + "41002", + "--quiet", + "--no-summary", + ], + ) + + assert result.exit_code == 0, result.output + assert [phase for phase, _ in calls] == ["enter", "exit"] + assert calls[0][1].model == "Qwen/Qwen3-8B" + assert captured["model"] == "openai/Qwen/Qwen3-8B" + assert captured["lm_kwargs"]["api_base"] == "http://stub:41002/v1" + + +def test_graph_lm_block_is_honoured(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + calls: list[Any] = [] + _stub_launcher(monkeypatch, calls) + _stub_dspy(monkeypatch, {}) + graph = tmp_path / "graph.yaml" + graph.write_text(_GRAPH_WITH_LM, encoding="utf-8") + + result = CliRunner().invoke( + _make_app(), + [ + str(graph), + "--checkpoint", + str(tmp_path / "ck.sqlite"), + "--quiet", + "--no-summary", + ], + ) + + assert result.exit_code == 0, result.output + spec = calls[0][1] + assert (spec.model, spec.port) == ("Qwen/Qwen3-8B", 41002) + + +def test_no_sglang_request_never_touches_the_launcher( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls: list[Any] = [] + _stub_launcher(monkeypatch, calls) + + result = CliRunner().invoke( + _make_app(), + [ + str(SAMPLE_GRAPH), + "--checkpoint", + str(tmp_path / "ck.sqlite"), + "--quiet", + "--no-summary", + ], + ) + + assert result.exit_code == 0, result.output + assert calls == [] + + +@pytest.mark.parametrize( + "conflicting", + [["--lm-url", "http://localhost:11434/v1"], ["--lm-model", "gpt-oss:20b"]], +) +def test_lm_flags_conflict_with_sglang( + conflicting: list[str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls: list[Any] = [] + _stub_launcher(monkeypatch, calls) + + result = CliRunner().invoke( + _make_app(), + [ + str(SAMPLE_GRAPH), + "--checkpoint", + str(tmp_path / "ck.sqlite"), + "--sglang-model", + "Qwen/Qwen3-8B", + *conflicting, + "--quiet", + "--no-summary", + ], + ) + + assert result.exit_code != 0 + assert calls == [] + + +def test_inspect_does_not_start_a_server(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + calls: list[Any] = [] + _stub_launcher(monkeypatch, calls) + graph = tmp_path / "graph.yaml" + graph.write_text(_GRAPH_WITH_LM, encoding="utf-8") + + result = CliRunner().invoke(_make_app(), [str(graph), "--inspect"]) + + assert result.exit_code == 0, result.output + assert calls == [] + + +# -------------------------------------------------------------------------- +# operator-only fields on a graph-declared block +# -------------------------------------------------------------------------- + + +def test_graph_declared_args_are_refused() -> None: + declared = SGLangServer(model="phi-4", args=["--trust-remote-code"]) + + with pytest.raises(typer.BadParameter, match="--sglang-arg"): + _resolve(_ir(declared)) + + +def test_graph_declared_args_are_allowed_once_restated_as_flags() -> None: + declared = SGLangServer(model="phi-4", args=["--trust-remote-code"]) + + spec = _resolve(_ir(declared), args=["--attention-backend", "triton"]) + + assert spec is not None + assert spec.args == ["--attention-backend", "triton"] + + +def test_graph_declared_non_loopback_host_is_refused() -> None: + declared = SGLangServer(model="phi-4", host="evil.example.com") + + with pytest.raises(typer.BadParameter, match="--sglang-host"): + _resolve(_ir(declared)) + + +def test_graph_declared_non_loopback_host_is_allowed_once_restated() -> None: + declared = SGLangServer(model="phi-4", host="evil.example.com") + + spec = _resolve(_ir(declared), host="evil.example.com") + + assert spec is not None + assert spec.host == "evil.example.com" + + +@pytest.mark.parametrize("host", ["127.0.0.1", "127.0.0.5", "::1", "localhost"]) +def test_graph_declared_loopback_hosts_pass(host: str) -> None: + spec = _resolve(_ir(SGLangServer(model="phi-4", host=host))) + + assert spec is not None + assert spec.host == host + + +def test_flag_declared_spec_may_reach_anywhere() -> None: + """Flags are the operator speaking; no restriction applies to them.""" + spec = _resolve(_ir(), model="phi-4", host="10.0.0.9", args=["--trust-remote-code"]) + + assert spec is not None + assert (spec.host, spec.args) == ("10.0.0.9", ["--trust-remote-code"]) + + +@pytest.mark.parametrize("graph_yaml", [_GRAPH_WITH_LM_ARGS, _GRAPH_WITH_REMOTE_LM]) +def test_cli_refuses_the_untrusted_halves_of_a_graph_block( + graph_yaml: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls: list[Any] = [] + _stub_launcher(monkeypatch, calls) + graph = tmp_path / "graph.yaml" + graph.write_text(graph_yaml, encoding="utf-8") + + result = CliRunner().invoke( + _make_app(), + [str(graph), "--checkpoint", str(tmp_path / "ck.sqlite"), "--quiet", "--no-summary"], + ) + + assert result.exit_code != 0 + assert calls == [] diff --git a/tests/unit/test_authoring_compiler.py b/tests/unit/test_authoring_compiler.py index 8ba61502..2701c263 100644 --- a/tests/unit/test_authoring_compiler.py +++ b/tests/unit/test_authoring_compiler.py @@ -145,3 +145,40 @@ def test_authoring_clips_renders_rules() -> None: assert any("r-judge-fail" in line and "goto work" in line for line in lines) assert any('(verdict (value "pass"))' in line and "halt" in line for line in lines) + + +def test_lm_block_lowers_to_an_ir_endpoint_spec() -> None: + ir = compile_authoring( + _doc( + lm={ + "provider": "sglang", + "model": "Qwen/Qwen3-8B", + "port": 41002, + "args": ["--attention-backend", "triton"], + } + ) + ) + + assert ir.lm is not None + assert (ir.lm.model, ir.lm.port) == ("Qwen/Qwen3-8B", 41002) + assert ir.lm.args == ["--attention-backend", "triton"] + assert ir.lm.host == "127.0.0.1" # default + + +def test_no_lm_block_leaves_the_endpoint_unbound() -> None: + assert compile_authoring(_doc()).lm is None + + +@pytest.mark.parametrize( + ("block", "match"), + [ + ("sglang", "must be a mapping"), + ({"provider": "sglang"}, "model"), + ({"model": "m", "bogus": 1}, "bogus"), + ({"provider": "vllm", "model": "m"}, "provider"), + ({"model": "m", "port": "http"}, "port"), + ], +) +def test_loud_lm_block_errors(block: Any, match: str) -> None: + with pytest.raises(IRValidationError, match=match): + compile_authoring(_doc(lm=block)) diff --git a/tests/unit/test_errors_walker.py b/tests/unit/test_errors_walker.py index d2abf177..f2681224 100644 --- a/tests/unit/test_errors_walker.py +++ b/tests/unit/test_errors_walker.py @@ -64,6 +64,9 @@ # and step/plan contract violations. "RLNodeError", "SimulationError", + # LM endpoint launcher (stargraph.lm.sglang) — StargraphRuntimeError + # subclass covering attach/spawn/readiness failures of a declared server. + "LMServerError", # foundry assembler — raised by stargraph.skills.foundry.assemble when no # runnable graph spine landed. StargraphRuntimeError-rooted so a single # ``except StargraphRuntimeError`` catches it alongside engine runtime failures. diff --git a/tests/unit/test_lm_sglang.py b/tests/unit/test_lm_sglang.py new file mode 100644 index 00000000..5f010961 --- /dev/null +++ b/tests/unit/test_lm_sglang.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for :mod:`stargraph.lm.sglang` -- probe, attach, argv, teardown. + +No SGLang and no sockets here: :func:`served_models` is monkeypatched so the +attach/spawn decision is exercised in isolation. The real spawn -> ready -> +teardown path runs against a stub HTTP server in +``tests/integration/test_lm_sglang_spawn.py``. +""" + +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING, NoReturn + +import pytest + +from stargraph.errors import LMServerError +from stargraph.ir import SGLangServer +from stargraph.lm import sglang as sg + +if TYPE_CHECKING: + from collections.abc import Callable + from pathlib import Path + +pytestmark = pytest.mark.unit + + +def _spec(**overrides: object) -> SGLangServer: + return SGLangServer.model_validate({"model": "Qwen/Qwen3-8B", **overrides}) + + +def _probe(result: list[str] | None) -> Callable[..., list[str] | None]: + """Stand in for :func:`served_models` with a fixed answer.""" + + def _stub(_url: str, **_kwargs: object) -> list[str] | None: + return result + + return _stub + + +def _never_spawn(*_args: object, **_kwargs: object) -> NoReturn: + pytest.fail("attached run must not spawn") + + +def _fake_spawn(*_args: object, **_kwargs: object) -> str: + return "proc" + + +def _noop_ready(*_args: object, **_kwargs: object) -> None: + return None + + +def test_base_url_is_openai_compatible() -> None: + assert sg.base_url(_spec(host="10.0.0.4", port=41002)) == "http://10.0.0.4:41002/v1" + + +def test_launch_argv_passes_model_port_and_extra_args() -> None: + argv = sg._launch_argv(_spec(port=41002, args=["--attention-backend", "triton"])) # pyright: ignore[reportPrivateUsage] + + assert argv[:3] == [sys.executable, "-m", "sglang.launch_server"] + assert argv[3:5] == ["--model-path", "Qwen/Qwen3-8B"] + assert "--port" in argv and argv[argv.index("--port") + 1] == "41002" + assert argv[-2:] == ["--attention-backend", "triton"] + + +def test_served_models_returns_none_when_nothing_answers() -> None: + # Port 1 is privileged and unbound: the probe must report "no server" + # rather than raising, so the caller spawns one. + assert sg.served_models("http://127.0.0.1:1/v1", timeout=0.25) is None + + +def test_attaches_to_running_server_without_spawning(monkeypatch: pytest.MonkeyPatch) -> None: + spec = _spec(port=41002) + monkeypatch.setattr(sg, "served_models", _probe([spec.model])) + monkeypatch.setattr(sg, "_spawn", _never_spawn) + messages: list[str] = [] + + with sg.sglang_server(spec, echo=messages.append) as url: + assert url == "http://127.0.0.1:41002/v1" + + assert any("attached" in m for m in messages) + + +def test_running_server_with_other_model_is_loud(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(sg, "served_models", _probe(["some/other-model"])) + + with pytest.raises(LMServerError, match="serves"), sg.sglang_server(_spec()): + pytest.fail("must not yield against the wrong model") + + +def test_spawned_server_is_terminated_on_exit(monkeypatch: pytest.MonkeyPatch) -> None: + spec = _spec() + terminated: list[object] = [] + monkeypatch.setattr(sg, "served_models", _probe(None)) + monkeypatch.setattr(sg, "_spawn", _fake_spawn) + monkeypatch.setattr(sg, "_await_ready", _noop_ready) + monkeypatch.setattr(sg, "_terminate", terminated.append) + + with sg.sglang_server(spec) as url: + assert url == sg.base_url(spec) + assert terminated == [] + + assert terminated == ["proc"] + + +def test_spawned_server_is_terminated_when_the_body_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + terminated: list[object] = [] + monkeypatch.setattr(sg, "served_models", _probe(None)) + monkeypatch.setattr(sg, "_spawn", _fake_spawn) + monkeypatch.setattr(sg, "_await_ready", _noop_ready) + monkeypatch.setattr(sg, "_terminate", terminated.append) + + with pytest.raises(RuntimeError, match="boom"), sg.sglang_server(_spec()): + raise RuntimeError("boom") + + assert terminated == ["proc"] + + +def test_log_tail_reports_last_lines(tmp_path: Path) -> None: + log = tmp_path / "sglang.log" + log.write_text("\n".join(f"line {i}" for i in range(50)), encoding="utf-8") + + tail = sg._log_tail(log) # pyright: ignore[reportPrivateUsage] + + assert tail.splitlines()[0] == "line 30" + assert tail.splitlines()[-1] == "line 49" From 4f7cc0174a0330c5fe7bc6f6cf516b1d1f16b79a Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Mon, 17 Aug 2026 20:40:03 +0000 Subject: [PATCH 02/11] =?UTF-8?q?examples:=20sglang-qa.yaml=20=E2=80=94=20?= =?UTF-8?q?a=20graph=20that=20carries=20its=20own=20LM=20endpoint?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the runnable example for the `lm:` block: the graph names the model it wants and `stargraph run` resolves the endpoint before the first node, instead of the caller supplying --lm-url/--lm-model. Writing it surfaced a regression in the previous commit. `_configure_lm` used to run before `build_node_registry`; folding it into the `_lm_endpoint` context manager moved it *after*, and `kind: dspy` validates that an LM is configured while the node is being constructed, not when it runs. Every `kind: dspy` graph was therefore unrunnable -- with a declared `lm:` block *and* with --lm-url/--lm-model, which is a regression of shipped behaviour. The endpoint context now opens before the registry is built and still spans the whole run, so a spawned server is torn down only once the loop is done or has raised. The existing sglang tests missed it because none of them build a registry containing a dspy node; the graphs are all `kind: echo`. `test_the_lm_is_configured_before_the_node_registry_is_built` pins the call order directly rather than relying on the one example that happens to use a dspy node. The example's golden test starts a loopback OpenAI-compatible stub server and re-points the declared block with `--sglang-port`, which drives the production *attach* branch -- a server already serving the requested model is used as-is. SGLang is GPU-only, so the spawn branch stays covered by tests/integration/test_lm_sglang_spawn.py; no engine code is monkeypatched here, and the assertion on the stub's exact answer proves the resolved base URL plus model really configured the LM the node ran against. Signed-off-by: Sean Mauk --- examples/README.md | 8 +- examples/sglang-qa.yaml | 36 +++++++ src/stargraph/cli/run.py | 69 ++++++------ tests/integration/test_examples.py | 164 ++++++++++++++++++++++++++++- tests/unit/cli/test_run_sglang.py | 54 ++++++++++ 5 files changed, 297 insertions(+), 34 deletions(-) create mode 100644 examples/sglang-qa.yaml diff --git a/examples/README.md b/examples/README.md index f89693b0..ffe7ecb3 100644 --- a/examples/README.md +++ b/examples/README.md @@ -9,6 +9,7 @@ golden test (`tests/integration/test_examples.py`) and must reach | `hello.yaml` | Smallest graph: 2 nodes, 2 rules, 1 state field, no LLM | `stargraph run examples/hello.yaml --inputs message=hello` | | `pipeline.yaml` | Three-step linear routing, one Fathom rule per hop | `stargraph run examples/pipeline.yaml --inputs message=hello` | | `research-bot.yaml` | Authoring format (no `ir_version`): react + judge + verdict-routed feedback loop in ~20 lines | `stargraph run examples/research-bot.yaml --lm-url --lm-model --inputs question="..."` | +| `sglang-qa.yaml` | The `lm:` block: the graph carries its own SGLang endpoint, so no `--lm-url`/`--lm-model` | `stargraph run examples/sglang-qa.yaml --inputs question="..."` | Trace rule firings without executing nodes: @@ -19,8 +20,11 @@ stargraph run examples/hello.yaml --inspect ## What's intentionally *not* here Most examples use `echo`/`halt` nodes, so they stay self-contained and fast -(`research-bot.yaml` is the exception: it demonstrates LLM nodes and needs -`--lm-url`/`--lm-model`; its golden test drives it with a scripted stub LM). +(`research-bot.yaml` and `sglang-qa.yaml` are the exceptions: both demonstrate +LLM nodes. `research-bot.yaml` needs `--lm-url`/`--lm-model` and its golden test +drives it with a scripted stub LM; `sglang-qa.yaml` declares its own endpoint and +its golden test attaches to a loopback stub server, exercising the real attach +path without a GPU). For the features that need more wiring, read the full graphs under [`demos/`](../demos/): diff --git a/examples/sglang-qa.yaml b/examples/sglang-qa.yaml new file mode 100644 index 00000000..3f5cca0f --- /dev/null +++ b/examples/sglang-qa.yaml @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# examples/sglang-qa.yaml — a graph that carries its own LM endpoint. +# +# The `lm:` block is the point: instead of the caller supplying +# --lm-url/--lm-model, the graph names the model it wants and `stargraph run` +# resolves the endpoint before the first node executes. If a server is already +# serving that model on that port it attaches and leaves it running; otherwise +# it launches `python -m sglang.launch_server`, waits for the endpoint to +# answer, and terminates it when the run ends. +# +# Run it (SGLang installed, weights downloadable): +# stargraph run examples/sglang-qa.yaml --inputs question="what routes stargraph?" +# +# Point it at a server you already have, without editing this file: +# stargraph run examples/sglang-qa.yaml --sglang-port 41010 --inputs question="..." +# +# See the lowered IR (the `lm:` block becomes an SGLangServer model, and is +# deliberately excluded from the structural graph hash — an endpoint is an +# environment binding, not topology): +# stargraph compile examples/sglang-qa.yaml +id: sglang-qa +lm: + provider: sglang # the only provider today + model: Qwen/Qwen2.5-0.5B-Instruct + port: 30000 # sglang's own default + startup_timeout_s: 600 # first run downloads weights +state: + question: str + answer: str +nodes: + # Every key except `kind` becomes the node's config, so this is a plain + # dspy.Predict over the declared signature — bound to the endpoint above. + ask: {kind: dspy, signature: "question -> answer"} +routes: + ask: done diff --git a/src/stargraph/cli/run.py b/src/stargraph/cli/run.py index 09044864..3a147989 100644 --- a/src/stargraph/cli/run.py +++ b/src/stargraph/cli/run.py @@ -512,47 +512,54 @@ def cmd( else: initial_values = parse_inputs(inputs or [], ir.state_schema) initial_state = g.state_schema(**initial_values) - try: - node_registry = build_node_registry(ir.nodes, ir_dir=graph.parent.resolve()) - except StargraphError as e: - # Registry failures are graph-authoring mistakes (unknown kind, bad - # module ref, missing sub-IR) — surface them as parameter errors. - raise typer.BadParameter(str(e)) from e - run = GraphRun( - run_id=run_id, - graph=g, - initial_state=initial_state, - node_registry=node_registry, - checkpointer=checkpointer, - fathom=build_ir_routing(ir, g.state_schema), - ) - console = Console() progress = ProgressPrinter(console, quiet=quiet, verbose=verbose) hitl: HITLHandler | None = None if non_interactive else HITLHandler(console) - async def _bootstrap_and_drive() -> RunSummary: - await checkpointer.bootstrap() - try: - if live_broker: - from stargraph.serve.lifecycle import broker_lifespan - - async with broker_lifespan(): - return await _drive_interactive(run, audit_sink, progress, hitl, console) - return await _drive_interactive(run, audit_sink, progress, hitl, console) - finally: - await checkpointer.close() - if audit_sink is not None: - await audit_sink.close() - def _echo(message: str) -> None: if not quiet: console.print(message, style="dim", highlight=False, markup=False) try: - # The endpoint outlives the whole run: a spawned sglang server is torn - # down only once the loop is done (or has raised). + # The endpoint is resolved *before* the node registry is built: a + # ``kind: dspy`` node checks for a configured LM as it is constructed, + # so a graph-declared ``lm:`` block (like --lm-url/--lm-model) has to + # be live by then. It outlives the whole run as well -- a spawned + # sglang server is torn down only once the loop is done (or has + # raised). with _lm_endpoint(sglang_spec, lm_url, lm_model, lm_key, lm_timeout, _echo): + try: + node_registry = build_node_registry(ir.nodes, ir_dir=graph.parent.resolve()) + except StargraphError as e: + # Registry failures are graph-authoring mistakes (unknown kind, + # bad module ref, missing sub-IR) — surface them as parameter + # errors. + raise typer.BadParameter(str(e)) from e + run = GraphRun( + run_id=run_id, + graph=g, + initial_state=initial_state, + node_registry=node_registry, + checkpointer=checkpointer, + fathom=build_ir_routing(ir, g.state_schema), + ) + + async def _bootstrap_and_drive() -> RunSummary: + await checkpointer.bootstrap() + try: + if live_broker: + from stargraph.serve.lifecycle import broker_lifespan + + async with broker_lifespan(): + return await _drive_interactive( + run, audit_sink, progress, hitl, console + ) + return await _drive_interactive(run, audit_sink, progress, hitl, console) + finally: + await checkpointer.close() + if audit_sink is not None: + await audit_sink.close() + summary = asyncio.run(_bootstrap_and_drive()) except KeyboardInterrupt: console.print("[yellow]cancelled[/yellow]") diff --git a/tests/integration/test_examples.py b/tests/integration/test_examples.py index 4a2cfc72..4127cee6 100644 --- a/tests/integration/test_examples.py +++ b/tests/integration/test_examples.py @@ -9,6 +9,13 @@ from __future__ import annotations import json +import socket +import subprocess +import sys +import textwrap +import time +import urllib.error +import urllib.request from pathlib import Path import pytest @@ -22,7 +29,7 @@ # driven below under a scripted DummyLM instead (CliRunner shares the process, # so a dspy.context around invoke() is the stub seam). Live runs pass # --lm-url/--lm-model. -_LM_EXAMPLES = {"research-bot.yaml"} +_LM_EXAMPLES = {"research-bot.yaml", "sglang-qa.yaml"} EXAMPLE_GRAPHS = sorted(p for p in EXAMPLES_DIR.glob("*.yaml") if p.name not in _LM_EXAMPLES) @@ -107,3 +114,158 @@ def test_research_bot_loop_with_scripted_lm(runner: CliRunner, tmp_path: Path) - assert state["answer"] == "draft-2 names CLIPS" # second round's answer won # The judge's round-1 rationale was re-injected into the round-2 brief. assert "too vague, name the engine" in state["brief"] + + +# --------------------------------------------------------------------------- # +# sglang-qa.yaml -- the graph carries its own endpoint via the `lm:` block # +# --------------------------------------------------------------------------- # + +_OPENAI_STUB = textwrap.dedent( + ''' + """OpenAI-compatible stub: GET /v1/models + POST /v1/chat/completions. + + Speaks just enough for the attach probe to recognise the model and for + DSPy's chat adapter to parse one field back out. + """ + import json + import sys + from http.server import BaseHTTPRequestHandler, HTTPServer + + MODEL, PORT, ANSWER = sys.argv[1], int(sys.argv[2]), sys.argv[3] + + + class Handler(BaseHTTPRequestHandler): + def _send(self, payload): + body = json.dumps(payload).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): # noqa: N802 + self._send({"data": [{"id": MODEL}]}) + + def do_POST(self): # noqa: N802 + self.rfile.read(int(self.headers.get("content-length", 0) or 0)) + content = f"[[ ## answer ## ]]\\n{ANSWER}\\n\\n[[ ## completed ## ]]" + self._send( + { + "id": "chatcmpl-stub", + "object": "chat.completion", + "model": MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": content}, + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ) + + def log_message(self, *_args): # keep the captured log quiet + return + + + HTTPServer(("127.0.0.1", PORT), Handler).serve_forever() + ''' +) + +_STUB_ANSWER = "Fathom, a CLIPS rules engine, decides every transition." + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _declared_model(graph: Path) -> str: + """The ``lm.model`` the example declares -- the stub must report exactly it. + + Read from the YAML rather than duplicated here so a rename of the model in + the example cannot leave this test attaching to something else. + """ + for line in graph.read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if stripped.startswith("model:"): + return stripped.split(":", 1)[1].split("#")[0].strip() + raise AssertionError(f"no lm.model declared in {graph}") + + +@pytest.fixture +def openai_stub(tmp_path: Path) -> object: + """Serve one model id on a free loopback port; yield ``(port, model)``. + + SGLang is GPU-only, so the example's *spawn* path cannot run in CI. Its + *attach* path can: a server already serving the requested model is left + alone and used as-is, which is the production branch taken here -- no + launch argv is stubbed and no engine code is monkeypatched. + """ + graph = EXAMPLES_DIR / "sglang-qa.yaml" + model = _declared_model(graph) + port = _free_port() + script = tmp_path / "openai_stub.py" + script.write_text(_OPENAI_STUB, encoding="utf-8") + proc = subprocess.Popen( + [sys.executable, str(script), model, str(port), _STUB_ANSWER], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + try: + with urllib.request.urlopen( + f"http://127.0.0.1:{port}/v1/models", timeout=1 + ) as response: + if response.status == 200: + break + except (urllib.error.URLError, OSError, TimeoutError): + time.sleep(0.05) + else: + raise AssertionError("stub server never became ready") + yield port, model + finally: + proc.terminate() + proc.wait(timeout=10) + + +@pytest.mark.integration +def test_sglang_qa_attaches_to_a_running_server( + openai_stub: tuple[int, str], runner: CliRunner, tmp_path: Path +) -> None: + """sglang-qa.yaml binds its LM from the graph, not from --lm-url/--lm-model. + + ``--sglang-port`` re-points the declared block field-by-field, which is how + an operator aims the example at a server they already have. Reaching + ``status=done`` with the stub's answer in state proves the whole chain: + the block lowered to an ``SGLangServer``, the endpoint resolved by + attaching, and the derived base URL + model configured the DSPy LM that + the ``ask`` node ran against. + """ + port, _model = openai_stub + result = runner.invoke( + app, + [ + "run", + str(EXAMPLES_DIR / "sglang-qa.yaml"), + "--sglang-port", + str(port), + "--checkpoint", + str(tmp_path / "ck.sqlite"), + "--inputs", + "question=what routes stargraph?", + "--quiet", + "--summary-json", + ], + ) + + assert result.exit_code == 0, result.output + lines = [ln for ln in result.stdout.splitlines() if ln.strip().startswith("{")] + assert lines, f"no JSON summary in output: {result.stdout!r}" + payload = json.loads(lines[-1]) + assert payload["status"] == "done" + assert payload["state_summary"]["answer"] == _STUB_ANSWER diff --git a/tests/unit/cli/test_run_sglang.py b/tests/unit/cli/test_run_sglang.py index 8f4887bb..abb8e3f7 100644 --- a/tests/unit/cli/test_run_sglang.py +++ b/tests/unit/cli/test_run_sglang.py @@ -341,3 +341,57 @@ def test_cli_refuses_the_untrusted_halves_of_a_graph_block( assert result.exit_code != 0 assert calls == [] + + +def test_the_lm_is_configured_before_the_node_registry_is_built( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Endpoint first, nodes second -- ``kind: dspy`` needs a live LM to build. + + :func:`stargraph.nodes.dspy.dspy_node_from_config` raises "no LM + configured" while it is *constructing* the node, not when the node runs. + Resolving the endpoint after the registry was built therefore made every + ``kind: dspy`` graph unrunnable -- with a declared ``lm:`` block *and* with + ``--lm-url``/``--lm-model``. Asserting the call order pins the constraint + directly, so a future reshuffle of this function fails here rather than in + the one example that happens to use a dspy node. + """ + import stargraph.cli.run as run_mod + + order: list[str] = [] + _stub_launcher(monkeypatch, []) + + import dspy # pyright: ignore[reportMissingTypeStubs] + + def _fake_lm(*_args: Any, **_kwargs: Any) -> object: + return object() + + def _record_configure(**_kwargs: Any) -> None: + order.append("configure-lm") + + monkeypatch.setattr(dspy, "LM", _fake_lm) + monkeypatch.setattr(dspy, "configure", _record_configure) + + real_build = run_mod.build_node_registry + + def _record_build(*args: Any, **kwargs: Any) -> Any: + order.append("build-nodes") + return real_build(*args, **kwargs) + + monkeypatch.setattr(run_mod, "build_node_registry", _record_build) + + result = CliRunner().invoke( + _make_app(), + [ + str(SAMPLE_GRAPH), + "--checkpoint", + str(tmp_path / "ck.sqlite"), + "--sglang-model", + "Qwen/Qwen3-8B", + "--quiet", + "--no-summary", + ], + ) + + assert result.exit_code == 0, result.output + assert order == ["configure-lm", "build-nodes"] From aa952b1081531b6f7b685fd3cbf516c782a9d6b3 Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Mon, 17 Aug 2026 21:17:20 +0000 Subject: [PATCH 03/11] feat(lm): preflight hardware, sglang runtime and weights before spawning Spawning an SGLang server only works when three separate things line up: the machine has an accelerator, the interpreter we spawn has an sglang build *for that accelerator*, and the weights are on disk. Left to itself sglang reports all three as the same opaque subprocess death, part-way through a startup timeout. `stargraph.lm.hardware` checks each one before the spawn: * Hardware comes from the vendor tools (`nvidia-smi`, `rocm-smi`, `xpu-smi`, `npu-smi`), never from torch. `torch.cuda.is_available()` answers "was this torch built with CUDA", not "does this box have a GPU" -- a CPU-only wheel on a two-L40S machine says False, which is exactly the case that needs an install command. * The runtime is probed inside the interpreter that would be spawned, and a mismatch is reported with the command for that platform. Install commands are copied verbatim from the SGLang install docs, including the CUDA 12.9 wheel set for drivers older than r580. `--install-runtime` runs them; without it, nothing is installed and the run stops. Kernel drivers are only ever reported. ROCm, XPU, Ascend NPU and Apple Metal ship as platform builds rather than plain wheels, so those get their doc page instead of a pip line that would fail. * Weights are fetched before the server starts, so `startup_timeout_s` measures server boot rather than racing a multi-gigabyte download. * The model format is validated, not rewritten: a GGUF repo is refused with the servable repo named (that is llama.cpp's format), and an FP8 checkpoint on pre-sm_89 hardware warns. A graph runs the weights it declares -- replay depends on it. The attach branch preflights nothing: that server is already up, and it is not ours to diagnose. Also switches the example to LiquidAI/LFM2.5-1.2B-Instruct. Signed-off-by: Sean Mauk --- docs/how-to/authoring-format.md | 22 + docs/reference/cli.md | 8 +- examples/sglang-qa.yaml | 5 +- src/stargraph/cli/run.py | 20 +- src/stargraph/lm/hardware.py | 580 ++++++++++++++++++++++ src/stargraph/lm/sglang.py | 12 + tests/integration/test_lm_sglang_spawn.py | 19 + tests/unit/cli/test_run_sglang.py | 44 ++ tests/unit/test_lm_hardware.py | 404 +++++++++++++++ tests/unit/test_lm_sglang.py | 18 + 10 files changed, 1127 insertions(+), 5 deletions(-) create mode 100644 src/stargraph/lm/hardware.py create mode 100644 tests/unit/test_lm_hardware.py diff --git a/docs/how-to/authoring-format.md b/docs/how-to/authoring-format.md index a0adbc6e..53657d2c 100644 --- a/docs/how-to/authoring-format.md +++ b/docs/how-to/authoring-format.md @@ -89,6 +89,28 @@ default-deny capability gate): The block is not part of the graph hash: an endpoint is an environment binding, not topology. +Before it launches anything, `stargraph run` checks that the machine and the +interpreter agree: + +- **Hardware** is read from the vendor tools (`nvidia-smi`, `rocm-smi`, + `xpu-smi`, `npu-smi`), never from torch. `torch.cuda.is_available()` answers + "was this torch built with CUDA", not "does this box have a GPU" -- a + CPU-only wheel on a two-GPU machine says `False`. +- **The runtime** is probed inside the interpreter that would be spawned. A + missing sglang, or a torch built for the wrong accelerator, is reported with + the install command for *that* platform. `--install-runtime` runs it; + without the flag nothing is installed and the run stops. Kernel drivers are + never touched -- a missing or too-old CUDA/ROCm driver can only be reported. + SGLang publishes plain wheels for NVIDIA only, so ROCm, XPU, Ascend NPU and + Apple Metal are reported with a pointer to their platform page rather than a + command that would not work. +- **The weights** are fetched before the server starts, so `startup_timeout_s` + measures server boot rather than racing a multi-gigabyte download. +- **The format** is validated, not rewritten. A GGUF repo is refused (that is + llama.cpp's format; sglang serves safetensors) with the servable repo named, + and an FP8 checkpoint on pre-sm_89 hardware warns. The graph always runs the + weights it declares. + ### `routes` Declaration order is the default flow: with no rule firing, execution diff --git a/docs/reference/cli.md b/docs/reference/cli.md index b68c7d76..5a01baea 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -82,9 +82,15 @@ on `failed`. | `--sglang-port PORT` | int | `30000` | SGLang port. | | `--sglang-arg ARG` | str (repeatable) | _(empty)_ | Extra argv passed through to `sglang.launch_server` verbatim. | | `--sglang-timeout SEC` | int | `600` | Seconds to wait for a launched SGLang server to answer. | +| `--install-runtime` | flag | `false` | Install the sglang build matching the detected accelerator before launching. | `--quiet` and `--verbose` are mutually exclusive. `--lm-url` and -`--lm-model` must be supplied together (or neither). +`--lm-model` must be supplied together (or neither). A spawned SGLang server is +preflighted first: hardware is detected from the vendor tools, the spawn +interpreter's sglang/torch build is checked against it, and the weights are +fetched before the startup clock starts. Without `--install-runtime` a runtime +that cannot serve is reported with the exact install command and the run stops; +nothing is ever installed implicitly, and kernel drivers are never touched. The `--sglang-*` flags bind the run to a local [SGLang](https://docs.sglang.ai/) server, and derive `--lm-url` / diff --git a/examples/sglang-qa.yaml b/examples/sglang-qa.yaml index 3f5cca0f..a10f4111 100644 --- a/examples/sglang-qa.yaml +++ b/examples/sglang-qa.yaml @@ -9,7 +9,8 @@ # it launches `python -m sglang.launch_server`, waits for the endpoint to # answer, and terminates it when the run ends. # -# Run it (SGLang installed, weights downloadable): +# Run it (weights are fetched on first use; add --install-runtime to have +# stargraph install the sglang build matching your GPU): # stargraph run examples/sglang-qa.yaml --inputs question="what routes stargraph?" # # Point it at a server you already have, without editing this file: @@ -22,7 +23,7 @@ id: sglang-qa lm: provider: sglang # the only provider today - model: Qwen/Qwen2.5-0.5B-Instruct + model: LiquidAI/LFM2.5-1.2B-Instruct port: 30000 # sglang's own default startup_timeout_s: 600 # first run downloads weights state: diff --git a/src/stargraph/cli/run.py b/src/stargraph/cli/run.py index 3a147989..d8835009 100644 --- a/src/stargraph/cli/run.py +++ b/src/stargraph/cli/run.py @@ -205,6 +205,7 @@ def _lm_endpoint( lm_key: str, lm_timeout: int, echo: Callable[[str], None], + install_runtime: bool = False, ) -> Generator[None]: """Hold the run's LM endpoint open: boot/attach sglang, then configure dspy. @@ -217,7 +218,9 @@ def _lm_endpoint( if spec is not None: from stargraph.lm.sglang import sglang_server - lm_url = stack.enter_context(sglang_server(spec, echo=echo)) + lm_url = stack.enter_context( + sglang_server(spec, echo=echo, install_runtime=install_runtime) + ) lm_model = spec.model _configure_lm(lm_url, lm_model, lm_key, lm_timeout) yield @@ -429,6 +432,17 @@ def cmd( help="Seconds to wait for a launched SGLang server to answer (default: 600).", ), ] = None, + install_runtime: Annotated[ + bool, + typer.Option( + "--install-runtime", + help=( + "Install the sglang build matching the detected accelerator before " + "launching. Without it, a runtime that cannot serve is reported with " + "the exact command instead of being repaired." + ), + ), + ] = False, ) -> None: """Run a Stargraph graph end-to-end (FR-8 POC). @@ -527,7 +541,9 @@ def _echo(message: str) -> None: # be live by then. It outlives the whole run as well -- a spawned # sglang server is torn down only once the loop is done (or has # raised). - with _lm_endpoint(sglang_spec, lm_url, lm_model, lm_key, lm_timeout, _echo): + with _lm_endpoint( + sglang_spec, lm_url, lm_model, lm_key, lm_timeout, _echo, install_runtime + ): try: node_registry = build_node_registry(ir.nodes, ir_dir=graph.parent.resolve()) except StargraphError as e: diff --git a/src/stargraph/lm/hardware.py b/src/stargraph/lm/hardware.py new file mode 100644 index 00000000..4f60655a --- /dev/null +++ b/src/stargraph/lm/hardware.py @@ -0,0 +1,580 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Detect the accelerator, check the sglang runtime against it, fetch weights. + +``stargraph run`` can spawn an SGLang server (:mod:`stargraph.lm.sglang`), and +that only works when three things line up: the machine has an accelerator, the +interpreter we spawn has an sglang build *for that accelerator*, and the +weights are on disk. Each has its own failure mode, and left to itself sglang +reports all three as the same opaque subprocess death. + +The split that matters here is **hardware vs runtime**. ``torch.cuda.is_available()`` +answers "was this torch built with CUDA", not "does this box have a GPU" -- a +CPU-only wheel on a machine with two L40S reports ``False``. So hardware is +detected from the vendor tools (``nvidia-smi``, ``rocm-smi``, ``xpu-smi``, +``npu-smi``), and the runtime is probed separately inside the target +interpreter. A mismatch between the two is the interesting case, and the one +that produces an actionable install command. + +What this module will *not* do: + +* **Install kernel drivers.** A missing or too-old CUDA/ROCm driver is a + system-level, root-owned concern. It is reported, never repaired. +* **Install anything without being asked.** Repair runs only under + ``--install-runtime``; otherwise the plan is printed and the run stops. +* **Rewrite the model.** A graph names the weights it runs (replay depends on + it). An unservable format is refused with the servable equivalent named, not + silently swapped. +""" + +from __future__ import annotations + +import json +import platform +import shutil +import subprocess +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Literal + +from stargraph.errors import LMServerError + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from stargraph.ir import SGLangServer + +__all__ = [ + "Accelerator", + "InstallPlan", + "Runtime", + "check_model_format", + "detect_accelerator", + "ensure_runtime", + "ensure_weights", + "plan_runtime_install", + "probe_runtime", +] + +Vendor = Literal["nvidia", "amd", "intel", "ascend", "apple", "none"] + +_PROBE_TIMEOUT_S = 20.0 +_INSTALL_TIMEOUT_S = 3600.0 + +# CUDA 13 wheels need a r580+ driver (13.0 minimum is 580.65.06 on Linux); an +# older driver has to stay on the CUDA 12 wheel set. +_CUDA13_MIN_DRIVER_MAJOR = 580 + + +@dataclass(frozen=True, slots=True) +class Accelerator: + """What the *machine* has, as reported by the vendor tooling.""" + + vendor: Vendor + devices: tuple[str, ...] = () + arch: str = "" + """Compute capability (NVIDIA, e.g. ``8.9``) or gfx target (AMD).""" + driver: str = "" + + @property + def count(self) -> int: + return len(self.devices) + + def describe(self) -> str: + if self.vendor == "none": + return "no accelerator detected" + if not self.devices: + return self.vendor + name = self.devices[0] + suffix = f" x{self.count}" if self.count > 1 else "" + arch = f" ({self.arch})" if self.arch else "" + driver = f", driver {self.driver}" if self.driver else "" + return f"{name}{suffix}{arch}{driver}" + + +@dataclass(frozen=True, slots=True) +class Runtime: + """What the *interpreter we would spawn* has installed.""" + + python: str + sglang: str | None = None + torch: str | None = None + torch_cuda: str | None = None + """``torch.version.cuda`` -- set on CUDA builds.""" + torch_hip: str | None = None + """``torch.version.hip`` -- set on ROCm builds.""" + torch_xpu: bool = False + device_count: int = 0 + + @property + def backend(self) -> str: + """The accelerator family this torch build can actually drive.""" + if self.torch_hip: + return "rocm" + if self.torch_cuda: + return "cuda" + if self.torch_xpu: + return "xpu" + return "cpu" + + def describe(self) -> str: + if self.torch is None: + return "torch not installed" + return f"torch {self.torch} ({self.backend} build), {self.device_count} device(s) visible" + + +@dataclass(frozen=True, slots=True) +class InstallPlan: + """How to repair ``Runtime`` for the detected :class:`Accelerator`.""" + + reason: str + """Why the current runtime cannot serve -- shown to the operator verbatim.""" + commands: tuple[tuple[str, ...], ...] = () + """argv sequences to run, in order. Empty when repair is not automatable.""" + manual: str = "" + """Doc pointer for platforms with no documented pip path.""" + notes: tuple[str, ...] = field(default=()) + + @property + def automatable(self) -> bool: + return bool(self.commands) + + +def _run(argv: Sequence[str], *, timeout: float = _PROBE_TIMEOUT_S) -> str | None: + """stdout of ``argv``, or ``None`` if the binary is absent or it failed. + + Every hardware probe is best-effort: a missing vendor tool means "not this + vendor", never an error. + """ + if shutil.which(argv[0]) is None: + return None + try: + completed = subprocess.run( + list(argv), capture_output=True, text=True, timeout=timeout, check=False + ) + except (OSError, subprocess.SubprocessError): + return None + if completed.returncode != 0: + return None + return completed.stdout + + +def _detect_nvidia() -> Accelerator | None: + out = _run( + [ + "nvidia-smi", + "--query-gpu=name,compute_cap,driver_version", + "--format=csv,noheader", + ] + ) + if not out or not out.strip(): + return None + names: list[str] = [] + arch = "" + driver = "" + for line in out.strip().splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) < 3: + continue + names.append(parts[0]) + arch, driver = parts[1], parts[2] + if not names: + return None + return Accelerator(vendor="nvidia", devices=tuple(names), arch=arch, driver=driver) + + +def _detect_amd() -> Accelerator | None: + out = _run(["rocm-smi", "--showproductname", "--csv"]) + if not out or not out.strip(): + return None + names = [ + line.split(",")[1].strip() + for line in out.strip().splitlines()[1:] + if len(line.split(",")) > 1 and line.split(",")[1].strip() + ] + if not names: + return None + info = _run(["rocminfo"]) or "" + arch = next( + ( + tok + for line in info.splitlines() + if "gfx" in line + for tok in line.split() + if "gfx" in tok + ), + "", + ) + return Accelerator(vendor="amd", devices=tuple(names), arch=arch) + + +def _detect_intel() -> Accelerator | None: + out = _run(["xpu-smi", "discovery"]) + if not out or "Device" not in out: + return None + names = [ + line.split("|", 2)[-1].strip() for line in out.splitlines() if "Device Name" in line + ] or ["Intel XPU"] + return Accelerator(vendor="intel", devices=tuple(names)) + + +def _detect_ascend() -> Accelerator | None: + out = _run(["npu-smi", "info"]) + if not out or "NPU" not in out: + return None + return Accelerator(vendor="ascend", devices=("Ascend NPU",)) + + +def detect_accelerator() -> Accelerator: + """What this machine has, from the vendor tools -- never from torch. + + Order is by how unambiguous the probe is; a box with two vendors' tools + installed is vanishingly rare next to the cost of guessing wrong. + """ + for probe in (_detect_nvidia, _detect_amd, _detect_intel, _detect_ascend): + found = probe() + if found is not None: + return found + if platform.system() == "Darwin" and platform.machine() in {"arm64", "aarch64"}: + return Accelerator(vendor="apple", devices=(f"Apple silicon ({platform.machine()})",)) + return Accelerator(vendor="none") + + +_RUNTIME_PROBE = """ +import json, importlib.util +out = {"sglang": None, "torch": None, "torch_cuda": None, "torch_hip": None, + "torch_xpu": False, "device_count": 0} +spec = importlib.util.find_spec("sglang") +if spec is not None: + try: + from importlib.metadata import version + out["sglang"] = version("sglang") + except Exception: + out["sglang"] = "unknown" +if importlib.util.find_spec("torch") is not None: + try: + import torch + out["torch"] = torch.__version__ + out["torch_cuda"] = torch.version.cuda + out["torch_hip"] = torch.version.hip + out["torch_xpu"] = bool(getattr(torch, "xpu", None) and torch.xpu.is_available()) + out["device_count"] = torch.cuda.device_count() if torch.cuda.is_available() else 0 + except Exception: + pass +print(json.dumps(out)) +""" + + +def probe_runtime(python: str | None = None) -> Runtime: + """What ``python`` (default: the interpreter that would be spawned) has. + + ``sglang`` is located with :func:`importlib.util.find_spec` rather than + imported -- importing it pulls in torch and CUDA context on a path whose + only job is to answer "is it there". + """ + interpreter = python or sys.executable + out = _run([interpreter, "-c", _RUNTIME_PROBE]) + if not out: + return Runtime(python=interpreter) + try: + data = json.loads(out) + except json.JSONDecodeError: + return Runtime(python=interpreter) + return Runtime( + python=interpreter, + sglang=data["sglang"], + torch=data["torch"], + torch_cuda=data["torch_cuda"], + torch_hip=data["torch_hip"], + torch_xpu=bool(data["torch_xpu"]), + device_count=int(data["device_count"]), + ) + + +def _installer(python: str) -> tuple[str, ...]: + """``uv pip install`` targeting ``python`` when uv is around, else pip.""" + if shutil.which("uv") is not None: + return ("uv", "pip", "install", "--python", python) + return (python, "-m", "pip", "install") + + +def _nvidia_plan(accel: Accelerator, runtime: Runtime, install: tuple[str, ...]) -> InstallPlan: + """CUDA 13 wheels on a r580+ driver; the pinned CUDA 12 set below that.""" + driver_major = 0 + if accel.driver: + head = accel.driver.split(".")[0] + driver_major = int(head) if head.isdigit() else 0 + reason = ( + f"sglang is not installed for {runtime.python}" + if runtime.sglang is None + else f"torch in {runtime.python} is a {runtime.backend} build, " + f"but {accel.describe()} needs cuda" + ) + if driver_major >= _CUDA13_MIN_DRIVER_MAJOR: + return InstallPlan( + reason=reason, + commands=((*install, "--prerelease=allow", "sglang"),), + notes=(f"driver {accel.driver} supports the CUDA 13 wheels",), + ) + return InstallPlan( + reason=reason, + commands=( + (*install, "--prerelease=allow", "sglang"), + ( + *install, + "--force-reinstall", + "torch==2.13.0", + "torchaudio==2.11.0", + "torchvision", + "--index-url", + "https://download.pytorch.org/whl/cu129", + ), + ( + *install, + "--force-reinstall", + "sglang-kernel", + "--index-url", + "https://docs.sglang.ai/whl/cu129/", + ), + ( + *install, + "--force-reinstall", + "sgl-deep-gemm", + "--index-url", + "https://docs.sglang.ai/whl/cu129/", + "--no-deps", + ), + ), + notes=( + f"driver {accel.driver or 'unknown'} predates r{_CUDA13_MIN_DRIVER_MAJOR}, " + "so the CUDA 12.9 wheel set is used", + ), + ) + + +_INSTALL_DOC = "https://docs.sglang.io/docs/get-started/install" + +_MANUAL_PLATFORMS: dict[Vendor, str] = { + "amd": "https://docs.sglang.io/docs/hardware-platforms/amd_gpu", + "intel": "https://docs.sglang.io/docs/hardware-platforms/xpu", + "ascend": "https://docs.sglang.io/docs/hardware-platforms/ascend-npus/getting-started/installation", + "apple": "https://docs.sglang.io/docs/hardware-platforms/apple_metal", +} + + +def plan_runtime_install(accel: Accelerator, runtime: Runtime) -> InstallPlan | None: + """How to make ``runtime`` able to serve on ``accel`` -- ``None`` if it already can. + + Only NVIDIA has an install path SGLang documents as plain wheels. ROCm, + XPU, Ascend NPU and Apple Metal ship through platform-specific pyprojects, + docker images, or a source build against the Apple toolchain, so those + return a plan that names the page instead of pretending a ``pip install`` + would work. + """ + if runtime.sglang is not None and _backend_matches(accel, runtime): + return None + + install = _installer(runtime.python) + if accel.vendor == "nvidia": + return _nvidia_plan(accel, runtime, install) + if accel.vendor in _MANUAL_PLATFORMS: + return InstallPlan( + reason=( + f"sglang is not installed for {runtime.python}" + if runtime.sglang is None + else f"torch in {runtime.python} is a {runtime.backend} build, " + f"but {accel.describe()} needs {_expected_backend(accel)}" + ), + manual=_MANUAL_PLATFORMS[accel.vendor], + notes=("sglang ships this platform as its own build, not a plain wheel",), + ) + # No accelerator found. sglang can serve on CPU, so this is a warning the + # caller decides about -- not a plan. + if runtime.sglang is None: + return InstallPlan( + reason=f"sglang is not installed for {runtime.python}", + commands=((*install, "sglang"),), + notes=("no accelerator detected; a CPU install serves, slowly",), + ) + return None + + +def _expected_backend(accel: Accelerator) -> str: + return {"nvidia": "cuda", "amd": "rocm", "intel": "xpu", "ascend": "npu"}.get( + accel.vendor, "cpu" + ) + + +def _backend_matches(accel: Accelerator, runtime: Runtime) -> bool: + """Does this torch build drive this hardware? + + ``device_count`` is the tiebreaker: a CUDA-built torch that sees zero + devices is as unusable as a CPU one, and that is what a container missing + ``--gpus all`` looks like. + """ + expected = _expected_backend(accel) + if expected == "cpu": + return True + if runtime.backend != expected: + return False + return runtime.device_count > 0 or expected == "npu" + + +def check_model_format(model: str, accel: Accelerator) -> str | None: + """Refuse formats sglang cannot serve; warn on quantization the GPU lacks. + + Returns a warning string, or ``None`` when nothing is worth saying. Raises + :class:`~stargraph.errors.LMServerError` for a model that cannot work at + all -- GGUF being the trap, since most on-device model families publish a + GGUF repo beside the safetensors one and it is the format LM Studio and + llama.cpp take. + """ + tail = model.rsplit("/", 1)[-1] + if tail.lower().endswith("-gguf") or ".gguf" in tail.lower(): + servable = model[: -len("-GGUF")] if tail.lower().endswith("-gguf") else model + raise LMServerError( + f"{model!r} is a GGUF repo; sglang serves safetensors, not GGUF", + hint=f"use {servable!r} (or serve the GGUF with llama.cpp and pass --lm-url)", + model=model, + ) + if "fp8" in tail.lower() and accel.vendor == "nvidia" and accel.arch: + try: + capability = float(accel.arch) + except ValueError: + return None + if capability < 8.9: + return ( + f"{model} is an FP8 checkpoint but {accel.describe()} is " + f"sm_{accel.arch.replace('.', '')}; FP8 needs sm_89 or newer, " + "so sglang will fall back or fail" + ) + return None + + +def ensure_runtime( + spec: SGLangServer, + *, + python: str | None = None, + install: bool = False, + echo: Callable[[str], None] | None = None, +) -> None: + """Check hardware + runtime before a spawn; optionally repair the runtime. + + Raises :class:`~stargraph.errors.LMServerError` when the runtime cannot + serve and ``install`` is false, so the operator sees the exact command + rather than a subprocess that dies during startup. + """ + + def _say(message: str) -> None: + if echo is not None: + echo(message) + + accel = detect_accelerator() + _say(f"hardware: {accel.describe()}") + warning = check_model_format(spec.model, accel) + if warning is not None: + _say(f"warning: {warning}") + + runtime = probe_runtime(python) + _say( + f"runtime: {runtime.describe()}" + (f", sglang {runtime.sglang}" if runtime.sglang else "") + ) + plan = plan_runtime_install(accel, runtime) + if plan is None: + return + + if not plan.automatable: + raise LMServerError( + plan.reason, + hint=f"sglang documents this platform at {plan.manual}", + notes="; ".join(plan.notes), + ) + rendered = "\n".join(" " + " ".join(command) for command in plan.commands) + if not install: + raise LMServerError( + plan.reason, + hint=f"re-run with --install-runtime, or install it yourself:\n{rendered}", + notes="; ".join(plan.notes), + ) + + for note in plan.notes: + _say(f"note: {note}") + for index, command in enumerate(plan.commands, start=1): + _say(f"installing ({index}/{len(plan.commands)}): {' '.join(command)}") + completed = subprocess.run(list(command), timeout=_INSTALL_TIMEOUT_S, check=False) + if completed.returncode != 0: + raise LMServerError( + f"runtime install failed: {' '.join(command)}", + hint="install it manually, then re-run without --install-runtime", + exit_code=str(completed.returncode), + ) + repaired = probe_runtime(python) + if plan_runtime_install(accel, repaired) is not None: + raise LMServerError( + f"runtime install completed but the runtime still cannot serve ({repaired.describe()})", + hint=f"sglang documents this platform at {_INSTALL_DOC}", + python=repaired.python, + ) + _say(f"runtime ready: {repaired.describe()}") + + +_WEIGHTS_PROBE = """ +import json, sys +model = sys.argv[1] +local_only = sys.argv[2] == "cached" +try: + from huggingface_hub import snapshot_download +except Exception: + print(json.dumps({"status": "no-hub"})) + raise SystemExit(0) +try: + path = snapshot_download(model, local_files_only=local_only) +except Exception as exc: + print(json.dumps({"status": "miss", "error": type(exc).__name__})) + raise SystemExit(0) +print(json.dumps({"status": "ok", "path": path})) +""" + + +def ensure_weights( + spec: SGLangServer, + *, + python: str | None = None, + echo: Callable[[str], None] | None = None, +) -> None: + """Fetch the weights before the server starts, so the boot timeout is a boot timeout. + + sglang downloads on first use inside its own startup, which means a + multi-gigabyte fetch is racing ``startup_timeout_s`` -- the run dies + part-way through a download that would have succeeded. Pulling the + snapshot first separates "weights are arriving" from "the server is + wedged". + + Best-effort by design: a local path, an unreachable hub, or an absent + ``huggingface_hub`` all fall through to sglang's own behaviour rather than + blocking the run. + """ + + def _say(message: str) -> None: + if echo is not None: + echo(message) + + if Path(spec.model).exists(): + return + interpreter = python or sys.executable + cached = _run([interpreter, "-c", _WEIGHTS_PROBE, spec.model, "cached"]) + if cached and json.loads(cached).get("status") == "ok": + _say(f"weights: {spec.model} already in the local hub cache") + return + _say(f"weights: fetching {spec.model} (not in the local hub cache)") + fetched = _run( + [interpreter, "-c", _WEIGHTS_PROBE, spec.model, "fetch"], timeout=_INSTALL_TIMEOUT_S + ) + if not fetched: + _say("weights: prefetch unavailable; sglang will download during startup") + return + result = json.loads(fetched) + if result.get("status") == "ok": + _say(f"weights: ready at {result['path']}") + else: + _say("weights: prefetch did not complete; sglang will download during startup") diff --git a/src/stargraph/lm/sglang.py b/src/stargraph/lm/sglang.py index 75c7317d..5f4f0dfd 100644 --- a/src/stargraph/lm/sglang.py +++ b/src/stargraph/lm/sglang.py @@ -41,6 +41,7 @@ import httpx from stargraph.errors import LMServerError +from stargraph.lm.hardware import ensure_runtime, ensure_weights if TYPE_CHECKING: from collections.abc import Callable, Generator @@ -195,6 +196,7 @@ def sglang_server( *, log_path: Path | None = None, echo: Callable[[str], None] | None = None, + install_runtime: bool = False, ) -> Generator[str]: """Yield the base URL of a server for ``spec``, booting one if needed. @@ -203,6 +205,13 @@ def sglang_server( its process group) when the block ends. ``log_path`` receives the spawned server's stdout+stderr (default: a temp file); ``echo`` gets one-line progress messages. + + The spawn branch runs :func:`~stargraph.lm.hardware.ensure_runtime` and + :func:`~stargraph.lm.hardware.ensure_weights` first, so an unusable runtime + is reported as such (with the install command for the detected hardware, + run for you under ``install_runtime``) and the weights are on disk before + ``startup_timeout_s`` starts counting. The attach branch skips both: that + server is already up, and it is not ours to diagnose. """ url = base_url(spec) existing = served_models(url) @@ -218,6 +227,9 @@ def sglang_server( yield url return + ensure_runtime(spec, install=install_runtime, echo=echo) + ensure_weights(spec, echo=echo) + if log_path is None: fd, name = tempfile.mkstemp(prefix=f"sglang-{spec.port}-", suffix=".log") os.close(fd) diff --git a/tests/integration/test_lm_sglang_spawn.py b/tests/integration/test_lm_sglang_spawn.py index 94f86ce6..31aefbde 100644 --- a/tests/integration/test_lm_sglang_spawn.py +++ b/tests/integration/test_lm_sglang_spawn.py @@ -24,6 +24,25 @@ pytestmark = pytest.mark.integration + +@pytest.fixture(autouse=True) +def _skip_preflight(monkeypatch: pytest.MonkeyPatch) -> None: # pyright: ignore[reportUnusedFunction] + """Bypass the hardware/runtime preflight for the lifecycle tests. + + These drive a stub launcher on a box that deliberately has no sglang + installed, which is exactly what :func:`stargraph.lm.hardware.ensure_runtime` + exists to refuse. The preflight is covered on its own in + ``tests/unit/test_lm_hardware.py``; here it would only assert that the dev + venv is a dev venv. + """ + + def _noop(*_args: object, **_kwargs: object) -> None: + return None + + monkeypatch.setattr(sg, "ensure_runtime", _noop) + monkeypatch.setattr(sg, "ensure_weights", _noop) + + _STUB = textwrap.dedent( ''' """Minimal OpenAI-compatible stub: GET /v1/models -> one model id.""" diff --git a/tests/unit/cli/test_run_sglang.py b/tests/unit/cli/test_run_sglang.py index abb8e3f7..f4637f61 100644 --- a/tests/unit/cli/test_run_sglang.py +++ b/tests/unit/cli/test_run_sglang.py @@ -395,3 +395,47 @@ def _record_build(*args: Any, **kwargs: Any) -> Any: assert result.exit_code == 0, result.output assert order == ["configure-lm", "build-nodes"] + + +def _capture_launcher(monkeypatch: pytest.MonkeyPatch, seen: dict[str, Any]) -> None: + import stargraph.lm.sglang as sglang_mod + + @contextlib.contextmanager + def _fake(spec: SGLangServer, **kwargs: Any) -> Generator[str]: + seen.update(kwargs) + seen["spec"] = spec + yield "http://stub:41002/v1" + + monkeypatch.setattr(sglang_mod, "sglang_server", _fake) + + +@pytest.mark.parametrize("flag", [True, False]) +def test_install_runtime_flag_reaches_the_launcher( + flag: bool, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """``--install-runtime`` is the only thing that lets a run mutate the environment. + + Parametrized rather than asserted one-way: a wiring bug that hard-codes + ``True`` is exactly as bad as one that hard-codes ``False``, and only the + false case catches the first. + """ + seen: dict[str, Any] = {} + _capture_launcher(monkeypatch, seen) + _stub_dspy(monkeypatch, {}) + + argv = [ + str(SAMPLE_GRAPH), + "--checkpoint", + str(tmp_path / "ck.sqlite"), + "--sglang-model", + "LiquidAI/LFM2.5-1.2B-Instruct", + "--quiet", + "--no-summary", + ] + if flag: + argv.append("--install-runtime") + + result = CliRunner().invoke(_make_app(), argv) + + assert result.exit_code == 0, result.output + assert seen["install_runtime"] is flag diff --git a/tests/unit/test_lm_hardware.py b/tests/unit/test_lm_hardware.py new file mode 100644 index 00000000..e4da4e83 --- /dev/null +++ b/tests/unit/test_lm_hardware.py @@ -0,0 +1,404 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Accelerator detection, runtime planning, format refusal (``stargraph.lm.hardware``). + +The point of the module under test is that *hardware* and *runtime* are +separate questions, so the tests keep them separate too: detection is driven +by canned vendor-tool output, and planning is driven by explicit +:class:`Runtime` values. Nothing here needs a GPU, and nothing installs +anything -- the one test that exercises the install path asserts on the argv +it would have run. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import pytest + +from stargraph.errors import LMServerError +from stargraph.ir import SGLangServer +from stargraph.lm import hardware as hw + +if TYPE_CHECKING: + from collections.abc import Sequence + +pytestmark = pytest.mark.unit + +# Real `nvidia-smi --query-gpu=name,compute_cap,driver_version --format=csv,noheader` +# output from a two-L40S box. +_NVIDIA_SMI = "NVIDIA L40S, 8.9, 580.173.02\nNVIDIA L40S, 8.9, 580.173.02\n" +_ROCM_SMI = "device,Card series\ncard0,Instinct MI300X\n" + + +def _linux() -> str: + return "Linux" + + +def _darwin() -> str: + return "Darwin" + + +def _arm64() -> str: + return "arm64" + + +def _fake_run( + monkeypatch: pytest.MonkeyPatch, table: dict[str, str | None] +) -> list[tuple[str, ...]]: + """Route :func:`hardware._run` through ``table`` keyed by argv[0]; record calls.""" + seen: list[tuple[str, ...]] = [] + + def _run(argv: Sequence[str], *, timeout: float = 0.0) -> str | None: + del timeout + seen.append(tuple(argv)) + return table.get(argv[0]) + + monkeypatch.setattr(hw, "_run", _run) + return seen + + +# --------------------------------------------------------------------------- # +# Detection # +# --------------------------------------------------------------------------- # + + +def test_nvidia_is_detected_with_arch_and_driver(monkeypatch: pytest.MonkeyPatch) -> None: + _fake_run(monkeypatch, {"nvidia-smi": _NVIDIA_SMI}) + accel = hw.detect_accelerator() + assert (accel.vendor, accel.count, accel.arch, accel.driver) == ( + "nvidia", + 2, + "8.9", + "580.173.02", + ) + assert accel.describe() == "NVIDIA L40S x2 (8.9), driver 580.173.02" + + +def test_amd_is_detected(monkeypatch: pytest.MonkeyPatch) -> None: + _fake_run(monkeypatch, {"rocm-smi": _ROCM_SMI, "rocminfo": " Name: gfx942"}) + accel = hw.detect_accelerator() + assert (accel.vendor, accel.devices, accel.arch) == ("amd", ("Instinct MI300X",), "gfx942") + + +def test_no_vendor_tool_means_no_accelerator(monkeypatch: pytest.MonkeyPatch) -> None: + """The CPU-only case must not be reported as a vendor with zero devices.""" + _fake_run(monkeypatch, {}) + monkeypatch.setattr(hw.platform, "system", _linux) + accel = hw.detect_accelerator() + assert accel.vendor == "none" + assert accel.describe() == "no accelerator detected" + + +def test_apple_silicon_is_named_rather_than_probed(monkeypatch: pytest.MonkeyPatch) -> None: + _fake_run(monkeypatch, {}) + monkeypatch.setattr(hw.platform, "system", _darwin) + monkeypatch.setattr(hw.platform, "machine", _arm64) + assert hw.detect_accelerator().vendor == "apple" + + +def test_detection_does_not_consult_torch(monkeypatch: pytest.MonkeyPatch) -> None: + """Hardware comes from the vendor tools, never from whatever torch was installed. + + This is the bug the module exists for: a CPU-only torch wheel on a box with + two L40S answers ``torch.cuda.is_available() == False``. If detection ever + starts asking torch, this test sees ``python`` in the probe argv. + """ + seen = _fake_run(monkeypatch, {"nvidia-smi": _NVIDIA_SMI}) + assert hw.detect_accelerator().vendor == "nvidia" + assert all("python" not in argv[0] for argv in seen), seen + + +# --------------------------------------------------------------------------- # +# Runtime probe # +# --------------------------------------------------------------------------- # + + +def test_probe_reports_this_interpreter_honestly() -> None: + """No stubbing: the dev venv genuinely has no sglang, and torch is CPU-only. + + Asserting against the real interpreter keeps the probe script itself under + test -- a syntax error or renamed key in ``_RUNTIME_PROBE`` fails here. + """ + runtime = hw.probe_runtime() + assert runtime.sglang is None + assert runtime.backend in {"cpu", "cuda", "rocm", "xpu"} + + +def test_probe_of_a_broken_interpreter_degrades(monkeypatch: pytest.MonkeyPatch) -> None: + _fake_run(monkeypatch, {}) + runtime = hw.probe_runtime("/nonexistent/python") + assert runtime == hw.Runtime(python="/nonexistent/python") + + +# --------------------------------------------------------------------------- # +# Planning # +# --------------------------------------------------------------------------- # + +_NVIDIA = hw.Accelerator(vendor="nvidia", devices=("NVIDIA L40S",), arch="8.9", driver="580.173.02") +_NVIDIA_OLD_DRIVER = hw.Accelerator( + vendor="nvidia", devices=("NVIDIA A100",), arch="8.0", driver="550.54.15" +) +_WORKING_CUDA = hw.Runtime( + python="/venv/bin/python", sglang="0.5.0", torch="2.13.0", torch_cuda="13.0", device_count=2 +) + + +def _uv_on_path(_name: str) -> str: + return "/usr/bin/uv" + + +def _nvidia_box() -> hw.Accelerator: + return _NVIDIA + + +def _cuda_runtime(_python: str | None = None) -> hw.Runtime: + return _WORKING_CUDA + + +def test_a_working_cuda_runtime_needs_no_plan() -> None: + """Anti-vacuity: planning must be able to say "nothing to do".""" + assert hw.plan_runtime_install(_NVIDIA, _WORKING_CUDA) is None + + +def test_cpu_torch_on_an_nvidia_box_is_planned_as_cuda13(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(hw.shutil, "which", _uv_on_path) + runtime = hw.Runtime(python="/venv/bin/python", sglang="0.5.0", torch="2.11.0+cpu") + plan = hw.plan_runtime_install(_NVIDIA, runtime) + assert plan is not None + assert plan.automatable + assert plan.commands == ( + ("uv", "pip", "install", "--python", "/venv/bin/python", "--prerelease=allow", "sglang"), + ) + assert "cpu build" in plan.reason + + +def test_a_pre_r580_driver_gets_the_cuda12_wheel_set(monkeypatch: pytest.MonkeyPatch) -> None: + """CUDA 13 wheels require r580+; an older driver must not be handed them.""" + monkeypatch.setattr(hw.shutil, "which", _uv_on_path) + plan = hw.plan_runtime_install(_NVIDIA_OLD_DRIVER, hw.Runtime(python="/venv/bin/python")) + assert plan is not None + joined = [" ".join(command) for command in plan.commands] + # verbatim from https://docs.sglang.io/get_started/install.html (CUDA 12 path) + assert any("whl/cu129" in command for command in joined), joined + assert any("torch==2.13.0 torchaudio==2.11.0 torchvision" in command for command in joined), ( + joined + ) + assert any("sglang-kernel" in command for command in joined), joined + assert any("sgl-deep-gemm" in command and "--no-deps" in command for command in joined), joined + + +def test_cuda_torch_that_sees_no_devices_still_needs_a_plan( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A container started without --gpus all: right wheel, no devices.""" + monkeypatch.setattr(hw.shutil, "which", _uv_on_path) + runtime = hw.Runtime( + python="/venv/bin/python", sglang="0.5.0", torch="2.13.0", torch_cuda="13.0", device_count=0 + ) + assert hw.plan_runtime_install(_NVIDIA, runtime) is not None + + +def test_amd_is_planned_but_not_automated() -> None: + """ROCm ships as its own build, so a pip command would be a lie.""" + accel = hw.Accelerator(vendor="amd", devices=("Instinct MI300X",), arch="gfx942") + plan = hw.plan_runtime_install(accel, hw.Runtime(python="/venv/bin/python")) + assert plan is not None + assert not plan.automatable + assert "amd_gpu" in plan.manual + + +def test_apple_silicon_gets_the_metal_page_not_a_pip_command() -> None: + """sglang runs on Metal/MLX, but only from a source build against Xcode.""" + accel = hw.Accelerator(vendor="apple", devices=("Apple silicon (arm64)",)) + plan = hw.plan_runtime_install(accel, hw.Runtime(python="/venv/bin/python")) + assert plan is not None + assert not plan.automatable + assert "apple_metal" in plan.manual + + +def test_apple_silicon_with_sglang_already_installed_needs_no_plan() -> None: + accel = hw.Accelerator(vendor="apple", devices=("Apple silicon (arm64)",)) + runtime = hw.Runtime(python="/venv/bin/python", sglang="0.5.0", torch="2.13.0") + assert hw.plan_runtime_install(accel, runtime) is None + + +def test_a_cpu_box_with_sglang_installed_is_fine() -> None: + """sglang serves on CPU; "no GPU" is not by itself a failure.""" + runtime = hw.Runtime(python="/venv/bin/python", sglang="0.5.0", torch="2.13.0") + assert hw.plan_runtime_install(hw.Accelerator(vendor="none"), runtime) is None + + +# --------------------------------------------------------------------------- # +# Model format # +# --------------------------------------------------------------------------- # + + +def test_gguf_is_refused_with_the_servable_repo_named() -> None: + with pytest.raises(LMServerError) as excinfo: + hw.check_model_format("LiquidAI/LFM2.5-1.2B-Instruct-GGUF", _NVIDIA) + message = str(excinfo.value) + assert "GGUF" in message + assert "LiquidAI/LFM2.5-1.2B-Instruct'" in message + + +def test_a_safetensors_repo_passes_quietly() -> None: + assert hw.check_model_format("LiquidAI/LFM2.5-1.2B-Instruct", _NVIDIA) is None + + +def test_fp8_on_pre_ada_hardware_warns() -> None: + warning = hw.check_model_format("some-org/model-FP8", _NVIDIA_OLD_DRIVER) + assert warning is not None + assert "FP8" in warning + + +def test_fp8_on_ada_is_not_warned_about() -> None: + """Mutation guard: the FP8 check must key on capability, not on the name.""" + assert hw.check_model_format("some-org/model-FP8", _NVIDIA) is None + + +# --------------------------------------------------------------------------- # +# ensure_runtime # +# --------------------------------------------------------------------------- # + +_SPEC = SGLangServer(provider="sglang", model="LiquidAI/LFM2.5-1.2B-Instruct", port=30000) + + +def _plan_for(monkeypatch: pytest.MonkeyPatch, plan: hw.InstallPlan | None) -> None: + monkeypatch.setattr(hw, "detect_accelerator", _nvidia_box) + monkeypatch.setattr(hw, "probe_runtime", _cuda_runtime) + + def _planned(_accel: hw.Accelerator, _runtime: hw.Runtime) -> hw.InstallPlan | None: + return plan + + monkeypatch.setattr(hw, "plan_runtime_install", _planned) + + +def test_without_the_flag_the_command_is_reported_not_run( + monkeypatch: pytest.MonkeyPatch, +) -> None: + plan = hw.InstallPlan(reason="torch is a cpu build", commands=(("uv", "pip", "install", "x"),)) + _plan_for(monkeypatch, plan) + ran: list[tuple[object, ...]] = [] + + def _record(*args: object, **kwargs: object) -> None: + ran.append((args, kwargs)) + + monkeypatch.setattr(hw.subprocess, "run", _record) + + with pytest.raises(LMServerError) as excinfo: + hw.ensure_runtime(_SPEC, install=False) + + assert "uv pip install x" in str(excinfo.value) + assert ran == [], "nothing may be installed without --install-runtime" + + +def test_with_the_flag_the_plan_is_executed_in_order(monkeypatch: pytest.MonkeyPatch) -> None: + plan = hw.InstallPlan( + reason="sglang is not installed", + commands=(("uv", "pip", "install", "sglang"), ("uv", "pip", "install", "torch")), + ) + monkeypatch.setattr(hw, "detect_accelerator", _nvidia_box) + monkeypatch.setattr(hw, "probe_runtime", _cuda_runtime) + plans = iter([plan, None]) # second call re-checks after installing + + def _next_plan(_accel: hw.Accelerator, _runtime: hw.Runtime) -> hw.InstallPlan | None: + return next(plans) + + monkeypatch.setattr(hw, "plan_runtime_install", _next_plan) + + ran: list[list[str]] = [] + + class _Completed: + returncode = 0 + + def _run(argv: list[str], **_kwargs: Any) -> _Completed: + ran.append(argv) + return _Completed() + + monkeypatch.setattr(hw.subprocess, "run", _run) + hw.ensure_runtime(_SPEC, install=True) + + assert ran == [["uv", "pip", "install", "sglang"], ["uv", "pip", "install", "torch"]] + + +def test_a_failed_install_is_loud(monkeypatch: pytest.MonkeyPatch) -> None: + plan = hw.InstallPlan(reason="sglang missing", commands=(("uv", "pip", "install", "sglang"),)) + _plan_for(monkeypatch, plan) + + class _Failed: + returncode = 1 + + def _failed(*_args: object, **_kwargs: object) -> _Failed: + return _Failed() + + monkeypatch.setattr(hw.subprocess, "run", _failed) + with pytest.raises(LMServerError, match="runtime install failed"): + hw.ensure_runtime(_SPEC, install=True) + + +def test_a_usable_runtime_installs_nothing(monkeypatch: pytest.MonkeyPatch) -> None: + _plan_for(monkeypatch, None) + ran: list[tuple[object, ...]] = [] + + def _record(*args: object, **kwargs: object) -> None: + ran.append((args, kwargs)) + + monkeypatch.setattr(hw.subprocess, "run", _record) + hw.ensure_runtime(_SPEC, install=True) + assert ran == [] + + +def test_an_unservable_model_is_refused_before_anything_is_probed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(hw, "detect_accelerator", _nvidia_box) + + def _too_late(_python: str | None = None) -> hw.Runtime: + pytest.fail("the model format must be refused before anything is probed") + + monkeypatch.setattr(hw, "probe_runtime", _too_late) + spec = SGLangServer(provider="sglang", model="LiquidAI/LFM2.5-1.2B-Instruct-GGUF", port=30000) + with pytest.raises(LMServerError, match="GGUF"): + hw.ensure_runtime(spec, install=False) + + +# --------------------------------------------------------------------------- # +# ensure_weights # +# --------------------------------------------------------------------------- # + + +def test_cached_weights_are_not_refetched(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[tuple[str, ...]] = [] + + def _run(argv: Sequence[str], **_kwargs: Any) -> str: + calls.append(tuple(argv)) + return '{"status": "ok"}' + + monkeypatch.setattr(hw, "_run", _run) + said: list[str] = [] + hw.ensure_weights(_SPEC, echo=said.append) + assert len(calls) == 1, "a cache hit must not trigger a second, downloading probe" + assert "already in the local hub cache" in said[0] + + +def test_a_cache_miss_triggers_a_fetch(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[tuple[str, ...]] = [] + + def _run(argv: Sequence[str], **_kwargs: Any) -> str: + calls.append(tuple(argv)) + return '{"status": "miss"}' if argv[-1] == "cached" else '{"status": "ok", "path": "/w"}' + + monkeypatch.setattr(hw, "_run", _run) + said: list[str] = [] + hw.ensure_weights(_SPEC, echo=said.append) + assert [argv[-1] for argv in calls] == ["cached", "fetch"] + assert any("ready at /w" in message for message in said), said + + +def test_a_local_model_path_is_left_alone(monkeypatch: pytest.MonkeyPatch, tmp_path: Any) -> None: + def _never(*_args: object, **_kwargs: object) -> str | None: + pytest.fail("a local model path must not be looked up on the hub") + + monkeypatch.setattr(hw, "_run", _never) + spec = SGLangServer(provider="sglang", model=str(tmp_path), port=30000) + hw.ensure_weights(spec) diff --git a/tests/unit/test_lm_sglang.py b/tests/unit/test_lm_sglang.py index 5f010961..5af2ea29 100644 --- a/tests/unit/test_lm_sglang.py +++ b/tests/unit/test_lm_sglang.py @@ -25,6 +25,24 @@ pytestmark = pytest.mark.unit +@pytest.fixture(autouse=True) +def _skip_preflight(monkeypatch: pytest.MonkeyPatch) -> None: # pyright: ignore[reportUnusedFunction] + """Bypass the hardware/runtime preflight for the lifecycle tests. + + These drive a stub launcher on a box that deliberately has no sglang + installed, which is exactly what :func:`stargraph.lm.hardware.ensure_runtime` + exists to refuse. The preflight is covered on its own in + ``tests/unit/test_lm_hardware.py``; here it would only assert that the dev + venv is a dev venv. + """ + + def _noop(*_args: object, **_kwargs: object) -> None: + return None + + monkeypatch.setattr(sg, "ensure_runtime", _noop) + monkeypatch.setattr(sg, "ensure_weights", _noop) + + def _spec(**overrides: object) -> SGLangServer: return SGLangServer.model_validate({"model": "Qwen/Qwen3-8B", **overrides}) From a9b30d0840af39e44d330b6e23687dc1bb034055 Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Mon, 17 Aug 2026 21:36:01 +0000 Subject: [PATCH 04/11] fix(lm): repair a wrong-backend torch instead of reinstalling sglang over it `--install-runtime` installed sglang and then failed its own verify with "runtime install completed but the runtime still cannot serve (torch 2.11.0+cpu)". The plan was wrong, not the check. Installing sglang cannot fix a CPU torch: `2.11.0+cpu` satisfies sglang's `torch==2.11.0` pin, because PEP 440 ignores the local version segment. The resolver is satisfied, the CPU wheel stays, and the box's GPUs stay invisible to that interpreter -- which is exactly the state a venv gets into when something pinned torch to the PyTorch CPU index (this repo does, deliberately, to keep CI wheels small). The plan now distinguishes the two failures: * sglang missing -> install sglang, and stop there. Its torch pin is not readable until it lands. * sglang present, torch built for the wrong accelerator -> `--force-reinstall` torch/torchaudio/torchvision off the CUDA index, pinned to the versions the *installed* sglang requires (read from its metadata, not hard-coded here -- sglang moves that pin every release). * sglang present, CUDA torch, zero devices, current driver -> not an install problem at all. Reported, with the container/driver cause named, instead of reinstalling wheels that are already correct. Repair therefore runs in rounds (install, re-probe, re-plan) with a no-progress guard: a round whose plan repeats stops the loop and hands over the commands that did not take, rather than retrying them. Index selection follows the driver as before: cu130 on r580+, the cu129 wheel set plus sglang-kernel/sgl-deep-gemm below it. Signed-off-by: Sean Mauk --- docs/how-to/authoring-format.md | 6 + docs/reference/cli.md | 5 + src/stargraph/lm/hardware.py | 229 +++++++++++++++++++++----------- tests/unit/test_lm_hardware.py | 187 ++++++++++++++++++++++++-- 4 files changed, 338 insertions(+), 89 deletions(-) diff --git a/docs/how-to/authoring-format.md b/docs/how-to/authoring-format.md index 53657d2c..0fb3a04a 100644 --- a/docs/how-to/authoring-format.md +++ b/docs/how-to/authoring-format.md @@ -101,6 +101,12 @@ interpreter agree: the install command for *that* platform. `--install-runtime` runs it; without the flag nothing is installed and the run stops. Kernel drivers are never touched -- a missing or too-old CUDA/ROCm driver can only be reported. +- **A CPU torch is repaired explicitly**, because installing sglang does not + do it: `2.11.0+cpu` satisfies sglang's `torch==2.11.0` pin (PEP 440 ignores + the local version), so the resolver is happy and the CPU wheel stays. + Repair therefore runs in rounds -- install sglang, re-probe, then + force-reinstall torch off the CUDA index at the pin *that* sglang resolved + to. A round that changes nothing stops the loop rather than retrying. SGLang publishes plain wheels for NVIDIA only, so ROCm, XPU, Ascend NPU and Apple Metal are reported with a pointer to their platform page rather than a command that would not work. diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 5a01baea..dad01af2 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -91,6 +91,11 @@ interpreter's sglang/torch build is checked against it, and the weights are fetched before the startup clock starts. Without `--install-runtime` a runtime that cannot serve is reported with the exact install command and the run stops; nothing is ever installed implicitly, and kernel drivers are never touched. +With the flag, repair runs in rounds: sglang first, then -- if the torch beside +it is a CPU (or otherwise mismatched) build -- a `--force-reinstall` off the +CUDA index at the pin sglang resolved to. Installing sglang alone does not fix +that torch: `2.11.0+cpu` satisfies `torch==2.11.0`, so the resolver leaves it +where it is. The `--sglang-*` flags bind the run to a local [SGLang](https://docs.sglang.ai/) server, and derive `--lm-url` / diff --git a/src/stargraph/lm/hardware.py b/src/stargraph/lm/hardware.py index 4f60655a..7654a628 100644 --- a/src/stargraph/lm/hardware.py +++ b/src/stargraph/lm/hardware.py @@ -60,6 +60,8 @@ _PROBE_TIMEOUT_S = 20.0 _INSTALL_TIMEOUT_S = 3600.0 +# sglang first, then the torch its metadata pins; a third is slack, not a plan. +_MAX_INSTALL_ROUNDS = 3 # CUDA 13 wheels need a r580+ driver (13.0 minimum is 580.65.06 on Linux); an # older driver has to stay on the CUDA 12 wheel set. @@ -105,6 +107,9 @@ class Runtime: """``torch.version.hip`` -- set on ROCm builds.""" torch_xpu: bool = False device_count: int = 0 + pinned_torch: str | None = None + """``torch==`` pin of the *installed* sglang -- what a repair must match.""" + pinned_torchaudio: str | None = None @property def backend(self) -> str: @@ -243,14 +248,19 @@ def detect_accelerator() -> Accelerator: _RUNTIME_PROBE = """ import json, importlib.util out = {"sglang": None, "torch": None, "torch_cuda": None, "torch_hip": None, - "torch_xpu": False, "device_count": 0} + "torch_xpu": False, "device_count": 0, "pinned_torch": None, + "pinned_torchaudio": None} spec = importlib.util.find_spec("sglang") if spec is not None: try: - from importlib.metadata import version + from importlib.metadata import requires, version out["sglang"] = version("sglang") + for req in requires("sglang") or []: + name, _, rest = req.partition("==") + if name.strip() in ("torch", "torchaudio") and rest: + out["pinned_" + name.strip()] = rest.split(";")[0].strip() except Exception: - out["sglang"] = "unknown" + out["sglang"] = out["sglang"] or "unknown" if importlib.util.find_spec("torch") is not None: try: import torch @@ -288,6 +298,8 @@ def probe_runtime(python: str | None = None) -> Runtime: torch_hip=data["torch_hip"], torch_xpu=bool(data["torch_xpu"]), device_count=int(data["device_count"]), + pinned_torch=data.get("pinned_torch"), + pinned_torchaudio=data.get("pinned_torchaudio"), ) @@ -298,58 +310,102 @@ def _installer(python: str) -> tuple[str, ...]: return (python, "-m", "pip", "install") +def _driver_major(accel: Accelerator) -> int: + head = accel.driver.split(".")[0] if accel.driver else "" + return int(head) if head.isdigit() else 0 + + +def _torch_repair(runtime: Runtime, install: tuple[str, ...], index: str) -> tuple[str, ...]: + """Force the torch family back onto ``index``, at the pins sglang resolved to. + + The versions come from the installed sglang's own metadata rather than a + number written down here: sglang moves its torch pin every release, and a + stale hard-coded pin would fight the resolver instead of satisfying it. + """ + packages = [f"torch=={runtime.pinned_torch}" if runtime.pinned_torch else "torch"] + if runtime.pinned_torchaudio: + packages.append(f"torchaudio=={runtime.pinned_torchaudio}") + packages.append("torchvision") + return (*install, "--force-reinstall", *packages, "--index-url", index) + + +_CUDA13_INDEX = "https://download.pytorch.org/whl/cu130" +_CUDA12_INDEX = "https://download.pytorch.org/whl/cu129" +_SGLANG_CU129_WHL = "https://docs.sglang.ai/whl/cu129/" + + def _nvidia_plan(accel: Accelerator, runtime: Runtime, install: tuple[str, ...]) -> InstallPlan: - """CUDA 13 wheels on a r580+ driver; the pinned CUDA 12 set below that.""" - driver_major = 0 - if accel.driver: - head = accel.driver.split(".")[0] - driver_major = int(head) if head.isdigit() else 0 - reason = ( - f"sglang is not installed for {runtime.python}" - if runtime.sglang is None - else f"torch in {runtime.python} is a {runtime.backend} build, " - f"but {accel.describe()} needs cuda" + """CUDA 13 wheels on an r580+ driver; the CUDA 12.9 wheel set below that. + + Two failures live here and they need different commands. A *missing* + sglang is one ``uv pip install`` away. A torch built for the wrong + accelerator is not: ``2.11.0+cpu`` satisfies sglang's ``torch==2.11.0`` + pin (PEP 440 ignores the local segment), so installing sglang next to it + resolves cleanly and leaves the CPU wheel exactly where it was. That case + needs an explicit ``--force-reinstall`` off the CUDA index, which is also + why :func:`ensure_runtime` re-plans after each round: the torch pin to + repair to is only knowable once sglang is installed. + """ + driver = _driver_major(accel) + cuda13 = driver >= _CUDA13_MIN_DRIVER_MAJOR + index = _CUDA13_INDEX if cuda13 else _CUDA12_INDEX + driver_note = ( + f"driver {accel.driver} supports the CUDA 13 wheels" + if cuda13 + else f"driver {accel.driver or 'unknown'} predates r{_CUDA13_MIN_DRIVER_MAJOR}, " + "so the CUDA 12.9 wheel set is used" ) - if driver_major >= _CUDA13_MIN_DRIVER_MAJOR: + + if runtime.sglang is None: return InstallPlan( - reason=reason, + reason=f"sglang is not installed for {runtime.python}", commands=((*install, "--prerelease=allow", "sglang"),), - notes=(f"driver {accel.driver} supports the CUDA 13 wheels",), + notes=(driver_note,), ) - return InstallPlan( - reason=reason, - commands=( - (*install, "--prerelease=allow", "sglang"), - ( - *install, - "--force-reinstall", - "torch==2.13.0", - "torchaudio==2.11.0", - "torchvision", - "--index-url", - "https://download.pytorch.org/whl/cu129", + + if runtime.backend == "cuda" and runtime.device_count == 0 and cuda13: + # Right wheel, current driver, no devices: nothing pip can install + # fixes this -- it is a container started without --gpus, or a driver + # that is not loaded. + return InstallPlan( + reason=( + f"torch in {runtime.python} is a cuda build but sees 0 devices, " + f"while {accel.describe()} is present" ), - ( - *install, - "--force-reinstall", - "sglang-kernel", - "--index-url", - "https://docs.sglang.ai/whl/cu129/", + manual=_INSTALL_DOC, + notes=( + "the GPUs are not visible to that interpreter -- a container " + "started without --gpus all, or a driver that is not loaded", ), + ) + + reason = ( + f"torch in {runtime.python} is a {runtime.backend} build, but {accel.describe()} needs cuda" + if runtime.backend != "cuda" + else f"torch in {runtime.python} is a cuda build that sees 0 devices on " + f"driver {accel.driver or 'unknown'}" + ) + if runtime.pinned_torch and runtime.torch: + reason += ( + f"; {runtime.torch} already satisfies sglang's torch==" + f"{runtime.pinned_torch} pin, so installing sglang alone leaves it in place" + ) + commands: list[tuple[str, ...]] = [_torch_repair(runtime, install, index)] + if not cuda13: + commands.append( + (*install, "--force-reinstall", "sglang-kernel", "--index-url", _SGLANG_CU129_WHL) + ) + commands.append( ( *install, "--force-reinstall", "sgl-deep-gemm", "--index-url", - "https://docs.sglang.ai/whl/cu129/", + _SGLANG_CU129_WHL, "--no-deps", - ), - ), - notes=( - f"driver {accel.driver or 'unknown'} predates r{_CUDA13_MIN_DRIVER_MAJOR}, " - "so the CUDA 12.9 wheel set is used", - ), - ) + ) + ) + return InstallPlan(reason=reason, commands=tuple(commands), notes=(driver_note,)) _INSTALL_DOC = "https://docs.sglang.io/docs/get-started/install" @@ -463,6 +519,12 @@ def ensure_runtime( Raises :class:`~stargraph.errors.LMServerError` when the runtime cannot serve and ``install`` is false, so the operator sees the exact command rather than a subprocess that dies during startup. + + Repair runs in rounds, because one round cannot see the next problem: a + box with no sglang gets sglang, and only then is its torch pin readable -- + which is what a wrong-backend torch has to be force-reinstalled to. The + loop stops as soon as a re-probe plans nothing, and refuses to spin when a + round leaves the runtime asking for the same commands again. """ def _say(message: str) -> None: @@ -479,43 +541,58 @@ def _say(message: str) -> None: _say( f"runtime: {runtime.describe()}" + (f", sglang {runtime.sglang}" if runtime.sglang else "") ) - plan = plan_runtime_install(accel, runtime) - if plan is None: - return - if not plan.automatable: - raise LMServerError( - plan.reason, - hint=f"sglang documents this platform at {plan.manual}", - notes="; ".join(plan.notes), - ) - rendered = "\n".join(" " + " ".join(command) for command in plan.commands) - if not install: - raise LMServerError( - plan.reason, - hint=f"re-run with --install-runtime, or install it yourself:\n{rendered}", - notes="; ".join(plan.notes), - ) - - for note in plan.notes: - _say(f"note: {note}") - for index, command in enumerate(plan.commands, start=1): - _say(f"installing ({index}/{len(plan.commands)}): {' '.join(command)}") - completed = subprocess.run(list(command), timeout=_INSTALL_TIMEOUT_S, check=False) - if completed.returncode != 0: + previous: tuple[tuple[str, ...], ...] | None = None + for _round in range(_MAX_INSTALL_ROUNDS): + plan = plan_runtime_install(accel, runtime) + if plan is None: + if previous is not None: + _say(f"runtime ready: {runtime.describe()}") + return + if not plan.automatable: raise LMServerError( - f"runtime install failed: {' '.join(command)}", - hint="install it manually, then re-run without --install-runtime", - exit_code=str(completed.returncode), + plan.reason, + hint=f"sglang documents this platform at {plan.manual or _INSTALL_DOC}", + notes="; ".join(plan.notes), ) - repaired = probe_runtime(python) - if plan_runtime_install(accel, repaired) is not None: - raise LMServerError( - f"runtime install completed but the runtime still cannot serve ({repaired.describe()})", - hint=f"sglang documents this platform at {_INSTALL_DOC}", - python=repaired.python, - ) - _say(f"runtime ready: {repaired.describe()}") + rendered = "\n".join(" " + " ".join(command) for command in plan.commands) + if not install: + raise LMServerError( + plan.reason, + hint=f"re-run with --install-runtime, or install it yourself:\n{rendered}", + notes="; ".join(plan.notes), + ) + if plan.commands == previous: + # The round ran and changed nothing the plan cares about. Running + # it again would loop, so hand the operator the state and the + # commands that did not take. + raise LMServerError( + f"runtime install ran but did not take: {plan.reason}", + hint=f"install it by hand and re-run without --install-runtime:\n{rendered}", + python=runtime.python, + ) + + for note in plan.notes: + _say(f"note: {note}") + for index, command in enumerate(plan.commands, start=1): + _say(f"installing ({index}/{len(plan.commands)}): {' '.join(command)}") + completed = subprocess.run(list(command), timeout=_INSTALL_TIMEOUT_S, check=False) + if completed.returncode != 0: + raise LMServerError( + f"runtime install failed: {' '.join(command)}", + hint="install it manually, then re-run without --install-runtime", + exit_code=str(completed.returncode), + ) + previous = plan.commands + runtime = probe_runtime(python) + _say(f"runtime: {runtime.describe()}") + + raise LMServerError( + f"runtime still cannot serve after {_MAX_INSTALL_ROUNDS} install rounds " + f"({runtime.describe()})", + hint=f"sglang documents this platform at {_INSTALL_DOC}", + python=runtime.python, + ) _WEIGHTS_PROBE = """ diff --git a/tests/unit/test_lm_hardware.py b/tests/unit/test_lm_hardware.py index e4da4e83..c9d0ab4c 100644 --- a/tests/unit/test_lm_hardware.py +++ b/tests/unit/test_lm_hardware.py @@ -11,6 +11,7 @@ from __future__ import annotations +import sys from typing import TYPE_CHECKING, Any import pytest @@ -114,14 +115,19 @@ def test_detection_does_not_consult_torch(monkeypatch: pytest.MonkeyPatch) -> No def test_probe_reports_this_interpreter_honestly() -> None: - """No stubbing: the dev venv genuinely has no sglang, and torch is CPU-only. + """No stubbing: the probe runs against the interpreter running the tests. Asserting against the real interpreter keeps the probe script itself under test -- a syntax error or renamed key in ``_RUNTIME_PROBE`` fails here. + Which packages happen to be in the venv is not asserted: a developer who + installs sglang (or a CUDA torch) must not turn this red. """ runtime = hw.probe_runtime() - assert runtime.sglang is None + assert runtime.python == sys.executable assert runtime.backend in {"cpu", "cuda", "rocm", "xpu"} + if runtime.sglang is not None: + # sglang always pins torch exactly; that pin is what a repair targets. + assert runtime.pinned_torch is not None def test_probe_of_a_broken_interpreter_degrades(monkeypatch: pytest.MonkeyPatch) -> None: @@ -160,42 +166,131 @@ def test_a_working_cuda_runtime_needs_no_plan() -> None: assert hw.plan_runtime_install(_NVIDIA, _WORKING_CUDA) is None -def test_cpu_torch_on_an_nvidia_box_is_planned_as_cuda13(monkeypatch: pytest.MonkeyPatch) -> None: +_CPU_TORCH_WITH_SGLANG = hw.Runtime( + python="/venv/bin/python", + sglang="0.5.17", + torch="2.11.0+cpu", + pinned_torch="2.11.0", + pinned_torchaudio="2.11.0", +) + + +def test_a_missing_sglang_is_installed_before_anything_is_said_about_torch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Round one is sglang alone -- its torch pin is unreadable until it lands.""" monkeypatch.setattr(hw.shutil, "which", _uv_on_path) - runtime = hw.Runtime(python="/venv/bin/python", sglang="0.5.0", torch="2.11.0+cpu") - plan = hw.plan_runtime_install(_NVIDIA, runtime) + plan = hw.plan_runtime_install(_NVIDIA, hw.Runtime(python="/venv/bin/python")) assert plan is not None - assert plan.automatable assert plan.commands == ( ("uv", "pip", "install", "--python", "/venv/bin/python", "--prerelease=allow", "sglang"), ) - assert "cpu build" in plan.reason + + +def test_cpu_torch_beside_an_installed_sglang_is_force_reinstalled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The bug this exists for: ``2.11.0+cpu`` satisfies ``torch==2.11.0``. + + Installing sglang next to a CPU wheel resolves cleanly and leaves it in + place, so the plan has to name torch explicitly, off the CUDA index, at + the pin sglang itself resolved to. + """ + monkeypatch.setattr(hw.shutil, "which", _uv_on_path) + plan = hw.plan_runtime_install(_NVIDIA, _CPU_TORCH_WITH_SGLANG) + assert plan is not None + assert plan.commands == ( + ( + "uv", + "pip", + "install", + "--python", + "/venv/bin/python", + "--force-reinstall", + "torch==2.11.0", + "torchaudio==2.11.0", + "torchvision", + "--index-url", + "https://download.pytorch.org/whl/cu130", + ), + ) + assert "satisfies sglang's torch==2.11.0 pin" in plan.reason + + +def test_the_torch_pin_comes_from_sglang_not_from_a_constant( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Mutation guard: a hard-coded pin would fight the installed sglang.""" + monkeypatch.setattr(hw.shutil, "which", _uv_on_path) + runtime = hw.Runtime( + python="/venv/bin/python", + sglang="9.9.9", + torch="3.4.0+cpu", + pinned_torch="3.4.0", + pinned_torchaudio="3.4.0", + ) + plan = hw.plan_runtime_install(_NVIDIA, runtime) + assert plan is not None + assert "torch==3.4.0" in " ".join(plan.commands[0]) def test_a_pre_r580_driver_gets_the_cuda12_wheel_set(monkeypatch: pytest.MonkeyPatch) -> None: """CUDA 13 wheels require r580+; an older driver must not be handed them.""" monkeypatch.setattr(hw.shutil, "which", _uv_on_path) - plan = hw.plan_runtime_install(_NVIDIA_OLD_DRIVER, hw.Runtime(python="/venv/bin/python")) + plan = hw.plan_runtime_install(_NVIDIA_OLD_DRIVER, _CPU_TORCH_WITH_SGLANG) assert plan is not None joined = [" ".join(command) for command in plan.commands] - # verbatim from https://docs.sglang.io/get_started/install.html (CUDA 12 path) + # the CUDA 12 shape from https://docs.sglang.io/docs/get-started/install assert any("whl/cu129" in command for command in joined), joined - assert any("torch==2.13.0 torchaudio==2.11.0 torchvision" in command for command in joined), ( + assert any("torch==2.11.0 torchaudio==2.11.0 torchvision" in command for command in joined), ( joined ) assert any("sglang-kernel" in command for command in joined), joined assert any("sgl-deep-gemm" in command and "--no-deps" in command for command in joined), joined -def test_cuda_torch_that_sees_no_devices_still_needs_a_plan( +def test_the_cuda13_path_carries_no_cuda12_kernel_steps(monkeypatch: pytest.MonkeyPatch) -> None: + """Mutation guard: the kernel wheels are the old-driver path, not both.""" + monkeypatch.setattr(hw.shutil, "which", _uv_on_path) + plan = hw.plan_runtime_install(_NVIDIA, _CPU_TORCH_WITH_SGLANG) + assert plan is not None + assert not any("cu129" in " ".join(command) for command in plan.commands) + + +def test_cuda_torch_that_sees_no_devices_is_reported_not_reinstalled( monkeypatch: pytest.MonkeyPatch, ) -> None: - """A container started without --gpus all: right wheel, no devices.""" + """A container started without --gpus all: right wheel, current driver, no devices. + + No wheel fixes that, so it must not be planned as an install. + """ monkeypatch.setattr(hw.shutil, "which", _uv_on_path) runtime = hw.Runtime( python="/venv/bin/python", sglang="0.5.0", torch="2.13.0", torch_cuda="13.0", device_count=0 ) - assert hw.plan_runtime_install(_NVIDIA, runtime) is not None + plan = hw.plan_runtime_install(_NVIDIA, runtime) + assert plan is not None + assert not plan.automatable + assert "sees 0 devices" in plan.reason + + +def test_cuda_torch_seeing_no_devices_on_an_old_driver_is_rebuilt_for_cuda12( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The same symptom with a pre-r580 driver *is* fixable: wrong wheel set.""" + monkeypatch.setattr(hw.shutil, "which", _uv_on_path) + runtime = hw.Runtime( + python="/venv/bin/python", + sglang="0.5.17", + torch="2.11.0", + torch_cuda="13.0", + device_count=0, + pinned_torch="2.11.0", + ) + plan = hw.plan_runtime_install(_NVIDIA_OLD_DRIVER, runtime) + assert plan is not None + assert plan.automatable + assert any("cu129" in " ".join(command) for command in plan.commands) def test_amd_is_planned_but_not_automated() -> None: @@ -402,3 +497,69 @@ def _never(*_args: object, **_kwargs: object) -> str | None: monkeypatch.setattr(hw, "_run", _never) spec = SGLangServer(provider="sglang", model=str(tmp_path), port=30000) hw.ensure_weights(spec) + + +def test_repair_runs_in_rounds_because_round_one_cannot_see_round_two( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Install sglang, re-probe, then repair the torch its metadata pins. + + This is the user-visible bug the rounds exist for: one pass installed + sglang, the CPU torch beside it satisfied the pin, and the run died with + "install completed but the runtime still cannot serve". + """ + monkeypatch.setattr(hw, "detect_accelerator", _nvidia_box) + monkeypatch.setattr(hw.shutil, "which", _uv_on_path) + probes = iter([hw.Runtime(python="/venv/bin/python"), _CPU_TORCH_WITH_SGLANG, _WORKING_CUDA]) + + def _next_probe(_python: str | None = None) -> hw.Runtime: + return next(probes) + + monkeypatch.setattr(hw, "probe_runtime", _next_probe) + + ran: list[list[str]] = [] + + class _Completed: + returncode = 0 + + def _run(argv: list[str], **_kwargs: Any) -> _Completed: + ran.append(argv) + return _Completed() + + monkeypatch.setattr(hw.subprocess, "run", _run) + hw.ensure_runtime(_SPEC, install=True) + + assert len(ran) == 2, ran + assert ran[0][-1] == "sglang" + assert "torch==2.11.0" in ran[1] + assert "https://download.pytorch.org/whl/cu130" in ran[1] + + +def test_a_round_that_changes_nothing_stops_instead_of_looping( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An install that does not take must not be run again and again.""" + monkeypatch.setattr(hw, "detect_accelerator", _nvidia_box) + monkeypatch.setattr(hw.shutil, "which", _uv_on_path) + + def _stuck(_python: str | None = None) -> hw.Runtime: + return _CPU_TORCH_WITH_SGLANG + + monkeypatch.setattr(hw, "probe_runtime", _stuck) + + ran: list[list[str]] = [] + + class _Completed: + returncode = 0 + + def _run(argv: list[str], **_kwargs: Any) -> _Completed: + ran.append(argv) + return _Completed() + + monkeypatch.setattr(hw.subprocess, "run", _run) + + with pytest.raises(LMServerError) as excinfo: + hw.ensure_runtime(_SPEC, install=True) + + assert "did not take" in str(excinfo.value) + assert len(ran) == 1, "the same command must not be retried" From 4433a5805709d3b0327e8944cedd7dc3cf099d25 Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Mon, 17 Aug 2026 21:42:03 +0000 Subject: [PATCH 05/11] feat(cli): --sglang-python serves from the venv that has sglang stargraph spawns `python -m sglang.launch_server` from its own interpreter, which is the wrong one whenever sglang lives somewhere else -- the common case, since a venv that runs stargraph may deliberately pin a CPU torch (this repo does) while the venv that serves has the CUDA stack. `--sglang-python` moves the whole spawn: preflight, weight fetch and launch all target the named interpreter, so the environment stargraph runs in never has to become the environment sglang serves from. A venv directory is accepted and resolved to its `bin/python`, and a path that is not an executable interpreter is refused at the flag rather than as an ENOENT from a subprocess after the graph has compiled. Flag only, deliberately: naming the interpreter to execute is argv into a subprocess, the same operator-only class as `lm.args` and a non-loopback `lm.host` that `_check_graph_declared` already refuses. A graph that could choose the interpreter would be choosing what code runs. Signed-off-by: Sean Mauk --- docs/how-to/authoring-format.md | 5 ++ docs/reference/cli.md | 8 +++ src/stargraph/cli/run.py | 48 +++++++++++++++- src/stargraph/lm/sglang.py | 27 ++++++--- tests/integration/test_lm_sglang_spawn.py | 6 +- tests/unit/cli/test_run_sglang.py | 69 +++++++++++++++++++++++ tests/unit/test_lm_sglang.py | 49 ++++++++++++++++ 7 files changed, 199 insertions(+), 13 deletions(-) diff --git a/docs/how-to/authoring-format.md b/docs/how-to/authoring-format.md index 0fb3a04a..0afdd936 100644 --- a/docs/how-to/authoring-format.md +++ b/docs/how-to/authoring-format.md @@ -110,6 +110,11 @@ interpreter agree: SGLang publishes plain wheels for NVIDIA only, so ROCm, XPU, Ascend NPU and Apple Metal are reported with a pointer to their platform page rather than a command that would not work. +- **The interpreter** is stargraph's own unless `--sglang-python` names + another one (a venv directory works). Preflight, weight fetch and launch all + move together, so a graph can run from a CPU-only venv and still serve from + the venv that has a CUDA sglang. It is a flag, never an `lm:` key -- naming + the interpreter to execute is operator-only. - **The weights** are fetched before the server starts, so `startup_timeout_s` measures server boot rather than racing a multi-gigabyte download. - **The format** is validated, not rewritten. A GGUF repo is refused (that is diff --git a/docs/reference/cli.md b/docs/reference/cli.md index dad01af2..b0ef2834 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -82,6 +82,7 @@ on `failed`. | `--sglang-port PORT` | int | `30000` | SGLang port. | | `--sglang-arg ARG` | str (repeatable) | _(empty)_ | Extra argv passed through to `sglang.launch_server` verbatim. | | `--sglang-timeout SEC` | int | `600` | Seconds to wait for a launched SGLang server to answer. | +| `--sglang-python PATH` | str | _(this interpreter)_ | Interpreter (or venv directory) to serve from. | | `--install-runtime` | flag | `false` | Install the sglang build matching the detected accelerator before launching. | `--quiet` and `--verbose` are mutually exclusive. `--lm-url` and @@ -97,6 +98,13 @@ CUDA index at the pin sglang resolved to. Installing sglang alone does not fix that torch: `2.11.0+cpu` satisfies `torch==2.11.0`, so the resolver leaves it where it is. +`--sglang-python` points the whole spawn -- preflight, weight fetch and launch +-- at another interpreter, so the venv running stargraph never has to become +the venv serving the model. A venv directory is accepted and resolved to its +`bin/python`. There is deliberately no `lm:` key for it: the interpreter is +argv into a subprocess, the same class of operator-only value as `args` and a +non-loopback `host`. + The `--sglang-*` flags bind the run to a local [SGLang](https://docs.sglang.ai/) server, and derive `--lm-url` / `--lm-model` from it — so they conflict with those two flags. Before the diff --git a/src/stargraph/cli/run.py b/src/stargraph/cli/run.py index d8835009..deda7ad8 100644 --- a/src/stargraph/cli/run.py +++ b/src/stargraph/cli/run.py @@ -31,6 +31,7 @@ import asyncio import contextlib import ipaddress +import os from pathlib import Path from typing import TYPE_CHECKING, Annotated, Any @@ -122,6 +123,26 @@ def _is_loopback(host: str) -> bool: return False +def _validated_python(python: str | None) -> str | None: + """Resolve ``--sglang-python`` to an executable interpreter, or fail early. + + A venv directory is accepted and resolved to its ``bin/python``: that is + what an operator has in hand ("the venv where sglang lives"), and the + alternative is a subprocess that dies with a bare ENOENT much later, after + the graph has already been compiled. + """ + if python is None: + return None + candidate = Path(python).expanduser() + if candidate.is_dir(): + candidate = candidate / "bin" / "python" + if not candidate.is_file() or not os.access(candidate, os.X_OK): + raise typer.BadParameter( + f"--sglang-python {python!r} is not an executable interpreter (looked at {candidate})" + ) + return str(candidate.resolve()) + + def _check_graph_declared( declared: SGLangServer, *, args_from_flags: bool, host_from_flags: str | None ) -> None: @@ -206,6 +227,7 @@ def _lm_endpoint( lm_timeout: int, echo: Callable[[str], None], install_runtime: bool = False, + sglang_python: str | None = None, ) -> Generator[None]: """Hold the run's LM endpoint open: boot/attach sglang, then configure dspy. @@ -219,7 +241,12 @@ def _lm_endpoint( from stargraph.lm.sglang import sglang_server lm_url = stack.enter_context( - sglang_server(spec, echo=echo, install_runtime=install_runtime) + sglang_server( + spec, + echo=echo, + install_runtime=install_runtime, + python=sglang_python, + ) ) lm_model = spec.model _configure_lm(lm_url, lm_model, lm_key, lm_timeout) @@ -432,6 +459,16 @@ def cmd( help="Seconds to wait for a launched SGLang server to answer (default: 600).", ), ] = None, + sglang_python: Annotated[ + str | None, + typer.Option( + "--sglang-python", + help=( + "Interpreter to serve from (default: the one running stargraph). " + "Point it at the venv that has sglang installed." + ), + ), + ] = None, install_runtime: Annotated[ bool, typer.Option( @@ -542,7 +579,14 @@ def _echo(message: str) -> None: # sglang server is torn down only once the loop is done (or has # raised). with _lm_endpoint( - sglang_spec, lm_url, lm_model, lm_key, lm_timeout, _echo, install_runtime + sglang_spec, + lm_url, + lm_model, + lm_key, + lm_timeout, + _echo, + install_runtime, + _validated_python(sglang_python), ): try: node_registry = build_node_registry(ir.nodes, ir_dir=graph.parent.resolve()) diff --git a/src/stargraph/lm/sglang.py b/src/stargraph/lm/sglang.py index 5f4f0dfd..ea701433 100644 --- a/src/stargraph/lm/sglang.py +++ b/src/stargraph/lm/sglang.py @@ -79,10 +79,10 @@ def served_models(url: str, *, timeout: float = _PROBE_TIMEOUT_S) -> list[str] | return [str(entry.get("id", "")) for entry in data] -def _launch_argv(spec: SGLangServer) -> list[str]: +def _launch_argv(spec: SGLangServer, python: str | None = None) -> list[str]: """Argv for the launch subprocess (monkeypatched in tests).""" return [ - sys.executable, + python or sys.executable, "-m", "sglang.launch_server", "--model-path", @@ -104,8 +104,10 @@ def _log_tail(log_path: Path) -> str: return "\n".join(lines[-_LOG_TAIL_LINES:]) -def _spawn(spec: SGLangServer, log_path: Path) -> subprocess.Popen[bytes]: - argv = _launch_argv(spec) +def _spawn( + spec: SGLangServer, log_path: Path, python: str | None = None +) -> subprocess.Popen[bytes]: + argv = _launch_argv(spec, python) handle = log_path.open("wb") try: return subprocess.Popen( @@ -117,7 +119,7 @@ def _spawn(spec: SGLangServer, log_path: Path) -> subprocess.Popen[bytes]: except OSError as exc: raise LMServerError( f"failed to launch sglang: {exc}", - hint=f"is sglang installed for {sys.executable}? `pip install sglang`", + hint=f"is sglang installed for {python or sys.executable}? `pip install sglang`", argv=" ".join(argv), ) from exc finally: @@ -197,6 +199,7 @@ def sglang_server( log_path: Path | None = None, echo: Callable[[str], None] | None = None, install_runtime: bool = False, + python: str | None = None, ) -> Generator[str]: """Yield the base URL of a server for ``spec``, booting one if needed. @@ -212,6 +215,14 @@ def sglang_server( run for you under ``install_runtime``) and the weights are on disk before ``startup_timeout_s`` starts counting. The attach branch skips both: that server is already up, and it is not ours to diagnose. + + ``python`` names the interpreter to serve from, defaulting to the one + running stargraph. It is the answer to "my sglang lives in a different + venv": the preflight, the weights fetch and the launch all move there + together, so the environment stargraph runs in never has to become the + environment sglang serves from. It is an operator-supplied value only -- + :func:`stargraph.cli.run._check_graph_declared` keeps a graph from + choosing which interpreter gets executed. """ url = base_url(spec) existing = served_models(url) @@ -227,8 +238,8 @@ def sglang_server( yield url return - ensure_runtime(spec, install=install_runtime, echo=echo) - ensure_weights(spec, echo=echo) + ensure_runtime(spec, python=python, install=install_runtime, echo=echo) + ensure_weights(spec, python=python, echo=echo) if log_path is None: fd, name = tempfile.mkstemp(prefix=f"sglang-{spec.port}-", suffix=".log") @@ -236,7 +247,7 @@ def sglang_server( log_path = Path(name) if echo is not None: echo(f"starting sglang ({spec.model}) on {url}; output -> {log_path}") - proc = _spawn(spec, log_path) + proc = _spawn(spec, log_path, python) try: _await_ready(proc, spec, url, log_path, echo) yield url diff --git a/tests/integration/test_lm_sglang_spawn.py b/tests/integration/test_lm_sglang_spawn.py index 31aefbde..db2ae91c 100644 --- a/tests/integration/test_lm_sglang_spawn.py +++ b/tests/integration/test_lm_sglang_spawn.py @@ -93,7 +93,7 @@ def test_spawns_waits_for_ready_then_tears_down( port = _free_port() spec = SGLangServer(model="stub/model", port=port, startup_timeout_s=30) - def _stub_argv(spec: SGLangServer) -> list[str]: + def _stub_argv(spec: SGLangServer, _python: str | None = None) -> list[str]: return [sys.executable, str(stub), spec.model, str(spec.port)] monkeypatch.setattr(sg, "_launch_argv", _stub_argv) @@ -114,7 +114,7 @@ def test_launch_that_dies_reports_exit_code_and_log( spec = SGLangServer(model="stub/model", port=_free_port(), startup_timeout_s=30) - def _dying_argv(_spec: SGLangServer) -> list[str]: + def _dying_argv(_spec: SGLangServer, _python: str | None = None) -> list[str]: return [sys.executable, "-c", 'import sys; print("CUDA go boom"); sys.exit(3)'] monkeypatch.setattr(sg, "_launch_argv", _dying_argv) @@ -135,7 +135,7 @@ def test_startup_timeout_is_loud(tmp_path: Path, monkeypatch: pytest.MonkeyPatch spec = SGLangServer(model="stub/model", port=_free_port(), startup_timeout_s=1) - def _hanging_argv(_spec: SGLangServer) -> list[str]: + def _hanging_argv(_spec: SGLangServer, _python: str | None = None) -> list[str]: return [sys.executable, "-c", "import time; time.sleep(60)"] monkeypatch.setattr(sg, "_launch_argv", _hanging_argv) diff --git a/tests/unit/cli/test_run_sglang.py b/tests/unit/cli/test_run_sglang.py index f4637f61..03aabeff 100644 --- a/tests/unit/cli/test_run_sglang.py +++ b/tests/unit/cli/test_run_sglang.py @@ -11,6 +11,7 @@ from __future__ import annotations import contextlib +import sys from pathlib import Path from typing import TYPE_CHECKING, Any @@ -439,3 +440,71 @@ def test_install_runtime_flag_reaches_the_launcher( assert result.exit_code == 0, result.output assert seen["install_runtime"] is flag + + +def _run_with( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, *extra: str +) -> tuple[Any, dict[str, Any]]: + seen: dict[str, Any] = {} + _capture_launcher(monkeypatch, seen) + _stub_dspy(monkeypatch, {}) + result = CliRunner().invoke( + _make_app(), + [ + str(SAMPLE_GRAPH), + "--checkpoint", + str(tmp_path / "ck.sqlite"), + "--sglang-model", + "LiquidAI/LFM2.5-1.2B-Instruct", + "--quiet", + "--no-summary", + *extra, + ], + ) + return result, seen + + +def test_sglang_python_reaches_the_launcher( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The venv that has sglang is rarely the venv that has stargraph.""" + result, seen = _run_with(monkeypatch, tmp_path, "--sglang-python", sys.executable) + + assert result.exit_code == 0, result.output + assert seen["python"] == str(Path(sys.executable).resolve()) + + +def test_a_venv_directory_is_resolved_to_its_interpreter( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """What an operator has in hand is the venv, not the path to bin/python.""" + venv = tmp_path / "sglang-venv" + (venv / "bin").mkdir(parents=True) + interpreter = venv / "bin" / "python" + interpreter.symlink_to(sys.executable) + + result, seen = _run_with(monkeypatch, tmp_path, "--sglang-python", str(venv)) + + assert result.exit_code == 0, result.output + assert seen["python"] == str(interpreter.resolve()) + + +def test_a_bad_interpreter_is_refused_before_the_graph_runs( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Fail on the flag, not on an ENOENT from a subprocess minutes later.""" + result, seen = _run_with(monkeypatch, tmp_path, "--sglang-python", str(tmp_path / "nope")) + + assert result.exit_code != 0 + assert "not an executable interpreter" in result.output + assert seen == {}, "nothing may be launched once the flag is refused" + + +def test_no_flag_means_the_launcher_picks_the_default_interpreter( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Mutation guard: the default must stay None, not this process's path.""" + result, seen = _run_with(monkeypatch, tmp_path) + + assert result.exit_code == 0, result.output + assert seen["python"] is None diff --git a/tests/unit/test_lm_sglang.py b/tests/unit/test_lm_sglang.py index 5af2ea29..09b5e08f 100644 --- a/tests/unit/test_lm_sglang.py +++ b/tests/unit/test_lm_sglang.py @@ -81,6 +81,55 @@ def test_launch_argv_passes_model_port_and_extra_args() -> None: assert argv[-2:] == ["--attention-backend", "triton"] +def test_launch_argv_serves_from_the_named_interpreter() -> None: + """``--sglang-python``: the venv that has sglang need not be stargraph's own.""" + argv = sg._launch_argv(_spec(), "/opt/sglang-venv/bin/python") # pyright: ignore[reportPrivateUsage] + + assert argv[:3] == ["/opt/sglang-venv/bin/python", "-m", "sglang.launch_server"] + + +def test_the_preflight_and_the_launch_target_the_same_interpreter( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """A preflight that checked a different interpreter than the one we launch is a lie.""" + seen: dict[str, object] = {} + + def _runtime(_spec: SGLangServer, **kwargs: object) -> None: + seen["runtime_python"] = kwargs.get("python") + + def _weights(_spec: SGLangServer, **kwargs: object) -> None: + seen["weights_python"] = kwargs.get("python") + + monkeypatch.setattr(sg, "ensure_runtime", _runtime) + monkeypatch.setattr(sg, "ensure_weights", _weights) + + def _nothing_listening(_url: str) -> list[str] | None: + return None + + def _argv(_spec: SGLangServer, python: str | None = None) -> list[str]: + seen["argv_python"] = python + return ["true"] + + monkeypatch.setattr(sg, "served_models", _nothing_listening) + monkeypatch.setattr(sg, "_launch_argv", _argv) + + def _boom(*_args: object, **_kwargs: object) -> None: + raise LMServerError("stop after the spawn decision") + + monkeypatch.setattr(sg, "_await_ready", _boom) + with ( + pytest.raises(LMServerError), + sg.sglang_server( + _spec(port=41999), log_path=tmp_path / "log", python="/opt/sglang-venv/bin/python" + ), + ): + pass + + assert seen["runtime_python"] == "/opt/sglang-venv/bin/python" + assert seen["weights_python"] == "/opt/sglang-venv/bin/python" + assert seen["argv_python"] == "/opt/sglang-venv/bin/python" + + def test_served_models_returns_none_when_nothing_answers() -> None: # Port 1 is privileged and unbound: the probe must report "no server" # rather than raising, so the caller spawns one. From 3b3eeda7d57c2a2280b060e8bf2688eaa41284a1 Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Mon, 17 Aug 2026 21:48:18 +0000 Subject: [PATCH 06/11] fix(lm): give the spawned server the PATH a venv activation would have A launched sglang died mid-startup with `sglang exited with code -9`, which reads like an OOM kill and is not one. The child's log had the real cause: FileNotFoundError: [Errno 2] No such file or directory: 'ninja' ... flashinfer/jit/cpp_ext.py, in run_ninja Launching an interpreter by absolute path runs its packages but not its scripts -- nothing puts that venv's `bin/` on PATH, which activation (or `uv run`) would have done. sglang depends on `ninja` and flashinfer shells out to it to JIT-compile attention kernels during startup, so the console script is installed and unreachable at the same time. The failure surfaces as a signal from a grandchild process, minutes into a run. The child now inherits PATH with the interpreter's directory prepended, which is what activation does and nothing more. Already-present entries are not duplicated. Found by running the example end-to-end in a throwaway venv on a real GPU box; caught by a test that spawns a child and reads back the PATH it saw, because asserting on the helper alone left `env=` free to be deleted. Signed-off-by: Sean Mauk --- src/stargraph/lm/sglang.py | 20 +++++++++++++ tests/integration/test_lm_sglang_spawn.py | 34 ++++++++++++++++++++--- tests/unit/test_lm_sglang.py | 25 ++++++++++++++++- 3 files changed, 74 insertions(+), 5 deletions(-) diff --git a/src/stargraph/lm/sglang.py b/src/stargraph/lm/sglang.py index ea701433..9afbcf56 100644 --- a/src/stargraph/lm/sglang.py +++ b/src/stargraph/lm/sglang.py @@ -104,6 +104,25 @@ def _log_tail(log_path: Path) -> str: return "\n".join(lines[-_LOG_TAIL_LINES:]) +def _child_env(python: str | None) -> dict[str, str]: + """The environment a venv activation would have given the server. + + Launching an interpreter by absolute path runs its packages but not its + scripts: nothing puts the venv's ``bin/`` on ``PATH``, which activation + (or ``uv run``) would have done. sglang notices at the worst moment -- + flashinfer JIT-compiles its attention kernels during startup and shells + out to ``ninja``, a console script sglang itself depends on. Without it + the child dies mid-startup and the parent reports SIGKILL, which reads + like an OOM and is not one. + """ + env = os.environ.copy() + bindir = str(Path(python or sys.executable).resolve().parent) + path = env.get("PATH", "") + if bindir not in path.split(os.pathsep): + env["PATH"] = os.pathsep.join([bindir, path]) if path else bindir + return env + + def _spawn( spec: SGLangServer, log_path: Path, python: str | None = None ) -> subprocess.Popen[bytes]: @@ -115,6 +134,7 @@ def _spawn( stdout=handle, stderr=subprocess.STDOUT, start_new_session=True, + env=_child_env(python), ) except OSError as exc: raise LMServerError( diff --git a/tests/integration/test_lm_sglang_spawn.py b/tests/integration/test_lm_sglang_spawn.py index db2ae91c..8f514c59 100644 --- a/tests/integration/test_lm_sglang_spawn.py +++ b/tests/integration/test_lm_sglang_spawn.py @@ -9,9 +9,11 @@ from __future__ import annotations +import os import socket +import sys import textwrap -from typing import TYPE_CHECKING +from pathlib import Path import pytest @@ -19,9 +21,6 @@ from stargraph.ir import SGLangServer from stargraph.lm import sglang as sg -if TYPE_CHECKING: - from pathlib import Path - pytestmark = pytest.mark.integration @@ -145,3 +144,30 @@ def _hanging_argv(_spec: SGLangServer, _python: str | None = None) -> list[str]: sg.sglang_server(spec, log_path=tmp_path / "sglang.log"), ): pytest.fail("must not yield before the endpoint answers") + + +def test_the_spawned_child_really_sees_the_venv_bin_on_path( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """End-to-end on the env, not just on the helper that builds it. + + The unit test covers what ``_child_env`` returns; this covers that the + value reaches :class:`subprocess.Popen`. Dropping ``env=`` was invisible + to every other test in the suite, and it is the difference between a + server that starts and one that dies in flashinfer's JIT looking for + ``ninja``. + """ + log_path = tmp_path / "child.log" + + def _print_path_argv(_spec: SGLangServer, python: str | None = None) -> list[str]: + return [python or sys.executable, "-c", "import os; print(os.environ['PATH'])"] + + monkeypatch.setattr(sg, "_launch_argv", _print_path_argv) + monkeypatch.setenv("PATH", "/usr/bin") + + spec = SGLangServer(model="stub/model", port=_free_port(), startup_timeout_s=30) + proc = sg._spawn(spec, log_path, sys.executable) # pyright: ignore[reportPrivateUsage] + proc.wait(timeout=30) + + seen = log_path.read_text().strip().split(os.pathsep) + assert seen[0] == str(Path(sys.executable).resolve().parent) diff --git a/tests/unit/test_lm_sglang.py b/tests/unit/test_lm_sglang.py index 09b5e08f..994fa040 100644 --- a/tests/unit/test_lm_sglang.py +++ b/tests/unit/test_lm_sglang.py @@ -9,7 +9,9 @@ from __future__ import annotations +import os import sys +from pathlib import Path from typing import TYPE_CHECKING, NoReturn import pytest @@ -20,7 +22,6 @@ if TYPE_CHECKING: from collections.abc import Callable - from pathlib import Path pytestmark = pytest.mark.unit @@ -130,6 +131,28 @@ def _boom(*_args: object, **_kwargs: object) -> None: assert seen["argv_python"] == "/opt/sglang-venv/bin/python" +def test_the_child_gets_the_interpreters_bin_on_path() -> None: + """Console scripts sglang shells out to (``ninja``) live in the venv's bin/. + + Launching by absolute interpreter path does not put that directory on + PATH -- activation does. flashinfer's JIT then cannot find ninja, the + child dies during startup, and the parent reports a bare SIGKILL. + """ + env = sg._child_env("/opt/sglang-venv/bin/python") # pyright: ignore[reportPrivateUsage] + + assert env["PATH"].split(os.pathsep)[0] == "/opt/sglang-venv/bin" + + +def test_a_bin_already_on_path_is_not_duplicated(monkeypatch: pytest.MonkeyPatch) -> None: + """Mutation guard: prepending unconditionally would grow PATH every run.""" + bindir = str(Path(sys.executable).resolve().parent) + monkeypatch.setenv("PATH", os.pathsep.join([bindir, "/usr/bin"])) + + env = sg._child_env(None) # pyright: ignore[reportPrivateUsage] + + assert env["PATH"].split(os.pathsep).count(bindir) == 1 + + def test_served_models_returns_none_when_nothing_answers() -> None: # Port 1 is privileged and unbound: the probe must report "no server" # rather than raising, so the caller spawns one. From 928aa5869fb9ffc0e4453d5592a1d8018be3cc26 Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Mon, 17 Aug 2026 21:51:59 +0000 Subject: [PATCH 07/11] fix(lm): keep the venv's bin on PATH by not resolving the interpreter symlink The previous commit prepended the wrong directory. A venv's `bin/python` is a symlink into the base install, so `.resolve().parent` lands in `~/.local/share/uv/python/.../bin` -- where none of the venv's console scripts are. The end-to-end run failed identically after the "fix": ninja sat in the venv's bin, unreachable. Paths are now made absolute and never resolved, in both places that touch an interpreter path: the spawn's PATH and `--sglang-python` (which would otherwise hand sglang an interpreter outside the venv the operator named, with none of its packages). The tests missed this because both sides of the assertion resolved. They now build a venv-shaped symlink and assert the venv's own bin, which fails against `.resolve()`. Signed-off-by: Sean Mauk --- src/stargraph/cli/run.py | 8 ++++++-- src/stargraph/lm/sglang.py | 5 ++++- tests/integration/test_lm_sglang_spawn.py | 11 +++++++++-- tests/unit/cli/test_run_sglang.py | 6 ++++-- tests/unit/test_lm_sglang.py | 18 ++++++++++++++++++ 5 files changed, 41 insertions(+), 7 deletions(-) diff --git a/src/stargraph/cli/run.py b/src/stargraph/cli/run.py index deda7ad8..48d132f7 100644 --- a/src/stargraph/cli/run.py +++ b/src/stargraph/cli/run.py @@ -126,10 +126,14 @@ def _is_loopback(host: str) -> bool: def _validated_python(python: str | None) -> str | None: """Resolve ``--sglang-python`` to an executable interpreter, or fail early. - A venv directory is accepted and resolved to its ``bin/python``: that is + A venv directory is accepted and mapped to its ``bin/python``: that is what an operator has in hand ("the venv where sglang lives"), and the alternative is a subprocess that dies with a bare ENOENT much later, after the graph has already been compiled. + + The path is made absolute but never resolved. A venv's ``bin/python`` is + a symlink into the base install; following it hands back an interpreter + outside the venv the operator named, with none of its packages. """ if python is None: return None @@ -140,7 +144,7 @@ def _validated_python(python: str | None) -> str | None: raise typer.BadParameter( f"--sglang-python {python!r} is not an executable interpreter (looked at {candidate})" ) - return str(candidate.resolve()) + return str(candidate.absolute()) def _check_graph_declared( diff --git a/src/stargraph/lm/sglang.py b/src/stargraph/lm/sglang.py index 9afbcf56..cee5df63 100644 --- a/src/stargraph/lm/sglang.py +++ b/src/stargraph/lm/sglang.py @@ -116,7 +116,10 @@ def _child_env(python: str | None) -> dict[str, str]: like an OOM and is not one. """ env = os.environ.copy() - bindir = str(Path(python or sys.executable).resolve().parent) + # Deliberately not resolved: a venv's ``bin/python`` is a symlink to the + # base interpreter, so resolving it lands in the base install's bin and + # silently misses every console script the venv actually has. + bindir = str(Path(python or sys.executable).absolute().parent) path = env.get("PATH", "") if bindir not in path.split(os.pathsep): env["PATH"] = os.pathsep.join([bindir, path]) if path else bindir diff --git a/tests/integration/test_lm_sglang_spawn.py b/tests/integration/test_lm_sglang_spawn.py index 8f514c59..691566f9 100644 --- a/tests/integration/test_lm_sglang_spawn.py +++ b/tests/integration/test_lm_sglang_spawn.py @@ -165,9 +165,16 @@ def _print_path_argv(_spec: SGLangServer, python: str | None = None) -> list[str monkeypatch.setattr(sg, "_launch_argv", _print_path_argv) monkeypatch.setenv("PATH", "/usr/bin") + # Through a venv-shaped symlink, because that is what a real venv is and + # what makes the difference between the venv's bin and the base install's. + bindir = tmp_path / "venv" / "bin" + bindir.mkdir(parents=True) + interpreter = bindir / "python" + interpreter.symlink_to(sys.executable) + spec = SGLangServer(model="stub/model", port=_free_port(), startup_timeout_s=30) - proc = sg._spawn(spec, log_path, sys.executable) # pyright: ignore[reportPrivateUsage] + proc = sg._spawn(spec, log_path, str(interpreter)) # pyright: ignore[reportPrivateUsage] proc.wait(timeout=30) seen = log_path.read_text().strip().split(os.pathsep) - assert seen[0] == str(Path(sys.executable).resolve().parent) + assert seen[0] == str(bindir) diff --git a/tests/unit/cli/test_run_sglang.py b/tests/unit/cli/test_run_sglang.py index 03aabeff..9f94a9d8 100644 --- a/tests/unit/cli/test_run_sglang.py +++ b/tests/unit/cli/test_run_sglang.py @@ -471,7 +471,7 @@ def test_sglang_python_reaches_the_launcher( result, seen = _run_with(monkeypatch, tmp_path, "--sglang-python", sys.executable) assert result.exit_code == 0, result.output - assert seen["python"] == str(Path(sys.executable).resolve()) + assert seen["python"] == sys.executable def test_a_venv_directory_is_resolved_to_its_interpreter( @@ -486,7 +486,9 @@ def test_a_venv_directory_is_resolved_to_its_interpreter( result, seen = _run_with(monkeypatch, tmp_path, "--sglang-python", str(venv)) assert result.exit_code == 0, result.output - assert seen["python"] == str(interpreter.resolve()) + # The venv's own path, not the symlink target: resolving it would hand + # sglang the base interpreter, which has none of the venv's packages. + assert seen["python"] == str(interpreter) def test_a_bad_interpreter_is_refused_before_the_graph_runs( diff --git a/tests/unit/test_lm_sglang.py b/tests/unit/test_lm_sglang.py index 994fa040..811234d1 100644 --- a/tests/unit/test_lm_sglang.py +++ b/tests/unit/test_lm_sglang.py @@ -131,6 +131,24 @@ def _boom(*_args: object, **_kwargs: object) -> None: assert seen["argv_python"] == "/opt/sglang-venv/bin/python" +def test_a_venv_symlink_keeps_the_venvs_bin_not_the_base_installs(tmp_path: Path) -> None: + """A venv's ``bin/python`` is a symlink; resolving it loses the venv. + + This is the failure this whole helper exists to prevent, one level down: + prepending the *resolved* parent puts the base interpreter's bin on PATH, + where none of the venv's console scripts (``ninja``) live. Every console + script is in the venv's own bin, next to the symlink. + """ + bindir = tmp_path / "venv" / "bin" + bindir.mkdir(parents=True) + interpreter = bindir / "python" + interpreter.symlink_to(sys.executable) + + env = sg._child_env(str(interpreter)) # pyright: ignore[reportPrivateUsage] + + assert env["PATH"].split(os.pathsep)[0] == str(bindir) + + def test_the_child_gets_the_interpreters_bin_on_path() -> None: """Console scripts sglang shells out to (``ninja``) live in the venv's bin/. From 2804d482791c4459fb25e25fd8ab44f3d94ba212 Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Mon, 17 Aug 2026 22:04:50 +0000 Subject: [PATCH 08/11] test(examples): prove the stub was called, and keep DSPy's disk cache out of it The attach test asserted only that the stub's answer landed in state. DSPy caches completions on disk across processes (`~/.dspy_cache`, keyed by model id + prompt + params), so that assertion could be satisfied with no server involved at all -- and worse, the test writes an entry that a later *real* run of the same example, with the same model and question, is served instead of calling the GPU. That is exactly what happened while verifying this branch: a run against a live sglang server on two L40S returned the stub's sentence and reported `0 llm calls`. The stub now appends to a hits file on every completion request, the test asserts it was reached, and the fixture disables DSPy's disk and memory caches so the run under test cannot be answered from -- or poison -- a cache shared with the developer's own runs. Signed-off-by: Sean Mauk --- tests/integration/test_examples.py | 28 ++++++++++++++++++----- tests/integration/test_lm_sglang_spawn.py | 5 +++- 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/tests/integration/test_examples.py b/tests/integration/test_examples.py index 4127cee6..54ee432b 100644 --- a/tests/integration/test_examples.py +++ b/tests/integration/test_examples.py @@ -131,7 +131,7 @@ def test_research_bot_loop_with_scripted_lm(runner: CliRunner, tmp_path: Path) - import sys from http.server import BaseHTTPRequestHandler, HTTPServer - MODEL, PORT, ANSWER = sys.argv[1], int(sys.argv[2]), sys.argv[3] + MODEL, PORT, ANSWER, HITS = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] class Handler(BaseHTTPRequestHandler): @@ -148,6 +148,8 @@ def do_GET(self): # noqa: N802 def do_POST(self): # noqa: N802 self.rfile.read(int(self.headers.get("content-length", 0) or 0)) + with open(HITS, "a", encoding="utf-8") as handle: + handle.write(self.path + "\\n") content = f"[[ ## answer ## ]]\\n{ANSWER}\\n\\n[[ ## completed ## ]]" self._send( { @@ -197,20 +199,31 @@ def _declared_model(graph: Path) -> str: @pytest.fixture def openai_stub(tmp_path: Path) -> object: - """Serve one model id on a free loopback port; yield ``(port, model)``. + """Serve one model id on a free loopback port; yield ``(port, model, hits)``. SGLang is GPU-only, so the example's *spawn* path cannot run in CI. Its *attach* path can: a server already serving the requested model is left alone and used as-is, which is the production branch taken here -- no launch argv is stubbed and no engine code is monkeypatched. + + ``hits`` is the file the stub appends to on every completion request. It + exists because DSPy caches responses on disk across processes, keyed by + model id + prompt + params: without proof the stub was reached, this test + passes on a cache entry written by an earlier run (or writes one that a + later *real* run against the same model and question serves instead of + calling the GPU). The cache is disabled here for the same reason. """ + import dspy + + dspy.configure_cache(enable_disk_cache=False, enable_memory_cache=False) graph = EXAMPLES_DIR / "sglang-qa.yaml" model = _declared_model(graph) port = _free_port() + hits = tmp_path / "stub-hits.txt" script = tmp_path / "openai_stub.py" script.write_text(_OPENAI_STUB, encoding="utf-8") proc = subprocess.Popen( - [sys.executable, str(script), model, str(port), _STUB_ANSWER], + [sys.executable, str(script), model, str(port), _STUB_ANSWER, str(hits)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) @@ -227,7 +240,7 @@ def openai_stub(tmp_path: Path) -> object: time.sleep(0.05) else: raise AssertionError("stub server never became ready") - yield port, model + yield port, model, hits finally: proc.terminate() proc.wait(timeout=10) @@ -235,7 +248,7 @@ def openai_stub(tmp_path: Path) -> object: @pytest.mark.integration def test_sglang_qa_attaches_to_a_running_server( - openai_stub: tuple[int, str], runner: CliRunner, tmp_path: Path + openai_stub: tuple[int, str, Path], runner: CliRunner, tmp_path: Path ) -> None: """sglang-qa.yaml binds its LM from the graph, not from --lm-url/--lm-model. @@ -246,7 +259,7 @@ def test_sglang_qa_attaches_to_a_running_server( attaching, and the derived base URL + model configured the DSPy LM that the ``ask`` node ran against. """ - port, _model = openai_stub + port, _model, hits = openai_stub result = runner.invoke( app, [ @@ -269,3 +282,6 @@ def test_sglang_qa_attaches_to_a_running_server( payload = json.loads(lines[-1]) assert payload["status"] == "done" assert payload["state_summary"]["answer"] == _STUB_ANSWER + # The answer proves nothing on its own -- DSPy would serve it from disk + # cache with no server involved at all. + assert hits.exists() and hits.read_text().strip(), "the stub was never called" diff --git a/tests/integration/test_lm_sglang_spawn.py b/tests/integration/test_lm_sglang_spawn.py index 691566f9..19429204 100644 --- a/tests/integration/test_lm_sglang_spawn.py +++ b/tests/integration/test_lm_sglang_spawn.py @@ -13,7 +13,7 @@ import socket import sys import textwrap -from pathlib import Path +from typing import TYPE_CHECKING import pytest @@ -21,6 +21,9 @@ from stargraph.ir import SGLangServer from stargraph.lm import sglang as sg +if TYPE_CHECKING: + from pathlib import Path + pytestmark = pytest.mark.integration From f1b7f93a7f537ddf49e7284bc9e86b578203cda7 Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Mon, 17 Aug 2026 22:59:20 +0000 Subject: [PATCH 09/11] fix(cli): count LM calls off the client, and name cache hits The run summary's "llm calls" number was fed by ToolCallEvent, so it counted tool calls and reported zero for a graph whose only work was an LM call. A `kind: dspy` node calls its LM directly and publishes nothing on the bus, so the count cannot come from there: it now comes from the DSPy client's own history, read at summary time. DSPy's disk cache stays enabled -- it is what makes a re-run cheap -- but a cached run reached no server and must not read like one that did. The history says which completions came back with `cache_hit`, so the summary names them: done in 72ms (0 steps, 1 llm calls, 1 cached) That line is from a run whose endpoint was a dead port. Before this change the same run printed "0 llm calls" and was indistinguishable from a GPU-served one -- which is exactly how a cached answer was mistaken for a live one while verifying the sglang example. The tool counter keeps its meaning under its own name (`tool_call_count`, JSON only); `llm_call_count` and the new `llm_cache_hits` join it in --summary-json. Signed-off-by: Sean Mauk --- docs/reference/cli.md | 13 +++++ src/stargraph/cli/_progress.py | 28 +++++++--- src/stargraph/cli/_summary.py | 8 ++- src/stargraph/cli/run.py | 30 ++++++++++- tests/integration/cli/test_run_progress.py | 7 ++- tests/integration/cli/test_run_summary.py | 38 ++++++++++++- tests/integration/test_examples.py | 9 +++- tests/unit/cli/test_run_lm_usage.py | 63 ++++++++++++++++++++++ 8 files changed, 180 insertions(+), 16 deletions(-) create mode 100644 tests/unit/cli/test_run_lm_usage.py diff --git a/docs/reference/cli.md b/docs/reference/cli.md index b0ef2834..4671c254 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -125,6 +125,19 @@ non-loopback `host` are operator-only and must be re-stated as `--sglang-arg` / `--sglang-host`. Neither is part of the graph hash: like `--lm-url`, an endpoint is an environment binding, not topology. +The end-of-run summary counts LM calls off the DSPy client, not off the event +bus -- a `kind: dspy` node calls its LM directly and publishes no event. DSPy's +disk cache (`~/.dspy_cache`) is on by default and is left on, so the line names +cache hits when there are any: + +``` +✓ done in 1.2s (1 steps, 1 llm calls, 1 cached) +``` + +A cached run reached no server. `--summary-json` carries the same three numbers +as `llm_call_count`, `llm_cache_hits` and `tool_call_count` (tool calls are +counted from `ToolCallEvent`, separately from LM calls). + **Examples** ```bash diff --git a/src/stargraph/cli/_progress.py b/src/stargraph/cli/_progress.py index 5534533c..536fe927 100644 --- a/src/stargraph/cli/_progress.py +++ b/src/stargraph/cli/_progress.py @@ -12,8 +12,10 @@ * :class:`TransitionEvent` boundaries close the prior node line and open the next; sentinel nodes (``__start__``, ``__end__``) are excluded from rendering and from the step count. -* :class:`ToolCallEvent` increments the LLM-call counter and is buffered - on the in-flight node so the closing line names the tool(s) used. +* :class:`ToolCallEvent` increments the tool-call counter and is buffered + on the in-flight node so the closing line names the tool(s) used. It does + *not* count as an LLM call: LM calls leave no event on the bus, so they are + read from the LM client afterwards (:func:`stargraph.cli.run._lm_usage`). * :class:`ToolResultEvent` adds ``result.usage.total_tokens`` to the running total when present; in ``verbose`` mode the result payload is also dumped under the node line. @@ -52,9 +54,13 @@ class ProgressStats: """Aggregated counters captured across a run.""" step_count: int - llm_call_count: int + tool_call_count: int total_tool_tokens: int node_durations_ms: dict[str, int] + llm_call_count: int = 0 + """Completions the LM client recorded -- cache hits included.""" + llm_cache_hits: int = 0 + """How many of those were served from DSPy's cache, without a request.""" @dataclass @@ -84,7 +90,7 @@ def __init__( self._verbose = verbose self._current: _NodeInflight | None = None self._step_counter = 0 - self._llm_calls = 0 + self._tool_calls = 0 self._tool_tokens = 0 self._durations: dict[str, int] = {} self._final_state: dict[str, Any] | None = None @@ -105,12 +111,20 @@ async def consume(self, bus: EventBus) -> None: self.finalize(ev.ts) return - def stats(self) -> ProgressStats: + def stats(self, *, llm_calls: int = 0, llm_cache_hits: int = 0) -> ProgressStats: + """Counters for the run; LM usage is supplied by the caller. + + The bus carries no LM event -- a DSPy node calls its client directly -- + so the two LM numbers come from the client's own history rather than + from anything seen here. + """ return ProgressStats( step_count=self._step_counter, - llm_call_count=self._llm_calls, + tool_call_count=self._tool_calls, total_tool_tokens=self._tool_tokens, node_durations_ms=dict(self._durations), + llm_call_count=llm_calls, + llm_cache_hits=llm_cache_hits, ) def final_state_dict(self) -> dict[str, Any] | None: @@ -167,7 +181,7 @@ def _on_transition(self, ev: Any) -> None: ) def _on_tool_call(self, ev: Any) -> None: - self._llm_calls += 1 + self._tool_calls += 1 if self._current is not None: self._current.tool_calls.append(ev.tool_name) diff --git a/src/stargraph/cli/_summary.py b/src/stargraph/cli/_summary.py index b8e2ac53..8897ed08 100644 --- a/src/stargraph/cli/_summary.py +++ b/src/stargraph/cli/_summary.py @@ -130,6 +130,8 @@ def render( "duration_ms": duration_ms, "step_count": stats.step_count, "llm_call_count": stats.llm_call_count, + "llm_cache_hits": stats.llm_cache_hits, + "tool_call_count": stats.tool_call_count, "total_tool_tokens": stats.total_tool_tokens, "node_durations_ms": dict(stats.node_durations_ms), "artifacts": artifact_relpaths, @@ -143,10 +145,14 @@ def render( # Human summary. status_mark = "✓" if summary.status == "done" else "✗" status_color = "green" if summary.status == "done" else "red" + # Cache hits are named only when there are any: a run served from + # DSPy's disk cache reached no server, and reading "1 llm calls" for + # it would be indistinguishable from a live one. + cached = f", {stats.llm_cache_hits} cached" if stats.llm_cache_hits else "" self._console.print( f"\n[{status_color}]{status_mark} {summary.status}[/{status_color}] " f"in {_fmt_duration(duration_ms)} " - f"({stats.step_count} steps, {stats.llm_call_count} llm calls)" + f"({stats.step_count} steps, {stats.llm_call_count} llm calls{cached})" ) if artifact_relpaths: diff --git a/src/stargraph/cli/run.py b/src/stargraph/cli/run.py index 48d132f7..9db03b9e 100644 --- a/src/stargraph/cli/run.py +++ b/src/stargraph/cli/run.py @@ -32,8 +32,9 @@ import contextlib import ipaddress import os +import sys from pathlib import Path -from typing import TYPE_CHECKING, Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any, cast import anyio import typer @@ -257,6 +258,30 @@ def _lm_endpoint( yield +def _lm_usage() -> tuple[int, int]: + """``(completions, cache hits)`` the configured LM recorded for this run. + + Read off the client instead of the bus: a ``kind: dspy`` node calls its LM + directly and publishes no event, so the progress printer never sees one. + Cache hits are counted apart from the total because DSPy's disk cache is on + by default -- a run answered entirely from it contacted no server, and must + not read like a run that did. + + Zero on both counts when dspy was never imported (no LM node ran) or when + history is disabled: both mean "nothing observable", not "nothing happened". + """ + module = sys.modules.get("dspy") + if module is None: + return (0, 0) + settings: Any = getattr(module, "settings", None) + history = cast( + "list[dict[str, Any]]", + getattr(getattr(settings, "lm", None), "history", None) or [], + ) + hits = sum(1 for entry in history if getattr(entry.get("response"), "cache_hit", False)) + return (len(history), hits) + + def _build_audit_sink(log_file: Path) -> AuditSink: """Choose chained vs legacy sink for ``log_file`` (chain-write, dual-read). @@ -630,6 +655,7 @@ async def _bootstrap_and_drive() -> RunSummary: raise typer.Exit(code=130) from None if not no_summary: + llm_calls, llm_cache_hits = _lm_usage() # Reconstruct final state model from the ResultEvent's snapshot. final_state_dict = progress.final_state_dict() or {} try: @@ -643,7 +669,7 @@ async def _bootstrap_and_drive() -> RunSummary: renderer.render( summary=summary, final_state=final_state, - stats=progress.stats(), + stats=progress.stats(llm_calls=llm_calls, llm_cache_hits=llm_cache_hits), artifacts_dir=artifacts_dir, run_id=run.run_id, checkpoint=ckpt_path, diff --git a/tests/integration/cli/test_run_progress.py b/tests/integration/cli/test_run_progress.py index 2caac517..387a5ebe 100644 --- a/tests/integration/cli/test_run_progress.py +++ b/tests/integration/cli/test_run_progress.py @@ -25,7 +25,7 @@ def _now() -> datetime: @pytest.mark.integration @pytest.mark.anyio -async def test_renders_node_lines_and_counts_llm_calls() -> None: +async def test_renders_node_lines_and_counts_tool_calls() -> None: out = io.StringIO() console = Console(file=out, force_terminal=False, width=120) printer = ProgressPrinter(console) @@ -121,7 +121,10 @@ async def _produce() -> None: s = printer.stats() assert s.step_count == 2 # nodes a, b (sentinels excluded) - assert s.llm_call_count == 1 + # A tool call is a tool call: the bus carries no LM event, so the printer + # has nothing to say about LM calls and leaves that count to the caller. + assert s.tool_call_count == 1 + assert s.llm_call_count == 0 @pytest.mark.integration diff --git a/tests/integration/cli/test_run_summary.py b/tests/integration/cli/test_run_summary.py index b64a8a0f..9a25d9d9 100644 --- a/tests/integration/cli/test_run_summary.py +++ b/tests/integration/cli/test_run_summary.py @@ -43,12 +43,14 @@ def _summary(status: str = "done", duration_ms: int = 1234) -> RunSummary: ) -def _stats() -> ProgressStats: +def _stats(llm_calls: int = 1, cache_hits: int = 0) -> ProgressStats: return ProgressStats( step_count=3, - llm_call_count=1, + tool_call_count=2, total_tool_tokens=42, node_durations_ms={"a": 12, "b": 4200}, + llm_call_count=llm_calls, + llm_cache_hits=cache_hits, ) @@ -78,6 +80,36 @@ def test_human_summary_writes_artifacts_to_disk(tmp_path: Path) -> None: assert "misc" in text # state field listed +def _render_counters(tmp_path: Path, stats: ProgressStats) -> str: + out = StringIO() + console = Console(file=out, force_terminal=False, width=120) + SummaryRenderer(console).render( + summary=_summary(), + final_state=_State(), + stats=stats, + artifacts_dir=tmp_path, + run_id="r-test", + checkpoint=tmp_path / "ck.sqlite", + ) + return out.getvalue() + + +def test_a_cached_run_says_so(tmp_path: Path) -> None: + """DSPy's disk cache is on by default; a run it answered contacted nobody. + + Without this the line is identical to a live run's, which is how a cached + answer gets mistaken for a served one. + """ + text = _render_counters(tmp_path, _stats(llm_calls=2, cache_hits=2)) + assert "2 llm calls, 2 cached" in text + + +def test_a_live_run_carries_no_cache_note(tmp_path: Path) -> None: + text = _render_counters(tmp_path, _stats(llm_calls=2, cache_hits=0)) + assert "2 llm calls)" in text + assert "cached" not in text + + def test_verifier_results_rendered(tmp_path: Path) -> None: out = StringIO() console = Console(file=out, force_terminal=False, width=120) @@ -128,6 +160,8 @@ def test_json_mode_emits_machine_readable(tmp_path: Path) -> None: assert payload["duration_ms"] == 500 assert payload["step_count"] == 3 assert payload["llm_call_count"] == 1 + assert payload["llm_cache_hits"] == 0 + assert payload["tool_call_count"] == 2 assert payload["artifacts"] == ["a.py"] assert payload["verifier_results"] == [ {"kind": "static", "passed": True, "duration_ms": 0, "findings": []} diff --git a/tests/integration/test_examples.py b/tests/integration/test_examples.py index 54ee432b..bdc04011 100644 --- a/tests/integration/test_examples.py +++ b/tests/integration/test_examples.py @@ -213,9 +213,11 @@ def openai_stub(tmp_path: Path) -> object: later *real* run against the same model and question serves instead of calling the GPU). The cache is disabled here for the same reason. """ - import dspy + import dspy # pyright: ignore[reportMissingTypeStubs] - dspy.configure_cache(enable_disk_cache=False, enable_memory_cache=False) + dspy.configure_cache( # pyright: ignore[reportUnknownMemberType] + enable_disk_cache=False, enable_memory_cache=False + ) graph = EXAMPLES_DIR / "sglang-qa.yaml" model = _declared_model(graph) port = _free_port() @@ -285,3 +287,6 @@ def test_sglang_qa_attaches_to_a_running_server( # The answer proves nothing on its own -- DSPy would serve it from disk # cache with no server involved at all. assert hits.exists() and hits.read_text().strip(), "the stub was never called" + # ... and the summary must say the same thing the stub's log does. + assert payload["llm_call_count"] == 1 + assert payload["llm_cache_hits"] == 0 diff --git a/tests/unit/cli/test_run_lm_usage.py b/tests/unit/cli/test_run_lm_usage.py new file mode 100644 index 00000000..474b36f8 --- /dev/null +++ b/tests/unit/cli/test_run_lm_usage.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``_lm_usage`` -- reading LM calls (and cache hits) off the client. + +A ``kind: dspy`` node calls its LM directly, so no LM event ever reaches the +bus and the progress printer cannot count one. The summary therefore asks the +configured client afterwards. DSPy's disk cache is on by default, so the count +that matters is really two counts: completions, and how many of those came +back from cache without a request. +""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any + +from stargraph.cli.run import _lm_usage # pyright: ignore[reportPrivateUsage] + +if TYPE_CHECKING: + import pytest + + +def _fake_dspy(history: list[dict[str, Any]] | None) -> SimpleNamespace: + lm = None if history is None else SimpleNamespace(history=history) + return SimpleNamespace(settings=SimpleNamespace(lm=lm)) + + +def _entry(*, cache_hit: bool) -> dict[str, Any]: + response: SimpleNamespace = SimpleNamespace() + if cache_hit: + response.cache_hit = True + return {"response": response} + + +def test_a_run_that_never_imported_dspy_reports_nothing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delitem(sys.modules, "dspy", raising=False) + assert _lm_usage() == (0, 0) + + +def test_no_configured_lm_reports_nothing(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "dspy", _fake_dspy(None)) # pyright: ignore[reportArgumentType] + assert _lm_usage() == (0, 0) + + +def test_every_completion_is_counted(monkeypatch: pytest.MonkeyPatch) -> None: + history = [_entry(cache_hit=False), _entry(cache_hit=False)] + monkeypatch.setitem(sys.modules, "dspy", _fake_dspy(history)) # pyright: ignore[reportArgumentType] + assert _lm_usage() == (2, 0) + + +def test_cached_completions_are_counted_apart_from_the_total( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The whole point: a cached run must be distinguishable from a live one. + + Both runs record the same number of completions; only the second number + says whether a server was ever contacted. + """ + history = [_entry(cache_hit=True), _entry(cache_hit=False), _entry(cache_hit=True)] + monkeypatch.setitem(sys.modules, "dspy", _fake_dspy(history)) # pyright: ignore[reportArgumentType] + assert _lm_usage() == (3, 2) From 8d92770744d1492d723a5eae8642e4316a1e6674 Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Tue, 18 Aug 2026 01:12:10 +0000 Subject: [PATCH 10/11] chore: regenerate uv.lock for the release version bump release-please bumps `pyproject.toml`, `__init__.py`, `openapi.json` and the manifest, but not `uv.lock`, so the lock's own `stargraph` entry sat two releases behind (0.5.3 against a 0.7.0 project). A stale project version leaves the lock unresolved-clean, so every `uv run` re-resolved and rewrote it -- and since the pre-commit ruff and pyright hooks shell through `uv run`, pre-commit failed on the modified lockfile while ruff itself reported no findings. `UV_FROZEN=1` worked around it one commit at a time. Regenerated with plain `uv lock`: no package moved. The diff is the project version plus the marker annotations current uv writes for dependencies it can already resolve (`more-itertools` under keyring, the torch/torchvision split); no url or hash changed. `uv lock --check` is clean and `uv sync --locked` audits the existing environment unchanged. The next release bump will drift it again until release-please regenerates the lock itself. Signed-off-by: Sean Mauk --- uv.lock | 62 ++++++++++++++++++++++++++++----------------------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/uv.lock b/uv.lock index aaf5ee29..58f0edc6 100644 --- a/uv.lock +++ b/uv.lock @@ -1673,7 +1673,7 @@ name = "jaraco-classes" version = "3.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "more-itertools" }, + { name = "more-itertools", marker = "platform_machine != 's390x'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/06/c0/ed4a27bc5571b99e3cff68f8a9fa5b56ff7df1c2251cc715a652ddd26402/jaraco.classes-3.4.0.tar.gz", hash = "sha256:47a024b51d0239c0dd8c8540c6c7f484be3b8fcf0b2d85c13825780d3b3f3acd", size = 11780, upload-time = "2024-03-31T07:27:36.643Z" } wheels = [ @@ -1694,7 +1694,7 @@ name = "jaraco-functools" version = "4.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "more-itertools" }, + { name = "more-itertools", marker = "platform_machine != 's390x'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/27/056e0638a86749374d6f57d0b0db39f29509cce9313cf91bdc0ac4d91084/jaraco_functools-4.4.0.tar.gz", hash = "sha256:da21933b0417b89515562656547a77b4931f98176eb173644c0d35032a33d6bb", size = 19943, upload-time = "2025-12-21T09:29:43.6Z" } wheels = [ @@ -1856,12 +1856,12 @@ name = "keyring" version = "25.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "jaraco-classes" }, - { name = "jaraco-context" }, - { name = "jaraco-functools" }, - { name = "jeepney", marker = "sys_platform == 'linux'" }, - { name = "pywin32-ctypes", marker = "sys_platform == 'win32'" }, - { name = "secretstorage", marker = "sys_platform == 'linux'" }, + { name = "jaraco-classes", marker = "platform_machine != 's390x'" }, + { name = "jaraco-context", marker = "platform_machine != 's390x'" }, + { name = "jaraco-functools", marker = "platform_machine != 's390x'" }, + { name = "jeepney", marker = "platform_machine != 's390x' and sys_platform == 'linux'" }, + { name = "pywin32-ctypes", marker = "platform_machine != 's390x' and sys_platform == 'win32'" }, + { name = "secretstorage", marker = "platform_machine != 's390x' and sys_platform == 'linux'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/43/4b/674af6ef2f97d56f0ab5153bf0bfa28ccb6c3ed4d1babf4305449668807b/keyring-25.7.0.tar.gz", hash = "sha256:fe01bd85eb3f8fb3dd0405defdeac9a5b4f6f0439edbb3149577f244a2e8245b", size = 63516, upload-time = "2025-11-16T16:26:09.482Z" } wheels = [ @@ -4323,8 +4323,8 @@ name = "secretstorage" version = "3.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cryptography" }, - { name = "jeepney" }, + { name = "cryptography", marker = "platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, + { name = "jeepney", marker = "platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } wheels = [ @@ -4517,7 +4517,7 @@ wheels = [ [[package]] name = "stargraph" -version = "0.5.3" +version = "0.7.0" source = { editable = "." } dependencies = [ { name = "aiosqlite" }, @@ -5034,13 +5034,13 @@ resolution-markers = [ "python_full_version < '3.14' and platform_machine != 's390x' and sys_platform == 'darwin'", ] dependencies = [ - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx" }, - { name = "setuptools" }, - { name = "sympy" }, - { name = "typing-extensions" }, + { name = "filelock", marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, + { name = "fsspec", marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, + { name = "jinja2", marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, + { name = "networkx", marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, + { name = "setuptools", marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, + { name = "sympy", marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, + { name = "typing-extensions", marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:33449899ce5496c1b84b4853179d94fd102028ae1407314d9fb956bb79e70d09", upload-time = "2026-07-08T12:26:23Z" }, @@ -5065,13 +5065,13 @@ resolution-markers = [ "python_full_version < '3.14' and platform_machine == 's390x' and sys_platform == 'darwin'", ] dependencies = [ - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx" }, - { name = "setuptools" }, - { name = "sympy" }, - { name = "typing-extensions" }, + { name = "filelock", marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, + { name = "fsspec", marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, + { name = "jinja2", marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, + { name = "networkx", marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, + { name = "setuptools", marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, + { name = "sympy", marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, + { name = "typing-extensions", marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp313-cp313-linux_s390x.whl", hash = "sha256:966d020354f465672dc7dd10d3a5c6cd17d7eb48620aa1d265b48a1f78f06898", upload-time = "2026-07-08T19:29:30Z" }, @@ -5102,9 +5102,9 @@ resolution-markers = [ "python_full_version < '3.14' and platform_machine != 's390x' and sys_platform == 'darwin'", ] dependencies = [ - { name = "numpy" }, - { name = "pillow" }, - { name = "torch", version = "2.13.0", source = { registry = "https://download.pytorch.org/whl/cpu" } }, + { name = "numpy", marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, + { name = "pillow", marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, + { name = "torch", version = "2.13.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "platform_machine != 's390x' and sys_platform == 'darwin'" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torchvision-0.28.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:d483b4aa3f5237569053f749cd1a2b5bb548ca456e40461a5dd087f21149d123", upload-time = "2026-07-08T12:26:40Z" }, @@ -5129,9 +5129,9 @@ resolution-markers = [ "python_full_version < '3.14' and platform_machine == 's390x' and sys_platform == 'darwin'", ] dependencies = [ - { name = "numpy" }, - { name = "pillow" }, - { name = "torch", version = "2.13.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" } }, + { name = "numpy", marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, + { name = "pillow", marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, + { name = "torch", version = "2.13.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "platform_machine == 's390x' or sys_platform != 'darwin'" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torchvision-0.28.0%2Bcpu-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:879ae6d4e2e3651582fb7187eafd535601cb5d019595d47e2c874262a000e88e", upload-time = "2026-07-08T12:26:39Z" }, From 1d105c4639168efd01d9dfde7168971762ca030a Mon Sep 17 00:00:00 2001 From: Sean Mauk Date: Tue, 18 Aug 2026 14:02:52 +0000 Subject: [PATCH 11/11] docs(changelog): entry for the IR lm block (NFR-2 gate) changelog-check guards src/stargraph/ir/ and src/stargraph/schemas/; this branch adds SGLangServer + IRDocument.lm and regenerates both schemas, so the gate wants the entry that says what changed and what it means for the hash. Signed-off-by: Sean Mauk --- CHANGELOG.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 70f6b3e3..ea8e160f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -81,6 +81,18 @@ ## [Unreleased] +### Added (sglang LM block, #199) +- IR gained `SGLangServer` and an optional `IRDocument.lm` block: a graph can + declare the LM endpoint it wants (`provider: sglang`, `model`, `host`, + `port`, `args`, `startup_timeout_s`) instead of the caller passing + `--lm-url`/`--lm-model`. `stargraph run` resolves it before the first node — + attaching to a server already serving that model, else spawning + `python -m sglang.launch_server` and terminating it at run end. The block is + **excluded from the structural graph hash**: an endpoint is an environment + binding, not topology. `args` and a non-loopback `host` are operator-only — + a graph-declared value is refused unless re-stated as `--sglang-arg` / + `--sglang-host`. JSON schemas + openapi.json regenerated. + ### Added (batteries-included, #189) - IR `RuleSpec.when` now accepts mapping sugar alongside the raw CLIPS string: `{node: , : , ...}` compiles to