mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-08 10:47:52 +02:00
On stream interrupt cleanup any (potentially) pending tasks
This commit is contained in:
Nuno Campos
+267
-229
@@ -253,127 +253,144 @@ class Pregel(
|
||||
input_keys: Optional[Union[str, Sequence[str]]] = None,
|
||||
output_keys: Optional[Union[str, Sequence[str]]] = None,
|
||||
) -> Iterator[Union[dict[str, Any], Any]]:
|
||||
if config["recursion_limit"] < 1:
|
||||
raise ValueError("recursion_limit must be at least 1")
|
||||
# assign defaults
|
||||
if output_keys is None:
|
||||
output_keys = [chan for chan in self.channels if chan not in self.hidden]
|
||||
else:
|
||||
validate_keys(output_keys, self.channels)
|
||||
if input_keys is None:
|
||||
input_keys = self.input
|
||||
else:
|
||||
validate_keys(input_keys, self.channels)
|
||||
# copy nodes to ignore mutations during execution
|
||||
processes = {**self.nodes}
|
||||
# get checkpoint from saver, or create an empty one
|
||||
checkpoint = self.checkpointer.get(config) if self.checkpointer else None
|
||||
checkpoint = checkpoint or empty_checkpoint()
|
||||
# create channels from checkpoint
|
||||
with ChannelsManager(
|
||||
self.channels, checkpoint
|
||||
) as channels, get_executor_for_config(config) as executor:
|
||||
# map inputs to channel updates
|
||||
_apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
deque(w for c in input for w in map_input(input_keys, c)),
|
||||
config,
|
||||
0,
|
||||
)
|
||||
|
||||
read = partial(_read_channel, channels)
|
||||
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
# computation proceeds in steps, while there are channel updates
|
||||
# channel updates from step N are only visible in step N+1
|
||||
# channels are guaranteed to be immutable for the duration of the step,
|
||||
# with channel updates applied only at the transition between steps
|
||||
for step in range(config["recursion_limit"] + 1):
|
||||
next_tasks = _prepare_next_tasks(checkpoint, processes, channels)
|
||||
|
||||
# if no more tasks, we're done
|
||||
if not next_tasks:
|
||||
break
|
||||
elif step == config["recursion_limit"]:
|
||||
raise GraphRecursionError(
|
||||
f"Recursion limit of {config['recursion_limit']} reached"
|
||||
"without hitting a stop condition. You can increase the limit"
|
||||
"by setting the `recursion_limit` config key."
|
||||
)
|
||||
|
||||
if self.debug:
|
||||
print_step_start(step, next_tasks)
|
||||
|
||||
# collect all writes to channels, without applying them yet
|
||||
pending_writes = deque[tuple[str, Any]]()
|
||||
|
||||
# prepare tasks with config
|
||||
tasks_w_config = [
|
||||
(
|
||||
proc,
|
||||
input,
|
||||
patch_config(
|
||||
config,
|
||||
run_name=name,
|
||||
callbacks=run_manager.get_child(f"graph:step:{step}"),
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_SEND: pending_writes.extend,
|
||||
CONFIG_KEY_READ: read,
|
||||
},
|
||||
),
|
||||
)
|
||||
for proc, input, name in next_tasks
|
||||
try:
|
||||
if config["recursion_limit"] < 1:
|
||||
raise ValueError("recursion_limit must be at least 1")
|
||||
# assign defaults
|
||||
if output_keys is None:
|
||||
output_keys = [
|
||||
chan for chan in self.channels if chan not in self.hidden
|
||||
]
|
||||
|
||||
# execute tasks, and wait for one to fail or all to finish.
|
||||
# each task is independent from all other concurrent tasks
|
||||
done, inflight = concurrent.futures.wait(
|
||||
[
|
||||
executor.submit(proc.invoke, input, config)
|
||||
for proc, input, config in tasks_w_config
|
||||
],
|
||||
return_when=concurrent.futures.FIRST_EXCEPTION,
|
||||
timeout=self.step_timeout,
|
||||
else:
|
||||
validate_keys(output_keys, self.channels)
|
||||
if input_keys is None:
|
||||
input_keys = self.input
|
||||
else:
|
||||
validate_keys(input_keys, self.channels)
|
||||
# copy nodes to ignore mutations during execution
|
||||
processes = {**self.nodes}
|
||||
# get checkpoint from saver, or create an empty one
|
||||
checkpoint = self.checkpointer.get(config) if self.checkpointer else None
|
||||
checkpoint = checkpoint or empty_checkpoint()
|
||||
# create channels from checkpoint
|
||||
with ChannelsManager(
|
||||
self.channels, checkpoint
|
||||
) as channels, get_executor_for_config(config) as executor:
|
||||
# map inputs to channel updates
|
||||
_apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
deque(w for c in input for w in map_input(input_keys, c)),
|
||||
config,
|
||||
0,
|
||||
)
|
||||
|
||||
# interrupt on failure or timeout
|
||||
_interrupt_or_proceed(done, inflight, step)
|
||||
read = partial(_read_channel, channels)
|
||||
|
||||
# apply writes to channels
|
||||
_apply_writes(checkpoint, channels, pending_writes, config, step + 1)
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
# computation proceeds in steps, while there are channel updates
|
||||
# channel updates from step N are only visible in step N+1
|
||||
# channels are guaranteed to be immutable for the duration of the step,
|
||||
# with channel updates applied only at the transition between steps
|
||||
for step in range(config["recursion_limit"] + 1):
|
||||
next_tasks = _prepare_next_tasks(checkpoint, processes, channels)
|
||||
|
||||
if self.debug:
|
||||
print_checkpoint(step, channels)
|
||||
# if no more tasks, we're done
|
||||
if not next_tasks:
|
||||
break
|
||||
elif step == config["recursion_limit"]:
|
||||
raise GraphRecursionError(
|
||||
f"Recursion limit of {config['recursion_limit']} reached"
|
||||
"without hitting a stop condition. You can increase the limit"
|
||||
"by setting the `recursion_limit` config key."
|
||||
)
|
||||
|
||||
# yield current value and checkpoint view
|
||||
if step_output := map_output(output_keys, pending_writes, channels):
|
||||
yield step_output
|
||||
# we can detect updates when output is multiple channels (ie. dict)
|
||||
if not isinstance(output_keys, str):
|
||||
# if view was updated, apply writes to channels
|
||||
_apply_writes_from_view(checkpoint, channels, step_output)
|
||||
if self.debug:
|
||||
print_step_start(step, next_tasks)
|
||||
|
||||
# save end of step checkpoint
|
||||
# collect all writes to channels, without applying them yet
|
||||
pending_writes = deque[tuple[str, Any]]()
|
||||
|
||||
# prepare tasks with config
|
||||
tasks_w_config = [
|
||||
(
|
||||
proc,
|
||||
input,
|
||||
patch_config(
|
||||
config,
|
||||
run_name=name,
|
||||
callbacks=run_manager.get_child(f"graph:step:{step}"),
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_SEND: pending_writes.extend,
|
||||
CONFIG_KEY_READ: read,
|
||||
},
|
||||
),
|
||||
)
|
||||
for proc, input, name in next_tasks
|
||||
]
|
||||
|
||||
futures = [
|
||||
executor.submit(proc.invoke, input, config)
|
||||
for proc, input, config in tasks_w_config
|
||||
]
|
||||
|
||||
# execute tasks, and wait for one to fail or all to finish.
|
||||
# each task is independent from all other concurrent tasks
|
||||
done, inflight = concurrent.futures.wait(
|
||||
futures,
|
||||
return_when=concurrent.futures.FIRST_EXCEPTION,
|
||||
timeout=self.step_timeout,
|
||||
)
|
||||
|
||||
# interrupt on failure or timeout
|
||||
_interrupt_or_proceed(done, inflight, step)
|
||||
|
||||
# apply writes to channels
|
||||
_apply_writes(
|
||||
checkpoint, channels, pending_writes, config, step + 1
|
||||
)
|
||||
|
||||
if self.debug:
|
||||
print_checkpoint(step, channels)
|
||||
|
||||
# yield current value and checkpoint view
|
||||
if step_output := map_output(output_keys, pending_writes, channels):
|
||||
yield step_output
|
||||
# we can detect updates when output is multiple channels (ie. dict)
|
||||
if not isinstance(output_keys, str):
|
||||
# if view was updated, apply writes to channels
|
||||
_apply_writes_from_view(checkpoint, channels, step_output)
|
||||
|
||||
# save end of step checkpoint
|
||||
if (
|
||||
self.checkpointer is not None
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_STEP
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
self.checkpointer.put(config, checkpoint)
|
||||
|
||||
# interrupt if any channel written to is in interrupt list
|
||||
if any(
|
||||
chan for chan, _ in pending_writes if chan in self.interrupt
|
||||
):
|
||||
break
|
||||
|
||||
# save end of run checkpoint
|
||||
if (
|
||||
self.checkpointer is not None
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_STEP
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_RUN
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
self.checkpointer.put(config, checkpoint)
|
||||
|
||||
# interrupt if any channel written to is in interrupt list
|
||||
if any(chan for chan, _ in pending_writes if chan in self.interrupt):
|
||||
break
|
||||
|
||||
# save end of run checkpoint
|
||||
if (
|
||||
self.checkpointer is not None
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_RUN
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
self.checkpointer.put(config, checkpoint)
|
||||
finally:
|
||||
# cancel any pending tasks when generator is interrupted
|
||||
try:
|
||||
futures
|
||||
except NameError:
|
||||
return
|
||||
for task in futures:
|
||||
task.cancel()
|
||||
|
||||
async def _atransform(
|
||||
self,
|
||||
@@ -384,139 +401,160 @@ class Pregel(
|
||||
input_keys: Optional[Union[str, Sequence[str]]] = None,
|
||||
output_keys: Optional[Union[str, Sequence[str]]] = None,
|
||||
) -> AsyncIterator[Union[dict[str, Any], Any]]:
|
||||
if config["recursion_limit"] < 1:
|
||||
raise ValueError("recursion_limit must be at least 1")
|
||||
# if running from astream_log() run each proc with streaming
|
||||
do_stream = next(
|
||||
(
|
||||
h
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, LogStreamCallbackHandler)
|
||||
),
|
||||
None,
|
||||
)
|
||||
# assign defaults
|
||||
if output_keys is None:
|
||||
output_keys = [chan for chan in self.channels if chan not in self.hidden]
|
||||
else:
|
||||
validate_keys(output_keys, self.channels)
|
||||
if input_keys is None:
|
||||
input_keys = self.input
|
||||
else:
|
||||
validate_keys(input_keys, self.channels)
|
||||
# copy nodes to ignore mutations during execution
|
||||
processes = {**self.nodes}
|
||||
# get checkpoint from saver, or create an empty one
|
||||
checkpoint = await self.checkpointer.aget(config) if self.checkpointer else None
|
||||
checkpoint = checkpoint or empty_checkpoint()
|
||||
# create channels from checkpoint
|
||||
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
|
||||
# map inputs to channel updates
|
||||
_apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
deque([w async for c in input for w in map_input(input_keys, c)]),
|
||||
config,
|
||||
0,
|
||||
try:
|
||||
if config["recursion_limit"] < 1:
|
||||
raise ValueError("recursion_limit must be at least 1")
|
||||
# if running from astream_log() run each proc with streaming
|
||||
do_stream = next(
|
||||
(
|
||||
h
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, LogStreamCallbackHandler)
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
read = partial(_read_channel, channels)
|
||||
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
# computation proceeds in steps, while there are channel updates
|
||||
# channel updates from step N are only visible in step N+1,
|
||||
# channels are guaranteed to be immutable for the duration of the step,
|
||||
# channel updates being applied only at the transition between steps
|
||||
for step in range(config["recursion_limit"] + 1):
|
||||
next_tasks = _prepare_next_tasks(checkpoint, processes, channels)
|
||||
|
||||
# if no more tasks, we're done
|
||||
if not next_tasks:
|
||||
break
|
||||
elif step == config["recursion_limit"]:
|
||||
raise GraphRecursionError(
|
||||
f"Recursion limit of {config['recursion_limit']} reached"
|
||||
"without hitting a stop condition. You can increase the limit"
|
||||
"by setting the `recursion_limit` config key."
|
||||
)
|
||||
|
||||
if self.debug:
|
||||
print_step_start(step, next_tasks)
|
||||
|
||||
# collect all writes to channels, without applying them yet
|
||||
pending_writes = deque[tuple[str, Any]]()
|
||||
|
||||
# prepare tasks with config
|
||||
tasks_w_config = [
|
||||
(
|
||||
proc,
|
||||
input,
|
||||
patch_config(
|
||||
config,
|
||||
run_name=name,
|
||||
callbacks=run_manager.get_child(f"graph:step:{step}"),
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_SEND: pending_writes.extend,
|
||||
CONFIG_KEY_READ: read,
|
||||
},
|
||||
),
|
||||
)
|
||||
for proc, input, name in next_tasks
|
||||
# assign defaults
|
||||
if output_keys is None:
|
||||
output_keys = [
|
||||
chan for chan in self.channels if chan not in self.hidden
|
||||
]
|
||||
|
||||
# execute tasks, and wait for one to fail or all to finish.
|
||||
# each task is independent from all other concurrent tasks
|
||||
done, inflight = await asyncio.wait(
|
||||
[
|
||||
asyncio.create_task(_aconsume(proc.astream(input, config)))
|
||||
for proc, input, config in tasks_w_config
|
||||
]
|
||||
if do_stream
|
||||
else [
|
||||
asyncio.create_task(proc.ainvoke(input, config))
|
||||
for proc, input, config in tasks_w_config
|
||||
],
|
||||
return_when=asyncio.FIRST_EXCEPTION,
|
||||
timeout=self.step_timeout,
|
||||
else:
|
||||
validate_keys(output_keys, self.channels)
|
||||
if input_keys is None:
|
||||
input_keys = self.input
|
||||
else:
|
||||
validate_keys(input_keys, self.channels)
|
||||
# copy nodes to ignore mutations during execution
|
||||
processes = {**self.nodes}
|
||||
# get checkpoint from saver, or create an empty one
|
||||
checkpoint = (
|
||||
await self.checkpointer.aget(config) if self.checkpointer else None
|
||||
)
|
||||
checkpoint = checkpoint or empty_checkpoint()
|
||||
# create channels from checkpoint
|
||||
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
|
||||
# map inputs to channel updates
|
||||
_apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
deque([w async for c in input for w in map_input(input_keys, c)]),
|
||||
config,
|
||||
0,
|
||||
)
|
||||
|
||||
# interrupt on failure or timeout
|
||||
_interrupt_or_proceed(done, inflight, step)
|
||||
read = partial(_read_channel, channels)
|
||||
|
||||
# apply writes to channels
|
||||
_apply_writes(checkpoint, channels, pending_writes, config, step + 1)
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
# computation proceeds in steps, while there are channel updates
|
||||
# channel updates from step N are only visible in step N+1,
|
||||
# channels are guaranteed to be immutable for the duration of the step,
|
||||
# channel updates being applied only at the transition between steps
|
||||
for step in range(config["recursion_limit"] + 1):
|
||||
next_tasks = _prepare_next_tasks(checkpoint, processes, channels)
|
||||
|
||||
if self.debug:
|
||||
print_checkpoint(step, channels)
|
||||
# if no more tasks, we're done
|
||||
if not next_tasks:
|
||||
break
|
||||
elif step == config["recursion_limit"]:
|
||||
raise GraphRecursionError(
|
||||
f"Recursion limit of {config['recursion_limit']} reached"
|
||||
"without hitting a stop condition. You can increase the limit"
|
||||
"by setting the `recursion_limit` config key."
|
||||
)
|
||||
|
||||
# yield current value and checkpoint view
|
||||
if step_output := map_output(output_keys, pending_writes, channels):
|
||||
yield step_output
|
||||
# we can detect updates when output is multiple channels (ie. dict)
|
||||
if not isinstance(output_keys, str):
|
||||
# if view was updated, apply writes to channels
|
||||
_apply_writes_from_view(checkpoint, channels, step_output)
|
||||
if self.debug:
|
||||
print_step_start(step, next_tasks)
|
||||
|
||||
# save end of step checkpoint
|
||||
# collect all writes to channels, without applying them yet
|
||||
pending_writes = deque[tuple[str, Any]]()
|
||||
|
||||
# prepare tasks with config
|
||||
tasks_w_config = [
|
||||
(
|
||||
proc,
|
||||
input,
|
||||
patch_config(
|
||||
config,
|
||||
run_name=name,
|
||||
callbacks=run_manager.get_child(f"graph:step:{step}"),
|
||||
configurable={
|
||||
# deque.extend is thread-safe
|
||||
CONFIG_KEY_SEND: pending_writes.extend,
|
||||
CONFIG_KEY_READ: read,
|
||||
},
|
||||
),
|
||||
)
|
||||
for proc, input, name in next_tasks
|
||||
]
|
||||
|
||||
futures = (
|
||||
[
|
||||
asyncio.create_task(_aconsume(proc.astream(input, config)))
|
||||
for proc, input, config in tasks_w_config
|
||||
]
|
||||
if do_stream
|
||||
else [
|
||||
asyncio.create_task(proc.ainvoke(input, config))
|
||||
for proc, input, config in tasks_w_config
|
||||
]
|
||||
)
|
||||
|
||||
# execute tasks, and wait for one to fail or all to finish.
|
||||
# each task is independent from all other concurrent tasks
|
||||
done, inflight = await asyncio.wait(
|
||||
futures,
|
||||
return_when=asyncio.FIRST_EXCEPTION,
|
||||
timeout=self.step_timeout,
|
||||
)
|
||||
|
||||
# interrupt on failure or timeout
|
||||
_interrupt_or_proceed(done, inflight, step)
|
||||
|
||||
# apply writes to channels
|
||||
_apply_writes(
|
||||
checkpoint, channels, pending_writes, config, step + 1
|
||||
)
|
||||
|
||||
if self.debug:
|
||||
print_checkpoint(step, channels)
|
||||
|
||||
# yield current value and checkpoint view
|
||||
if step_output := map_output(output_keys, pending_writes, channels):
|
||||
yield step_output
|
||||
# we can detect updates when output is multiple channels (ie. dict)
|
||||
if not isinstance(output_keys, str):
|
||||
# if view was updated, apply writes to channels
|
||||
_apply_writes_from_view(checkpoint, channels, step_output)
|
||||
|
||||
# save end of step checkpoint
|
||||
if (
|
||||
self.checkpointer is not None
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_STEP
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
await self.checkpointer.aput(config, checkpoint)
|
||||
|
||||
# interrupt if any channel written to is in interrupt list
|
||||
if any(
|
||||
chan for chan, _ in pending_writes if chan in self.interrupt
|
||||
):
|
||||
break
|
||||
|
||||
# save end of run checkpoint
|
||||
if (
|
||||
self.checkpointer is not None
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_STEP
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_RUN
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
await self.checkpointer.aput(config, checkpoint)
|
||||
|
||||
# interrupt if any channel written to is in interrupt list
|
||||
if any(chan for chan, _ in pending_writes if chan in self.interrupt):
|
||||
break
|
||||
|
||||
# save end of run checkpoint
|
||||
if (
|
||||
self.checkpointer is not None
|
||||
and self.checkpointer.at == CheckpointAt.END_OF_RUN
|
||||
):
|
||||
checkpoint = create_checkpoint(checkpoint, channels)
|
||||
await self.checkpointer.aput(config, checkpoint)
|
||||
finally:
|
||||
# cancel any pending tasks when generator is interrupted
|
||||
try:
|
||||
futures
|
||||
except NameError:
|
||||
return
|
||||
for task in futures:
|
||||
task.cancel()
|
||||
|
||||
def invoke(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user