mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-21 23:22:27 +02:00
40ab009c62
## Summary
Adds cooperative drain support for Pregel runs so a graph can be asked
to stop at the next superstep boundary, persist its checkpoint, and
surface a resumable terminal exception.
- New `RunControl` (in `langgraph.runtime`) — a thread-safe handle whose
`request_drain(reason="shutdown")` sets a single flag.
- New `GraphDrained(GraphBubbleUp)` exception (in `langgraph.errors`)
raised when a run exits early due to drain. Carries the `reason` string.
- New `control: RunControl | None` kwarg on `invoke` / `ainvoke` /
`stream` / `astream` / `stream_v2` / `astream_v2`. Wired through to
`Runtime.control`, so nodes can read `runtime.control.drain_requested` /
`drain_reason` and even call `request_drain()` from inside a node.
- Stream transformers learn `"drained"` as a terminal `SubgraphStatus`.
The intended use is hooking SIGTERM (or any external supervisor signal)
to `control.request_drain("sigterm")` so an in-flight graph run can stop
cleanly and be resumed later from the saved checkpoint.
## Semantics: cooperative, between-superstep
`request_drain()` flips a flag. The Pregel loop checks it at the top of
each `tick()`, **after** the previous superstep's writes have been
applied and checkpointed. It never preempts work that is already
running.
| Scenario | Behavior |
|---|---|
| Node mid-execution (blocking I/O, sleeps, etc.) | Runs to completion.
Drain takes effect on the next superstep. |
| Node with a retry policy currently retrying | Retry loop runs to
exhaustion or success (drain is not checked between retries). Drain
takes effect on the next superstep. |
| Functional API: `@entrypoint` with pending `@task` futures |
Entrypoint and all dispatched tasks complete; drain takes effect after
the entrypoint returns. |
| Graph naturally finishes on the same tick where drain was requested
(no more tasks) | Treated as `done`; returns normally. **No
`GraphDrained` is raised.** The caller can inspect
`control.drain_requested` afterwards to distinguish a
drained-but-completed run from a normal one. |
| More tasks remain | Raises `GraphDrained(reason)`. The checkpoint of
the last completed superstep is saved (also under `durability="exit"`).
Resume with `invoke(None, config)` / `ainvoke(None, config)`. |
| Subgraph requests drain | `GraphDrained` bubbles up through the parent
loop and stops it at its own next superstep boundary; the parent's
checkpoint is saved and resumable. |
Drain does **not** cancel asyncio tasks or kill threads. Pair it with a
graceful timeout + `task.cancel()` (or process exit) if you need a hard
upper bound — see `test_drain_then_cancel_after_graceful_timeout` for
the recommended pattern.
## Usage
```python
from langgraph.runtime import RunControl
from langgraph.errors import GraphDrained
control = RunControl()
# In a signal handler, supervisor, etc.:
# control.request_drain("sigterm")
try:
result = graph.invoke(input, config, control=control)
if control.drain_requested:
# finished naturally on the same tick where drain was requested
...
except GraphDrained as e:
# checkpoint saved; resume later with the same config
log.info("graph drained: %s", e.reason)
```
## Test plan
- [x] Sync + async drain stops the next superstep
(`test_run_control_request_drain_stops_future_steps[_async]`)
- [x] Drain on the terminal step finishes normally
(`test_drain_requested_in_terminal_step_finishes_normally[_async]`)
- [x] `durability=\"exit\"` persists a resumable checkpoint on drain
(`test_drain_with_exit_durability_persists_resume_checkpoint`)
- [x] Subgraph drain bubbles up and parent resumes correctly
(`test_drain_from_subgraph_can_resume_parent`)
- [x] External thread / task triggering drain mid-run
(`test_external_drain_concurrent_sync` / `_async`)
- [x] Drain + hard cancel after graceful timeout
(`test_drain_then_cancel_after_graceful_timeout`)
- [x] Functional API: in-flight `@task` futures still resolve after
`request_drain()`
(`test_request_drain_allows_inflight_[a]call_scheduling`)
- [x] `control` kwarg wired through `stream_v2`
(`test_stream_v2_accepts_control_for_drain`)
- [x] `Runtime.merge` preserves `control`
(`test_merge_runtime_preserves_run_control`)
---------
Co-authored-by: Quanzheng Long <long@langchain.dev>
Co-authored-by: Will Fu-Hinthorn <will@langchain.dev>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>