-
feat: allow graph to graceful shutdown/drain by request (#7274)
发布于
2026-04-29 22:23:31 +00:00 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(inlanggraph.runtime) — a thread-safe handle whose
request_drain(reason="shutdown")sets a single flag. - New
GraphDrained(GraphBubbleUp)exception (inlanggraph.errors)
raised when a run exits early due to drain. Carries thereasonstring. - New
control: RunControl | Nonekwarg oninvoke/ainvoke/
stream/astream/stream_v2/astream_v2. Wired through to
Runtime.control, so nodes can readruntime.control.drain_requested/
drain_reasonand even callrequest_drain()from inside a node. - Stream transformers learn
"drained"as a terminalSubgraphStatus.
The intended use is hooking SIGTERM (or any external supervisor signal)
tocontrol.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
eachtick(), 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: @entrypointwith pending@taskfuturesEntrypoint 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. **NoGraphDrainedis raised.** The caller can inspectcontrol.drain_requestedafterwards to distinguish adrained-but-completed run from a normal one. More tasks remain Raises GraphDrained(reason). The checkpoint ofthe last completed superstep is saved (also under durability="exit").Resume with invoke(None, config)/ainvoke(None, config).Subgraph requests drain GraphDrainedbubbles up through the parentloop 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 — seetest_drain_then_cancel_after_graceful_timeoutfor
the recommended pattern.Usage
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
- Sync + async drain stops the next superstep
(test_run_control_request_drain_stops_future_steps[_async]) - Drain on the terminal step finishes normally
(test_drain_requested_in_terminal_step_finishes_normally[_async]) durability=\"exit\"persists a resumable checkpoint on drain
(test_drain_with_exit_durability_persists_resume_checkpoint)- Subgraph drain bubbles up and parent resumes correctly
(test_drain_from_subgraph_can_resume_parent) - External thread / task triggering drain mid-run
(test_external_drain_concurrent_sync/_async) - Drain + hard cancel after graceful timeout
(test_drain_then_cancel_after_graceful_timeout) - Functional API: in-flight
@taskfutures still resolve after
request_drain()
(test_request_drain_allows_inflight_[a]call_scheduling) controlkwarg wired throughstream_v2
(test_stream_v2_accepts_control_for_drain)Runtime.mergepreservescontrol
(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下载附件
- New