diff --git a/CHANGELOG.md b/CHANGELOG.md index 70f6b3e..ea8e160 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 diff --git a/docs/how-to/authoring-format.md b/docs/how-to/authoring-format.md index 5ec485a..0afdd93 100644 --- a/docs/how-to/authoring-format.md +++ b/docs/how-to/authoring-format.md @@ -51,6 +51,77 @@ 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. + +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. +- **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. +- **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 + 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 @@ -111,6 +182,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 732a37b..4671c25 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -77,9 +77,66 @@ 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. | +| `--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 -`--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. +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. + +`--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 +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. + +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** @@ -92,6 +149,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 5ba220b..11527c0 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 b362b5f..7d68037 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/examples/README.md b/examples/README.md index f89693b..ffe7ecb 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 0000000..a10f411 --- /dev/null +++ b/examples/sglang-qa.yaml @@ -0,0 +1,37 @@ +# 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 (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: +# 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: LiquidAI/LFM2.5-1.2B-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/authoring.py b/src/stargraph/authoring.py index 211ff12..cdee310 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/_progress.py b/src/stargraph/cli/_progress.py index 5534533..536fe92 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 b8e2ac5..8897ed0 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 083ab6c..9db03b9 100644 --- a/src/stargraph/cli/run.py +++ b/src/stargraph/cli/run.py @@ -30,8 +30,11 @@ import asyncio 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 @@ -54,7 +57,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 +72,8 @@ ) if TYPE_CHECKING: + from collections.abc import Callable, Generator + from stargraph.checkpoint.protocol import RunSummary __all__ = ["cmd", "node_kinds"] @@ -76,6 +81,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 +99,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 +114,174 @@ 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 _validated_python(python: str | None) -> str | None: + """Resolve ``--sglang-python`` to an executable interpreter, or fail early. + + 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 + 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.absolute()) + + +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], + install_runtime: bool = False, + sglang_python: str | None = 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, + install_runtime=install_runtime, + python=sglang_python, + ) + ) + lm_model = spec.model + _configure_lm(lm_url, lm_model, lm_key, lm_timeout) + 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). @@ -273,6 +451,64 @@ 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, + 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( + "--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). @@ -287,7 +523,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 +535,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 @@ -340,46 +592,70 @@ 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 + def _echo(message: str) -> None: + if not quiet: + console.print(message, style="dim", highlight=False, markup=False) + + try: + # 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, + install_runtime, + _validated_python(sglang_python), + ): + 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(): + async with broker_lifespan(): + return await _drive_interactive( + run, audit_sink, progress, hitl, console + ) 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() + finally: + await checkpointer.close() + if audit_sink is not None: + await audit_sink.close() - try: - summary = asyncio.run(_bootstrap_and_drive()) + summary = asyncio.run(_bootstrap_and_drive()) except KeyboardInterrupt: console.print("[yellow]cancelled[/yellow]") 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: @@ -393,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/src/stargraph/errors/__init__.py b/src/stargraph/errors/__init__.py index 58904d9..452b4f2 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 1ee1703..6b53ff6 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 ea39595..59b7d1b 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 0a91c96..ab411be 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 0000000..d874658 --- /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/hardware.py b/src/stargraph/lm/hardware.py new file mode 100644 index 0000000..7654a62 --- /dev/null +++ b/src/stargraph/lm/hardware.py @@ -0,0 +1,657 @@ +# 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 +# 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. +_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 + 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: + """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, "pinned_torch": None, + "pinned_torchaudio": None} +spec = importlib.util.find_spec("sglang") +if spec is not None: + try: + 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"] = out["sglang"] or "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"]), + pinned_torch=data.get("pinned_torch"), + pinned_torchaudio=data.get("pinned_torchaudio"), + ) + + +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 _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 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 runtime.sglang is None: + return InstallPlan( + reason=f"sglang is not installed for {runtime.python}", + commands=((*install, "--prerelease=allow", "sglang"),), + notes=(driver_note,), + ) + + 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" + ), + 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", + _SGLANG_CU129_WHL, + "--no-deps", + ) + ) + return InstallPlan(reason=reason, commands=tuple(commands), notes=(driver_note,)) + + +_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. + + 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: + 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 "") + ) + + 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( + plan.reason, + hint=f"sglang documents this platform at {plan.manual or _INSTALL_DOC}", + 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), + ) + 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 = """ +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 new file mode 100644 index 0000000..cee5df6 --- /dev/null +++ b/src/stargraph/lm/sglang.py @@ -0,0 +1,278 @@ +# 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 +from stargraph.lm.hardware import ensure_runtime, ensure_weights + +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, python: str | None = None) -> list[str]: + """Argv for the launch subprocess (monkeypatched in tests).""" + return [ + python or 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 _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() + # 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 + return env + + +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( + argv, + stdout=handle, + stderr=subprocess.STDOUT, + start_new_session=True, + env=_child_env(python), + ) + except OSError as exc: + raise LMServerError( + f"failed to launch sglang: {exc}", + hint=f"is sglang installed for {python or 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, + install_runtime: bool = False, + python: str | 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. + + 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. + + ``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) + 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 + + 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") + 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, python) + 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 94a76a3..888e382 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 bf66a74..55440bc 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/cli/test_run_progress.py b/tests/integration/cli/test_run_progress.py index 2caac51..387a5eb 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 b64a8a0..9a25d9d 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 4a2cfc7..bdc0401 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,179 @@ 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, HITS = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] + + + 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)) + with open(HITS, "a", encoding="utf-8") as handle: + handle.write(self.path + "\\n") + 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, 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 # pyright: ignore[reportMissingTypeStubs] + + 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() + 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, str(hits)], + 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, hits + finally: + proc.terminate() + proc.wait(timeout=10) + + +@pytest.mark.integration +def test_sglang_qa_attaches_to_a_running_server( + 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. + + ``--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, hits = 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 + # 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/integration/test_lm_sglang_spawn.py b/tests/integration/test_lm_sglang_spawn.py new file mode 100644 index 0000000..1942920 --- /dev/null +++ b/tests/integration/test_lm_sglang_spawn.py @@ -0,0 +1,183 @@ +# 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 os +import socket +import sys +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 + + +@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.""" + 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, _python: str | None = None) -> 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, _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) + 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, _python: str | None = None) -> 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") + + +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") + + # 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, str(interpreter)) # pyright: ignore[reportPrivateUsage] + proc.wait(timeout=30) + + seen = log_path.read_text().strip().split(os.pathsep) + assert seen[0] == str(bindir) 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 0000000..474b36f --- /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) diff --git a/tests/unit/cli/test_run_sglang.py b/tests/unit/cli/test_run_sglang.py new file mode 100644 index 0000000..9f94a9d --- /dev/null +++ b/tests/unit/cli/test_run_sglang.py @@ -0,0 +1,512 @@ +# 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 +import sys +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 == [] + + +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"] + + +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 + + +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"] == sys.executable + + +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 + # 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( + 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_authoring_compiler.py b/tests/unit/test_authoring_compiler.py index 8ba6150..2701c26 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 d2abf17..f268122 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_hardware.py b/tests/unit/test_lm_hardware.py new file mode 100644 index 0000000..c9d0ab4 --- /dev/null +++ b/tests/unit/test_lm_hardware.py @@ -0,0 +1,565 @@ +# 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 + +import sys +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 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.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: + _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 + + +_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) + plan = hw.plan_runtime_install(_NVIDIA, hw.Runtime(python="/venv/bin/python")) + assert plan is not None + assert plan.commands == ( + ("uv", "pip", "install", "--python", "/venv/bin/python", "--prerelease=allow", "sglang"), + ) + + +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, _CPU_TORCH_WITH_SGLANG) + assert plan is not None + joined = [" ".join(command) for command in plan.commands] + # 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.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_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, 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 + ) + 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: + """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) + + +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" diff --git a/tests/unit/test_lm_sglang.py b/tests/unit/test_lm_sglang.py new file mode 100644 index 0000000..811234d --- /dev/null +++ b/tests/unit/test_lm_sglang.py @@ -0,0 +1,236 @@ +# 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 os +import sys +from pathlib import Path +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 + +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}) + + +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_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_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/. + + 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. + 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" diff --git a/uv.lock b/uv.lock index aaf5ee2..58f0edc 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" },