发布

  • feat: allow graph to graceful shutdown/drain by request (#7274)

    frostbyte_neo 发布于 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 (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

    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 @task futures still resolve after
      request_drain()
      (test_request_drain_allows_inflight_[a]call_scheduling)
    • control kwarg wired through stream_v2
      (test_stream_v2_accepts_control_for_drain)
    • 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

    下载附件