Add concurrent generation and cancellation on the vLLM backends, and a per-request provider mode

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
 

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

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

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" }

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",

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.