From 8899c9cfa6769fbdb0e3b082ea17e7c705aae112 Mon Sep 17 00:00:00 2001 From: Erik Miehling Date: Sun, 25 Oct 2026 22:50:34 -0400 Subject: [PATCH] Add per-request generation on the vLLM backends and an opt-in per-request provider mode The offline vLLM backend runs an AsyncLLM engine on an event loop it owns on a daemon thread, and its session submits one engine request per item. Concurrent generate() calls therefore share one running batch, and each request returns when its own generation finishes. The engine is built under the LLM_CLASS usage context, which keeps the offline scheduler defaults. SteeringPipeline gains supports_concurrent_generate and a prompt lock: each call's prompt phase runs under the lock, and the dispatch runs concurrently only when the backend lacks IN_PROCESS_TORCH and no decoding driver is enabled. The reserved cancel_token keyword cancels a call from another thread (the offline engine aborts the request; the in-process and serve sessions stop between items). ProviderOptions.max_concurrency selects per-request mode in the Inspect provider: each request is its own pipeline call, seeded from the base seed and the request, and a cancelled request aborts its generation. The dispatch mode is recorded in run provenance. The version is 0.5.7. Signed-off-by: Erik Miehling --- AGENTS.md | 44 +- docs/concepts/steering_pipelines.md | 26 + docs/tutorials/evaluate_steering_pipelines.md | 27 + pyproject.toml | 2 +- .../algorithms/core/execution/__init__.py | 3 + .../algorithms/core/execution/backend.py | 25 +- .../algorithms/core/execution/cancellation.py | 74 +++ .../core/execution/session_utils.py | 8 +- .../algorithms/core/steering_pipeline.py | 497 +++++++++++++----- steerability/backends/huggingface/session.py | 22 + steerability/backends/vllm/__init__.py | 8 +- steerability/backends/vllm/backend.py | 213 ++++++-- steerability/backends/vllm/engine_thread.py | 169 ++++++ steerability/backends/vllm/session.py | 171 +++++- steerability/evaluation/batching.py | 286 +++++++--- steerability/evaluation/provider.py | 61 ++- steerability/evaluation/runner.py | 13 +- tests/controls/test_pasta.py | 2 +- tests/controls/test_phased_driver_batching.py | 10 +- .../controls/test_phased_token_boundaries.py | 2 +- tests/core/test_backend_execution.py | 2 +- tests/core/test_cancellation.py | 148 ++++++ tests/core/test_concurrent_generate.py | 384 ++++++++++++++ tests/core/test_driver_rollout_anchor.py | 2 +- tests/core/test_engine_thread.py | 178 +++++++ tests/core/test_model_access.py | 4 +- tests/core/test_staged_steer.py | 4 +- tests/core/test_vllm_backend_construction.py | 82 ++- tests/core/test_vllm_engine.py | 120 +++++ tests/core/test_vllm_plugin_engine.py | 53 ++ tests/core/test_vllm_release.py | 32 ++ tests/core/test_vllm_serve_backend.py | 119 ++++- tests/evaluation/conftest.py | 3 + tests/evaluation/test_batching.py | 270 +++++++++- tests/evaluation/test_provider.py | 57 ++ tests/evaluation/test_runner.py | 22 + tests/utils/runtime_helpers.py | 4 +- uv.lock | 2 +- 38 files changed, 2791 insertions(+), 358 deletions(-) create mode 100644 steerability/algorithms/core/execution/cancellation.py create mode 100644 steerability/backends/vllm/engine_thread.py create mode 100644 tests/core/test_cancellation.py create mode 100644 tests/core/test_concurrent_generate.py create mode 100644 tests/core/test_engine_thread.py diff --git a/AGENTS.md b/AGENTS.md index 68de2417..dbc13158 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -240,6 +240,10 @@ Behaviors that differ from bare Hugging Face usage: `add_generation_prompt`, `return_dict`). The toolkit does not interpret its contents; keys are model-family specific (e.g. `enable_thinking`). Because it rides inside `gen_kwargs`, thinking-on and thinking-off runs get distinct configuration identities in sweeps. +- `cancel_token` is a reserved key inside `gen_kwargs` (a `CancelToken` from `algorithms/core/execution`, else + `TypeError`). Cancelling the token from another thread stops the call at the granularity of the session that runs + it (see Execution backends), and the call raises `GenerationCancelled`. A decoding driver's rollouts receive the + token through their `SteeredSession`. - Token ids are returned as generated on every backend (stop text and any token-boundary overrun stay in the ids); decoded continuation text is truncated at the first stop-string occurrence by one client-side rule. - `generate(..., return_output=True)` returns an `Output` object (or list of them) with fields `output_ids`, @@ -267,6 +271,12 @@ Behaviors that differ from bare Hugging Face usage: state control; the Inspect model provider batches concurrent requests when it is `True` and serializes them otherwise. With state hooks in process, the driver path raises `ValueError` for a batch of more than one prompt when the driver sets `whole_batch_rollouts = False` or the call is seeded under `seed_scope="item"`. +- `pipeline.supports_concurrent_generate` is `True` when the backend lacks `IN_PROCESS_TORCH` (the vLLM backends) and + no enabled decoding driver is present. `generate()` accepts concurrent calls from several threads on any pipeline. + Each call's prompt phase (input controls, templating, tokenization, item assembly) runs under the pipeline's + prompt lock, and the dispatch runs concurrently with other calls only when `supports_concurrent_generate` is + `True`; otherwise the calls run one at a time. `steer()`, `release_backends()`, and `compute_logprobs()` hold the + lock for their whole duration and must not run concurrently with `generate()`. - `pipeline.compute_logprobs(input_ids, ref_output_ids=...)` scores reference tokens teacher-forced with the full steering applied; output controls with `include_in_scoring=False` are excluded from scoring. - Controls with a `tokenizer` attribute left as `None` get the pipeline tokenizer injected automatically. @@ -310,6 +320,23 @@ pipeline = SteeringPipeline( forcing `VLLM_HOOK_WORKER=unified` for `hook_plugin` boots. `serve_environment` returns the same policy as a fresh mapping for a `vllm serve` process. The model-runner constraint is owned by the vLLM-Hook plugin (which pins the legacy runner or supports V2), not the toolkit. +- The offline engine is vLLM's `AsyncLLM`, constructed from `AsyncEngineArgs(model=..., **engine_kwargs)` on an + event loop that the backend owns on a daemon thread (`backends/vllm/engine_thread.py`). The session submits one + engine request per item, and the engine admits each request as it arrives. Concurrent `generate()` calls therefore + share one running batch, and each request returns when it finishes. The vLLM-Hook plugin's patch of `AsyncLLM.generate` + performs each request's admission, spec staging, capture attachment, and cleanup. `release()` cancels in-flight + requests (their callers raise `EngineClosedError`), shuts the engine down on its loop, and joins the thread. The + engine is built under vLLM's `LLM_CLASS` usage context, which keeps the offline scheduler defaults (`max_num_seqs`, + `max_num_batched_tokens`). Per-request lifecycle logging defaults off (`enable_log_requests=False`), and the + periodic stats logging keeps vLLM's default (on, which reports preemptions); an explicit entry in `engine_kwargs` + takes precedence for either. +- Sessions honor a `cancel_token` at different granularity. The offline vLLM session aborts the call's unfinished + engine requests. The in-process session checks the token before each decode and lets a running `model.generate` + finish, and the serve session checks it before each HTTP request and lets a request already sent complete. +- On `vllm-serve`, a single-prompt call issues one HTTP request on the caller's thread, and the server batches + requests from concurrent calls. `request_timeout` must cover one full generation, `max_retries` applies to + transport errors only, and `max_concurrency` bounds the fan-out inside one call rather than the number of + concurrent calls. Launching and supervising the server stays with the caller. - Structural controls train on the staged model and serve their artifacts (checkpoint or LoRA) on vLLM backends. An engine serves the last checkpoint and one LoRA adapter, and `check()` fails a second adapter; on `vllm-serve`, a checkpoint is served only when its path is the spec's `model`. @@ -493,6 +520,15 @@ dispatch carries `seed_scope` from `ProviderOptions` (default `"dispatch"`), so the Hugging Face backend, and bitwise reproducibility of stochastic sampling is not preserved under concurrency (see the `steerability/evaluation/batching.py` module docstring for the full contract). +`ProviderOptions(max_concurrency=N)` selects per-request mode on a pipeline whose `supports_concurrent_generate` is +`True`. Each Inspect request is then dispatched at once as its own single-prompt `pipeline.generate()` call in a +worker thread, with at most `N` in flight, and `max_connections()` returns `N`. A seeded request samples under +`derive_item_seed(seed, f"request:{digest}", 0)`, where `digest` hashes the request's prompt and per-sample runtime +kwargs. The request therefore reproduces regardless of concurrency and arrival order. A cancelled Inspect request cancels its +`CancelToken`, which aborts its engine request. On any other pipeline the provider logs at info level and stays in +batching mode. The mode is recorded as `dispatch_mode` in each result row's provenance, and `max_concurrency` is part +of the recorded protocol. + There is no results checkpoint: the `.eval` logs under `save_dir/inspect_logs/` are the store, and `eval_set` resumes each (config, trial, suite) cell from them at sample granularity. Because `eval_set` matches task identity only, the runner records the rest of the protocol (toolkit version, provider options and `seed_scope`, generate @@ -578,8 +614,8 @@ own in the common case. Required hooks per category: at every scoring step. `decode` returns `[B * n, T + L]` ids, where `n` is the caller's `num_return_sequences` (default 1). Each row is the padded prompt row as received followed by one candidate's continuation, right-padded, in row-major, candidate-minor order. The pipeline always passes `session=` (a `SteeredSession` containing the - generation's steering entries); drivers issue their rollouts through it, and `model` is None on backends without a - loaded model. `session_generate_items(session, rows, eos_token_ids=(), **gen_kwargs)` (in + generation's steering entries and its `cancel_token`); drivers issue their rollouts through it, and `model` is None + on backends without a loaded model. `session_generate_items(session, rows, eos_token_ids=(), **gen_kwargs)` (in `core/execution/session_utils.py`) runs one session call over unpadded rows and returns one `(continuation_ids, finish_reason)` per row; pass the generation config's eos ids as `eos_token_ids` so that an eos sharing the pad id is kept. @@ -855,6 +891,10 @@ Rules that hold regardless of task: during `steer()` never changes a `config_id` or a `.spipe` fit digest. 17. `Output.finish_reasons` has one entry per row of `output_ids`, enforced at construction together with the `[batch, seq]` shape of `output_ids`. +18. Concurrent `generate()` calls serialize their prompt phase under the pipeline's prompt lock and dispatch + concurrently only when `supports_concurrent_generate` is True; a decoding driver's `decode` never runs + concurrently with another generation of the same pipeline. A cancelled `CancelToken` stops its generation at the + session's stated granularity and surfaces as `GenerationCancelled`. ## Pointers diff --git a/docs/concepts/steering_pipelines.md b/docs/concepts/steering_pipelines.md index 7897adc2..27d2207a 100644 --- a/docs/concepts/steering_pipelines.md +++ b/docs/concepts/steering_pipelines.md @@ -154,6 +154,28 @@ with SteeringPipeline(controls=[caa], backend="vllm") as pipeline: # the engine is shut down on exit ``` +### Concurrent generation + +A pipeline accepts `generate()` calls from several threads at once. Each call's prompt phase (input controls, chat +templating, tokenization) runs under a lock owned by the pipeline, since input controls keep state between calls. On +the vLLM backends the dispatch then runs concurrently with other calls. The offline engine admits each request into +its running batch as soon as there is room, and a request returns as soon as its own generation finishes. The +`supports_concurrent_generate` property reports whether this applies. It is `False` on the Hugging Face backend +(one loaded model) and under an enabled decoding driver, and concurrent calls on such a pipeline run one at a time. + +A call can be cancelled from another thread by passing a `CancelToken` (via `cancel_token`) and cancelling it. On the +offline vLLM engine this aborts the call's requests, which frees their slots in the running batch. The in-process and +serve backends stop between items and let a running decode or HTTP request finish. In each case the call raises +`GenerationCancelled`. + +```python +from steerability.algorithms.core.execution import CancelToken + +token = CancelToken() +# on another thread: token.cancel() +response = pipeline.generate(text="...", max_new_tokens=64, cancel_token=token) +``` + ### Evaluation The `SteeringEval` runner forwards its `backend` and `fit` arguments to the pipelines it builds and checks support over @@ -189,6 +211,10 @@ spec = BackendSpec( ) ``` +Note that `request_timeout` (default 120 seconds) bounds each HTTP request and must cover one full generation. +The `max_retries` option applies to transport errors only, and `max_concurrency` bounds the requests issued in +parallel by one `generate()` call rather than the number of concurrent calls. + Serving activation interventions through the vLLM-Hook plugin additionally requires the plugin in the serving environment, `VLLM_HOOK_WORKER=unified` and eager execution on the server, and `hook_plugin: True` on the spec. Artifacts reach the server either through `artifact_dir`, the server's registry directory (its diff --git a/docs/tutorials/evaluate_steering_pipelines.md b/docs/tutorials/evaluate_steering_pipelines.md index bcba0933..f33765f9 100644 --- a/docs/tutorials/evaluate_steering_pipelines.md +++ b/docs/tutorials/evaluate_steering_pipelines.md @@ -41,6 +41,7 @@ where the `ProviderOptions` dataclass contains the provider's configuration: - `on_unsupported_param`: the policy for `GenerateConfig` parameters the pipeline cannot honor (`"raise"` by default or `"warn"`) - `seed_scope`: how a seeded dispatch applies its seed to the rows of a batch (see below) +- `max_concurrency`: selects per-request mode on the vLLM backends (see below) The provider decides how to deliver prompts to the pipeline when it is constructed. With a chat-templated tokenizer, prompts are dispatched as `messages=` and every input control participates as it does in @@ -92,6 +93,32 @@ batch compositions, since padded-batch numerics differ from single-item numerics Trial-to-trial variation under sampling is therefore measured rather than eliminated, which is the role of `num_trials` and the per-metric standard error. +### Per-request mode + +On the vLLM backends we can instead set `max_concurrency` in `ProviderOptions`. The provider then dispatches each +request at once as its own pipeline call, with at most `max_concurrency` calls in flight. The engine admits each +request into its running batch as soon as there is room. This means that a short generation returns without +waiting for the longest generation of a batch. Per-request mode applies when the pipeline's `supports_concurrent_generate` is +`True`, i.e., on a vLLM backend without an enabled decoding driver. For other pipelines the provider logs at info +level and keeps batching. The mode in effect is recorded as `dispatch_mode` in the run provenance. + +```python +options = ProviderOptions(max_concurrency=64) +model = as_inspect_model(pipeline, options=options) +``` + +The engine's `max_num_seqs` (an `engine_kwargs` entry) bounds the running batch, and the engine queues the requests +beyond it. This means that `max_concurrency` should be at least `max_num_seqs` to keep the batch full. Note that +`max_num_seqs` should be sized against the KV cache budget (`gpu_memory_utilization`, `max_model_len`), since a +running batch that outgrows the cache preempts sequences. The engine's periodic stats log reports preemptions. It is +on by default and can be turned off with `disable_log_stats=True` in `engine_kwargs`. + +Per-request mode changes the reproducibility contract. A seeded request samples under a seed derived from the base +seed and the request itself (its prompt and per-sample runtime kwargs). This means that a seeded request reproduces +regardless of concurrency and arrival order, and two requests with the same prompt and per-sample kwargs sample the +same stream. The `seed_scope` setting has no effect in this mode, since every dispatch is one request. A request +that Inspect cancels also aborts its generation on the engine. + ## Suites and the runner An `InspectSuite` specifies a set of tasks evaluated together. `SteeringEval` runs each configuration diff --git a/pyproject.toml b/pyproject.toml index 2ab463d2..04dd5214 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "steerability" -version = "0.5.6" +version = "0.5.7" description = "Steerability Toolkit" readme = "README.md" license = { text = "Apache-2.0" } diff --git a/steerability/algorithms/core/execution/__init__.py b/steerability/algorithms/core/execution/__init__.py index 9c8b105b..c81a47e4 100644 --- a/steerability/algorithms/core/execution/__init__.py +++ b/steerability/algorithms/core/execution/__init__.py @@ -13,6 +13,7 @@ capabilities_for_spec, resolve_backend_class, ) +from steerability.algorithms.core.execution.cancellation import CancelToken, GenerationCancelled from steerability.algorithms.core.execution.contracts import ( Alternative, BackendCapabilities, @@ -71,6 +72,7 @@ "Backend", "BackendCapabilities", "BackendSpec", + "CancelToken", "Capability", "CaptureKinds", "ConstraintEntry", @@ -79,6 +81,7 @@ "as_constraint_source", "CaptureResult", "CheckpointArtifact", + "GenerationCancelled", "GenerationItem", "GenerationParams", "HookEntry", diff --git a/steerability/algorithms/core/execution/backend.py b/steerability/algorithms/core/execution/backend.py index d9d84743..dd172ed3 100644 --- a/steerability/algorithms/core/execution/backend.py +++ b/steerability/algorithms/core/execution/backend.py @@ -14,6 +14,7 @@ import torch +from steerability.algorithms.core.execution.cancellation import CancelToken from steerability.algorithms.core.execution.contracts import ( BackendCapabilities, Capability, @@ -42,6 +43,9 @@ class SteeringSession(Protocol): holds only per-operation state. Session-contract facts, provided by every backend and therefore never capability atoms, include token-id prompts, stop rules, minimum tokens, multiple candidates, seeded sampling, prompt-logprob scoring, and the model layout. + + `generate` and `score` accept a `cancel_token`. A session raises `GenerationCancelled` once + it observes the token cancelled, at the granularity its docstring states. """ @property @@ -53,6 +57,8 @@ def generate( self, items: Sequence[GenerationItem], params: GenerationParams, + *, + cancel_token: CancelToken | None = None, ) -> list[ItemResult]: """Generate one result per item, in item order.""" ... @@ -61,6 +67,8 @@ def score( self, items: Sequence[ScoringItem], params: GenerationParams, + *, + cancel_token: CancelToken | None = None, ) -> torch.Tensor: """Teacher-forced log-probabilities of each item's reference tokens, shape `[num_items, ref_len]`.""" @@ -88,20 +96,25 @@ class SteeredSession: generation's hooks ambiently for the span of the driver's decode; on spec-consuming backends the injected entries are the generation's lowered interventions with prompt-relative scopes rewritten to absolute positions at the generation's original prompt - boundary, so re-prefilled continuation tokens are steered at their original positions. + boundary, so re-prefilled continuation tokens are steered at their original positions. The + wrapper also forwards the generation's `CancelToken` on every `generate` and `score`, and a + driver's rollouts therefore stop with their generation. Attributes: inner: The wrapped backend session. state_entries: Entries injected ahead of each item's own. + cancel_token: The generation's cancellation token, forwarded to the wrapped session, or + None. generated_tokens: Running non-pad token count over every `ItemResult.output.output_ids` this wrapper has returned from `generate`, across all of a driver's rollouts. The pad id is read from the wrapped session's tokenizer; `numel()` counts every position when there is none. The pipeline reads this total after the driver's `decode` returns. """ - def __init__(self, inner, state_entries: tuple = ()): + def __init__(self, inner, state_entries: tuple = (), cancel_token: CancelToken | None = None): self.inner = inner self.state_entries = tuple(state_entries) + self.cancel_token = cancel_token self.generated_tokens = 0 @property @@ -124,7 +137,9 @@ def _inject(self, item): def generate(self, items, params): """Generate with the wrapper's entries injected into every item, accumulating the non-pad token count of every returned output into `generated_tokens`.""" - results = self.inner.generate([self._inject(item) for item in items], params) + results = self.inner.generate( + [self._inject(item) for item in items], params, cancel_token=self.cancel_token, + ) pad_token_id = getattr(getattr(self.inner, "tokenizer", None), "pad_token_id", None) for result in results: output_ids = result.output.output_ids @@ -136,7 +151,9 @@ def generate(self, items, params): def score(self, items, params): """Score with the wrapper's entries injected into every item.""" - return self.inner.score([self._inject(item) for item in items], params) + return self.inner.score( + [self._inject(item) for item in items], params, cancel_token=self.cancel_token, + ) def capture(self, prompts, layers, mode, location="layer_output"): """Capture through the wrapped session under the wrapper's entries. diff --git a/steerability/algorithms/core/execution/cancellation.py b/steerability/algorithms/core/execution/cancellation.py new file mode 100644 index 00000000..e040bea8 --- /dev/null +++ b/steerability/algorithms/core/execution/cancellation.py @@ -0,0 +1,74 @@ +"""Cooperative cancellation for one generation call: `CancelToken` and `GenerationCancelled`. + +A caller creates one token per `SteeringPipeline.generate()` call (the reserved `cancel_token` +keyword) and cancels it from any thread. The pipeline forwards the token to the session that runs +the call, and each session honors it at the granularity its docstring states. +""" +import logging +import threading +from collections.abc import Callable + +logger = logging.getLogger(__name__) + + +class GenerationCancelled(RuntimeError): + """A generation stopped because its `CancelToken` was cancelled before it completed.""" + + +class CancelToken: + """A thread-safe, one-way cancellation signal for one generation call. + + `cancel()` is idempotent and runs every registered callback once, on the cancelling thread. A + callback registered after cancellation runs at once, on the registering thread. Callback + exceptions are logged and do not propagate. Sessions honor the token as their docstrings state: + the offline vLLM session aborts the call's engine requests, and the in-process and serve + sessions stop between items and let a running decode or request finish. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._cancelled = False + self._callbacks: list[Callable[[], None]] = [] + + def cancel(self) -> None: + """Cancel the token and run the registered callbacks; later calls do nothing.""" + with self._lock: + if self._cancelled: + return + self._cancelled = True + callbacks, self._callbacks = self._callbacks, [] + for callback in callbacks: + self._run(callback) + + @property + def cancelled(self) -> bool: + """Whether `cancel()` has been called.""" + return self._cancelled + + def on_cancel(self, callback: Callable[[], None]) -> None: + """Register a zero-argument callback to run once when the token is cancelled. + + Args: + callback: The callable to run. It runs at once when the token is already cancelled. + """ + with self._lock: + if not self._cancelled: + self._callbacks.append(callback) + return + self._run(callback) + + def raise_if_cancelled(self) -> None: + """Raise `GenerationCancelled` when the token is cancelled. + + Raises: + GenerationCancelled: If `cancel()` has been called. + """ + if self._cancelled: + raise GenerationCancelled("The generation was cancelled through its CancelToken.") + + @staticmethod + def _run(callback: Callable[[], None]) -> None: + try: + callback() + except Exception: + logger.warning("A CancelToken callback raised; continuing.", exc_info=True) diff --git a/steerability/algorithms/core/execution/session_utils.py b/steerability/algorithms/core/execution/session_utils.py index a116193e..799004fe 100644 --- a/steerability/algorithms/core/execution/session_utils.py +++ b/steerability/algorithms/core/execution/session_utils.py @@ -237,15 +237,15 @@ def _require_rollouts(self) -> None: "higher." ) - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): """Generate through the venue session; requires `ModelAccess.ROLLOUTS` or higher.""" self._require_rollouts() - return self.inner.generate(items, params) + return self.inner.generate(items, params, cancel_token=cancel_token) - def score(self, items, params): + def score(self, items, params, *, cancel_token=None): """Score through the venue session; requires `ModelAccess.ROLLOUTS` or higher.""" self._require_rollouts() - return self.inner.score(items, params) + return self.inner.score(items, params, cancel_token=cancel_token) def capture(self, prompts, layers, mode, location="layer_output"): """Capture through the venue session; requires `ModelAccess.CAPTURE` or higher.""" diff --git a/steerability/algorithms/core/steering_pipeline.py b/steerability/algorithms/core/steering_pipeline.py index 123460db..6022760a 100644 --- a/steerability/algorithms/core/steering_pipeline.py +++ b/steerability/algorithms/core/steering_pipeline.py @@ -3,6 +3,7 @@ """ import contextlib import logging +import threading import time import traceback import warnings @@ -18,6 +19,7 @@ from steerability.algorithms.core.execution.access import ModelAccess, PlannedFit, PlannedStep, SteerPlan from steerability.algorithms.core.execution.backend import SteeredSession, capabilities_for_spec, resolve_backend_class +from steerability.algorithms.core.execution.cancellation import CancelToken from steerability.algorithms.core.execution.contracts import ( BackendCapabilities, Capability, @@ -70,7 +72,7 @@ resolve_token_prompt, ) from steerability.algorithms.input_control.base import InputControl -from steerability.algorithms.output_control.base import OutputControl +from steerability.algorithms.output_control.base import DecodingDriver, OutputControl from steerability.algorithms.state_control.base import StateControl from steerability.algorithms.structural_control.base import StructuralControl from steerability.utils.tokenization import ensure_pad_token, to_left_pad @@ -93,6 +95,62 @@ def _split_generated_tokens(total: int | None, num_rows: int) -> list[int | None return [base + (1 if i < remainder else 0) for i in range(num_rows)] +@dataclass(frozen=True, slots=True) +class _DriverDecode: + """The decoding-driver inputs of one generation, assembled in its prompt phase. + + Attributes: + driver: The enabled `DecodingDriver`. + logits_processors: The composed logits-processor stack passed to `decode`. + stopping_criteria: The composed stopping-criteria stack passed to `decode`. + state_entries: The generation's state-control entries. + rollout_entries: The entries the `SteeredSession` injects into each rollout; empty when the + session hosts the state entries as in-process hooks for the span of the decode. + hooks_in_process: Whether the session hosts the state entries as in-process hooks. + runtime_kwargs: The call's runtime kwargs, passed to `decode`. + """ + + driver: DecodingDriver + logits_processors: Any + stopping_criteria: Any + state_entries: tuple + rollout_entries: tuple + hooks_in_process: bool + runtime_kwargs: dict + + +@dataclass(frozen=True, slots=True) +class _PreparedGeneration: + """The assembled inputs of one `generate()` call, built in its prompt phase. + + Attributes: + backend: The inference backend that runs the call. + steered_input_ids: The left-packed steered prompt ids, `[batch, seq_len]`. + steered_attention_mask: The attention mask matching `steered_input_ids`. + params: The call's normalized generation parameters. + items: One `GenerationItem` per prompt row on the default decode path; empty under a driver. + driver_decode: The decoding-driver inputs, or None on the default decode path. + decode_text: Whether the return is decoded to text. + is_single: Whether the caller passed one prompt, which removes the batch dimension of the + return. + return_output: Whether the return is `Output` objects. + return_full_sequence: Whether the returned ids include the prompt. + cancel_token: The call's cancellation token, or None. + """ + + backend: Any + steered_input_ids: torch.Tensor + steered_attention_mask: torch.Tensor + params: GenerationParams + items: tuple[GenerationItem, ...] + driver_decode: _DriverDecode | None + decode_text: bool + is_single: bool + return_output: bool + return_full_sequence: bool + cancel_token: CancelToken | None + + @dataclass(slots=True, eq=False, weakref_slot=True) class SteeringPipeline: """Main steering pipeline for applying various control methods to Hugging Face causal language models. @@ -156,6 +214,11 @@ class SteeringPipeline: unless a preloaded `model=` was passed. - On engine backends, `model` is non-None only while the staged in-process model exists during the steer phase; the stage is freed before the engine boots. + - `generate()` accepts concurrent calls from several threads. Each call's prompt phase runs + under the pipeline's prompt lock, and the dispatch runs concurrently with other calls when + `supports_concurrent_generate` is True; otherwise the calls run one at a time. + `steer()`, `release_backends()`, and `compute_logprobs()` must not run concurrently with + `generate()`. A call is cancelled from another thread through its `cancel_token`. For the state category, list order in `controls` defines the composition surface. List order sets `steer()` order, hook registration order, and execution order for hooks on the same module. PyTorch @@ -207,6 +270,8 @@ class SteeringPipeline: _is_steered: bool = field(default=False, init=False, repr=False) _prompt_warnings: PromptWarnings = field(default_factory=PromptWarnings, init=False, repr=False) + # reentrant so that a control calling back into the pipeline during steer() cannot deadlock + _prompt_lock: threading.RLock = field(default_factory=threading.RLock, init=False, repr=False) def __post_init__(self) -> None: @@ -365,6 +430,27 @@ def supports_batching(self) -> bool: return not any(control.enabled for control in self.state_controls) return True + @property + def supports_concurrent_generate(self) -> bool: + """Return True if concurrent `generate()` calls on this pipeline dispatch concurrently. + + The pipeline's backend must lack `Capability.IN_PROCESS_TORCH`, which excludes the Hugging + Face backend (one loaded model, with hooks registered per item), and no enabled + `DecodingDriver` may be present, since a driver's `decode` runs client-side code with + per-call state. The backend's capabilities are its negotiated capabilities when the + pipeline holds a constructed backend (after `steer()`), and are read from its spec + otherwise. Reading the property never constructs a backend. The result does not depend + on the arguments of a `generate()` call. When it is False, concurrent calls are accepted and + run one at a time. + """ + spec = self._resolve_backend_spec(self.backend) + with self._prompt_lock: + backend = self._backends.get(spec) + capabilities = backend.negotiated_capabilities() if backend is not None else capabilities_for_spec(spec) + if Capability.IN_PROCESS_TORCH in capabilities.atoms: + return False + return resolve_decoding_driver(self.output_controls) is None + def _inject_tokenizer(self) -> None: """Attach the pipeline tokenizer to every control exposing an unset `tokenizer`.""" if self.tokenizer is None: @@ -471,14 +557,17 @@ def release_backends(self) -> None: Subsequent operations construct fresh backends against the same specs. Lowered intervention entries and staged artifacts persist, so a released pipeline remains usable at the cost of re-booting engines on next use. Release is idempotent. Engine-owning - backends shut down deterministically. + backends shut down deterministically. Release must not run concurrently with + `generate()`; a generation whose engine requests are still in flight raises + `EngineClosedError`. """ - backends, self._backends = self._backends, {} - for backend in backends.values(): - try: - backend.release() - except Exception: - logger.warning("Backend release failed", exc_info=True) + with self._prompt_lock: + backends, self._backends = self._backends, {} + for backend in backends.values(): + try: + backend.release() + except Exception: + logger.warning("Backend release failed", exc_info=True) def __enter__(self) -> "SteeringPipeline": """Return the pipeline for use as a context manager.""" @@ -714,7 +803,9 @@ def steer(self, **steer_kwargs) -> None: retried `steer()` starts new engines. The stage frames in the error's traceback are cleared of their local variables, which frees the stage weights even while the caller keeps a reference to the error. A post-mortem debugger shows those frames without their - local variables. A repeated call on an already-steered pipeline is a no-op. + local variables. A repeated call on an already-steered pipeline is a no-op. The steer + holds the pipeline's prompt lock for its whole duration and must not run concurrently + with `generate()`. Warns: UserWarning: If two or more enabled controls declare the same `RUNTIME_KWARGS_SCHEMA` @@ -733,6 +824,11 @@ def steer(self, **steer_kwargs) -> None: ModuleNotFoundError: If a configured backend kind requires an optional dependency that is not installed (e.g. the `vllm` extra). """ + with self._prompt_lock: + self._steer_locked(steer_kwargs) + + def _steer_locked(self, steer_kwargs: dict) -> None: + """Run `steer()` with the prompt lock held.""" if self._is_steered: return @@ -1236,7 +1332,7 @@ def generate( input_ids: Token prompt as a 1-D/2-D integer tensor, `list[int]`, or `list[list[int]]`. **gen_kwargs: Generation parameters in `model.generate` vocabulary, normalized through `GenerationParams` and executed by the inference backend's session - (unlisted keys pass through in process and raise on API backends). Two keys are + (unlisted keys pass through in process and raise on API backends). Three keys are reserved and consumed here rather than forwarded to the backend: - `return_full_sequence: bool` to include the prompt in the returned token IDs. @@ -1251,6 +1347,18 @@ def generate( does not interpret the mapping; keys are model-family specific (e.g. `enable_thinking`), and models whose templates expose no such switch are unaffected. + - `cancel_token: CancelToken` to cancel the call from another thread. The + session that runs the call honors it at its stated granularity (the + offline vLLM session aborts the call's engine requests, and the in-process + and serve sessions stop between items), and a decoding driver's rollouts + receive it through their `SteeredSession`. A call whose token is cancelled + before it completes raises `GenerationCancelled`. + + Concurrent calls from several threads are supported. Each call's prompt phase (input + controls, templating, tokenization, and the assembly of the session's items) runs under the + pipeline's prompt lock. The dispatch then runs concurrently with other calls when + `supports_concurrent_generate` is True. Otherwise each call holds the lock until it returns, + and concurrent calls run one at a time. Returns: See dispatch table above. @@ -1259,8 +1367,9 @@ def generate( RuntimeError: If `steer()` has not yet been called. TypeError: If no prompt source or more than one is provided, if a source fails validation, if `attention_mask` is paired with `text=`/`messages=`, if - `chat_template_kwargs` is paired with `text=`/`input_ids=`, or if - `chat_template_kwargs` is not a mapping. + `chat_template_kwargs` is paired with `text=`/`input_ids=`, if + `chat_template_kwargs` is not a mapping, or if `cancel_token` is neither a + `CancelToken` nor None. ValueError: If a token tensor is not 1-D/2-D, nested token lists are ragged, a text/ chat sequence is empty, `chat_template_kwargs` contains a pipeline-owned template argument, or multiple candidates per prompt (`num_return_sequences`/`n` greater @@ -1269,6 +1378,7 @@ def generate( state controls run as in-process hooks, and either the driver sets `whole_batch_rollouts` to False or the call passes a seed with `seed_scope="item"` (the default). + GenerationCancelled: If `cancel_token` is cancelled before the call completes. """ if not self._is_steered: raise RuntimeError("Must call `.steer()` before `.generate()`.") @@ -1276,6 +1386,9 @@ def generate( runtime_kwargs = runtime_kwargs or {} return_full_sequence = bool(gen_kwargs.pop("return_full_sequence", False)) chat_template_kwargs = gen_kwargs.pop("chat_template_kwargs", None) + cancel_token = gen_kwargs.pop("cancel_token", None) + if cancel_token is not None and not isinstance(cancel_token, CancelToken): + raise TypeError(f"cancel_token must be a CancelToken or None; got {type(cancel_token).__name__}.") if chat_template_kwargs is not None: if not isinstance(chat_template_kwargs, Mapping): @@ -1320,42 +1433,49 @@ def generate( "input_ids= for a [batch * n, gen_len] token return." ) - # resolve the prompt tensors per modality - message_handled: set[int] = set() - if kind == "text": - prompt_input_ids, prompt_attention_mask, is_single = resolve_text_prompt( - payload, - input_controls=self.input_controls, - tokenizer=self.tokenizer, - warnings_state=self._prompt_warnings, - ) - elif kind == "messages": - prompt_input_ids, prompt_attention_mask, message_handled, is_single = resolve_messages_prompt( - payload, - runtime_kwargs, - input_controls=self.input_controls, - tokenizer=self.tokenizer, - chat_template_kwargs=chat_template_kwargs, - ) - else: # tokens - prompt_input_ids, prompt_attention_mask, is_single = resolve_token_prompt( - payload, - attention_mask, - input_controls=self.input_controls, - warnings_state=self._prompt_warnings, - ) + # the prompt phase runs under the prompt lock; the dispatch releases it only when the + # pipeline supports concurrent generation + with self._prompt_lock: + # resolve the prompt tensors per modality + message_handled: set[int] = set() + if kind == "text": + prompt_input_ids, prompt_attention_mask, is_single = resolve_text_prompt( + payload, + input_controls=self.input_controls, + tokenizer=self.tokenizer, + warnings_state=self._prompt_warnings, + ) + elif kind == "messages": + prompt_input_ids, prompt_attention_mask, message_handled, is_single = resolve_messages_prompt( + payload, + runtime_kwargs, + input_controls=self.input_controls, + tokenizer=self.tokenizer, + chat_template_kwargs=chat_template_kwargs, + ) + else: # tokens + prompt_input_ids, prompt_attention_mask, is_single = resolve_token_prompt( + payload, + attention_mask, + input_controls=self.input_controls, + warnings_state=self._prompt_warnings, + ) - return self._execute_generation( - prompt_input_ids=prompt_input_ids, - prompt_attention_mask=prompt_attention_mask, - message_handled=message_handled, - decode_text=(kind != "tokens"), - is_single=is_single, - runtime_kwargs=runtime_kwargs, - return_output=return_output, - return_full_sequence=return_full_sequence, - gen_kwargs=gen_kwargs, - ) + prepared = self._prepare_generation( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + message_handled=message_handled, + decode_text=(kind != "tokens"), + is_single=is_single, + runtime_kwargs=runtime_kwargs, + return_output=return_output, + return_full_sequence=return_full_sequence, + cancel_token=cancel_token, + gen_kwargs=gen_kwargs, + ) + if prepared.driver_decode is not None or not self.supports_concurrent_generate: + return self._dispatch_generation(prepared) + return self._dispatch_generation(prepared) @overload def generate_text( @@ -1365,6 +1485,7 @@ def generate_text( runtime_kwargs: dict | None = ..., return_output: Literal[False] = ..., return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> str: ... @overload @@ -1375,6 +1496,7 @@ def generate_text( runtime_kwargs: dict | None = ..., return_output: Literal[False] = ..., return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> list[str]: ... @overload @@ -1385,6 +1507,7 @@ def generate_text( runtime_kwargs: dict | None = ..., return_output: Literal[True], return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> Output | list[Output]: ... @overload @@ -1395,6 +1518,7 @@ def generate_text( runtime_kwargs: dict | None = ..., return_output: bool = ..., return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> str | list[str] | Output | list[Output]: ... @@ -1405,6 +1529,7 @@ def generate_text( runtime_kwargs: dict | None = None, return_output: bool = False, return_full_sequence: bool = False, + cancel_token: CancelToken | None = None, **gen_kwargs: Any, ) -> str | list[str] | Output | list[Output]: """Generate from plain-text prompts. @@ -1422,12 +1547,15 @@ def generate_text( return_full_sequence: If True, include the prompt in the returned token IDs. Returned ids exclude the prompt by default; do not slice the result by prompt length. + cancel_token: Token that cancels the call from another thread, as for `generate()`. **gen_kwargs: Generation parameters, as for `generate()`. Returns: `str`, `list[str]`, `Output`, or `list[Output]`. """ gen_kwargs["return_full_sequence"] = return_full_sequence + if cancel_token is not None: + gen_kwargs["cancel_token"] = cancel_token return self.generate( text=text, runtime_kwargs=runtime_kwargs, return_output=return_output, **gen_kwargs, ) @@ -1441,6 +1569,7 @@ def generate_messages( runtime_kwargs: dict | None = ..., return_output: Literal[False] = ..., return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> str: ... @overload @@ -1452,6 +1581,7 @@ def generate_messages( runtime_kwargs: dict | None = ..., return_output: Literal[False] = ..., return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> list[str]: ... @overload @@ -1463,6 +1593,7 @@ def generate_messages( runtime_kwargs: dict | None = ..., return_output: Literal[True], return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> Output | list[Output]: ... @overload @@ -1474,6 +1605,7 @@ def generate_messages( runtime_kwargs: dict | None = ..., return_output: bool = ..., return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> str | list[str] | Output | list[Output]: ... @@ -1485,6 +1617,7 @@ def generate_messages( runtime_kwargs: dict | None = None, return_output: bool = False, return_full_sequence: bool = False, + cancel_token: CancelToken | None = None, **gen_kwargs: Any, ) -> str | list[str] | Output | list[Output]: """Generate from chat prompts through the tokenizer's chat template. @@ -1509,12 +1642,15 @@ def generate_messages( return_full_sequence: If True, include the prompt in the returned token IDs. Returned ids exclude the prompt by default; do not slice the result by prompt length. + cancel_token: Token that cancels the call from another thread, as for `generate()`. **gen_kwargs: Generation parameters, as for `generate()`. Returns: `str`, `list[str]`, `Output`, or `list[Output]`. """ gen_kwargs["return_full_sequence"] = return_full_sequence + if cancel_token is not None: + gen_kwargs["cancel_token"] = cancel_token if chat_template_kwargs is not None: gen_kwargs["chat_template_kwargs"] = chat_template_kwargs return self.generate( @@ -1531,6 +1667,7 @@ def generate_tokens( runtime_kwargs: dict | None = ..., return_output: Literal[False] = ..., return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> torch.Tensor: ... @overload @@ -1542,6 +1679,7 @@ def generate_tokens( runtime_kwargs: dict | None = ..., return_output: Literal[True], return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> Output | list[Output]: ... @overload @@ -1553,6 +1691,7 @@ def generate_tokens( runtime_kwargs: dict | None = ..., return_output: bool = ..., return_full_sequence: bool = ..., + cancel_token: CancelToken | None = ..., **gen_kwargs: Any, ) -> torch.Tensor | Output | list[Output]: ... @@ -1564,6 +1703,7 @@ def generate_tokens( runtime_kwargs: dict | None = None, return_output: bool = False, return_full_sequence: bool = False, + cancel_token: CancelToken | None = None, **gen_kwargs: Any, ) -> torch.Tensor | Output | list[Output]: """Generate from already-tokenized prompts. @@ -1583,18 +1723,21 @@ def generate_tokens( return_full_sequence: If True, include the prompt in the returned token IDs. Returned ids exclude the prompt by default; do not slice the result by prompt length. + cancel_token: Token that cancels the call from another thread, as for `generate()`. **gen_kwargs: Generation parameters, as for `generate()`. Returns: `torch.Tensor`, `Output`, or `list[Output]`. """ gen_kwargs["return_full_sequence"] = return_full_sequence + if cancel_token is not None: + gen_kwargs["cancel_token"] = cancel_token return self.generate( input_ids=input_ids, attention_mask=attention_mask, runtime_kwargs=runtime_kwargs, return_output=return_output, **gen_kwargs, ) - def _execute_generation( + def _prepare_generation( self, prompt_input_ids: torch.Tensor, prompt_attention_mask: torch.Tensor | None, @@ -1605,21 +1748,20 @@ def _execute_generation( runtime_kwargs: dict, return_output: bool, return_full_sequence: bool, + cancel_token: CancelToken | None, gen_kwargs: dict, - ) -> str | list[str] | torch.Tensor | Output | list[Output]: - """Run generation from the prompt tensors and shape the return value. + ) -> _PreparedGeneration: + """Assemble one generation's inputs from the prompt tensors; the prompt phase of `generate()`. The token-level input-control chain runs first. The steered prompt tensors are then left-packed, since the sessions run batched forwards in that layout. The later steps (state hooks, per-item prompts, driver inputs, full-sequence returns) use the same layout. Output controls that export generation parameters are merged into the call's `GenerationParams`, and the state-control entries are collected for the backend. With - the default decoding driver, each prompt row becomes a `GenerationItem` that the - inference backend's session runs. An enabled `DecodingDriver` instead runs in the client - process and issues its rollouts through `session=`. In-process state hooks then stay - registered for the whole decode. The returned token IDs contain only the continuation - unless `return_full_sequence` is True. Decoded continuation text is truncated at the - first stop string. + the default decoding driver, each prompt row becomes a `GenerationItem` for the + inference backend's session. With an enabled `DecodingDriver`, the composed processor + and criteria stacks and the rollout entries are assembled for its `decode`. The caller + holds the prompt lock. Args: prompt_input_ids: Prompt token IDs [seq_len] or [batch, seq_len] before the token-level @@ -1635,12 +1777,12 @@ def _execute_generation( return_output: If True, return `Output` (single) or `list[Output]` (batched) regardless of `decode_text`. return_full_sequence: If True, include the prompt in the returned token IDs. + cancel_token: The call's cancellation token, or None. gen_kwargs: Generation parameters, normalized through `GenerationParams` and passed to the decoding driver or the session. Returns: - `str` or `list[str]` when `decode_text` is True, `torch.Tensor` when False, or - `Output`/`list[Output]` when `return_output` is True. + The `_PreparedGeneration` that `_dispatch_generation` runs. Raises: ValueError: If a batch of more than one row reaches a decoding driver while enabled @@ -1739,42 +1881,152 @@ def _execute_generation( model=self.model, **gen_kwargs ) - driver_generated_tokens: int | None = None - with backend.open_session() as session: - if decoding_driver is not None: - # client-side driver path: composed stacks, session-hosted hooks for the span - # of the decode, rollouts through a SteeredSession - logits_processors, stopping_criteria = compose_stacks( - self.output_controls, steered_input_ids, runtime_kwargs, - steered_attention_mask, gen_kwargs, skip_ids=skip_ids, + if decoding_driver is not None: + # client-side driver path: composed stacks, session-hosted hooks for the span of the + # decode, rollouts through a SteeredSession + logits_processors, stopping_criteria = compose_stacks( + self.output_controls, steered_input_ids, runtime_kwargs, + steered_attention_mask, gen_kwargs, skip_ids=skip_ids, + ) + params = GenerationParams.from_gen_kwargs(**gen_kwargs) + for contribution in lowered.values(): + params = merge_lowered_params(params, contribution) + # in process, the session hosts this generation's hooks for the whole decode, so + # rollouts through the session and auxiliary forwards on the loaded model are steered + # alike and the SteeredSession injects nothing (ambient hooks already cover its items); + # on spec-consuming backends the SteeredSession injects a rollout variant of each + # lowered entry whose prompt-relative scopes are rewritten to absolute positions at the + # generation's original prompt boundary + driver_rollout_entries: tuple = () + if state_entries and not hooks_in_process: + driver_rollout_entries = rollout_entries( + state_entries, steered_input_ids, steered_attention_mask, ) - params = GenerationParams.from_gen_kwargs(**gen_kwargs) - for contribution in lowered.values(): - params = merge_lowered_params(params, contribution) - # in process, the session hosts this generation's hooks for the whole decode, - # so rollouts through the session and auxiliary forwards on the live model are - # steered alike and the SteeredSession injects nothing (ambient hooks already - # cover its items); on spec-consuming backends the SteeredSession injects a - # rollout variant of each lowered entry whose prompt-relative scopes are - # rewritten to absolute positions at the generation's original prompt boundary - driver_rollout_entries: tuple = () - if state_entries and hooks_in_process: - applied = session.entries_applied(state_entries) + return _PreparedGeneration( + backend=backend, + steered_input_ids=steered_input_ids, + steered_attention_mask=steered_attention_mask, + params=params, + items=(), + driver_decode=_DriverDecode( + driver=decoding_driver, + logits_processors=logits_processors, + stopping_criteria=stopping_criteria, + state_entries=tuple(state_entries), + rollout_entries=tuple(driver_rollout_entries), + hooks_in_process=hooks_in_process, + runtime_kwargs=runtime_kwargs, + ), + decode_text=decode_text, + is_single=is_single, + return_output=return_output, + return_full_sequence=return_full_sequence, + cancel_token=cancel_token, + ) + + # default path: per-prompt items executed by the session; on backends hosting structured + # outputs natively, declarative constraints lower in place of their in-memory processors + constraint_sources: dict[int, ConstraintSource] = {} + processor_specs: dict[int, ProcessorSpecEntry] = {} + if not hooks_in_process: + constraint_sources = constraint_contributions(self.output_controls, runtime_kwargs) + processor_specs = processor_spec_contributions( + self.output_controls, runtime_kwargs, inference_capabilities, + ) + output_entries = collect_output_entries( + self.output_controls, steered_input_ids, runtime_kwargs, + attention_mask=steered_attention_mask, + skip_ids=skip_ids | frozenset(constraint_sources) | frozenset(processor_specs), + **gen_kwargs, + ) + if constraint_sources or processor_specs: + output_entries = output_entries + tuple( + ConstraintEntry(source=source) for source in constraint_sources.values() + ) + tuple(processor_specs.values()) + user_processors = gen_kwargs.pop("logits_processor", None) or [] + user_criteria = gen_kwargs.pop("stopping_criteria", None) or [] + params = GenerationParams.from_gen_kwargs(**gen_kwargs) + for contribution in lowered.values(): + params = merge_lowered_params(params, contribution) + if user_processors or user_criteria: + extra = dict(params.extra) + if user_processors: + extra["logits_processor"] = user_processors + if user_criteria: + extra["stopping_criteria"] = user_criteria + params = replace(params, extra=extra) + + items = tuple( + GenerationItem( + prompt=PreparedPrompt.from_token_ids( + steered_input_ids[i:i + 1], steered_attention_mask[i:i + 1], + ), + state_entries=state_entry_rows[i] if state_entry_rows is not None else state_entries, + output_entries=output_entries, + ) + for i in range(steered_input_ids.size(0)) + ) + return _PreparedGeneration( + backend=backend, + steered_input_ids=steered_input_ids, + steered_attention_mask=steered_attention_mask, + params=params, + items=items, + driver_decode=None, + decode_text=decode_text, + is_single=is_single, + return_output=return_output, + return_full_sequence=return_full_sequence, + cancel_token=cancel_token, + ) + + def _dispatch_generation( + self, prepared: _PreparedGeneration, + ) -> str | list[str] | torch.Tensor | Output | list[Output]: + """Run one prepared generation on its backend's session and shape the return value. + + With the default decoding driver, the session generates the prepared items. An enabled + `DecodingDriver` instead runs in the client process and issues its rollouts through a + `SteeredSession`, which forwards the call's `cancel_token`. In-process state hooks stay + registered for the whole decode. The returned token IDs contain only the continuation + unless `return_full_sequence` is True. Decoded continuation text is truncated at the first + stop string. The driver path runs under the prompt lock. + + Args: + prepared: The generation's inputs, from `_prepare_generation`. + + Returns: + `str` or `list[str]` when `decode_text` is True, `torch.Tensor` when False, or + `Output`/`list[Output]` when `return_output` is True. + + Raises: + GenerationCancelled: If the call's `cancel_token` is cancelled before it completes. + """ + steered_input_ids = prepared.steered_input_ids + steered_attention_mask = prepared.steered_attention_mask + params = prepared.params + cancel_token = prepared.cancel_token + return_full_sequence = prepared.return_full_sequence + driver_decode = prepared.driver_decode + + driver_generated_tokens: int | None = None + with prepared.backend.open_session() as session: + if driver_decode is not None: + if driver_decode.state_entries and driver_decode.hooks_in_process: + applied = session.entries_applied(driver_decode.state_entries) else: applied = contextlib.nullcontext() - if state_entries: - driver_rollout_entries = rollout_entries( - state_entries, steered_input_ids, steered_attention_mask, - ) - driver_session = SteeredSession(session, driver_rollout_entries) + driver_session = SteeredSession(session, driver_decode.rollout_entries, cancel_token=cancel_token) + if cancel_token is not None: + cancel_token.raise_if_cancelled() with applied: - full_output_ids = decoding_driver.decode( + full_output_ids = driver_decode.driver.decode( input_ids=steered_input_ids, attention_mask=steered_attention_mask, model=self.model, - logits_processors=logits_processors, - stopping_criteria=stopping_criteria, - runtime_kwargs=runtime_kwargs, + logits_processors=driver_decode.logits_processors, + stopping_criteria=driver_decode.stopping_criteria, + runtime_kwargs=driver_decode.runtime_kwargs, session=driver_session, **params.to_gen_kwargs(), ) @@ -1791,50 +2043,7 @@ def _execute_generation( tokenizer=self.tokenizer, ) else: - # default path: per-prompt items executed by the session; on backends hosting - # structured outputs natively, declarative constraints lower in place of their - # live processors - constraint_sources: dict[int, ConstraintSource] = {} - processor_specs: dict[int, ProcessorSpecEntry] = {} - if not hooks_in_process: - constraint_sources = constraint_contributions(self.output_controls, runtime_kwargs) - processor_specs = processor_spec_contributions( - self.output_controls, runtime_kwargs, inference_capabilities, - ) - output_entries = collect_output_entries( - self.output_controls, steered_input_ids, runtime_kwargs, - attention_mask=steered_attention_mask, - skip_ids=skip_ids | frozenset(constraint_sources) | frozenset(processor_specs), - **gen_kwargs, - ) - if constraint_sources or processor_specs: - output_entries = output_entries + tuple( - ConstraintEntry(source=source) for source in constraint_sources.values() - ) + tuple(processor_specs.values()) - user_processors = gen_kwargs.pop("logits_processor", None) or [] - user_criteria = gen_kwargs.pop("stopping_criteria", None) or [] - params = GenerationParams.from_gen_kwargs(**gen_kwargs) - for contribution in lowered.values(): - params = merge_lowered_params(params, contribution) - if user_processors or user_criteria: - extra = dict(params.extra) - if user_processors: - extra["logits_processor"] = user_processors - if user_criteria: - extra["stopping_criteria"] = user_criteria - params = replace(params, extra=extra) - - items = [ - GenerationItem( - prompt=PreparedPrompt.from_token_ids( - steered_input_ids[i:i + 1], steered_attention_mask[i:i + 1], - ), - state_entries=state_entry_rows[i] if state_entry_rows is not None else state_entries, - output_entries=output_entries, - ) - for i in range(steered_input_ids.size(0)) - ] - results = session.generate(items, params) + results = session.generate(list(prepared.items), params, cancel_token=cancel_token) new_tokens, reasons = self._assemble_item_outputs(results) num_candidates = params.n or 1 if return_full_sequence: @@ -1852,8 +2061,8 @@ def _execute_generation( # rows (remainder on the first row) so arm-level sums are preserved; the driverless path # leaves generated_tokens None and consumers fall back to the returned continuation count generated_per_row = _split_generated_tokens(driver_generated_tokens, num_rows) - if return_output: - if is_single: + if prepared.return_output: + if prepared.is_single: # one Output covers every candidate row, so it carries the full dispatch total return Output( output_ids=new_tokens, @@ -1871,14 +2080,14 @@ def _execute_generation( for i in range(num_rows) ] - if not decode_text: + if not prepared.decode_text: return returned_ids # text / chat → decode; decoded continuation text truncates at the first stop string decoded = self.tokenizer.decode(returned_ids, skip_special_tokens=True) if params.stop_strings and not return_full_sequence: decoded = [truncate_at_stop_strings(text, params.stop_strings) for text in decoded] - if is_single: + if prepared.is_single: return decoded[0] return decoded @@ -1926,7 +2135,9 @@ def compute_logprobs( When all pipeline controls support batching, the items share one set of control entries and score in a single pass (inputs are left-padded internally for correct positional alignment); otherwise each item is prepared and scored sequentially. Encoder-decoder - models score in process against the live model. + models score in process against the live model. Scoring holds the pipeline's prompt lock + for its whole duration. Concurrent calls therefore score one at a time and wait for the + prompt phase of a concurrent `generate()`. Args: input_ids: Input token IDs as list or tensor [seq_len] or [batch, seq_len] @@ -1945,6 +2156,20 @@ def compute_logprobs( UnsupportedPipelineError: If an enabled control is score-unsupported on the configured inference backend. """ + with self._prompt_lock: + return self._compute_logprobs_locked( + input_ids, attention_mask, ref_output_ids, runtime_kwargs, **forward_kwargs, + ) + + def _compute_logprobs_locked( + self, + input_ids: list[int] | torch.LongTensor, + attention_mask: torch.Tensor | None, + ref_output_ids: list[int] | torch.LongTensor | None, + runtime_kwargs: dict | None, + **forward_kwargs: Any, + ) -> torch.Tensor: + """Run `compute_logprobs()` with the prompt lock held.""" if not self._is_steered: raise RuntimeError("Must call `.steer()` before `.compute_logprobs()`.") if ref_output_ids is None: diff --git a/steerability/backends/huggingface/session.py b/steerability/backends/huggingface/session.py index 43e752d2..854d5fdb 100644 --- a/steerability/backends/huggingface/session.py +++ b/steerability/backends/huggingface/session.py @@ -7,6 +7,7 @@ import torch from transformers import LogitsProcessorList, PreTrainedModel, PreTrainedTokenizerBase, StoppingCriteriaList +from steerability.algorithms.core.execution.cancellation import CancelToken from steerability.algorithms.core.execution.contracts import UnsupportedOperationError from steerability.algorithms.core.execution.fanout import derive_item_seed from steerability.algorithms.core.execution.params import GenerationParams @@ -365,6 +366,8 @@ def generate( self, items: Sequence[GenerationItem], params: GenerationParams, + *, + cancel_token: CancelToken | None = None, ) -> list[ItemResult]: """Generate one result per item under the item's state and output entries. @@ -389,9 +392,13 @@ def generate( restart their count at the prefill of each call. The restart also applies to hooks that `entries_applied` keeps registered across calls. + A cancelled `cancel_token` is checked before each item's decode on the serial path and + once before a batched decode; a running `model.generate` call finishes. + Args: items: The generation items. params: Normalized generation parameters shared by all items. + cancel_token: Token checked before each decode, or None. Returns: One `ItemResult` per item, in item order. The `finish_reasons` of each result contain @@ -403,6 +410,7 @@ def generate( UnsupportedOperationError: If an item contains a state entry that is not a `HookEntry` or an output entry that is not a `StackEntry`. ValueError: If `params.stop_strings` is set and the session has no tokenizer. + GenerationCancelled: If `cancel_token` is cancelled before a decode starts. """ self._ensure_open() model = self.model @@ -431,6 +439,8 @@ def generate( and len(set(seeds)) == 1 ) if batchable: + if cancel_token is not None: + cancel_token.raise_if_cancelled() return self._generate_batched( items, params, gen_kwargs, user_processors, user_criteria, seeds[0], ) @@ -439,6 +449,8 @@ def generate( results: list[ItemResult] = [] for index, item in enumerate(items): + if cancel_token is not None: + cancel_token.raise_if_cancelled() input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt, drop_leading_pads=len(items) > 1) processors, criteria = self._compose_entry_stacks( item.output_entries, extra_processors=user_processors, extra_criteria=user_criteria, @@ -532,6 +544,8 @@ def score( self, items: Sequence[ScoringItem], params: GenerationParams, + *, + cancel_token: CancelToken | None = None, ) -> torch.Tensor: """Teacher-forced log-probabilities of each item's reference tokens. @@ -543,10 +557,13 @@ def score( batched forward pass (prompts right-padded to a common length, then left-packed together); otherwise items score serially. Stopping criteria never apply. Decoder-only models only; the pipeline's `compute_logprobs` serves encoder-decoder models in-process. + A cancelled `cancel_token` is checked before each item's forward pass on the serial path + and once before a batched pass. Args: items: The scoring items. Every item must carry the same reference length. params: `params.extra` passes through as forward keyword arguments. + cancel_token: Token checked before each forward pass, or None. Returns: Log probabilities of shape `[num_items, ref_len]`. @@ -554,6 +571,7 @@ def score( Raises: ValueError: If items carry differing reference lengths. UnsupportedOperationError: If the model is an encoder-decoder model. + GenerationCancelled: If `cancel_token` is cancelled before a forward pass starts. """ self._ensure_open() model = self.model @@ -572,10 +590,14 @@ def score( raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.") if len(items) > 1 and self._entries_identical(items): + if cancel_token is not None: + cancel_token.raise_if_cancelled() return self._score_batched(items, forward_kwargs) all_logprobs: list[torch.Tensor] = [] for item in items: + if cancel_token is not None: + cancel_token.raise_if_cancelled() input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt) input_ids, attention_mask = to_left_pad(input_ids, attention_mask) ref_output_ids = item.ref_output_ids diff --git a/steerability/backends/vllm/__init__.py b/steerability/backends/vllm/__init__.py index 88eab73f..79306182 100644 --- a/steerability/backends/vllm/__init__.py +++ b/steerability/backends/vllm/__init__.py @@ -4,9 +4,11 @@ This package imports cleanly without vLLM installed. The capability tables (`capabilities.py`) are static data used by `check()`; the strict parameter-rendering table and the request/response mapping helpers (`rendering.py`) are plain functions; the two backend classes and their -backend-side helpers live in `backend.py`, and the request sessions in `session.py`. Constructing -`VLLMBackend` requires the `vllm` optional dependency (it boots an engine); `VLLMServeBackend` -needs only a reachable vLLM server. +backend-side helpers live in `backend.py`, and the request sessions in `session.py`. The offline +backend runs its `AsyncLLM` engine on the event-loop thread in `engine_thread.py`, and its +sessions submit one engine request per item, which lets concurrent calls share the engine's +running batch. Constructing `VLLMBackend` requires the `vllm` optional dependency (it boots an +engine); `VLLMServeBackend` needs only a reachable vLLM server. """ from steerability.backends.vllm.backend import VLLMBackend, VLLMServeBackend from steerability.backends.vllm.capabilities import VLLM_BASELINE_CAPABILITIES diff --git a/steerability/backends/vllm/backend.py b/steerability/backends/vllm/backend.py index ab3f49fa..81c3c9d7 100644 --- a/steerability/backends/vllm/backend.py +++ b/steerability/backends/vllm/backend.py @@ -2,13 +2,15 @@ (`"vllm-serve"`). This module holds the backend-side helpers (the artifact uploader, layout and tokenizer -resolution, discovery reconciliation) and the two backend classes. It imports cleanly without -vLLM installed: constructing `VLLMBackend` requires the `vllm` optional dependency (it boots an -engine); `VLLMServeBackend` needs only a reachable vLLM server. Every `from vllm ...` and -`from vllm_hook_plugins ...` import stays inside a function or method body. +resolution, discovery reconciliation, the engine request coroutines) and the two backend classes. +It imports cleanly without vLLM installed: constructing `VLLMBackend` requires the `vllm` optional +dependency (it boots an engine); `VLLMServeBackend` needs only a reachable vLLM server. Every +`from vllm ...` and `from vllm_hook_plugins ...` import stays inside a function or method body. """ +import concurrent.futures import gc import hashlib +import itertools import json import logging import os @@ -16,6 +18,8 @@ import urllib.request import uuid from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any import torch @@ -33,6 +37,7 @@ from steerability.algorithms.core.internals.fingerprint import config_digest_dict, is_absent_chat_template_fingerprint from steerability.algorithms.core.internals.model_layout import config_head_geometry, text_config from steerability.backends.vllm.capabilities import _intersect_with_discovery, _reconcile_discovery, _vllm_capabilities +from steerability.backends.vllm.engine_thread import EngineThread from steerability.backends.vllm.environment import engine_boot_environment, engine_environment from steerability.backends.vllm.rendering import raise_for_spec_rejection from steerability.backends.vllm.session import VLLMOfflineSession, VLLMServeSession @@ -133,6 +138,65 @@ def _client_tokenizer(source: str, trust_remote_code: bool = False): return ensure_pad_token(tokenizer) +async def _construct_engine(model_ref: str, engine_kwargs: dict): + """An `AsyncLLM` built on the running loop, which the engine's own tasks then belong to. + + The `LLM_CLASS` usage context is required for the offline scheduler defaults (`max_num_seqs`, + `max_num_batched_tokens`); `AsyncLLM.from_engine_args` otherwise uses `ENGINE_CONTEXT`, for + which vLLM applies its smaller generic `SchedulerConfig` defaults. + """ + from vllm import AsyncEngineArgs + from vllm.usage.usage_lib import UsageContext + from vllm.v1.engine.async_llm import AsyncLLM + + return AsyncLLM.from_engine_args( + AsyncEngineArgs(model=model_ref, **engine_kwargs), usage_context=UsageContext.LLM_CLASS, + ) + + +async def _final_output(engine, prompt, sampling_params, request_id: str, lora_request=None): + """The finished `RequestOutput` of one `AsyncLLM.generate` call.""" + final = None + async for output in engine.generate(prompt, sampling_params, request_id, lora_request=lora_request): + final = output + return final + + +async def _collective_rpc(engine, method: str): + """The per-worker replies of one collective RPC on the engine.""" + return await engine.collective_rpc(method, None, (), None) + + +async def _shutdown(engine) -> None: + """Shut the engine down on the loop its tasks belong to.""" + engine.shutdown() + + +@dataclass(frozen=True, slots=True) +class RequestHandle: + """One engine request submitted through the backend's engine thread. + + Attributes: + request_id: The request id the engine was given. + future: The future that resolves with the request's finished `RequestOutput`. + """ + + request_id: str + future: concurrent.futures.Future + + def result(self, timeout: float | None = None) -> Any: + """Block until the request finishes and return its `RequestOutput`.""" + return self.future.result(timeout) + + def done(self) -> bool: + """Whether the request finished, failed, or was cancelled.""" + return self.future.done() + + def cancel(self) -> bool: + """Cancel the request, which aborts it on the engine; False when it already completed.""" + return self.future.cancel() + + class _DiscoveryNegotiation: """The negotiated capabilities of a vLLM backend that may hold a discovery payload.""" @@ -149,15 +213,28 @@ def negotiated_capabilities(self) -> BackendCapabilities: class VLLMBackend(_DiscoveryNegotiation, Backend): """The offline vLLM engine backend. - Boots one engine per backend instance from the spec (`engine_kwargs` option forwarded to - `vllm.LLM`); requires the `vllm` optional dependency. The last `CheckpointArtifact` overrides - the served model reference, and the `LoRAArtifact` attaches as a LoRA request on every - generation. - When the spec declares `hook_plugin`, the unified worker is selected via - `VLLM_HOOK_WORKER=unified`, and the discovery payload is fetched from the engine at - construction and kept on the backend. `negotiated_capabilities()` narrows the spec's - advertisement to it, while `capabilities_for_spec` gives the spec's advertisement without - constructing the backend. + Boots one vLLM `AsyncLLM` engine per backend instance from the spec (`engine_kwargs` option + forwarded to `AsyncEngineArgs`); requires the `vllm` optional dependency. The engine runs on an + event loop that the backend owns on a daemon thread, and every engine operation is a coroutine + submitted to that loop. The engine's scheduler admits requests from any thread as they + arrive, and concurrent session calls therefore share one running batch. On a + `hook_plugin` engine, the plugin's patch of `AsyncLLM.generate` stages each request's + intervention and capture specs on the workers, attaches its capture payload, and clears its + state when the request finishes or is aborted. + + The last `CheckpointArtifact` overrides the served model reference, and the `LoRAArtifact` + attaches as a LoRA request on every generation. When the spec declares `hook_plugin`, the + unified worker is selected via `VLLM_HOOK_WORKER=unified`, and the discovery payload is + fetched from the engine at construction and kept on the backend. `negotiated_capabilities()` + narrows the spec's advertisement to it, while `capabilities_for_spec` gives the spec's + advertisement without constructing the backend. + + The engine is built under vLLM's `LLM_CLASS` usage context, which keeps the scheduler defaults + (`max_num_seqs`, `max_num_batched_tokens`) of an offline `vllm.LLM` engine. Without it vLLM + applies its smaller generic scheduler defaults. Per-request lifecycle logging is off by default + (`enable_log_requests=False`), and the engine's periodic stats logging keeps vLLM's default + (`disable_log_stats=False`), which reports preemptions. An explicit `enable_log_requests` or + `disable_log_stats` entry in `engine_kwargs` takes precedence. """ def __init__(self, spec: BackendSpec, artifacts: Sequence[Artifact] = ()) -> None: @@ -165,10 +242,11 @@ def __init__(self, spec: BackendSpec, artifacts: Sequence[Artifact] = ()) -> Non raise ValueError(f"VLLMBackend requires a 'vllm' spec; got kind {spec.kind!r}.") self.spec = spec self._released = False + self._engine = None + self._engine_thread: EngineThread | None = None + self._lora_request = None require("vllm") - from vllm import LLM - checkpoint, lora = split_artifacts(artifacts) model_ref = checkpoint.path if checkpoint is not None else spec.model if model_ref is None: @@ -188,17 +266,21 @@ def __init__(self, spec: BackendSpec, artifacts: Sequence[Artifact] = ()) -> Non # worker hooks do not run under CUDA-graph replay; spec construction rejects an # explicit False, so this only fills the default engine_kwargs.setdefault("enforce_eager", True) + # per-request lifecycle logging defaults off; stats logging keeps vLLM's default; caller kwargs win + engine_kwargs.setdefault("enable_log_requests", False) - # the boot environment is scoped to this engine's construction and restored afterwards, - # so a later plugin-free engine in the same process is unaffected - forced, defaults = engine_boot_environment(bool(spec.get_option("hook_plugin"))) - with engine_environment(forced, defaults): - self._llm = LLM(model=model_ref, **engine_kwargs) - + self._request_counter = itertools.count() # the pipeline records a backend only when the constructor returns, so a failure from - # here on must release the engine before propagating + # here on must release the engine thread and any engine before propagating try: - self._lora_request = None + self._engine_thread = EngineThread(name=f"vllm-{spec.spec_hash[:8]}") + self._engine_thread.start() + # the boot environment is scoped to this engine's construction and restored + # afterwards, so a later plugin-free engine in the same process is unaffected + forced, defaults = engine_boot_environment(bool(spec.get_option("hook_plugin"))) + with engine_environment(forced, defaults): + self._engine = self._engine_thread.run(_construct_engine(model_ref, engine_kwargs)) + if lora is not None: from vllm.lora.request import LoRARequest @@ -228,17 +310,12 @@ def stage_artifacts(self, payloads) -> None: self._artifact_uploader.upload_payloads(payloads) def _fetch_discovery(self) -> dict | None: - payload = None - for target in (self._llm, getattr(self._llm, "llm_engine", None)): - rpc = getattr(target, "collective_rpc", None) - if callable(rpc): - try: - replies = rpc("hook_capabilities") - except Exception as error: - logger.warning("vLLM-Hook discovery failed: %s", error) - return None - payload = next((reply for reply in replies if reply), None) - break + try: + replies = self._engine_thread.run(_collective_rpc(self._engine, "hook_capabilities")) + except Exception as error: + logger.warning("vLLM-Hook discovery failed: %s", error) + return None + payload = next((reply for reply in replies or () if reply), None) if payload is None: logger.warning( "vLLM-Hook discovery returned no payload; is VLLM_HOOK_WORKER=unified active?" @@ -258,49 +335,73 @@ def open_session(self) -> "VLLMOfflineSession": Raises: RuntimeError: If the backend has been released. """ - self._require_llm() + self._require_engine() return VLLMOfflineSession(self) - def _require_llm(self): - """The live engine, or a `RuntimeError` when the backend has been released.""" - if self._llm is None: + def _require_engine(self): + """The loaded engine, or a `RuntimeError` when the backend has been released.""" + if self._engine is None: raise RuntimeError( "This VLLMBackend was released; construct a new backend (or a new " "SteeringPipeline operation, which reconstructs backends automatically)." ) - return self._llm + return self._engine + + def _submit_request(self, prompt, sampling_params, *, lora_request=None) -> RequestHandle: + """Submit one engine request through the engine thread and return its handle. + + The request's output kind is set to final-only, which makes the engine return one output + when the request finishes. Request ids combine the spec hash with a per-backend counter. + + Args: + prompt: The engine prompt (a `TokensPrompt`, with its `cache_salt` when salted). + sampling_params: The request's `SamplingParams`; its `output_kind` is overwritten. + lora_request: The LoRA request to serve the request under, or None. + + Returns: + The `RequestHandle` whose future resolves with the finished `RequestOutput`. + + Raises: + RuntimeError: If the backend has been released. + EngineClosedError: If the engine thread closes before the request is submitted. + """ + engine = self._require_engine() + from vllm.sampling_params import RequestOutputKind + + sampling_params.output_kind = RequestOutputKind.FINAL_ONLY + request_id = f"{self.spec.spec_hash[:8]}-{next(self._request_counter)}" + future = self._engine_thread.submit( + _final_output(engine, prompt, sampling_params, request_id, lora_request=lora_request) + ) + return RequestHandle(request_id=request_id, future=future) def release(self) -> None: """Shut the engine down explicitly and mark the backend unusable. + In-flight requests are cancelled first, and their waiters raise `EngineClosedError`. The + engine is then shut down on its event loop, and the engine thread is stopped and joined. Release is idempotent; after it, a new backend must be constructed. The distributed-state - teardown is process-global, so release assumes no other live vLLM engine in the process. + teardown is process-global, so release assumes no other loaded vLLM engine in the process. Ray-based executors are out of scope. Engine-touching calls on any still-open session raise - after release. + `RuntimeError` after release. """ if self._released: return self._released = True - llm = self._llm - self._llm = None + engine, self._engine = self._engine, None + engine_thread = self._engine_thread self._lora_request = None - for resolve in ( - lambda: getattr(llm, "shutdown", None), - lambda: getattr(getattr(llm, "llm_engine", None), "shutdown", None), - lambda: getattr( - getattr(getattr(llm, "llm_engine", None), "engine_core", None), "shutdown", None - ), - ): - shutdown = resolve() - if callable(shutdown): + if engine_thread is not None: + engine_thread.cancel_pending() + if engine is not None: try: - shutdown() + engine_thread.run(_shutdown(engine)) except Exception: - logger.warning("vLLM engine shutdown hop failed; continuing.", exc_info=True) - break + logger.warning("vLLM engine shutdown failed; continuing.", exc_info=True) + engine_thread.close() - del llm + del engine gc.collect() try: diff --git a/steerability/backends/vllm/engine_thread.py b/steerability/backends/vllm/engine_thread.py new file mode 100644 index 00000000..4853b295 --- /dev/null +++ b/steerability/backends/vllm/engine_thread.py @@ -0,0 +1,169 @@ +"""The event-loop thread that runs the offline vLLM engine's coroutines. + +`EngineThread` runs one asyncio event loop on a daemon thread. The offline backend constructs its +`AsyncLLM` engine on that loop and submits every engine operation to it as a coroutine, which lets +callers on any thread share one engine. The module has no vLLM dependency; the backend owns what runs on +the loop. +""" +import asyncio +import concurrent.futures +import logging +import threading +from collections.abc import Coroutine +from typing import Any, TypeVar + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + + +class EngineClosedError(RuntimeError): + """The engine thread was closed (the backend released) before a request completed.""" + + +class EngineThread: + """An asyncio event loop on a daemon thread, the single place engine coroutines run. + + `submit` and `run` accept coroutines from any thread. Every future `submit` returns is tracked + until it completes, which lets `cancel_pending` and `close` cancel the work still in flight. Once + closed, the thread accepts no further submissions. + + Args: + name: Thread name, for logs. + """ + + def __init__(self, *, name: str) -> None: + self._name = name + self._lock = threading.Lock() + self._loop: asyncio.AbstractEventLoop | None = None + self._thread: threading.Thread | None = None + self._pending: set[concurrent.futures.Future] = set() + self._closed = False + + @property + def closed(self) -> bool: + """Whether `close()` has been called.""" + return self._closed + + def start(self) -> None: + """Create the event loop and run it on a new daemon thread. + + Raises: + EngineClosedError: If the thread was closed. + RuntimeError: If the thread was already started. + """ + with self._lock: + if self._closed: + raise EngineClosedError(f"Engine thread {self._name} is closed.") + if self._thread is not None: + raise RuntimeError(f"Engine thread {self._name} is already started.") + loop = asyncio.new_event_loop() + ready = threading.Event() + + def run_loop() -> None: + asyncio.set_event_loop(loop) + loop.call_soon(ready.set) + loop.run_forever() + + thread = threading.Thread(target=run_loop, name=self._name, daemon=True) + thread.start() + ready.wait() + self._loop = loop + self._thread = thread + + def submit(self, coro: Coroutine[Any, Any, T]) -> concurrent.futures.Future: + """Schedule `coro` on the loop and return a future for its result. + + Cancelling the returned future cancels the coroutine on the loop, which then sees + `asyncio.CancelledError` at its current await. + + Args: + coro: The coroutine to run. + + Returns: + A `concurrent.futures.Future` that resolves with the coroutine's result or exception. + + Raises: + EngineClosedError: If the thread is closed or was never started; `coro` is closed + without running. + """ + with self._lock: + if self._closed or self._loop is None: + coro.close() + raise EngineClosedError(f"Engine thread {self._name} is closed.") + future = asyncio.run_coroutine_threadsafe(coro, self._loop) + self._pending.add(future) + future.add_done_callback(self._forget) + return future + + def run(self, coro: Coroutine[Any, Any, T]) -> T: + """Run `coro` on the loop and block the calling thread until it completes. + + Args: + coro: The coroutine to run. + + Returns: + The coroutine's result. + + Raises: + EngineClosedError: If the thread is closed, or is closed while `coro` runs. + Exception: The coroutine's own exception, re-raised on the calling thread. + """ + future = self.submit(coro) + try: + return future.result() + except concurrent.futures.CancelledError as error: + raise EngineClosedError(f"Engine thread {self._name} was closed while the call ran.") from error + + def cancel_pending(self) -> None: + """Cancel every submitted future that is not yet done.""" + with self._lock: + pending = list(self._pending) + for future in pending: + future.cancel() + + def close(self) -> None: + """Cancel pending work, stop the loop, join the thread, and close the loop. + + Close is idempotent and thread-safe. Afterwards `submit` and `run` raise + `EngineClosedError`. Must not be called from a coroutine running on the loop. + """ + with self._lock: + if self._closed: + return + self._closed = True + loop, thread = self._loop, self._thread + pending = list(self._pending) + for future in pending: + future.cancel() + if loop is None or thread is None: + return + if thread is threading.current_thread(): + raise RuntimeError(f"Engine thread {self._name} cannot close itself from its own loop.") + loop.call_soon_threadsafe(self._drain_and_stop, loop) + thread.join() + try: + loop.run_until_complete(loop.shutdown_asyncgens()) + except Exception: + logger.debug("Async generator shutdown on %s failed.", self._name, exc_info=True) + loop.close() + + @staticmethod + def _drain_and_stop(loop: asyncio.AbstractEventLoop) -> None: + """Cancel the tasks still on the loop and stop it once their cancellation has run.""" + tasks = [task for task in asyncio.all_tasks(loop) if not task.done()] + for task in tasks: + task.cancel() + if not tasks: + loop.stop() + return + + async def wait_then_stop() -> None: + await asyncio.wait(tasks, timeout=5.0) + loop.stop() + + loop.create_task(wait_then_stop()) + + def _forget(self, future: concurrent.futures.Future) -> None: + with self._lock: + self._pending.discard(future) diff --git a/steerability/backends/vllm/session.py b/steerability/backends/vllm/session.py index f3be9ef5..ca5338df 100644 --- a/steerability/backends/vllm/session.py +++ b/steerability/backends/vllm/session.py @@ -1,10 +1,11 @@ """The vLLM request sessions: the offline-engine session and the server session. -`_RequestSessionBase` holds the lifecycle and layout shared by both; `VLLMOfflineSession` -submits batched engine calls and serves plugin capture, and `VLLMServeSession` fans requests out -over the completions endpoint. This module imports cleanly without vLLM installed; the `vllm` -imports the sessions need stay method-local. +`_RequestSessionBase` holds the lifecycle and layout shared by both. `VLLMOfflineSession` submits +one engine request per item through the backend's engine thread and serves plugin capture, and +`VLLMServeSession` fans requests out over the completions endpoint. This module imports cleanly +without vLLM installed; the `vllm` imports the sessions need stay method-local. """ +import concurrent.futures import json import uuid from collections.abc import Sequence @@ -12,6 +13,7 @@ import torch +from steerability.algorithms.core.execution.cancellation import CancelToken, GenerationCancelled from steerability.algorithms.core.execution.contracts import CaptureKinds, UnsupportedOperationError from steerability.algorithms.core.execution.fanout import ( PartialBatchError, @@ -32,6 +34,7 @@ ) from steerability.algorithms.core.output import Output from steerability.backends.vllm.capabilities import _refuse_by_constraints, _refuse_by_engine_facts +from steerability.backends.vllm.engine_thread import EngineClosedError from steerability.backends.vllm.rendering import ( _load_safetensors_bytes, _split_item_entries, @@ -44,7 +47,7 @@ ) if TYPE_CHECKING: - from steerability.backends.vllm.backend import VLLMBackend, VLLMServeBackend + from steerability.backends.vllm.backend import RequestHandle, VLLMBackend, VLLMServeBackend class _RequestSessionBase: @@ -95,6 +98,48 @@ def layout(self) -> ModelFacts: ) return layout + @staticmethod + def _await_handles(handles: Sequence["RequestHandle"], cancel_token: CancelToken | None) -> list: + """Wait for every submitted engine request and return their outputs in submission order. + + With a token, cancelling it cancels the requests not yet done, which aborts them on the + engine. A request that fails or is interrupted cancels the remaining requests before the + error propagates. + + Args: + handles: The submitted requests, in item order. + cancel_token: The call's cancellation token, or None. + + Returns: + One finished `RequestOutput` per handle. + + Raises: + GenerationCancelled: If the token was cancelled before every request finished. + EngineClosedError: If the backend was released while a request was in flight. + """ + # the token's callback references this list, which is emptied once the call ends so that a + # long-lived token does not keep the finished requests alive + pending = list(handles) + if cancel_token is not None: + cancel_token.on_cancel(lambda: [handle.cancel() for handle in pending if not handle.done()]) + outputs = [] + try: + for handle in handles: + outputs.append(handle.result()) + except concurrent.futures.CancelledError as error: + for handle in handles: + handle.cancel() + if cancel_token is not None and cancel_token.cancelled: + raise GenerationCancelled("The generation was cancelled through its CancelToken.") from error + raise EngineClosedError("The backend was released while a request was in flight.") from error + except BaseException: + for handle in handles: + handle.cancel() + raise + finally: + pending.clear() + return outputs + def _item_seed(self, item: GenerationItem, params: GenerationParams, index: int) -> int | None: if item.seed is not None: return item.seed @@ -183,10 +228,25 @@ def capture( class VLLMOfflineSession(_RequestSessionBase): """Request session over the offline engine. - Token-id prompts submit as `TokensPrompt`s in one engine call with per-item sampling - parameters; the engine schedules the batch internally, so no client-side fan-out is needed. + Each item submits as one `TokensPrompt` engine request with its own sampling parameters, + through the backend's engine thread. The engine's scheduler admits each request as it arrives, + so the items of one call and the items of concurrent calls on any thread share one running + batch, and each request finishes independently. """ + def _submit_all(self, requests, lora_request=None) -> list["RequestHandle"]: + """Submit `(prompt, sampling_params)` pairs in order, cancelling the submitted requests when + a later submission fails.""" + handles: list = [] + try: + for prompt, sampling_params in requests: + handles.append(self._backend._submit_request(prompt, sampling_params, lora_request=lora_request)) + except BaseException: + for handle in handles: + handle.cancel() + raise + return handles + def capture( self, prompts: list[PreparedPrompt], @@ -271,7 +331,7 @@ def capture( ) if intervention is not None: extra_args["intervention_spec"] = intervention.to_wire() - engine_prompts = [] + requests = [] prompt_lens: list[int] = [] for prompt in prompts: resolved = prompt.resolve_token_ids(self.tokenizer) @@ -282,10 +342,9 @@ def capture( prompt_lens.append(len(ids)) engine_prompt = TokensPrompt(prompt_token_ids=ids) engine_prompt["cache_salt"] = uuid.uuid4().hex - engine_prompts.append(engine_prompt) - sampling = SamplingParams(max_tokens=1, temperature=0.0, extra_args=extra_args) + requests.append((engine_prompt, SamplingParams(max_tokens=1, temperature=0.0, extra_args=dict(extra_args)))) - request_outputs = self._backend._require_llm().generate(engine_prompts, sampling, use_tqdm=False) + request_outputs = self._await_handles(self._submit_all(requests), None) rows_per_layer: dict[int, list[torch.Tensor]] = {layer: [] for layer in layer_ids} for index, request_output in enumerate(request_outputs): @@ -329,20 +388,34 @@ def generate( self, items: Sequence[GenerationItem], params: GenerationParams, + *, + cancel_token: CancelToken | None = None, ) -> list[ItemResult]: """Generate one result per item through the engine. + Each item is submitted as one engine request, and the engine admits each request as it is + submitted. The call returns once every item's request has finished. Cancelling + `cancel_token` aborts the call's unfinished requests on the engine, which frees their + slots in the running batch. + Args: items: The generation items; state entries lower as intervention specs on - plugin-active backends, and no client-side hooks or live processors execute + plugin-active backends, and no client-side hooks or in-memory processors execute here. params: Normalized generation parameters shared by all items; unmapped `extra` keys - raise. + raise `ValueError`. + cancel_token: Token whose cancellation aborts the call's requests, or None. Returns: One `ItemResult` per item, in item order. + + Raises: + GenerationCancelled: If `cancel_token` is cancelled before every request finished. + EngineClosedError: If the backend is released while a request is in flight. """ self._ensure_open() + if cancel_token is not None: + cancel_token.raise_if_cancelled() if not items: return [] item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm") @@ -372,10 +445,8 @@ def generate( sampling.append(SamplingParams(**args)) self._generate_count += 1 - generate_kwargs: dict[str, Any] = {"use_tqdm": False} - if self._backend._lora_request is not None: - generate_kwargs["lora_request"] = self._backend._lora_request - request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs) + handles = self._submit_all(zip(prompts, sampling), lora_request=self._backend._lora_request) + request_outputs = self._await_handles(handles, cancel_token) results: list[ItemResult] = [] for index, request_output in enumerate(request_outputs): @@ -395,16 +466,20 @@ def score( self, items: Sequence[ScoringItem], params: GenerationParams, + *, + cancel_token: CancelToken | None = None, ) -> torch.Tensor: """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs. - Each item's prompt and reference concatenate into one token-id prompt submitted with - `prompt_logprobs=0`, and the reference positions' log-probabilities are read back. + Each item's prompt and reference concatenate into one token-id prompt submitted as one + engine request with `prompt_logprobs=0`, and the reference positions' log-probabilities + are read back. Cancelling `cancel_token` aborts the call's unfinished requests. Args: items: The scoring items. Every item must carry the same reference length. params: Must carry no `extra` keys; forward keyword arguments have no remote rendering. + cancel_token: Token whose cancellation aborts the call's requests, or None. Returns: Log probabilities of shape `[num_items, ref_len]` on CPU. @@ -412,8 +487,12 @@ def score( Raises: ValueError: If items carry differing reference lengths or `params.extra` is non-empty. + GenerationCancelled: If `cancel_token` is cancelled before every request finished. + EngineClosedError: If the backend is released while a request is in flight. """ self._ensure_open() + if cancel_token is not None: + cancel_token.raise_if_cancelled() if params.extra: raise ValueError( f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote " @@ -452,10 +531,8 @@ def score( prompts.append(prompt) sampling.append(SamplingParams(**args)) - generate_kwargs: dict[str, Any] = {"use_tqdm": False} - if self._backend._lora_request is not None: - generate_kwargs["lora_request"] = self._backend._lora_request - request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs) + handles = self._submit_all(zip(prompts, sampling), lora_request=self._backend._lora_request) + request_outputs = self._await_handles(handles, cancel_token) rows = [ extract_ref_logprobs(request_output.prompt_logprobs, ref_ids) for request_output, ref_ids in zip(request_outputs, ref_ids_per_item) @@ -466,31 +543,44 @@ def score( class VLLMServeSession(_RequestSessionBase): """Request session over a vLLM server's completions endpoint. - Items fan out concurrently under the backend's `max_concurrency`; transport failures retry - with exponential backoff, and a batch whose items partially fail raises `PartialBatchError` - carrying the successes and the re-issuable failures. + The items of one call fan out concurrently under the backend's `max_concurrency`, and a + single-item call issues its one HTTP request on the caller's thread. The server batches + requests from concurrent calls, and the caller bounds how many calls run at once. Transport + failures retry with exponential backoff (`max_retries` applies to transport errors only), and + a batch whose items partially fail raises `PartialBatchError` containing the successes and the + re-issuable failures. The backend's `request_timeout` bounds each HTTP request and must + cover one full generation. """ def generate( self, items: Sequence[GenerationItem], params: GenerationParams, + *, + cancel_token: CancelToken | None = None, ) -> list[ItemResult]: """Generate one result per item through the completions endpoint. + A cancelled `cancel_token` stops the call before each item's request starts; a request + already sent completes. + Args: items: The generation items; entries must be empty. params: Normalized generation parameters shared by all items; unmapped `extra` keys - raise. + raise `ValueError`. + cancel_token: Token checked before each item's request starts, or None. Returns: One `ItemResult` per item, in item order. Raises: + GenerationCancelled: If `cancel_token` is cancelled before an item's request starts. PartialBatchError: If some items failed after transport retries while others succeeded. """ self._ensure_open() + if cancel_token is not None: + cancel_token.raise_if_cancelled() if not items: return [] item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm-serve") @@ -503,6 +593,8 @@ def generate( def make_task(index: int): def task() -> ItemResult: + if cancel_token is not None: + cancel_token.raise_if_cancelled() body: dict[str, Any] = { "model": backend._served_model, "prompt": item_ids[index], @@ -545,6 +637,7 @@ def task() -> ItemResult: return task outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency) + _raise_if_cancelled(outcomes, cancel_token) failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)] results = [outcome for outcome in outcomes if not isinstance(outcome, Exception)] if failures: @@ -555,12 +648,18 @@ def score( self, items: Sequence[ScoringItem], params: GenerationParams, + *, + cancel_token: CancelToken | None = None, ) -> torch.Tensor: """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs. + A cancelled `cancel_token` stops the call before each item's request starts; a request + already sent completes. + Args: items: The scoring items. Every item must carry the same reference length. params: Must carry no `extra` keys. + cancel_token: Token checked before each item's request starts, or None. Returns: Log probabilities of shape `[num_items, ref_len]` on CPU. @@ -568,10 +667,13 @@ def score( Raises: ValueError: If items carry differing reference lengths or `params.extra` is non-empty. + GenerationCancelled: If `cancel_token` is cancelled before an item's request starts. PartialBatchError: If some items failed after transport retries while others succeeded. """ self._ensure_open() + if cancel_token is not None: + cancel_token.raise_if_cancelled() if params.extra: raise ValueError( f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote " @@ -595,6 +697,8 @@ def score( def make_task(index: int): def task() -> list[float]: + if cancel_token is not None: + cancel_token.raise_if_cancelled() body = { "model": backend._served_model, "prompt": [*prompt_ids[index], *ref_ids[index]], @@ -621,8 +725,21 @@ def task() -> list[float]: return task outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency) + _raise_if_cancelled(outcomes, cancel_token) failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)] if failures: successes = [outcome for outcome in outcomes if not isinstance(outcome, Exception)] raise PartialBatchError(successes, failures) return torch.tensor(outcomes, dtype=torch.float32) + + +def _raise_if_cancelled(outcomes: list, cancel_token: CancelToken | None) -> None: + """Raise `GenerationCancelled` when a fan-out task stopped on the call's cancelled token. + + A cancellation then surfaces as itself rather than inside a `PartialBatchError`. + """ + if cancel_token is None or not cancel_token.cancelled: + return + for outcome in outcomes: + if isinstance(outcome, GenerationCancelled): + raise outcome diff --git a/steerability/evaluation/batching.py b/steerability/evaluation/batching.py index f798487f..ab754241 100644 --- a/steerability/evaluation/batching.py +++ b/steerability/evaluation/batching.py @@ -1,13 +1,17 @@ -"""Lock-leader batching collator for the steering-pipeline model provider. +"""Request dispatch for the steering-pipeline model provider: the lock-leader collator (batching +mode) and the per-request dispatcher (per-request mode). -Inspect issues one async request per sample, many concurrently; a `SteeringPipeline` is synchronous -and permits one in-flight generation, but accepts batched calls. The collator turns concurrent -requests into batched `pipeline.generate()` calls with no background task and no timer. One -`anyio.Lock` guards pipeline access: every request enqueues its record and then acquires the lock; -whoever holds it (the leader) takes its record plus the compatible queued records, dispatches one -batched call in a worker thread, and completes every member. A waiter that wakes with its record -already served returns without dispatching. The lock holder is the only caller of -`pipeline.generate`, so the one-in-flight invariant holds by construction. +Inspect issues one async request per sample, many concurrently, while a `SteeringPipeline` is +synchronous. Both dispatchers admit requests the same way (validation, runtime-kwarg partitioning, +and inert-kwarg accounting) and differ in how admitted requests reach `pipeline.generate()`. + +The lock-leader collator turns concurrent requests into batched `pipeline.generate()` calls with no +background task and no timer, for a pipeline that runs one generation at a time. One `anyio.Lock` +guards pipeline access: every request enqueues its record and then acquires the lock; whoever holds +it (the leader) takes its record plus the compatible queued records, dispatches one batched call in +a worker thread, and completes every member. A waiter that wakes with its record already served +returns without dispatching. The lock holder is the only caller of `pipeline.generate`, so at most +one generation is in flight by construction. An uncontended lock acquisition yields once to the event loop, so requests admitted in the same scheduling window enqueue before the first leader takes its batch, and while a batch is in flight @@ -47,6 +51,26 @@ stochastic sampling is not preserved under concurrency under either scope; a bitwise-reproducible stochastic run requires a batch ceiling of 1 and Inspect `max_connections=1`. Greedy decoding sidesteps seed sensitivity, with the one caveat that padded-batch numerics can differ from single-item numerics on some kernels. + +Per-request mode applies to a pipeline whose `supports_concurrent_generate` is True, when the provider is +configured with `ProviderOptions.max_concurrency`. The per-request dispatcher issues one single-prompt +`pipeline.generate()` call per admitted request, each in its own worker thread, with at most `max_concurrency` +calls in flight. The pipeline serializes the calls' prompt phases and the engine admits each request into its +running batch. A request therefore returns as soon as its own generation finishes. There is no queue, no batch key, and +no poison isolation, since every dispatch is one request and its error is that request's own. + +In per-request mode a request's seed is a function of the base seed and the request. The dispatcher replaces a +seeded request's `seed` with `derive_item_seed(seed, f"request:{digest}", 0)`, where `digest` is the SHA-256 of the +canonical JSON of the request's prompt and its retained per-sample runtime kwargs. A seeded request therefore +reproduces independently of concurrency and arrival order, distinct prompts sample distinct streams, and two +requests with identical prompts and per-sample kwargs sample the same stream. `seed_scope` has no effect in +per-request mode, since every dispatch is one item. + +In per-request mode a cancelled Inspect request aborts its generation. Each request carries its own `CancelToken` +into `pipeline.generate()`, and the dispatcher cancels the token when the requesting task is cancelled. The task +returns at once, and the worker thread ends with `GenerationCancelled` once the session observes the token (the +offline vLLM session aborts the request on the engine). In batching mode one dispatch serves several records, and a +cancelled record does not cancel its dispatch. """ import hashlib import json @@ -60,6 +84,8 @@ import anyio import anyio.to_thread +from steerability.algorithms.core.execution.cancellation import CancelToken +from steerability.algorithms.core.execution.fanout import derive_item_seed from steerability.algorithms.core.identity import canonical_value from steerability.algorithms.core.output import Output @@ -71,7 +97,7 @@ @dataclass(eq=False, slots=True) class BatchRequest: - """One admitted generation request, queued for a batched dispatch. + """One admitted generation request. Attributes: prompt: One conversation (a list of chat-message dicts) on the messages path, or one @@ -84,6 +110,7 @@ class BatchRequest: done: Event set by the leader once the record holds its result. output: The record's `Output` on success, else None. error: The record's exception on failure, else None. + cancel_token: The request's cancellation token in per-request mode, else None. """ prompt: Any @@ -94,19 +121,15 @@ class BatchRequest: done: anyio.Event = field(default_factory=anyio.Event) output: Output | None = None error: BaseException | None = None + cancel_token: CancelToken | None = None -class LockLeaderCollator: - """Collate concurrent async requests into batched `pipeline.generate()` calls. - - See the module docstring for the protocol, the batch-key semantics, poison isolation, and the - reproducibility contract. `admit()` validates and builds a record without enqueueing; the - async `serve()` enqueues it and runs the leader protocol. +class _RequestAdmission: + """Request admission shared by the dispatchers: validation, runtime-kwarg partitioning by declared + scope, and inert-kwarg accounting. Args: - pipeline: The steered pipeline; the lock holder is its only caller. - max_batch_size: Dispatch ceiling; the caller passes the effective value (already clamped - to 1 when the pipeline does not support batching). + pipeline: The steered pipeline. prompt_path: `"messages"` to dispatch conversations via `messages=`, `"text"` to dispatch rendered strings via `text=`. declared_scopes: Runtime-kwarg names declared by the pipeline's enabled controls, mapped @@ -121,22 +144,17 @@ def __init__( self, pipeline: "SteeringPipeline", *, - max_batch_size: int, prompt_path: Literal["messages", "text"], declared_scopes: Mapping[str, str], static_runtime_kwargs: Mapping[str, Any], chat_template_kwargs: Mapping[str, Any] | None = None, ) -> None: self._pipeline = pipeline - self._max_batch_size = int(max_batch_size) self._prompt_path = prompt_path self._declared_scopes = dict(declared_scopes) self._static_runtime_kwargs = dict(static_runtime_kwargs) self._chat_template_kwargs = dict(chat_template_kwargs) if chat_template_kwargs is not None else None - self._lock = anyio.Lock() - self._queue: list[BatchRequest] = [] self._closed = False - self._batch_failures_logged: set[str] = set() self._inert_keys: set[str] = set() # static kwargs are partitioned once by declared scope; undeclared names are inert @@ -153,11 +171,11 @@ def __init__( @property def closed(self) -> bool: - """Whether the collator has been closed.""" + """Whether the dispatcher has been closed.""" return self._closed def close(self) -> None: - """Refuse new admissions; an in-flight batch completes and delivers its results.""" + """Refuse new admissions; an in-flight dispatch completes and delivers its results.""" self._closed = True @property @@ -166,7 +184,7 @@ def inert_runtime_kwargs(self) -> frozenset[str]: return frozenset(self._inert_keys) def _mark_inert(self, key: str, tier: str) -> None: - """Record a runtime kwarg no enabled control declares, logging it once per collator.""" + """Record a runtime kwarg no enabled control declares, logging it once per dispatcher.""" if key in self._inert_keys: return self._inert_keys.add(key) @@ -198,7 +216,7 @@ def admit( The admitted `BatchRequest`. Raises: - RuntimeError: If the collator is closed. + RuntimeError: If the dispatcher is closed. ValueError: If a per-sample key is also supplied statically, or is declared `"call"`-scoped by an enabled control. """ @@ -248,6 +266,88 @@ def _batch_key( ) return hashlib.sha256(payload.encode()).hexdigest() + def _run_batch(self, batch: list[BatchRequest]) -> list[Output]: + """Issue one `pipeline.generate` call for `batch`, returning one `Output` per record. + + Raises: + RuntimeError: If the call returns a different number of outputs than records. + """ + leader = batch[0] + runtime_kwargs = dict(self._static_call_kwargs) + for key, value in self._static_row_kwargs.items(): + runtime_kwargs[key] = [value] * len(batch) + for key in leader.per_sample_runtime_kwargs: + runtime_kwargs[key] = [member.per_sample_runtime_kwargs[key] for member in batch] + gen_kwargs = dict(leader.gen_kwargs) + if self._chat_template_kwargs is not None: + gen_kwargs["chat_template_kwargs"] = dict(self._chat_template_kwargs) + if leader.cancel_token is not None: + gen_kwargs["cancel_token"] = leader.cancel_token + if leader.num_choices > 1: + prompt_kwargs = ( + {"messages": leader.prompt} if self._prompt_path == "messages" else {"text": leader.prompt} + ) + output = self._pipeline.generate( + runtime_kwargs=runtime_kwargs, return_output=True, n=leader.num_choices, + **prompt_kwargs, **gen_kwargs, + ) + return [output] + prompts = [member.prompt for member in batch] + prompt_kwargs = {"messages": prompts} if self._prompt_path == "messages" else {"text": prompts} + outputs = list(self._pipeline.generate( + runtime_kwargs=runtime_kwargs, return_output=True, **prompt_kwargs, **gen_kwargs, + )) + if len(outputs) != len(batch): + raise RuntimeError( + f"pipeline.generate returned {len(outputs)} output(s) for a dispatch of {len(batch)} " + "prompt(s); one Output per prompt is required." + ) + return outputs + + +class LockLeaderCollator(_RequestAdmission): + """Collate concurrent async requests into batched `pipeline.generate()` calls. + + See the module docstring for the protocol, the batch-key semantics, poison isolation, and the + reproducibility contract. `admit()` validates and builds a record without enqueueing; the + async `serve()` enqueues it and runs the leader protocol. + + Args: + pipeline: The steered pipeline; the lock holder is its only caller. + max_batch_size: Dispatch ceiling; the caller passes the effective value (already clamped + to 1 when the pipeline does not support batching). + prompt_path: `"messages"` to dispatch conversations via `messages=`, `"text"` to dispatch + rendered strings via `text=`. + declared_scopes: Runtime-kwarg names declared by the pipeline's enabled controls, mapped + to their scope (`"row"` or `"call"`). + static_runtime_kwargs: Runtime kwargs applied to every dispatch: `"call"`-scoped values + pass through unchanged, `"row"`-scoped values are one row's value broadcast to every + row, undeclared names are inert. Shallow-copied and never mutated. + chat_template_kwargs: Forwarded to `apply_chat_template` on the messages path, or None. + """ + + def __init__( + self, + pipeline: "SteeringPipeline", + *, + max_batch_size: int, + prompt_path: Literal["messages", "text"], + declared_scopes: Mapping[str, str], + static_runtime_kwargs: Mapping[str, Any], + chat_template_kwargs: Mapping[str, Any] | None = None, + ) -> None: + super().__init__( + pipeline, + prompt_path=prompt_path, + declared_scopes=declared_scopes, + static_runtime_kwargs=static_runtime_kwargs, + chat_template_kwargs=chat_template_kwargs, + ) + self._max_batch_size = int(max_batch_size) + self._lock = anyio.Lock() + self._queue: list[BatchRequest] = [] + self._batch_failures_logged: set[str] = set() + async def serve(self, record: BatchRequest) -> Output: """Enqueue one admitted record and run the leader protocol until it holds its result. @@ -328,38 +428,106 @@ def _dispatch(self, batch: list[BatchRequest]) -> None: for member, output in zip(batch, outputs): member.output = output - def _run_batch(self, batch: list[BatchRequest]) -> list[Output]: - """Issue one `pipeline.generate` call for `batch`, returning one `Output` per record. + +def request_seed(seed: int, prompt: Any, per_sample_runtime_kwargs: Mapping[str, Any]) -> int: + """The seed of one seeded request in per-request mode, a function of the base seed and the request. + + Args: + seed: The request's base seed. + prompt: The request's prompt (a conversation or a rendered string). + per_sample_runtime_kwargs: The request's retained per-sample runtime kwargs. + + Returns: + `derive_item_seed(seed, f"request:{digest}", 0)`, where `digest` is the SHA-256 of the + canonical JSON of the prompt and the per-sample runtime kwargs. + """ + payload = json.dumps( + canonical_value({"prompt": prompt, "per_sample": dict(per_sample_runtime_kwargs)}), sort_keys=True, + ) + digest = hashlib.sha256(payload.encode()).hexdigest() + return derive_item_seed(seed, f"request:{digest}", 0) + + +class PerRequestDispatcher(_RequestAdmission): + """Dispatch each admitted request at once in its own worker thread, at most `max_concurrency` in flight. + + Each request becomes one single-prompt (or, for multiple candidates, one bare-conversation) + `pipeline.generate()` call with its own `CancelToken`. A seeded request's seed is replaced by + `request_seed(...)`. See the module docstring for the per-request contract. `admit()` validates + and builds a record; the async `serve()` dispatches it. + + Args: + pipeline: The steered pipeline; its `supports_concurrent_generate` is True. + max_concurrency: Maximum number of `pipeline.generate()` calls in flight (at least 1). + prompt_path: `"messages"` to dispatch conversations via `messages=`, `"text"` to dispatch + rendered strings via `text=`. + declared_scopes: Runtime-kwarg names declared by the pipeline's enabled controls, mapped + to their scope (`"row"` or `"call"`). + static_runtime_kwargs: Runtime kwargs applied to every dispatch: `"call"`-scoped values + pass through unchanged, `"row"`-scoped values are one row's value, undeclared names + are inert. Shallow-copied and never mutated. + chat_template_kwargs: Forwarded to `apply_chat_template` on the messages path, or None. + """ + + def __init__( + self, + pipeline: "SteeringPipeline", + *, + max_concurrency: int, + prompt_path: Literal["messages", "text"], + declared_scopes: Mapping[str, str], + static_runtime_kwargs: Mapping[str, Any], + chat_template_kwargs: Mapping[str, Any] | None = None, + ) -> None: + super().__init__( + pipeline, + prompt_path=prompt_path, + declared_scopes=declared_scopes, + static_runtime_kwargs=static_runtime_kwargs, + chat_template_kwargs=chat_template_kwargs, + ) + if int(max_concurrency) < 1: + raise ValueError(f"max_concurrency must be >= 1; got {max_concurrency}.") + self._limiter = anyio.CapacityLimiter(int(max_concurrency)) + + async def serve(self, record: BatchRequest) -> Output: + """Dispatch one admitted record in a worker thread and return its result. + + Cancelling the requesting task cancels the record's `CancelToken` and returns at once; the + worker thread ends when the session observes the token, and its result is discarded. + + Args: + record: A record built by `admit()`. + + Returns: + The record's `Output`. Raises: - RuntimeError: If the call returns a different number of outputs than records. + Exception: The record's own dispatch exception, re-raised on the requesting task. """ - leader = batch[0] - runtime_kwargs = dict(self._static_call_kwargs) - for key, value in self._static_row_kwargs.items(): - runtime_kwargs[key] = [value] * len(batch) - for key in leader.per_sample_runtime_kwargs: - runtime_kwargs[key] = [member.per_sample_runtime_kwargs[key] for member in batch] - gen_kwargs = dict(leader.gen_kwargs) - if self._chat_template_kwargs is not None: - gen_kwargs["chat_template_kwargs"] = dict(self._chat_template_kwargs) - if leader.num_choices > 1: - prompt_kwargs = ( - {"messages": leader.prompt} if self._prompt_path == "messages" else {"text": leader.prompt} - ) - output = self._pipeline.generate( - runtime_kwargs=runtime_kwargs, return_output=True, n=leader.num_choices, - **prompt_kwargs, **gen_kwargs, - ) - return [output] - prompts = [member.prompt for member in batch] - prompt_kwargs = {"messages": prompts} if self._prompt_path == "messages" else {"text": prompts} - outputs = list(self._pipeline.generate( - runtime_kwargs=runtime_kwargs, return_output=True, **prompt_kwargs, **gen_kwargs, - )) - if len(outputs) != len(batch): - raise RuntimeError( - f"pipeline.generate returned {len(outputs)} output(s) for a dispatch of {len(batch)} " - "prompt(s); one Output per prompt is required." + record.cancel_token = CancelToken() + try: + await anyio.to_thread.run_sync( + self._run_request, record, limiter=self._limiter, abandon_on_cancel=True, ) - return outputs + except anyio.get_cancelled_exc_class(): + record.cancel_token.cancel() + raise + if record.error is not None: + raise record.error + return record.output + + def _run_request(self, record: BatchRequest) -> None: + """Run one record's `pipeline.generate` call in the worker thread, filling its result slot. + + Never raises; the call's exception lands on the record. + """ + try: + record.cancel_token.raise_if_cancelled() + if record.gen_kwargs.get("seed") is not None: + record.gen_kwargs["seed"] = request_seed( + int(record.gen_kwargs["seed"]), record.prompt, record.per_sample_runtime_kwargs, + ) + record.output = self._run_batch([record])[0] + except Exception as error: + record.error = error diff --git a/steerability/evaluation/provider.py b/steerability/evaluation/provider.py index 99824b67..6f2c0036 100644 --- a/steerability/evaluation/provider.py +++ b/steerability/evaluation/provider.py @@ -4,7 +4,9 @@ through `pipeline.generate()`: on the messages path prompts enter as structured chat messages, so the pipeline owns chat templating and message-phase input controls fire exactly as in deployment; tokenizers without a chat template evaluate through the text path instead. Concurrent Inspect -requests collate into batched pipeline calls through the lock-leader collator in `batching`. +requests collate into batched pipeline calls through the lock-leader collator in `batching`, or, +in per-request mode (`ProviderOptions.max_concurrency` on a pipeline that supports concurrent +generation), each request is dispatched as its own pipeline call. The provider is generation-only by design. Logprob-based scoring (including prompt-logprob perplexity) is structurally unsupported, since a pipeline's decoding-time controls are not @@ -36,9 +38,11 @@ modelapi, ) +from steerability.algorithms.core.execution.spec import BackendSpec from steerability.algorithms.core.output import Output, truncate_at_stop_strings +from steerability.algorithms.core.utils.assembly import resolve_decoding_driver from steerability.algorithms.core.utils.controls import runtime_kwargs_schema -from steerability.evaluation.batching import LockLeaderCollator +from steerability.evaluation.batching import LockLeaderCollator, PerRequestDispatcher from steerability.utils.rendering import has_chat_template, join_message_contents from steerability.utils.thinking import ( DEFAULT_THINK_TAGS, @@ -125,7 +129,15 @@ class ProviderOptions: across prompts). With `"item"`, each row derives its own seed and the rows are decoded one at a time. The collator groups requests by arrival time, which makes per-row reproducibility unattainable under either setting. `"item"` therefore only makes decoding - serial. The setting has no effect on the vLLM backends. + serial. The setting has no effect on the vLLM backends or in per-request mode. + max_concurrency: Selects per-request mode when set (at least 1). When the pipeline's + `supports_concurrent_generate` is True, each request is dispatched at once as its own + single-prompt `pipeline.generate()` call in a worker thread, with at most + `max_concurrency` calls in flight, and `max_connections()` returns `max_concurrency`. + A seeded request then samples under a seed derived from the base seed and the request + (its prompt and per-sample runtime kwargs), and a cancelled request aborts its + generation. When the pipeline does not support concurrent generation, the provider + logs at info level and keeps batching mode. None (default) selects batching mode. """ runtime_kwargs: Mapping[str, Any] = field(default_factory=dict) @@ -137,6 +149,7 @@ class ProviderOptions: reasoning_split: Literal["auto", "text", "tokens"] = "auto" on_unsupported_param: Literal["raise", "warn"] = "raise" seed_scope: Literal["item", "dispatch"] = "dispatch" + max_concurrency: int | None = None def __post_init__(self) -> None: if not isinstance(self.runtime_kwargs, Mapping): @@ -161,6 +174,8 @@ def __post_init__(self) -> None: raise ValueError(f"on_unsupported_param must be 'raise' or 'warn'; got {self.on_unsupported_param!r}.") if self.seed_scope not in ("item", "dispatch"): raise ValueError(f"seed_scope must be 'item' or 'dispatch'; got {self.seed_scope!r}.") + if self.max_concurrency is not None and int(self.max_concurrency) < 1: + raise ValueError(f"max_concurrency must be >= 1 or None; got {self.max_concurrency}.") class SteeringPipelineModelAPI(ModelAPI): @@ -252,22 +267,41 @@ def __init__( ) self._effective_max_batch = effective_max_batch + self._per_request = self._options.max_concurrency is not None and pipeline.supports_concurrent_generate + if self._options.max_concurrency is not None and not self._per_request: + logger.info( + "max_concurrency=%d is set, but the pipeline does not support concurrent generation (%s); " + "requests are dispatched in batching mode.", + self._options.max_concurrency, _concurrency_blocker(pipeline), + ) + schema = runtime_kwargs_schema(pipeline.controls) declared_scopes = {name: entry["scope"] for name, entry in schema.items()} - self._collator = LockLeaderCollator( - pipeline, - max_batch_size=effective_max_batch, + common: dict[str, Any] = dict( prompt_path=self._prompt_path, declared_scopes=declared_scopes, static_runtime_kwargs=self._options.runtime_kwargs, chat_template_kwargs=self._options.chat_template_kwargs, ) + self._collator: LockLeaderCollator | PerRequestDispatcher + if self._per_request: + self._collator = PerRequestDispatcher( + pipeline, max_concurrency=int(self._options.max_concurrency), **common, + ) + else: + self._collator = LockLeaderCollator(pipeline, max_batch_size=effective_max_batch, **common) @property def prompt_path(self) -> Literal["messages", "text"]: """`"messages"` when the tokenizer has a chat template, `"text"` otherwise.""" return self._prompt_path + @property + def dispatch_mode(self) -> Literal["batched", "per_request"]: + """`"per_request"` when each request is dispatched as its own pipeline call, `"batched"` when + the lock-leader collator combines requests into batched calls.""" + return "per_request" if self._per_request else "batched" + @property def effective_max_batch(self) -> int: """The collator's dispatch ceiling after the batching clamp.""" @@ -525,7 +559,10 @@ def max_tokens(self) -> int | None: return self._options.default_max_tokens def max_connections(self) -> int: - """Advertised concurrency; equals the effective batch ceiling so batches self-fill.""" + """Advertised concurrency: `max_concurrency` in per-request mode, else the effective batch + ceiling so batches fill.""" + if self._per_request: + return int(self._options.max_concurrency) return self._effective_max_batch def connection_key(self) -> str: @@ -559,6 +596,16 @@ def steerability_provider() -> type[SteeringPipelineModelAPI]: return SteeringPipelineModelAPI +def _concurrency_blocker(pipeline: "SteeringPipeline") -> str: + """The condition that keeps `pipeline` from concurrent generation, for the fallback log.""" + driver = resolve_decoding_driver(getattr(pipeline, "output_controls", ())) + if driver is not None: + return f"decoding driver {type(driver).__name__} is enabled" + backend = getattr(pipeline, "backend", None) + kind = backend.kind if isinstance(backend, BackendSpec) else (backend or "huggingface") + return f"backend kind {kind!r} runs in process" + + def _count_non_pad(ids: torch.Tensor | None, pad_token_id: int | None) -> int: """Count non-pad positions in a token-id tensor (best effort where pad equals EOS).""" if ids is None: diff --git a/steerability/evaluation/runner.py b/steerability/evaluation/runner.py index 93fc1ad3..14e37923 100644 --- a/steerability/evaluation/runner.py +++ b/steerability/evaluation/runner.py @@ -165,8 +165,9 @@ def run(self) -> dict[str, list[dict]]: dict for fixed and baseline configurations. - `"suites"`: Mapping from suite name to that suite's flattened results, with log paths relative to `save_dir`. - - `"provenance"`: Versions, backend kind, `prompt_path`, the shared-base - fingerprint when one exists, and a timestamp. + - `"provenance"`: Versions, backend kind, `prompt_path`, `dispatch_mode` + (`"batched"` or `"per_request"`), the shared-base fingerprint when one exists, + and a timestamp. Raises: RuntimeError: If any configuration is unsupported and `on_unsupported="raise"` (one @@ -273,6 +274,13 @@ def run(self) -> dict[str, list[dict]]: with factory.steered(point.controls_factory()) as pipeline: logger.info("Steered %s in %.0fs.", label, time.monotonic() - steer_started) prompt_path = "messages" if has_chat_template(pipeline.tokenizer) else "text" + # the provider dispatches per request when max_concurrency is set and the + # pipeline supports concurrent generation + per_request = ( + getattr(self.provider_options, "max_concurrency", None) is not None + and pipeline.supports_concurrent_generate + ) + dispatch_mode = "per_request" if per_request else "batched" for trial_id in range(self.num_trials): trial_seed = ( derive_trial_seed(self.seed, point.config_id, trial_id) @@ -320,6 +328,7 @@ def run(self) -> dict[str, list[dict]]: **versions, "backend": factory.backend_kind, "prompt_path": prompt_path, + "dispatch_mode": dispatch_mode, "shared_base_fingerprint": factory.shared_base_fingerprint, "recorded_utc": datetime.datetime.now( datetime.timezone.utc diff --git a/tests/controls/test_pasta.py b/tests/controls/test_pasta.py index 91a9f786..c12dba54 100644 --- a/tests/controls/test_pasta.py +++ b/tests/controls/test_pasta.py @@ -182,7 +182,7 @@ def __init__(self, tokenizer, lift_map: dict[tuple[int, int], float]): self.current_score = 0.0 self.item_count = 0 - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): entries = items[0].state_entries if entries: keywords = entries[0].hooks["pre"][0]["hook_func"].keywords diff --git a/tests/controls/test_phased_driver_batching.py b/tests/controls/test_phased_driver_batching.py index 9b83a6a5..df482c12 100644 --- a/tests/controls/test_phased_driver_batching.py +++ b/tests/controls/test_phased_driver_batching.py @@ -81,9 +81,9 @@ def _count_session_calls(monkeypatch) -> list[int]: calls: list[int] = [] original = ExclusiveSession.generate - def counting(self, items, params): + def counting(self, items, params, *, cancel_token=None): calls.append(len(items)) - return original(self, items, params) + return original(self, items, params, cancel_token=cancel_token) monkeypatch.setattr(ExclusiveSession, "generate", counting) return calls @@ -276,9 +276,9 @@ def _record_session_calls(monkeypatch, recorder: RecordingTransform) -> list[tup calls: list[tuple[list[int], list[torch.Tensor]]] = [] original = ExclusiveSession.generate - def recording(self, items, params): + def recording(self, items, params, *, cancel_token=None): start = len(recorder.masks) - results = original(self, items, params) + results = original(self, items, params, cancel_token=cancel_token) calls.append(([item.prompt.token_ids.size(-1) for item in items], recorder.masks[start:])) return results @@ -467,7 +467,7 @@ def __init__(self, tokenizer, continuations, reasons): self._reasons = reasons self.seen = [] - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): self.seen.append((items, params)) width = max(len(ids) for ids in self._continuations) pad = self.tokenizer.pad_token_id diff --git a/tests/controls/test_phased_token_boundaries.py b/tests/controls/test_phased_token_boundaries.py index 8822bb52..75994524 100644 --- a/tests/controls/test_phased_token_boundaries.py +++ b/tests/controls/test_phased_token_boundaries.py @@ -182,7 +182,7 @@ def __init__(self, tokenizer, continuation: list[int], reason: str | None): self._continuation = continuation self._reason = reason - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): return [ ItemResult(index=index, output=Output( output_ids=torch.tensor([self._continuation]), adapted_input_ids=item.prompt.token_ids, diff --git a/tests/core/test_backend_execution.py b/tests/core/test_backend_execution.py index b3e33524..eae54ee8 100644 --- a/tests/core/test_backend_execution.py +++ b/tests/core/test_backend_execution.py @@ -1131,7 +1131,7 @@ def __init__(self, row_ids, pad_token_id): self._row_ids = row_ids self.tokenizer = type("_T", (), {"pad_token_id": pad_token_id})() - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): from steerability.algorithms.core.execution.payloads import ItemResult return [ diff --git a/tests/core/test_cancellation.py b/tests/core/test_cancellation.py new file mode 100644 index 00000000..3ff8514a --- /dev/null +++ b/tests/core/test_cancellation.py @@ -0,0 +1,148 @@ +"""Tests for `CancelToken` and `GenerationCancelled`, and for the in-process session's cancellation +granularity (the token is checked before each decode; a running `model.generate` finishes).""" +import logging +import threading + +import pytest +import torch + +from steerability.algorithms.core.execution import ( + CancelToken, + GenerationCancelled, + GenerationItem, + GenerationParams, + PreparedPrompt, + ScoringItem, +) +from steerability.backends.huggingface import HFBackend +from tests.utils.tiny_models import tiny_llama, wordlevel_tokenizer + + +class TestCancelToken: + + def test_cancel_is_idempotent_and_runs_each_callback_once(self): + token = CancelToken() + calls = [] + token.on_cancel(lambda: calls.append("first")) + token.on_cancel(lambda: calls.append("second")) + assert not token.cancelled + token.cancel() + token.cancel() + assert token.cancelled + assert calls == ["first", "second"] + + def test_callbacks_run_on_the_cancelling_thread(self): + token = CancelToken() + threads = [] + token.on_cancel(lambda: threads.append(threading.get_ident())) + canceller = threading.Thread(target=token.cancel) + canceller.start() + canceller.join(5) + assert threads == [canceller.ident] + + def test_callback_registered_after_cancellation_runs_at_once(self): + token = CancelToken() + token.cancel() + calls = [] + token.on_cancel(lambda: calls.append("late")) + assert calls == ["late"] + + def test_callback_exceptions_are_logged_and_do_not_propagate(self, caplog): + token = CancelToken() + calls = [] + + def failing(): + raise RuntimeError("callback failed") + + token.on_cancel(failing) + token.on_cancel(lambda: calls.append("after")) + with caplog.at_level(logging.WARNING, logger="steerability.algorithms.core.execution.cancellation"): + token.cancel() + assert calls == ["after"] + assert any("callback raised" in record.getMessage() for record in caplog.records) + + def test_raise_if_cancelled(self): + token = CancelToken() + token.raise_if_cancelled() + token.cancel() + with pytest.raises(GenerationCancelled, match="cancelled"): + token.raise_if_cancelled() + + def test_generation_cancelled_is_a_runtime_error(self): + assert issubclass(GenerationCancelled, RuntimeError) + + +@pytest.fixture +def exclusive_backend(): + model = tiny_llama() + tokenizer = wordlevel_tokenizer() + from steerability.algorithms.core.execution import BackendSpec + + spec = BackendSpec(kind="huggingface", model=None) + return HFBackend.adopt(spec, lambda: model, lambda: tokenizer) + + +def _items(count): + return [ + GenerationItem(prompt=PreparedPrompt.from_token_ids(torch.tensor([[0, 3 + index, 4]])), seed=index) + for index in range(count) + ] + + +class TestExclusiveSessionCancellation: + + def test_cancelled_token_raises_before_any_decode(self, exclusive_backend, monkeypatch): + model = exclusive_backend._model_provider() + decodes = [] + original = model.generate + monkeypatch.setattr(model, "generate", lambda *args, **kwargs: decodes.append(1) or original(*args, **kwargs)) + token = CancelToken() + token.cancel() + with exclusive_backend.open_session() as session: + with pytest.raises(GenerationCancelled): + session.generate(_items(1), GenerationParams(max_new_tokens=2, greedy=True), cancel_token=token) + assert decodes == [] + + def test_serial_path_stops_between_items(self, exclusive_backend, monkeypatch): + model = exclusive_backend._model_provider() + token = CancelToken() + decodes = [] + original = model.generate + + def decode_then_cancel(*args, **kwargs): + decodes.append(1) + result = original(*args, **kwargs) + token.cancel() + return result + + monkeypatch.setattr(model, "generate", decode_then_cancel) + with exclusive_backend.open_session() as session: + with pytest.raises(GenerationCancelled): + # distinct item seeds decode the items one at a time + session.generate(_items(3), GenerationParams(max_new_tokens=2, greedy=True), cancel_token=token) + assert decodes == [1] + + def test_batched_path_checks_once_before_the_decode(self, exclusive_backend): + token = CancelToken() + token.cancel() + items = [GenerationItem(prompt=PreparedPrompt.from_token_ids(torch.tensor([[0, 3, 4]]))) for _ in range(2)] + with exclusive_backend.open_session() as session: + with pytest.raises(GenerationCancelled): + session.generate(items, GenerationParams(max_new_tokens=2, greedy=True), cancel_token=token) + + def test_score_checks_the_token(self, exclusive_backend): + token = CancelToken() + token.cancel() + item = ScoringItem( + prompt=PreparedPrompt.from_token_ids(torch.tensor([[0, 3, 4]])), ref_output_ids=torch.tensor([[5]]), + ) + with exclusive_backend.open_session() as session: + with pytest.raises(GenerationCancelled): + session.score([item], GenerationParams(), cancel_token=token) + + def test_uncancelled_token_generates_normally(self, exclusive_backend): + with exclusive_backend.open_session() as session: + results = session.generate( + _items(2), GenerationParams(max_new_tokens=2, greedy=True), cancel_token=CancelToken(), + ) + assert len(results) == 2 diff --git a/tests/core/test_concurrent_generate.py b/tests/core/test_concurrent_generate.py new file mode 100644 index 00000000..e51bd88e --- /dev/null +++ b/tests/core/test_concurrent_generate.py @@ -0,0 +1,384 @@ +"""Concurrent `generate()` calls: the `supports_concurrent_generate` truth table, the prompt lock +(prompt phases never overlap; dispatches overlap only when concurrent generation is supported), and +the reserved `cancel_token` keyword (validation, abort on cancellation, and forwarding to a +driver's rollouts). + +Engine-kind pipelines run against a recording backend registered by monkeypatching +`resolve_backend_class`, so no vLLM installation or engine is needed. +""" +import threading +import time + +import pytest +import torch + +from steerability.algorithms.core.execution import ( + BackendSpec, + CancelToken, + GenerationCancelled, + ItemResult, + ModelAccess, + ModelFacts, + Requirements, +) +from steerability.algorithms.core.execution.session_utils import session_generate +from steerability.algorithms.core.output import Output +from steerability.algorithms.core.steering_pipeline import SteeringPipeline +from steerability.algorithms.input_control.base import InputControl +from steerability.algorithms.output_control.base import DecodingDriver +from tests.utils.tiny_models import tiny_llama, wordlevel_tokenizer + +VLLM_SPEC = BackendSpec(kind="vllm", model="tiny-engine-model") + + +class _RecordingSession: + """An engine session double that records each call's thread and wall-clock interval. + + A call blocks until its backend's `release` event is set (when one is set on the backend), + or until its `cancel_token` is cancelled, whichever comes first. + """ + + def __init__(self, backend): + self._backend = backend + + @property + def tokenizer(self): + return self._backend.tokenizer + + @property + def layout(self) -> ModelFacts: + return ModelFacts( + num_layers=2, hidden_size=16, num_attention_heads=2, head_dim=8, dtype="float32", + model_fingerprint=None, model_type="llama", fingerprint_grade="config", + ) + + def __enter__(self): + return self + + def __exit__(self, *exc): + return None + + def close(self): + return None + + def generate(self, items, params, *, cancel_token=None): + backend = self._backend + started = time.monotonic() + backend.tokens_seen.append(cancel_token) + if cancel_token is not None: + cancel_token.on_cancel(backend.abort_seen.set) + deadline = started + backend.hold_seconds + while time.monotonic() < deadline: + if cancel_token is not None and cancel_token.cancelled: + backend.calls.append((threading.get_ident(), started, time.monotonic(), "cancelled")) + raise GenerationCancelled("cancelled") + if backend.release is not None and backend.release.is_set(): + break + time.sleep(0.01) + backend.calls.append((threading.get_ident(), started, time.monotonic(), "done")) + return [ + ItemResult(index=index, output=Output( + output_ids=torch.tensor([[3, 1]]), adapted_input_ids=item.prompt.token_ids, finish_reasons=("eos",), + )) + for index, item in enumerate(items) + ] + + def score(self, items, params, *, cancel_token=None): + return torch.zeros(len(items), items[0].ref_output_ids.shape[-1]) + + +class _RecordingBackend: + """An engine-kind backend double; its sessions record calls (see `_RecordingSession`).""" + + hold_seconds = 0.3 + release: threading.Event | None = None + calls: list = [] + tokens_seen: list = [] + abort_seen = threading.Event() + + def __init__(self, spec, artifacts=()): + self.spec = spec + self.tokenizer = wordlevel_tokenizer() + self._discovery = None + + @classmethod + def capabilities_for_spec(cls, spec): + from steerability.backends.vllm import VLLMBackend + + return VLLMBackend.capabilities_for_spec(spec) + + def negotiated_capabilities(self): + return self.capabilities_for_spec(self.spec) + + def open_session(self): + return _RecordingSession(self) + + def stage_artifacts(self, payloads): + return None + + def release(self): + return None + + @classmethod + def reset(cls): + cls.hold_seconds = 0.3 + cls.release = None + cls.calls = [] + cls.tokens_seen = [] + cls.abort_seen = threading.Event() + + +@pytest.fixture +def recording_engine(monkeypatch): + import steerability.algorithms.core.execution.backend as backend_module + import steerability.algorithms.core.steering_pipeline as pipeline_module + + original = backend_module.resolve_backend_class + + def resolver(spec): + if spec.kind == "huggingface": + return original(spec) + return _RecordingBackend + + monkeypatch.setattr(backend_module, "resolve_backend_class", resolver) + monkeypatch.setattr(pipeline_module, "resolve_backend_class", resolver) + _RecordingBackend.reset() + yield _RecordingBackend + _RecordingBackend.reset() + + +class _OverlapProbe(InputControl): + """A token-level input control that records whether two `adapt` calls ever overlapped.""" + + Args = None + supports_batching = True + + def __init__(self): + super().__init__() + self.active = 0 + self.max_active = 0 + + def adapt(self, input_ids, runtime_kwargs=None): + self.active += 1 + self.max_active = max(self.max_active, self.active) + time.sleep(0.05) + self.active -= 1 + return input_ids + + def requirements(self): + return Requirements() + + def steer_access(self): + return ModelAccess.FACTS + + +class _ProbeDriver(DecodingDriver): + """A driver that rolls out once through the session it receives and records that session.""" + + Args = None + + def _configure(self): + self.seen_session = None + + def requirements(self): + return Requirements() + + def steer_access(self): + return ModelAccess.FACTS + + def decode(self, input_ids, attention_mask, model, logits_processors, stopping_criteria, + runtime_kwargs, session=None, **gen_kwargs): + self.seen_session = session + return session_generate(session, input_ids, attention_mask, **gen_kwargs) + + +def _engine_pipeline(controls=()): + pipeline = SteeringPipeline(controls=list(controls), backend=VLLM_SPEC, tokenizer=wordlevel_tokenizer()) + pipeline.steer() + return pipeline + + +def _generate_in_threads(pipeline, count=2, **gen_kwargs): + errors = [] + + def call(index): + try: + pipeline.generate(input_ids=torch.tensor([[0, 3 + index, 4]]), max_new_tokens=2, **gen_kwargs) + except BaseException as error: + errors.append(error) + + threads = [threading.Thread(target=call, args=(index,)) for index in range(count)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(30) + return errors + + +def _overlap(first, second) -> bool: + return first[1] < second[2] and second[1] < first[2] + + +class TestSupportsConcurrentGenerate: + """The truth table, read from the spec before `steer()` (no backend is constructed).""" + + def test_huggingface_backend_is_false(self): + pipeline = SteeringPipeline(model=tiny_llama(), tokenizer=wordlevel_tokenizer()) + assert pipeline.supports_concurrent_generate is False + + def test_vllm_spec_without_a_driver_is_true(self): + pipeline = SteeringPipeline(backend=VLLM_SPEC) + assert pipeline.supports_concurrent_generate is True + + def test_vllm_serve_spec_without_a_driver_is_true(self): + spec = BackendSpec(kind="vllm-serve", model="m", options={"base_url": "http://localhost:1"}) + assert SteeringPipeline(backend=spec).supports_concurrent_generate is True + + def test_enabled_driver_is_false(self): + pipeline = SteeringPipeline(controls=[_ProbeDriver()], backend=VLLM_SPEC) + assert pipeline.supports_concurrent_generate is False + + def test_disabled_driver_is_true(self): + driver = _ProbeDriver() + driver.enabled = False + pipeline = SteeringPipeline(controls=[driver], backend=VLLM_SPEC) + assert pipeline.supports_concurrent_generate is True + + def test_steered_pipeline_reads_the_negotiated_capabilities(self, recording_engine): + pipeline = _engine_pipeline() + assert pipeline.supports_concurrent_generate is True + + +class TestPromptLock: + + def test_dispatches_overlap_and_prompt_phases_do_not(self, recording_engine): + probe = _OverlapProbe() + pipeline = _engine_pipeline([probe]) + assert pipeline.supports_concurrent_generate + + errors = _generate_in_threads(pipeline, count=3) + + assert errors == [] + assert probe.max_active == 1 + calls = recording_engine.calls + assert len(calls) == 3 + assert len({thread for thread, *_ in calls}) == 3 + assert any(_overlap(first, second) for i, first in enumerate(calls) for second in calls[i + 1:]) + + def test_unsupported_pipeline_serializes_dispatches(self, recording_engine): + pipeline = _engine_pipeline([_ProbeDriver()]) + assert not pipeline.supports_concurrent_generate + + errors = _generate_in_threads(pipeline, count=2) + + assert errors == [] + first, second = recording_engine.calls + assert not _overlap(first, second) + + def test_in_process_pipeline_serializes_calls(self, monkeypatch): + from steerability.backends.huggingface import ExclusiveSession + + intervals = [] + original = ExclusiveSession.generate + + def recording(self, items, params, *, cancel_token=None): + started = time.monotonic() + time.sleep(0.2) + result = original(self, items, params, cancel_token=cancel_token) + intervals.append((None, started, time.monotonic())) + return result + + monkeypatch.setattr(ExclusiveSession, "generate", recording) + pipeline = SteeringPipeline(model=tiny_llama(), tokenizer=wordlevel_tokenizer()) + pipeline.steer() + errors = _generate_in_threads(pipeline, count=2, do_sample=False) + + assert errors == [] + assert len(intervals) == 2 + assert not _overlap(*intervals) + + def test_compute_logprobs_runs_under_the_lock(self, recording_engine): + pipeline = _engine_pipeline() + holder_started = threading.Event() + holder_release = threading.Event() + + def hold_lock(): + with pipeline._prompt_lock: + holder_started.set() + holder_release.wait(5) + + holder = threading.Thread(target=hold_lock) + holder.start() + assert holder_started.wait(5) + finished = threading.Event() + + def score(): + pipeline.compute_logprobs(torch.tensor([[0, 3, 4]]), ref_output_ids=torch.tensor([[5]])) + finished.set() + + scorer = threading.Thread(target=score) + scorer.start() + assert not finished.wait(0.3) + holder_release.set() + assert finished.wait(5) + holder.join(5) + scorer.join(5) + + +class TestCancelTokenOnGenerate: + + def test_wrong_type_raises_type_error(self, recording_engine): + pipeline = _engine_pipeline() + with pytest.raises(TypeError, match="cancel_token must be a CancelToken"): + pipeline.generate(input_ids=torch.tensor([[0, 3, 4]]), max_new_tokens=2, cancel_token=object()) + + def test_token_reaches_the_session(self, recording_engine): + pipeline = _engine_pipeline() + token = CancelToken() + pipeline.generate(input_ids=torch.tensor([[0, 3, 4]]), max_new_tokens=2, cancel_token=token) + assert recording_engine.tokens_seen == [token] + + def test_cancelling_aborts_the_session_call(self, recording_engine): + recording_engine.hold_seconds = 30.0 + pipeline = _engine_pipeline() + token = CancelToken() + errors = [] + + def call(): + try: + pipeline.generate(input_ids=torch.tensor([[0, 3, 4]]), max_new_tokens=2, cancel_token=token) + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=call) + started = time.monotonic() + thread.start() + time.sleep(0.2) + token.cancel() + thread.join(10) + + assert not thread.is_alive() + assert time.monotonic() - started < 5 + assert len(errors) == 1 and isinstance(errors[0], GenerationCancelled) + assert recording_engine.abort_seen.is_set() + + def test_cancelled_token_before_the_call_raises(self, recording_engine): + pipeline = _engine_pipeline() + token = CancelToken() + token.cancel() + with pytest.raises(GenerationCancelled): + pipeline.generate(input_ids=torch.tensor([[0, 3, 4]]), max_new_tokens=2, cancel_token=token) + + def test_steered_session_forwards_the_token_to_a_driver_rollout(self, recording_engine): + driver = _ProbeDriver() + pipeline = _engine_pipeline([driver]) + token = CancelToken() + pipeline.generate(input_ids=torch.tensor([[0, 3, 4]]), max_new_tokens=2, cancel_token=token) + assert driver.seen_session.cancel_token is token + assert recording_engine.tokens_seen == [token] + + def test_generate_methods_accept_the_token(self, recording_engine): + pipeline = _engine_pipeline() + token = CancelToken() + pipeline.generate_tokens(torch.tensor([[0, 3, 4]]), max_new_tokens=2, cancel_token=token) + assert recording_engine.tokens_seen == [token] diff --git a/tests/core/test_driver_rollout_anchor.py b/tests/core/test_driver_rollout_anchor.py index 7f4d09e7..95d23780 100644 --- a/tests/core/test_driver_rollout_anchor.py +++ b/tests/core/test_driver_rollout_anchor.py @@ -176,7 +176,7 @@ def test_steered_session_injects_rewritten_entries_per_item(self): captured = {} class _ProbeSession: - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): captured["items"] = items return [] diff --git a/tests/core/test_engine_thread.py b/tests/core/test_engine_thread.py new file mode 100644 index 00000000..000ace33 --- /dev/null +++ b/tests/core/test_engine_thread.py @@ -0,0 +1,178 @@ +"""Tests for `EngineThread`, the event-loop thread that runs the offline vLLM engine's coroutines. + +The class has no vLLM dependency, so these tests run everywhere: concurrent submission from several +threads onto one loop, result and exception delivery, cancellation reaching the coroutine, and the +close lifecycle. +""" +import asyncio +import threading +import time + +import pytest + +from steerability.backends.vllm.engine_thread import EngineClosedError, EngineThread + + +@pytest.fixture +def engine_thread(): + thread = EngineThread(name="test-engine") + thread.start() + yield thread + thread.close() + + +def test_submissions_from_several_threads_run_concurrently_on_one_loop(engine_thread): + running = 0 + peak = 0 + loops = set() + + async def work(value): + nonlocal running, peak + loops.add(asyncio.get_running_loop()) + running += 1 + peak = max(peak, running) + await asyncio.sleep(0.2) + running -= 1 + return value * 2 + + results = {} + + def caller(value): + results[value] = engine_thread.submit(work(value)).result(5) + + threads = [threading.Thread(target=caller, args=(value,)) for value in range(4)] + started = time.monotonic() + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert results == {0: 0, 1: 2, 2: 4, 3: 6} + assert len(loops) == 1 + assert peak == 4 + assert time.monotonic() - started < 0.6 # the four sleeps overlapped + + +def test_run_returns_the_result_or_raises_the_coroutine_error(engine_thread): + async def ok(): + return "done" + + async def fail(): + raise KeyError("missing") + + assert engine_thread.run(ok()) == "done" + with pytest.raises(KeyError, match="missing"): + engine_thread.run(fail()) + + +def test_cancelling_a_future_cancels_the_coroutine(engine_thread): + started = threading.Event() + observed = threading.Event() + + async def long_running(): + started.set() + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + observed.set() + raise + + future = engine_thread.submit(long_running()) + assert started.wait(5) + assert future.cancel() + assert observed.wait(5) + + +def test_cancel_pending_cancels_only_unfinished_futures(engine_thread): + async def quick(): + return 1 + + async def slow(): + await asyncio.sleep(30) + + done = engine_thread.submit(quick()) + assert done.result(5) == 1 + pending = engine_thread.submit(slow()) + engine_thread.cancel_pending() + assert pending.cancelled() or pending.exception(5) is not None + assert done.result() == 1 + assert not engine_thread.closed + + +def test_close_cancels_pending_work_and_is_idempotent(): + thread = EngineThread(name="closing-engine") + thread.start() + observed = threading.Event() + started = threading.Event() + + async def long_running(): + started.set() + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + observed.set() + raise + + future = thread.submit(long_running()) + assert started.wait(5) + thread.close() + thread.close() + assert thread.closed + assert future.cancelled() + assert observed.is_set() + assert not any(t.name == "closing-engine" and t.is_alive() for t in threading.enumerate()) + + +def test_submit_and_run_raise_after_close(): + thread = EngineThread(name="closed-engine") + thread.start() + thread.close() + + async def never(): + return None + + coroutine = never() + with pytest.raises(EngineClosedError): + thread.submit(coroutine) + assert coroutine.cr_frame is None # the refused coroutine was closed, not left pending + with pytest.raises(EngineClosedError): + thread.run(never()) + + +def test_run_raises_engine_closed_when_closed_under_the_call(engine_thread): + started = threading.Event() + + async def long_running(): + started.set() + await asyncio.sleep(30) + + errors = [] + + def caller(): + try: + engine_thread.run(long_running()) + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=caller) + thread.start() + assert started.wait(5) + engine_thread.close() + thread.join(5) + assert len(errors) == 1 and isinstance(errors[0], EngineClosedError) + + +def test_start_twice_raises(engine_thread): + with pytest.raises(RuntimeError, match="already started"): + engine_thread.start() + + +def test_submit_before_start_raises(): + thread = EngineThread(name="unstarted-engine") + + async def never(): + return None + + with pytest.raises(EngineClosedError): + thread.submit(never()) + thread.close() diff --git a/tests/core/test_model_access.py b/tests/core/test_model_access.py index bc3394f9..11074204 100644 --- a/tests/core/test_model_access.py +++ b/tests/core/test_model_access.py @@ -125,11 +125,11 @@ def layout(self): self.calls.append("layout") return "layout-sentinel" - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): self.calls.append("generate") return [] - def score(self, items, params): + def score(self, items, params, *, cancel_token=None): self.calls.append("score") return torch.zeros(0, 0) diff --git a/tests/core/test_staged_steer.py b/tests/core/test_staged_steer.py index dbf929e6..518630f6 100644 --- a/tests/core/test_staged_steer.py +++ b/tests/core/test_staged_steer.py @@ -81,10 +81,10 @@ def capture(self, prompts, layers, mode, location="layer_output"): mode=mode, location=location, ) - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): raise AssertionError("steer-phase mocks do not generate") - def score(self, items, params): + def score(self, items, params, *, cancel_token=None): raise AssertionError("steer-phase mocks do not score") diff --git a/tests/core/test_vllm_backend_construction.py b/tests/core/test_vllm_backend_construction.py index 52c2bf90..e3fe2d9a 100644 --- a/tests/core/test_vllm_backend_construction.py +++ b/tests/core/test_vllm_backend_construction.py @@ -1,4 +1,5 @@ -"""Engine-free test: a `VLLMBackend.__init__` failure after the boot releases the engine.""" +"""Engine-free tests: a `VLLMBackend.__init__` failure after the boot releases the engine and stops +its engine thread, and a released backend refuses new requests.""" import sys import types @@ -6,32 +7,91 @@ from steerability.algorithms.core.execution import BackendSpec from steerability.backends.vllm import VLLMBackend +from steerability.backends.vllm.engine_thread import EngineThread -class _FakeLLM: +class _FakeEngine: instances: list = [] - def __init__(self, model, **kwargs): + def __init__(self, model_ref, engine_kwargs): + self.model_ref = model_ref + self.engine_kwargs = engine_kwargs self.shutdown_calls = 0 - _FakeLLM.instances.append(self) + _FakeEngine.instances.append(self) def shutdown(self): self.shutdown_calls += 1 -def test_post_boot_failure_releases_engine(monkeypatch): - module = types.ModuleType("vllm") - module.LLM = _FakeLLM - monkeypatch.setitem(sys.modules, "vllm", module) - _FakeLLM.instances.clear() - # hermetic: no hub lookups, and the realistic failure (tokenizer resolution) raises +@pytest.fixture +def fake_engine(monkeypatch): + """Boot a fake engine through the engine thread, with no vLLM and no hub lookups.""" + monkeypatch.setitem(sys.modules, "vllm", types.ModuleType("vllm")) + _FakeEngine.instances.clear() + + async def construct(model_ref, engine_kwargs): + return _FakeEngine(model_ref, engine_kwargs) + + monkeypatch.setattr("steerability.backends.vllm.backend._construct_engine", construct) monkeypatch.setattr("steerability.backends.vllm.backend._reject_encoder_decoder", lambda *a, **k: None) + yield _FakeEngine + _FakeEngine.instances.clear() + +def test_post_boot_failure_releases_engine(fake_engine, monkeypatch): + # the realistic failure after the boot is tokenizer resolution def failing(source, trust_remote_code=False): raise OSError("no such tokenizer") monkeypatch.setattr("steerability.backends.vllm.backend._client_tokenizer", failing) + created = [] + original_init = EngineThread.__init__ + + def recording_init(self, *, name): + original_init(self, name=name) + created.append(self) + + monkeypatch.setattr(EngineThread, "__init__", recording_init) with pytest.raises(OSError, match="no such tokenizer"): VLLMBackend(BackendSpec(kind="vllm", model="tiny")) - (engine,) = _FakeLLM.instances + (engine,) = fake_engine.instances + assert engine.shutdown_calls == 1 + (engine_thread,) = created + assert engine_thread.closed + + +def test_engine_logging_defaults_and_caller_overrides(fake_engine, monkeypatch): + monkeypatch.setattr("steerability.backends.vllm.backend._client_tokenizer", lambda *a, **k: object()) + monkeypatch.setattr("steerability.backends.vllm.backend._config_layout", lambda *a, **k: None) + backend = VLLMBackend(BackendSpec(kind="vllm", model="tiny")) + overridden = VLLMBackend(BackendSpec( + kind="vllm", model="tiny", + options={"engine_kwargs": {"enable_log_requests": True, "disable_log_stats": True}}, + )) + try: + default_kwargs, override_kwargs = (engine.engine_kwargs for engine in fake_engine.instances) + assert default_kwargs["enable_log_requests"] is False + # stats logging keeps vLLM's default (on), which reports preemptions + assert "disable_log_stats" not in default_kwargs + assert override_kwargs["enable_log_requests"] is True + assert override_kwargs["disable_log_stats"] is True + finally: + backend.release() + overridden.release() + + +def test_released_backend_refuses_requests(fake_engine, monkeypatch): + monkeypatch.setattr("steerability.backends.vllm.backend._client_tokenizer", lambda *a, **k: object()) + monkeypatch.setattr("steerability.backends.vllm.backend._config_layout", lambda *a, **k: None) + backend = VLLMBackend(BackendSpec(kind="vllm", model="tiny")) + session = backend.open_session() + backend.release() + backend.release() + + (engine,) = fake_engine.instances assert engine.shutdown_calls == 1 + with pytest.raises(RuntimeError, match="was released"): + backend.open_session() + with pytest.raises(RuntimeError, match="was released"): + backend._submit_request({"prompt_token_ids": [1]}, types.SimpleNamespace(output_kind=None)) + assert session.closed is False diff --git a/tests/core/test_vllm_engine.py b/tests/core/test_vllm_engine.py index 725cc90d..ff27aa18 100644 --- a/tests/core/test_vllm_engine.py +++ b/tests/core/test_vllm_engine.py @@ -157,3 +157,123 @@ def parsed(label: str, text: str): assert parsed("engine", engine_text) is not None assert parsed("hf", hf_text) is not None assert engine_text == hf_text + + +def _greedy_item(tokenizer, text: str) -> GenerationItem: + encoded = tokenizer(text, return_tensors="pt") + return GenerationItem(prompt=PreparedPrompt.from_token_ids(encoded["input_ids"])) + + +class TestPerRequestEngine: + """Per-request admission through the engine thread: concurrent calls share the running batch, a + short request returns before a long one, and cancellation and release abort in-flight requests.""" + + def test_score_matches_the_hugging_face_reference(self, engine_backend): + tokenizer = _tokenizer() + model = AutoModelForCausalLM.from_pretrained(TINY_MODEL, dtype=torch.float32) + prompt_ids = tokenizer("The capital of France is", return_tensors="pt")["input_ids"] + ref = tokenizer(" Paris", add_special_tokens=False, return_tensors="pt")["input_ids"] + with torch.no_grad(): + logits = model(torch.cat([prompt_ids, ref], dim=1)).logits + positions = torch.arange(prompt_ids.size(1) - 1, prompt_ids.size(1) - 1 + ref.size(1)) + reference = torch.log_softmax(logits[0, positions], dim=-1).gather(-1, ref[0].unsqueeze(-1)).squeeze(-1) + + item = ScoringItem(prompt=PreparedPrompt.from_token_ids(prompt_ids), ref_output_ids=ref) + with engine_backend.open_session() as session: + scored = session.score([item], GenerationParams()) + # the engine and the in-process forward use different attention kernels, which differ by + # about 1e-3 in float32 + assert torch.allclose(scored[0], reference, atol=1e-2) + + def test_concurrent_calls_share_the_engine_and_the_short_one_returns_first(self, engine_backend): + import threading + import time + + tokenizer = engine_backend.tokenizer + short_item = _greedy_item(tokenizer, "The capital of France is") + long_item = _greedy_item(tokenizer, "Once upon a time") + short_params = GenerationParams(max_new_tokens=4, greedy=True) + long_params = GenerationParams(max_new_tokens=400, min_new_tokens=400, greedy=True) + + with engine_backend.open_session() as session: + serial_short = session.generate([short_item], short_params)[0].output.output_ids.tolist() + serial_long = session.generate([long_item], long_params)[0].output.output_ids.tolist() + + finished: dict[str, float] = {} + outputs: dict[str, list] = {} + + def run(name, item, params): + with engine_backend.open_session() as session: + outputs[name] = session.generate([item], params)[0].output.output_ids.tolist() + finished[name] = time.monotonic() + + long_thread = threading.Thread(target=run, args=("long", long_item, long_params)) + short_thread = threading.Thread(target=run, args=("short", short_item, short_params)) + long_thread.start() + time.sleep(0.05) + short_thread.start() + long_thread.join(120) + short_thread.join(120) + + assert finished["short"] < finished["long"] + assert outputs["short"] == serial_short + assert outputs["long"] == serial_long + + def test_cancelling_a_token_aborts_the_engine_request(self, engine_backend): + import threading + import time + + from steerability.algorithms.core.execution import CancelToken, GenerationCancelled + + item = _greedy_item(engine_backend.tokenizer, "Once upon a time") + params = GenerationParams(max_new_tokens=480, min_new_tokens=480, greedy=True) + token = CancelToken() + errors = [] + + def run(): + try: + with engine_backend.open_session() as session: + session.generate([item], params, cancel_token=token) + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=run) + thread.start() + time.sleep(0.3) + token.cancel() + thread.join(60) + + assert len(errors) == 1 and isinstance(errors[0], GenerationCancelled) + deadline = time.monotonic() + 10 + while engine_backend._engine.output_processor.get_num_unfinished_requests() and time.monotonic() < deadline: + time.sleep(0.05) + assert engine_backend._engine.output_processor.get_num_unfinished_requests() == 0 + + def test_pipeline_generate_with_a_cancelled_token_raises(self, engine_backend): + import threading + import time + + from steerability.algorithms.core.execution import CancelToken, GenerationCancelled + + pipeline = SteeringPipeline(backend=engine_backend.spec, tokenizer=_tokenizer()) + pipeline._backends[engine_backend.spec] = engine_backend + pipeline.steer() + assert pipeline.supports_concurrent_generate + token = CancelToken() + errors = [] + + def run(): + try: + pipeline.generate( + text="Once upon a time", max_new_tokens=480, min_new_tokens=480, do_sample=False, + cancel_token=token, + ) + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=run) + thread.start() + time.sleep(0.3) + token.cancel() + thread.join(60) + assert len(errors) == 1 and isinstance(errors[0], GenerationCancelled) diff --git a/tests/core/test_vllm_plugin_engine.py b/tests/core/test_vllm_plugin_engine.py index 2eacd295..3b00119f 100644 --- a/tests/core/test_vllm_plugin_engine.py +++ b/tests/core/test_vllm_plugin_engine.py @@ -391,3 +391,56 @@ def test_routed_decoding_end_to_end_on_engine(self, plugin_backend): pipeline.steer() text = pipeline.generate(text="the committee approved it", max_new_tokens=8, do_sample=False) assert isinstance(text, str) + + +class TestPerRequestPluginEngine: + """Capture and interventions through the per-request path: a spec-bearing request that shares the + running batch with a plain request is steered exactly as it is when it runs alone.""" + + def test_capture_returns_the_plugin_payload_per_request(self, plugin_backend): + prompts = [PreparedPrompt.from_text("The cat sat on the mat"), PreparedPrompt.from_text("Hello")] + with plugin_backend.open_session() as session: + result = session.capture(prompts, [0, 1], "all_tokens") + assert set(result.hidden) == {0, 1} + assert result.hidden[0].shape[0] == 2 + assert result.attention_mask.shape[0] == 2 + + def test_steered_and_plain_requests_share_a_batch(self, plugin_backend): + import threading + + from steerability.algorithms.state_control.caa.control import CAA + + hidden = plugin_backend._layout.hidden_size + vector = _steered_vector(TINY_MODEL, hidden, layers=[1]) + control = CAA(steering_vector=vector, layer_id=1, multiplier=6.0) + pipeline = SteeringPipeline(controls=[control], backend=plugin_backend.spec) + pipeline.model = AutoModelForCausalLM.from_pretrained(TINY_MODEL) + pipeline.tokenizer = _tokenizer() + pipeline._backends[plugin_backend.spec] = plugin_backend + pipeline.steer() + plain = SteeringPipeline(backend=plugin_backend.spec, tokenizer=_tokenizer()) + plain._backends[plugin_backend.spec] = plugin_backend + plain.steer() + + steered_prompt, plain_prompt = "The committee reviewed the plan", "The weather today is" + kwargs = dict(max_new_tokens=24, do_sample=False, return_output=True) + serial_steered = pipeline.generate(text=steered_prompt, **kwargs).output_ids[0].tolist() + serial_plain = plain.generate(text=plain_prompt, **kwargs).output_ids[0].tolist() + + outputs = {} + + def run(name, target, prompt): + outputs[name] = target.generate(text=prompt, **kwargs).output_ids[0].tolist() + + threads = [ + threading.Thread(target=run, args=("steered", pipeline, steered_prompt)), + threading.Thread(target=run, args=("plain", plain, plain_prompt)), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(120) + + assert outputs["steered"] == serial_steered + assert outputs["plain"] == serial_plain + assert serial_steered != plain.generate(text=steered_prompt, **kwargs).output_ids[0].tolist() diff --git a/tests/core/test_vllm_release.py b/tests/core/test_vllm_release.py index ddb8d3e2..bc76fd86 100644 --- a/tests/core/test_vllm_release.py +++ b/tests/core/test_vllm_release.py @@ -136,3 +136,35 @@ def test_pipeline_end_to_end_with_stopping_rules(): assert out.output_ids.shape[1] <= 6 finally: pipeline.release_backends() + + +def test_release_with_a_request_in_flight_fails_that_request(): + """Release while a long request is in flight: the waiter raises `EngineClosedError`, the engine + shuts down, and the backend refuses new sessions.""" + import threading + import time + + from steerability.backends.vllm.engine_thread import EngineClosedError + + backend = _boot_or_skip() + item = GenerationItem(prompt=PreparedPrompt.from_text("Once upon a time")) + params = GenerationParams(max_new_tokens=480, min_new_tokens=480, greedy=True) + errors = [] + + def run(): + try: + with backend.open_session() as session: + session.generate([item], params) + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=run) + thread.start() + time.sleep(0.5) + backend.release() + thread.join(60) + + assert len(errors) == 1 and isinstance(errors[0], EngineClosedError) + assert backend._engine_thread.closed + with pytest.raises(RuntimeError, match="was released"): + backend.open_session() diff --git a/tests/core/test_vllm_serve_backend.py b/tests/core/test_vllm_serve_backend.py index 18a264b2..45740f68 100644 --- a/tests/core/test_vllm_serve_backend.py +++ b/tests/core/test_vllm_serve_backend.py @@ -378,6 +378,57 @@ def test_layout_unresolvable_raises_on_access(self, fake_server): _ = session.layout +class TestServeSessionCancellation: + """The serve session checks the token before each item's request; a sent request completes.""" + + def test_cancelled_token_sends_no_request(self, fake_server): + from steerability.algorithms.core.execution import CancelToken, GenerationCancelled + + backend = VLLMServeBackend(_serve_spec()) + token = CancelToken() + token.cancel() + items = [GenerationItem(prompt=PreparedPrompt.from_token_ids([0, 3, 4]))] + with backend.open_session() as session: + with pytest.raises(GenerationCancelled): + session.generate(items, GenerationParams(max_new_tokens=4), cancel_token=token) + with pytest.raises(GenerationCancelled): + session.score( + [ScoringItem(prompt=PreparedPrompt.from_token_ids([0, 3]), ref_output_ids=torch.tensor([[5]]))], + GenerationParams(), cancel_token=token, + ) + assert not [path for path, _ in fake_server.requests if path == "/v1/completions"] + + def test_cancellation_between_items_stops_the_remaining_requests(self, fake_server): + from steerability.algorithms.core.execution import CancelToken, GenerationCancelled + + backend = VLLMServeBackend(_serve_spec(max_concurrency=1)) + token = CancelToken() + original = fake_server.handle + + def cancel_after_first(path, payload): + response = original(path, payload) + if path == "/v1/completions": + token.cancel() + return response + + fake_server.handle = cancel_after_first + items = [GenerationItem(prompt=PreparedPrompt.from_token_ids([0, 3 + index])) for index in range(3)] + with backend.open_session() as session: + with pytest.raises(GenerationCancelled): + session.generate(items, GenerationParams(max_new_tokens=4), cancel_token=token) + completions = [payload for path, payload in fake_server.requests if path == "/v1/completions"] + assert [payload["prompt"] for payload in completions] == [[0, 3]] + + def test_uncancelled_token_generates_normally(self, fake_server): + from steerability.algorithms.core.execution import CancelToken + + backend = VLLMServeBackend(_serve_spec()) + items = [GenerationItem(prompt=PreparedPrompt.from_token_ids([0, 3, 4]))] + with backend.open_session() as session: + results = session.generate(items, GenerationParams(max_new_tokens=4), cancel_token=CancelToken()) + assert results[0].output.output_ids.tolist() == [[9, 1]] + + class TestEncoderDecoderSpecRejection: def test_local_encoder_decoder_config_rejected_for_vllm_kinds(self, tmp_path): @@ -799,39 +850,44 @@ def test_gemma4_config_dir_reports_per_layer_head_dim_as_none(self, tmp_path): assert facts.model_type == "gemma4" -class _FakeCaptureLLM: - """An offline engine double that records each capture request and returns zero hidden states.""" +class _FakeCaptureEngine: + """An `AsyncLLM` double that records each capture request and returns zero hidden states.""" def __init__(self, hidden: int = 4): self.hidden = hidden self.extra_args: list[dict] = [] + self.request_ids: list[str] = [] - def generate(self, prompts, sampling, use_tqdm=False): + async def generate(self, prompt, sampling_params, request_id, lora_request=None): import json from types import SimpleNamespace import safetensors.torch - self.extra_args.append(dict(sampling.extra_args)) - layers = sampling.extra_args["capture"]["layers"] - outputs = [] - for prompt in prompts: - length = len(prompt["prompt_token_ids"]) - tensors = {f"layer_{layer}": torch.zeros(length, self.hidden) for layer in layers} - outputs.append(SimpleNamespace(captures=(json.dumps({"positions": {}}), safetensors.torch.save(tensors)))) - return outputs + self.extra_args.append(dict(sampling_params.extra_args)) + self.request_ids.append(request_id) + layers = sampling_params.extra_args["capture"]["layers"] + length = len(prompt["prompt_token_ids"]) + tensors = {f"layer_{layer}": torch.zeros(length, self.hidden) for layer in layers} + yield SimpleNamespace(captures=(json.dumps({"positions": {}}), safetensors.torch.save(tensors))) -def _offline_backend(llm, tmp_path): - """A `VLLMBackend` over a fake engine, built without booting vLLM.""" +def _offline_backend(engine, tmp_path): + """A `VLLMBackend` over a fake engine on a started engine thread, built without booting vLLM.""" + import itertools + from steerability.algorithms.core.execution import ModelFacts from steerability.backends.vllm import VLLMBackend from steerability.backends.vllm.backend import _ArtifactUploader + from steerability.backends.vllm.engine_thread import EngineThread backend = VLLMBackend.__new__(VLLMBackend) backend.spec = BackendSpec(kind="vllm", model="m", options={"hook_plugin": True, "artifact_dir": str(tmp_path)}) backend._released = False - backend._llm = llm + backend._engine = engine + backend._engine_thread = EngineThread(name="fake-engine") + backend._engine_thread.start() + backend._request_counter = itertools.count() backend._lora_request = None backend.tokenizer = wordlevel_tokenizer() backend._layout = ModelFacts( @@ -849,30 +905,39 @@ class TestOfflineSteeredCapture: def test_capture_requests_carry_the_entries_spec(self, tmp_path): pytest.importorskip("vllm") pytest.importorskip("vllm_hook_plugins") - llm = _FakeCaptureLLM() - backend = _offline_backend(llm, tmp_path) + engine = _FakeCaptureEngine() + backend = _offline_backend(engine, tmp_path) spec = _mini_spec() prompts = [PreparedPrompt.from_token_ids([0, 3, 4]), PreparedPrompt.from_token_ids([0, 5])] - with backend.open_session() as session: - unsteered = session.capture(prompts, [0], "last_token") - steered = session.capture(prompts, [0], "last_token", state_entries=(InterventionEntry(spec=spec),)) + try: + with backend.open_session() as session: + unsteered = session.capture(prompts, [0], "last_token") + steered = session.capture(prompts, [0], "last_token", state_entries=(InterventionEntry(spec=spec),)) + finally: + backend._engine_thread.close() assert unsteered.hidden[0].shape == steered.hidden[0].shape == (2, 4) - assert "intervention_spec" not in llm.extra_args[0] - assert llm.extra_args[1]["intervention_spec"] == spec.to_wire() - assert llm.extra_args[1]["capture"] == llm.extra_args[0]["capture"] + # one engine request per prompt: the unsteered call's two requests, then the steered call's + assert len(engine.extra_args) == 4 + assert all("intervention_spec" not in extra for extra in engine.extra_args[:2]) + assert all(extra["intervention_spec"] == spec.to_wire() for extra in engine.extra_args[2:]) + assert all(extra["capture"] == engine.extra_args[0]["capture"] for extra in engine.extra_args) + assert len(set(engine.request_ids)) == 4 def test_last_token_capture_needs_only_the_all_tokens_wire_mode(self, tmp_path): pytest.importorskip("vllm") pytest.importorskip("vllm_hook_plugins") - llm = _FakeCaptureLLM() - backend = _offline_backend(llm, tmp_path) + engine = _FakeCaptureEngine() + backend = _offline_backend(engine, tmp_path) backend._discovery["capture_kinds"] = {**backend._discovery["capture_kinds"], "modes": ["all_tokens"]} assert backend.capture_kinds.modes == frozenset({"all_tokens"}) - with backend.open_session() as session: - result = session.capture([PreparedPrompt.from_token_ids([0, 3, 4])], [0], "last_token") + try: + with backend.open_session() as session: + result = session.capture([PreparedPrompt.from_token_ids([0, 3, 4])], [0], "last_token") + finally: + backend._engine_thread.close() assert result.hidden[0].shape == (1, 4) - assert llm.extra_args[0]["capture"]["mode"] == "all_tokens" + assert engine.extra_args[0]["capture"]["mode"] == "all_tokens" def test_steered_session_forwards_its_entries(self): from steerability.algorithms.core.execution.backend import SteeredSession diff --git a/tests/evaluation/conftest.py b/tests/evaluation/conftest.py index 8eff7ac2..923df914 100644 --- a/tests/evaluation/conftest.py +++ b/tests/evaluation/conftest.py @@ -74,6 +74,7 @@ class StubSteeringPipeline: calls: One dict per `generate` invocation with the received arguments. fail_above_batch_size: When set, a dispatch with more prompts than this raises. gate: Optional `threading.Event` the first dispatch waits on before returning. + supports_concurrent_generate: The value the provider reads to choose per-request mode. """ def __init__( @@ -81,11 +82,13 @@ def __init__( *, tokenizer: StubTokenizer | None = None, supports_batching: bool = True, + supports_concurrent_generate: bool = False, controls: tuple = (), decode_texts: list[str] | None = None, ): self._is_steered = True self.supports_batching = supports_batching + self.supports_concurrent_generate = supports_concurrent_generate self.controls = controls self.tokenizer = tokenizer if tokenizer is not None else StubTokenizer(decode_texts=decode_texts) self.calls: list[dict[str, Any]] = [] diff --git a/tests/evaluation/test_batching.py b/tests/evaluation/test_batching.py index 33f2edef..10cce697 100644 --- a/tests/evaluation/test_batching.py +++ b/tests/evaluation/test_batching.py @@ -1,5 +1,10 @@ -"""Tests for the lock-leader collator: batching under deterministic gating, batch keys, runtime- -kwargs collation, poison isolation, cancellation, strict serialization, and loop reuse.""" +"""Tests for the request dispatchers. + +The lock-leader collator tests cover batching under deterministic gating, batch keys, runtime-kwargs +collation, poison isolation, cancellation, strict serialization, and loop reuse. The per-request +dispatcher tests cover one call per request, the concurrency limiter, error isolation, request seeds, and +cancellation of a request's generation. +""" import threading import anyio @@ -7,7 +12,8 @@ pytest.importorskip("inspect_ai") -from steerability.evaluation.batching import LockLeaderCollator +from steerability.algorithms.core.execution import CancelToken +from steerability.evaluation.batching import LockLeaderCollator, PerRequestDispatcher, request_seed from tests.evaluation.conftest import StubSteeringPipeline @@ -350,3 +356,261 @@ async def main(): with pytest.raises(RuntimeError, match="boom"): anyio.run(main) + + +def _dispatcher(pipeline=None, *, max_concurrency=4, declared_scopes=None, static_runtime_kwargs=None, + prompt_path="messages") -> tuple[PerRequestDispatcher, StubSteeringPipeline]: + pipeline = pipeline if pipeline is not None else StubSteeringPipeline(supports_concurrent_generate=True) + dispatcher = PerRequestDispatcher( + pipeline, + max_concurrency=max_concurrency, + prompt_path=prompt_path, + declared_scopes={"spans": "row"} if declared_scopes is None else declared_scopes, + static_runtime_kwargs=static_runtime_kwargs or {}, + ) + return dispatcher, pipeline + + +class _BlockingPipeline(StubSteeringPipeline): + """A stub whose `generate` records its interval and blocks until released or cancelled.""" + + def __init__(self, *, release: threading.Event | None = None, hold_seconds: float = 0.3): + super().__init__(supports_concurrent_generate=True) + self.release = release + self.hold_seconds = hold_seconds + self.intervals: list[tuple[float, float]] = [] + self.active = 0 + self.peak = 0 + self._counter = threading.Lock() + + def generate(self, **kwargs): + import time + + with self._counter: + self.active += 1 + self.peak = max(self.peak, self.active) + started = time.monotonic() + token = kwargs.get("cancel_token") + try: + deadline = started + self.hold_seconds + while time.monotonic() < deadline: + if token is not None and token.cancelled: + from steerability.algorithms.core.execution import GenerationCancelled + + raise GenerationCancelled("cancelled") + if self.release is not None and self.release.is_set(): + break + time.sleep(0.01) + return super().generate(**kwargs) + finally: + self.intervals.append((started, time.monotonic())) + with self._counter: + self.active -= 1 + + +def _serve_all(dispatcher, records): + results: list = [None] * len(records) + errors: list = [None] * len(records) + + async def serve(index): + try: + results[index] = await dispatcher.serve(records[index]) + except Exception as error: + errors[index] = error + + async def main(): + async with anyio.create_task_group() as tg: + for index in range(len(records)): + tg.start_soon(serve, index) + + anyio.run(main) + return results, errors + + +class TestPerRequestDispatcher: + def test_every_record_dispatches_as_its_own_single_prompt_call(self): + pipeline = _BlockingPipeline(hold_seconds=0.3) + dispatcher, _ = _dispatcher(pipeline, max_concurrency=8) + records = [_admit(dispatcher, [{"role": "user", "content": f"q{i}"}], per_sample={"spans": [f"s{i}"]}) + for i in range(4)] + results, errors = _serve_all(dispatcher, records) + + assert errors == [None] * 4 + assert all(result is not None for result in results) + assert len(pipeline.calls) == 4 + # a batch of one conversation per call; the dispatcher does not wait for other records + assert all(len(call["messages"]) == 1 for call in pipeline.calls) + assert pipeline.peak == 4 + dispatched = sorted(call["messages"][0][0]["content"] for call in pipeline.calls) + assert dispatched == ["q0", "q1", "q2", "q3"] + spans = sorted(call["runtime_kwargs"]["spans"][0][0] for call in pipeline.calls) + assert spans == ["s0", "s1", "s2", "s3"] # per-sample values arrive as one-element lists + + def test_limiter_bounds_the_calls_in_flight(self): + pipeline = _BlockingPipeline(hold_seconds=0.2) + dispatcher, _ = _dispatcher(pipeline, max_concurrency=2) + records = [_admit(dispatcher, [{"role": "user", "content": f"q{i}"}]) for i in range(5)] + _, errors = _serve_all(dispatcher, records) + assert errors == [None] * 5 + assert pipeline.peak == 2 + + def test_each_record_receives_its_own_error(self): + class PoisonPipeline(StubSteeringPipeline): + def generate(self, *, messages=None, **kwargs): + (conversation,) = messages + if conversation[0]["content"] == "poison": + raise ValueError("bad sample") + return super().generate(messages=messages, **kwargs) + + dispatcher, _ = _dispatcher(PoisonPipeline(supports_concurrent_generate=True)) + records = [_admit(dispatcher, [{"role": "user", "content": prompt}]) for prompt in ("ok", "poison", "ok2")] + results, errors = _serve_all(dispatcher, records) + assert results[0] is not None and results[2] is not None + assert isinstance(errors[1], ValueError) and "bad sample" in str(errors[1]) + assert errors[0] is None and errors[2] is None + + def test_static_and_per_sample_kwargs_partition_like_the_collator(self, caplog): + pipeline = StubSteeringPipeline(supports_concurrent_generate=True) + artifact = {"a": "b"} + with caplog.at_level("INFO", logger="steerability.evaluation.batching"): + dispatcher, _ = _dispatcher( + pipeline, + declared_scopes={"spans": "row", "canned_responses": "call"}, + static_runtime_kwargs={"canned_responses": artifact, "other": 1}, + ) + record = _admit(dispatcher, [{"role": "user", "content": "q"}], per_sample={"spans": ["x"], "unused": 2}) + _serve_all(dispatcher, [record]) + (call,) = pipeline.calls + assert call["runtime_kwargs"]["canned_responses"] is artifact + assert call["runtime_kwargs"]["spans"] == [["x"]] + assert "other" not in call["runtime_kwargs"] and "unused" not in call["runtime_kwargs"] + assert dispatcher.inert_runtime_kwargs == frozenset({"other", "unused"}) + with pytest.raises(ValueError, match="both per sample"): + _admit(dispatcher, [{"role": "user", "content": "q"}], per_sample={"canned_responses": {}}) + scoped, _ = _dispatcher(declared_scopes={"canned_responses": "call"}) + with pytest.raises(ValueError, match="declared 'call'-scoped"): + _admit(scoped, [{"role": "user", "content": "q"}], per_sample={"canned_responses": {}}) + + def test_static_row_scoped_value_is_one_row(self): + pipeline = StubSteeringPipeline(supports_concurrent_generate=True) + dispatcher, _ = _dispatcher(pipeline, static_runtime_kwargs={"spans": ["x"]}) + _serve_all(dispatcher, [_admit(dispatcher, [{"role": "user", "content": "q"}])]) + (call,) = pipeline.calls + assert call["runtime_kwargs"]["spans"] == [["x"]] + + def test_multi_candidate_request_dispatches_a_bare_conversation(self): + pipeline = StubSteeringPipeline(supports_concurrent_generate=True) + dispatcher, _ = _dispatcher(pipeline) + conversation = [{"role": "user", "content": "q"}] + results, errors = _serve_all(dispatcher, [_admit(dispatcher, conversation, num_choices=3)]) + assert errors == [None] + (call,) = pipeline.calls + assert call["messages"] == conversation + assert call["gen_kwargs"]["n"] == 3 + assert results[0].output_ids.shape[0] == 3 + + def test_closed_dispatcher_refuses_admission(self): + dispatcher, _ = _dispatcher() + dispatcher.close() + assert dispatcher.closed + with pytest.raises(RuntimeError, match="closed"): + _admit(dispatcher, [{"role": "user", "content": "q"}]) + + def test_max_concurrency_is_validated(self): + with pytest.raises(ValueError, match="max_concurrency"): + _dispatcher(max_concurrency=0) + + def test_a_cancel_token_reaches_pipeline_generate(self): + pipeline = StubSteeringPipeline(supports_concurrent_generate=True) + dispatcher, _ = _dispatcher(pipeline) + record = _admit(dispatcher, [{"role": "user", "content": "q"}]) + _serve_all(dispatcher, [record]) + (call,) = pipeline.calls + token = call["gen_kwargs"]["cancel_token"] + assert isinstance(token, CancelToken) + assert token is record.cancel_token and not token.cancelled + + def test_cancelling_the_serving_task_cancels_the_token(self): + import time + + pipeline = _BlockingPipeline(hold_seconds=30.0) + dispatcher, _ = _dispatcher(pipeline) + record = _admit(dispatcher, [{"role": "user", "content": "q"}]) + outcome = {} + + async def main(): + with anyio.move_on_after(0.3) as scope: + await dispatcher.serve(record) + outcome["cancelled"] = scope.cancelled_caught + outcome["returned_at"] = time.monotonic() + + started = time.monotonic() + anyio.run(main) + assert outcome["cancelled"] is True + assert outcome["returned_at"] - started < 5 # the task returned at once + assert record.cancel_token.cancelled + deadline = time.monotonic() + 5 + while not pipeline.intervals and time.monotonic() < deadline: + time.sleep(0.02) + assert pipeline.intervals # the worker thread observed the token and ended + assert record.output is None + + def test_reuse_across_two_anyio_run_loops(self): + dispatcher, pipeline = _dispatcher() + + async def one_request(tag): + return await dispatcher.serve(_admit(dispatcher, [{"role": "user", "content": tag}])) + + assert anyio.run(one_request, "first") is not None + assert anyio.run(one_request, "second") is not None + assert len(pipeline.calls) == 2 + + +class TestRequestSeeds: + @staticmethod + def _seed_of(prompt, *, seed=7, per_sample=None): + pipeline = StubSteeringPipeline(supports_concurrent_generate=True) + dispatcher, _ = _dispatcher(pipeline) + record = _admit(dispatcher, prompt, gen_kwargs={"max_new_tokens": 4, "seed": seed, "seed_scope": "dispatch"}, + per_sample=per_sample or {}) + _serve_all(dispatcher, [record]) + (call,) = pipeline.calls + return call["gen_kwargs"]["seed"] + + def test_seed_is_derived_from_the_base_seed_and_the_request(self): + prompt = [{"role": "user", "content": "q"}] + assert self._seed_of(prompt) == request_seed(7, prompt, {}) + assert self._seed_of(prompt) != 7 + + def test_seed_is_deterministic_and_shared_by_identical_requests(self): + prompt = [{"role": "user", "content": "q"}] + assert self._seed_of(prompt) == self._seed_of([{"role": "user", "content": "q"}]) + + def test_distinct_prompts_and_base_seeds_give_distinct_seeds(self): + first = self._seed_of([{"role": "user", "content": "a"}]) + second = self._seed_of([{"role": "user", "content": "b"}]) + third = self._seed_of([{"role": "user", "content": "a"}], seed=8) + assert len({first, second, third}) == 3 + + def test_per_sample_kwargs_enter_the_seed(self): + prompt = [{"role": "user", "content": "q"}] + assert self._seed_of(prompt, per_sample={"spans": ["x"]}) != self._seed_of(prompt, per_sample={"spans": ["y"]}) + + def test_unseeded_request_attaches_no_seed(self): + pipeline = StubSteeringPipeline(supports_concurrent_generate=True) + dispatcher, _ = _dispatcher(pipeline) + _serve_all(dispatcher, [_admit(dispatcher, [{"role": "user", "content": "q"}])]) + (call,) = pipeline.calls + assert "seed" not in call["gen_kwargs"] + + def test_seed_is_independent_of_arrival_order(self): + prompts = [[{"role": "user", "content": f"q{i}"}] for i in range(4)] + + def seeds_in(order): + pipeline = StubSteeringPipeline(supports_concurrent_generate=True) + dispatcher, _ = _dispatcher(pipeline) + records = [_admit(dispatcher, prompts[i], gen_kwargs={"max_new_tokens": 4, "seed": 3}) for i in order] + _serve_all(dispatcher, records) + return {call["messages"][0][0]["content"]: call["gen_kwargs"]["seed"] for call in pipeline.calls} + + assert seeds_in([0, 1, 2, 3]) == seeds_in([3, 1, 0, 2]) diff --git a/tests/evaluation/test_provider.py b/tests/evaluation/test_provider.py index f7f5fdc4..7ca75621 100644 --- a/tests/evaluation/test_provider.py +++ b/tests/evaluation/test_provider.py @@ -138,7 +138,64 @@ def test_hooks(self): assert api.tools_required() is False +class TestDispatchMode: + def test_default_is_batched(self): + api = _api(StubSteeringPipeline(supports_concurrent_generate=True), options=ProviderOptions(max_batch_size=3)) + assert api.dispatch_mode == "batched" + assert api.max_connections() == 3 + + def test_max_concurrency_on_a_supporting_pipeline_is_per_request(self): + from steerability.evaluation.batching import PerRequestDispatcher + + pipeline = StubSteeringPipeline(supports_concurrent_generate=True) + api = _api(pipeline, options=ProviderOptions(max_concurrency=64, max_batch_size=3)) + assert api.dispatch_mode == "per_request" + assert api.max_connections() == 64 + assert isinstance(api._collator, PerRequestDispatcher) + + def test_max_concurrency_on_an_unsupporting_pipeline_falls_back_and_logs(self, caplog): + pipeline = StubSteeringPipeline(supports_concurrent_generate=False) + with caplog.at_level("INFO", logger="steerability.evaluation.provider"): + api = _api(pipeline, options=ProviderOptions(max_concurrency=64, max_batch_size=3)) + assert api.dispatch_mode == "batched" + assert api.max_connections() == 3 + fallback = [ + record for record in caplog.records if "does not support concurrent generation" in record.getMessage() + ] + assert len(fallback) == 1 + + def test_fallback_log_names_the_in_process_backend(self, caplog): + from steerability.algorithms.core.steering_pipeline import SteeringPipeline + from tests.utils.tiny_models import tiny_llama, wordlevel_tokenizer + + tokenizer = wordlevel_tokenizer() + tokenizer.chat_template = CHAT_TEMPLATE + pipeline = SteeringPipeline(model=tiny_llama(), tokenizer=tokenizer) + pipeline.steer() + with caplog.at_level("INFO", logger="steerability.evaluation.provider"): + api = _api(pipeline, options=ProviderOptions(max_concurrency=8)) + assert api.dispatch_mode == "batched" + assert any("backend kind 'huggingface' runs in process" in record.getMessage() for record in caplog.records) + + def test_per_request_generate_dispatches_one_conversation_with_a_token(self): + from steerability.algorithms.core.execution import CancelToken + + pipeline = StubSteeringPipeline(supports_concurrent_generate=True) + api = _api(pipeline, options=ProviderOptions(max_concurrency=4)) + output = _generate(api, [ChatMessageUser(content="hello")]) + assert output.choices[0].message.text == "row-0" + (call,) = pipeline.calls + assert call["messages"] == [[{"role": "user", "content": "hello"}]] + assert isinstance(call["gen_kwargs"]["cancel_token"], CancelToken) + + class TestProviderOptionsValidation: + def test_max_concurrency_validation(self): + with pytest.raises(ValueError, match="max_concurrency"): + ProviderOptions(max_concurrency=0) + assert ProviderOptions().max_concurrency is None + assert ProviderOptions(max_concurrency=1).max_concurrency == 1 + def test_bad_values_raise(self): with pytest.raises(ValueError, match="max_batch_size"): ProviderOptions(max_batch_size=0) diff --git a/tests/evaluation/test_runner.py b/tests/evaluation/test_runner.py index 5f180efd..ec9b43e7 100644 --- a/tests/evaluation/test_runner.py +++ b/tests/evaluation/test_runner.py @@ -109,6 +109,7 @@ def test_run_shape_and_results_frame(self, tiny_base, tmp_path): assert baseline_run["suites"]["capability"]["stub/task"]["metrics"]["match/accuracy"] == 0.5 assert baseline_run["provenance"]["backend"] == "huggingface" assert baseline_run["provenance"]["prompt_path"] == "messages" + assert baseline_run["provenance"]["dispatch_mode"] == "batched" frame = runner.results() assert list(frame.columns) == [ @@ -203,6 +204,27 @@ def test_changed_protocol_is_refused(self, tiny_base, tmp_path): with pytest.raises(ValueError, match=r"different protocol \(hf_model_kwargs differ"): second.run() + def test_changed_max_concurrency_is_a_changed_protocol(self, tiny_base, tmp_path): + import json + + from steerability.evaluation.provider import ProviderOptions + + _runner({"baseline": []}, save_dir=tmp_path, provider_options=ProviderOptions()).run() + recorded = json.loads((tmp_path / "protocol.json").read_text())["provider_options"] + assert recorded["fields"]["max_concurrency"] is None + second = _runner({"baseline": []}, save_dir=tmp_path, provider_options=ProviderOptions(max_concurrency=16)) + with pytest.raises(ValueError, match=r"different protocol \(provider_options differ"): + second.run() + + def test_recorded_provider_options_contain_max_concurrency(self, tiny_base, tmp_path): + import json + + from steerability.evaluation.provider import ProviderOptions + + _runner({"baseline": []}, save_dir=tmp_path, provider_options=ProviderOptions(max_concurrency=16)).run() + protocol = json.loads((tmp_path / "protocol.json").read_text()) + assert protocol["provider_options"]["fields"]["max_concurrency"] == 16 + def test_logs_without_a_protocol_are_refused(self, tiny_base, tmp_path): log_dir = tmp_path / "inspect_logs" / "baseline" / "trial_0" / "capability" log_dir.mkdir(parents=True) diff --git a/tests/utils/runtime_helpers.py b/tests/utils/runtime_helpers.py index b68fc26d..a5fc8ad3 100644 --- a/tests/utils/runtime_helpers.py +++ b/tests/utils/runtime_helpers.py @@ -110,7 +110,7 @@ def __init__(self, fake_generate, tokenizer=None): self._fake_generate = fake_generate self.tokenizer = tokenizer - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): from steerability.algorithms.core.execution.payloads import ItemResult from steerability.algorithms.core.output import Output @@ -140,7 +140,7 @@ def script_session_generate(monkeypatch, fake_generate): """ from steerability.backends.huggingface import ExclusiveSession - def generate(self, items, params): + def generate(self, items, params, *, cancel_token=None): return ScriptedSession(fake_generate, tokenizer=self.tokenizer).generate(items, params) monkeypatch.setattr(ExclusiveSession, "generate", generate) diff --git a/uv.lock b/uv.lock index a687cf31..070a84b6 100644 --- a/uv.lock +++ b/uv.lock @@ -7179,7 +7179,7 @@ wheels = [ [[package]] name = "steerability" -version = "0.5.6" +version = "0.5.7" source = { editable = "." } dependencies = [ { name = "accelerate" },