mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 05:07:51 +02:00
lib: Separate out executor logic from stream methods (#1048)
* lib: Separate out executor logic from stream methods - first step towards splitting out tick and stream * Lint * Lint * Lint * Lint
This commit is contained in:
Nuno Campos
@@ -41,7 +41,6 @@ from langchain_core.runnables.config import (
|
||||
ensure_config,
|
||||
get_async_callback_manager_for_config,
|
||||
get_callback_manager_for_config,
|
||||
get_executor_for_config,
|
||||
merge_configs,
|
||||
patch_config,
|
||||
)
|
||||
@@ -94,6 +93,7 @@ from langgraph.pregel.debug import (
|
||||
print_step_tasks,
|
||||
print_step_writes,
|
||||
)
|
||||
from langgraph.pregel.executor import AsyncBackgroundExecutor, BackgroundExecutor
|
||||
from langgraph.pregel.io import (
|
||||
map_input,
|
||||
map_output_updates,
|
||||
@@ -836,7 +836,6 @@ class Pregel(
|
||||
run_id=config.get("run_id"),
|
||||
)
|
||||
try:
|
||||
bg: list[concurrent.futures.Future] = []
|
||||
if config["recursion_limit"] < 1:
|
||||
raise ValueError("recursion_limit must be at least 1")
|
||||
if self.checkpointer and not config.get("configurable"):
|
||||
@@ -880,29 +879,25 @@ class Pregel(
|
||||
|
||||
start = saved.metadata.get("step", -2) + 1 if saved else -1
|
||||
# create channels from checkpoint
|
||||
with ChannelsManager(
|
||||
with BackgroundExecutor(config) as submit, ChannelsManager(
|
||||
self.channels, checkpoint, config
|
||||
) as channels, get_executor_for_config(
|
||||
config
|
||||
) as executor, ManagedValuesManager(
|
||||
) as channels, ManagedValuesManager(
|
||||
self.managed_values_dict, config, self
|
||||
) as managed:
|
||||
|
||||
def put_writes(task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
|
||||
if self.checkpointer is not None:
|
||||
bg.append(
|
||||
executor.submit(
|
||||
self.checkpointer.put_writes,
|
||||
{
|
||||
**checkpoint_config,
|
||||
"configurable": {
|
||||
**checkpoint_config["configurable"],
|
||||
"thread_ts": checkpoint["id"],
|
||||
},
|
||||
submit(
|
||||
self.checkpointer.put_writes,
|
||||
{
|
||||
**checkpoint_config,
|
||||
"configurable": {
|
||||
**checkpoint_config["configurable"],
|
||||
"thread_ts": checkpoint["id"],
|
||||
},
|
||||
writes,
|
||||
task_id,
|
||||
)
|
||||
},
|
||||
writes,
|
||||
task_id,
|
||||
)
|
||||
|
||||
def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]:
|
||||
@@ -920,13 +915,11 @@ class Pregel(
|
||||
checkpoint, channels, metadata["step"]
|
||||
)
|
||||
# save it, without blocking
|
||||
bg.append(
|
||||
executor.submit(
|
||||
self.checkpointer.put,
|
||||
checkpoint_config,
|
||||
copy_checkpoint(checkpoint),
|
||||
metadata,
|
||||
)
|
||||
submit(
|
||||
self.checkpointer.put,
|
||||
checkpoint_config,
|
||||
copy_checkpoint(checkpoint),
|
||||
metadata,
|
||||
)
|
||||
# update checkpoint config
|
||||
checkpoint_config = {
|
||||
@@ -1059,7 +1052,7 @@ class Pregel(
|
||||
# each task is independent from all other concurrent tasks
|
||||
# yield updates/debug output as each task finishes
|
||||
futures = {
|
||||
executor.submit(run_with_retry, task, self.retry_policy): task
|
||||
submit(run_with_retry, task, self.retry_policy): task
|
||||
for task in next_tasks
|
||||
if not task.writes
|
||||
}
|
||||
@@ -1080,6 +1073,8 @@ class Pregel(
|
||||
else None
|
||||
),
|
||||
)
|
||||
if not done:
|
||||
break # timed out
|
||||
for fut in done:
|
||||
task = futures.pop(fut)
|
||||
if fut.exception() is not None:
|
||||
@@ -1182,19 +1177,6 @@ class Pregel(
|
||||
except BaseException as e:
|
||||
run_manager.on_chain_error(e)
|
||||
raise
|
||||
finally:
|
||||
# cancel any pending tasks when generator is interrupted
|
||||
try:
|
||||
for task in futures:
|
||||
task.cancel()
|
||||
except NameError:
|
||||
pass
|
||||
# wait for all background tasks to finish
|
||||
done, _ = concurrent.futures.wait(
|
||||
bg, return_when=concurrent.futures.ALL_COMPLETED
|
||||
)
|
||||
for task in done:
|
||||
task.result()
|
||||
|
||||
async def astream(
|
||||
self,
|
||||
@@ -1294,7 +1276,6 @@ class Pregel(
|
||||
)
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
bg: list[asyncio.Task] = []
|
||||
if config["recursion_limit"] < 1:
|
||||
raise ValueError("recursion_limit must be at least 1")
|
||||
if self.checkpointer and not config.get("configurable"):
|
||||
@@ -1342,7 +1323,7 @@ class Pregel(
|
||||
|
||||
start = saved.metadata.get("step", -2) + 1 if saved else -1
|
||||
# create channels from checkpoint
|
||||
async with AsyncChannelsManager(
|
||||
async with AsyncBackgroundExecutor() as submit, AsyncChannelsManager(
|
||||
self.channels, checkpoint, config
|
||||
) as channels, AsyncManagedValuesManager(
|
||||
self.managed_values_dict, config, self
|
||||
@@ -1350,20 +1331,17 @@ class Pregel(
|
||||
|
||||
def put_writes(task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
|
||||
if self.checkpointer is not None:
|
||||
bg.append(
|
||||
asyncio.create_task(
|
||||
self.checkpointer.aput_writes(
|
||||
{
|
||||
**checkpoint_config,
|
||||
"configurable": {
|
||||
**checkpoint_config["configurable"],
|
||||
"thread_ts": checkpoint["id"],
|
||||
},
|
||||
},
|
||||
writes,
|
||||
task_id,
|
||||
)
|
||||
)
|
||||
submit(
|
||||
self.checkpointer.aput_writes,
|
||||
{
|
||||
**checkpoint_config,
|
||||
"configurable": {
|
||||
**checkpoint_config["configurable"],
|
||||
"thread_ts": checkpoint["id"],
|
||||
},
|
||||
},
|
||||
writes,
|
||||
task_id,
|
||||
)
|
||||
|
||||
def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]:
|
||||
@@ -1381,13 +1359,13 @@ class Pregel(
|
||||
checkpoint, channels, metadata["step"]
|
||||
)
|
||||
# save it, without blocking
|
||||
bg.append(
|
||||
asyncio.create_task(
|
||||
self.checkpointer.aput(
|
||||
checkpoint_config, copy_checkpoint(checkpoint), metadata
|
||||
)
|
||||
)
|
||||
submit(
|
||||
self.checkpointer.aput,
|
||||
checkpoint_config,
|
||||
copy_checkpoint(checkpoint),
|
||||
metadata,
|
||||
)
|
||||
|
||||
# update checkpoint config
|
||||
checkpoint_config = {
|
||||
**checkpoint_config,
|
||||
@@ -1517,8 +1495,13 @@ class Pregel(
|
||||
# each task is independent from all other concurrent tasks
|
||||
# yield updates/debug output as each task finishes
|
||||
futures = {
|
||||
asyncio.create_task(
|
||||
arun_with_retry(task, self.retry_policy, do_stream)
|
||||
submit(
|
||||
arun_with_retry,
|
||||
task,
|
||||
self.retry_policy,
|
||||
do_stream,
|
||||
__name__=task.name,
|
||||
__cancel_on_exit__=True,
|
||||
): task
|
||||
for task in next_tasks
|
||||
if not task.writes
|
||||
@@ -1536,6 +1519,8 @@ class Pregel(
|
||||
max(0, end_time - loop.time()) if end_time else None
|
||||
),
|
||||
)
|
||||
if not done:
|
||||
break # timed out
|
||||
for fut in done:
|
||||
task = futures.pop(fut)
|
||||
if fut.exception() is not None:
|
||||
@@ -1569,7 +1554,7 @@ class Pregel(
|
||||
del fut, task
|
||||
|
||||
# panic on failure or timeout
|
||||
_panic_or_proceed(done, inflight, step)
|
||||
_panic_or_proceed(done, inflight, step, asyncio.TimeoutError)
|
||||
# don't keep futures around in memory longer than needed
|
||||
del done, inflight, futures
|
||||
|
||||
@@ -1640,22 +1625,8 @@ class Pregel(
|
||||
# set final channel values as run output
|
||||
await run_manager.on_chain_end(read_channels(channels, output_keys))
|
||||
except BaseException as e:
|
||||
await run_manager.on_chain_error(e)
|
||||
await asyncio.shield(run_manager.on_chain_error(e))
|
||||
raise
|
||||
finally:
|
||||
# cancel any pending tasks when generator is interrupted
|
||||
try:
|
||||
for task in futures:
|
||||
task.cancel()
|
||||
bg.append(task)
|
||||
except NameError:
|
||||
pass
|
||||
# wait for all background tasks to finish
|
||||
fut = asyncio.gather(*bg)
|
||||
# mark the exception as retrieved
|
||||
fut.add_done_callback(_mark_cancelled_as_seen)
|
||||
# wait for the future to finish, shielded from cancellation
|
||||
await asyncio.shield(fut)
|
||||
|
||||
def invoke(
|
||||
self,
|
||||
@@ -1773,6 +1744,7 @@ def _panic_or_proceed(
|
||||
done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
|
||||
inflight: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
|
||||
step: int,
|
||||
timeout_exc_cls: Type[Exception] = TimeoutError,
|
||||
) -> None:
|
||||
while done:
|
||||
# if any task failed
|
||||
@@ -1789,7 +1761,7 @@ def _panic_or_proceed(
|
||||
# cancel all pending tasks
|
||||
inflight.pop().cancel()
|
||||
# raise timeout error
|
||||
raise TimeoutError(f"Timed out at step {step}")
|
||||
raise timeout_exc_cls(f"Timed out at step {step}")
|
||||
|
||||
|
||||
def _should_interrupt(
|
||||
@@ -2167,10 +2139,3 @@ def _with_mode(mode: StreamMode, on: bool, iter: Iterator[Any]) -> Iterator[Any]
|
||||
yield (mode, chunk)
|
||||
else:
|
||||
yield from iter
|
||||
|
||||
|
||||
def _mark_cancelled_as_seen(fut: concurrent.futures.Future) -> None:
|
||||
try:
|
||||
fut.exception()
|
||||
except (asyncio.CancelledError, concurrent.futures.CancelledError):
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
from contextvars import copy_context
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
AsyncContextManager,
|
||||
Callable,
|
||||
Iterator,
|
||||
Optional,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
)
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langchain_core.runnables.config import get_executor_for_config
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
P = ParamSpec("P")
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class Submit(Protocol[P, T]):
|
||||
def __call__(
|
||||
self,
|
||||
fn: Callable[P, T],
|
||||
*args: P.args,
|
||||
__name__: Optional[str] = None,
|
||||
__cancel_on_exit__: bool = False,
|
||||
**kwargs: P.kwargs,
|
||||
) -> concurrent.futures.Future[T]:
|
||||
...
|
||||
|
||||
|
||||
@contextmanager
|
||||
def BackgroundExecutor(config: RunnableConfig) -> Iterator[Submit]:
|
||||
tasks: dict[concurrent.futures.Future, bool] = {}
|
||||
with get_executor_for_config(config) as executor:
|
||||
|
||||
def done(task: concurrent.futures.Future) -> None:
|
||||
try:
|
||||
task.result()
|
||||
except BaseException:
|
||||
pass
|
||||
else:
|
||||
tasks.pop(task)
|
||||
|
||||
def submit(
|
||||
fn: Callable[P, T],
|
||||
*args: P.args,
|
||||
__name__: Optional[str] = None, # currently not used in sync version
|
||||
__cancel_on_exit__: bool = False,
|
||||
**kwargs: P.kwargs,
|
||||
) -> concurrent.futures.Future:
|
||||
task = executor.submit(fn, *args, **kwargs)
|
||||
tasks[task] = __cancel_on_exit__
|
||||
task.add_done_callback(done)
|
||||
return task
|
||||
|
||||
try:
|
||||
yield submit
|
||||
finally:
|
||||
for task, cancel in tasks.items():
|
||||
if cancel:
|
||||
task.cancel()
|
||||
# executor waits for all tasks to finish on exit
|
||||
for task in tasks:
|
||||
# the first task to have raised an exception will be re-raised here
|
||||
task.result()
|
||||
|
||||
|
||||
class AsyncBackgroundExecutor(AsyncContextManager):
|
||||
def __init__(self) -> None:
|
||||
self.context_not_supported = sys.version_info < (3, 11)
|
||||
self.tasks: dict[asyncio.Task, bool] = {}
|
||||
self.sentinel = object()
|
||||
|
||||
def submit(
|
||||
self,
|
||||
fn: Callable[P, T],
|
||||
*args: P.args,
|
||||
__name__: Optional[str] = None,
|
||||
__cancel_on_exit__: bool = False,
|
||||
**kwargs: P.kwargs,
|
||||
) -> asyncio.Task[T]:
|
||||
coro = fn(*args, **kwargs)
|
||||
if self.context_not_supported:
|
||||
task = asyncio.create_task(coro, name=__name__)
|
||||
else:
|
||||
task = asyncio.create_task(coro, name=__name__, context=copy_context())
|
||||
self.tasks[task] = __cancel_on_exit__
|
||||
task.add_done_callback(self.done)
|
||||
return task
|
||||
|
||||
def done(self, task: asyncio.Task) -> None:
|
||||
try:
|
||||
task.result()
|
||||
except BaseException:
|
||||
pass
|
||||
else:
|
||||
self.tasks.pop(task)
|
||||
|
||||
async def __aenter__(self) -> Submit:
|
||||
return self.submit
|
||||
|
||||
async def exit(self) -> None:
|
||||
fut = asyncio.gather(*self.tasks, return_exceptions=True)
|
||||
try:
|
||||
rtns = await asyncio.shield(fut)
|
||||
finally:
|
||||
del self.tasks
|
||||
for rtn in rtns:
|
||||
# if this is ever changed to BaseException, need to ignore CancelledError
|
||||
if isinstance(rtn, Exception):
|
||||
raise rtn
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: Optional[type[BaseException]],
|
||||
exc_value: Optional[BaseException],
|
||||
traceback: Optional[TracebackType],
|
||||
) -> Optional[bool]:
|
||||
for task, cancel in self.tasks.items():
|
||||
if cancel:
|
||||
task.cancel(self.sentinel)
|
||||
# wait for all background tasks to finish, shielded from cancellation
|
||||
await asyncio.shield(self.exit())
|
||||
@@ -41,6 +41,8 @@ def default_retry_on(exc: Exception) -> bool:
|
||||
|
||||
|
||||
class RetryPolicy(NamedTuple):
|
||||
"""Configuration for retrying nodes."""
|
||||
|
||||
initial_interval: float = 0.5
|
||||
"""Amount of time that must elapse before the first retry occurs. In seconds."""
|
||||
backoff_factor: float = 2.0
|
||||
|
||||
@@ -220,7 +220,7 @@ async def test_step_timeout_on_stream_hang() -> None:
|
||||
graph = builder.compile()
|
||||
graph.step_timeout = 1
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
async for chunk in graph.astream(1, stream_mode="updates"):
|
||||
assert chunk == {"alittlewhile": {"alittlewhile": "1"}}
|
||||
await asyncio.sleep(0.6)
|
||||
@@ -247,7 +247,7 @@ async def test_cancel_graph_astream(
|
||||
try:
|
||||
|
||||
class State(TypedDict):
|
||||
value: int
|
||||
value: Annotated[int, operator.add]
|
||||
|
||||
class AwhileMaker:
|
||||
def __init__(self) -> None:
|
||||
@@ -270,20 +270,31 @@ async def test_cancel_graph_astream(
|
||||
return {"value": 2}
|
||||
|
||||
awhile = AwhileMaker()
|
||||
aparallelwhile = AwhileMaker()
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("awhile", awhile)
|
||||
builder.add_node("aparallelwhile", aparallelwhile)
|
||||
builder.add_node(alittlewhile)
|
||||
builder.add_edge(START, "alittlewhile")
|
||||
builder.add_edge(START, "aparallelwhile")
|
||||
builder.add_edge("alittlewhile", "awhile")
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
|
||||
# test interrupting astream
|
||||
got_event = False
|
||||
thread1: RunnableConfig = {"configurable": {"thread_id": 1}}
|
||||
async with aclosing(graph.astream({"value": 1}, thread1)) as stream:
|
||||
async for chunk in stream:
|
||||
assert chunk == {"alittlewhile": {"value": 2}}
|
||||
got_event = True
|
||||
break
|
||||
|
||||
assert got_event
|
||||
|
||||
# node aparallelwhile should start, but be cancelled
|
||||
assert aparallelwhile.started is True
|
||||
assert aparallelwhile.cancelled is True
|
||||
|
||||
# node "awhile" should never start
|
||||
assert awhile.started is False
|
||||
|
||||
@@ -292,7 +303,10 @@ async def test_cancel_graph_astream(
|
||||
state = await graph.aget_state(thread1)
|
||||
assert state is not None
|
||||
assert state.values == {"value": 1}
|
||||
assert state.next == ("alittlewhile",)
|
||||
assert state.next == (
|
||||
"aparallelwhile",
|
||||
"alittlewhile",
|
||||
)
|
||||
assert state.metadata == {"source": "loop", "step": 0, "writes": None}
|
||||
finally:
|
||||
if getattr(checkpointer, "__aexit__", None):
|
||||
@@ -366,9 +380,10 @@ async def test_cancel_graph_astream_events_v2(
|
||||
# did break
|
||||
assert got_event
|
||||
|
||||
# node "awhile" starts but is cancelled
|
||||
assert awhile.started is True
|
||||
assert awhile.cancelled is True
|
||||
# node "awhile" maybe starts (impl detail of astream_events)
|
||||
# if it does start, it must be cancelled
|
||||
if awhile.started:
|
||||
assert awhile.cancelled is True
|
||||
|
||||
# node "anotherwhile" should never start
|
||||
assert anotherwhile.started is False
|
||||
|
||||
Reference in New Issue
Block a user