WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-03 01:48:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

Finish impl

Browse Source
This commit is contained in:
Nuno Campos
2024-12-04 15:37:31 -08:00
parent 7d8205633d
commit 0461d45d76
4 changed files with 241 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/algo.py
+91 -16
View File
@@ -58,6 +58,7 @@ from langgraph.constants import (
)
from langgraph.errors import EmptyChannelError, InvalidUpdateError
from langgraph.managed.base import ManagedValueMapping
from langgraph.pregel.call import get_runnable_for_func
from langgraph.pregel.io import read_channel, read_channels
from langgraph.pregel.log import logger
from langgraph.pregel.manager import ChannelsManager
@@ -98,9 +99,15 @@ class PregelTaskWrites(NamedTuple):
class Call:
__slots__ = ("func", "input")
func: str | Callable
input: Any
def __init__(self, func: str | Callable, input: Any) -> None:
self.func = func
self.input = input
def should_interrupt(
checkpoint: Checkpoint,
@@ -184,7 +191,7 @@ def local_write(
"""Function injected under CONFIG_KEY_SEND in task config, to write to channels.
Validates writes and forwards them to `commit` function."""
for chan, value in writes:
if chan in (PUSH, TASKS):
if chan in (PUSH, TASKS) and value is not None:
if not isinstance(value, Send):
raise InvalidUpdateError(f"Expected Send, got {value}")
if value.node not in process_keys:
@@ -464,7 +471,87 @@ def prepare_single_task(
configurable = config.get(CONF, {})
parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
if task_path[0] == PUSH:
if task_path[0] == PUSH and isinstance(task_path[-1], Call):
# (PUSH, parent task path, idx of PUSH write, id of parent task, Call)
task_path_t = cast(tuple[str, tuple, int, str, Optional[Call]], task_path)
call = task_path_t[-1]
proc = get_runnable_for_func(call.func)
name = proc.name
if name is None:
raise ValueError("`call` functions must have a `__name__` attribute")
# create task id
triggers = [PUSH]
checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
name,
PUSH,
_tuple_str(task_path[1]),
str(task_path[2]),
)
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
metadata = {
"langgraph_step": step,
"langgraph_node": name,
"langgraph_triggers": triggers,
"langgraph_path": task_path[:3],
"langgraph_checkpoint_ns": task_checkpoint_ns,
}
if task_id_checksum is not None:
assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
if for_execution:
writes: deque[tuple[str, Any]] = deque()
return PregelExecutableTask(
name,
call.input,
proc,
writes,
patch_config(
merge_configs(config, {"metadata": metadata, "tags": proc.tags}),
run_name=name,
callbacks=(
manager.get_child(f"graph:step:{step}") if manager else None
),
configurable={
CONFIG_KEY_TASK_ID: task_id,
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write,
writes.extend,
processes.keys(),
),
CONFIG_KEY_READ: partial(
local_read,
step,
checkpoint,
channels,
managed,
PregelTaskWrites(task_path[:3], name, writes, triggers),
config,
),
CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)),
CONFIG_KEY_CHECKPOINTER: (
checkpointer or configurable.get(CONFIG_KEY_CHECKPOINTER)
),
CONFIG_KEY_CHECKPOINT_MAP: {
**configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
parent_ns: checkpoint["id"],
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
},
),
triggers,
None,
None,
task_id,
task_path[:3],
)
else:
return PregelTask(task_id, name, task_path[:3])
elif task_path[0] == PUSH:
if len(task_path) == 2: # TODO: remove branch in 1.0
# legacy SEND tasks, executed in superstep n+1
# (PUSH, idx of pending send)
@@ -498,13 +585,7 @@ def prepare_single_task(
elif len(task_path) >= 4:
# new PUSH tasks, executed in superstep n
# (PUSH, parent task path, idx of PUSH write, id of parent task)
task_path_t = cast(
Union[
tuple[str, tuple, int, str],
tuple[str, tuple, int, str, Optional[Call]],
],
task_path,
)
task_path_t = cast(tuple[str, tuple, int, str], task_path)
writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]]
if task_path_t[2] >= len(writes_for_path):
logger.warning(
@@ -513,12 +594,7 @@ def prepare_single_task(
return
packet = writes_for_path[task_path_t[2]][2]
if packet is None:
if len(task_path_t) == 5:
packet = task_path_t[4]
else:
# no packet to replay, this is a "call" task
return
# TODO handle Call packets
return
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending writes"
@@ -623,7 +699,6 @@ def prepare_single_task(
task_path[:3],
writers=proc.flat_writers,
)
else:
return PregelTask(task_id, packet.node, task_path[:3])
elif task_path[0] == PULL:
libs/langgraph/langgraph/pregel/call.py
+113
View File
@@ -0,0 +1,113 @@
import sys
import types
from langgraph.utils.runnable import RunnableCallable
"""
Utilities borrowed from cloudpickle.
https://github.com/cloudpipe/cloudpickle/blob/6220b0ce83ffee5e47e06770a1ee38ca9e47c850/cloudpickle/cloudpickle.py#L265
"""
def _getattribute(obj, name):
for subpath in name.split("."):
if subpath == "<locals>":
raise AttributeError(
"Can't get local attribute {!r} on {!r}".format(name, obj)
)
try:
parent = obj
obj = getattr(obj, subpath)
except AttributeError:
raise AttributeError(
"Can't get attribute {!r} on {!r}".format(name, obj)
) from None
return obj, parent
def _whichmodule(obj, name):
"""Find the module an object belongs to.
This function differs from ``pickle.whichmodule`` in two ways:
- it does not mangle the cases where obj's module is __main__ and obj was
not found in any module.
- Errors arising during module introspection are ignored, as those errors
are considered unwanted side effects.
"""
module_name = getattr(obj, "__module__", None)
if module_name is not None:
return module_name
# Protect the iteration by using a copy of sys.modules against dynamic
# modules that trigger imports of other modules upon calls to getattr or
# other threads importing at the same time.
for module_name, module in sys.modules.copy().items():
# Some modules such as coverage can inject non-module objects inside
# sys.modules
if (
module_name == "__main__"
or module_name == "__mp_main__"
or module is None
or not isinstance(module, types.ModuleType)
):
continue
try:
if _getattribute(module, name)[0] is obj:
return module_name
except Exception:
pass
return None
def _lookup_module_and_qualname(obj, name=None):
if name is None:
name = getattr(obj, "__qualname__", None)
if name is None: # pragma: no cover
# This used to be needed for Python 2.7 support but is probably not
# needed anymore. However we keep the __name__ introspection in case
# users of cloudpickle rely on this old behavior for unknown reasons.
name = getattr(obj, "__name__", None)
module_name = _whichmodule(obj, name)
if module_name is None:
# In this case, obj.__module__ is None AND obj was not found in any
# imported module. obj is thus treated as dynamic.
return None
if module_name == "__main__":
return None
# Note: if module_name is in sys.modules, the corresponding module is
# assumed importable at unpickling time. See #357
module = sys.modules.get(module_name, None)
if module is None:
# The main reason why obj's module would not be imported is that this
# module has been dynamically created, using for example
# types.ModuleType. The other possibility is that module was removed
# from sys.modules after obj was created/imported. But this case is not
# supported, as the standard pickle does not support it either.
return None
try:
obj2, parent = _getattribute(module, name)
except AttributeError:
# obj was not found inside the module it points to
return None
if obj2 is not obj:
return None
return module, name
def get_runnable_for_func(
func: types.FunctionType,
) -> RunnableCallable:
if func in CACHE:
return CACHE[func]
elif not _lookup_module_and_qualname(func):
return RunnableCallable(func)
else:
return CACHE.setdefault(func, RunnableCallable(func))
CACHE: dict[types.FunctionType, RunnableCallable] = {}
libs/langgraph/langgraph/pregel/retry.py
+4 -6
View File
@@ -37,9 +37,7 @@ def run_with_retry(
# clear any writes from previous attempts
task.writes.clear()
# run the task
task.proc.invoke(task.input, config)
# if successful, end
break
return task.proc.invoke(task.input, config)
except ParentCommand as exc:
ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
cmd = exc.args[0]
@@ -128,10 +126,10 @@ async def arun_with_retry(
if stream:
async for _ in task.proc.astream(task.input, config):
pass
# if successful, end
break
else:
await task.proc.ainvoke(task.input, config)
# if successful, end
break
return await task.proc.ainvoke(task.input, config)
except ParentCommand as exc:
ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
cmd = exc.args[0]
libs/langgraph/langgraph/pregel/runner.py
+33 -19
View File
@@ -66,20 +66,26 @@ class PregelRunner:
get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None,
) -> Iterator[None]:
def writer(
task: PregelExecutableTask, writes: Sequence[tuple[str, Any]]
) -> None:
task: PregelExecutableTask,
writes: Sequence[tuple[str, Any]],
*,
calls: Optional[Sequence[Call]] = None,
) -> Sequence[Optional[concurrent.futures.Future]]:
prev_length = len(task.writes)
# delegate to the underlying writer
task.config[CONF][CONFIG_KEY_SEND](writes)
# confirm no other concurrent writes were added
# TODO could use a lock here instead, if writes can come from many threads
assert len(task.writes) == prev_length + len(writes)
rtn: dict[int, Optional[concurrent.futures.Future]] = {}
for idx, w in enumerate(writes, start=prev_length):
# bail if not a PUSH write
if w[0] != PUSH:
continue
# schedule the next task, if the callback returns one
if next_task := self.schedule_task(task, idx):
if next_task := self.schedule_task(
task, idx, calls[idx - prev_length] if calls else None
):
# if the parent task was retried,
# the next task might already be running
if any(
@@ -87,18 +93,26 @@ class PregelRunner:
):
continue
# schedule the next task
futures[
self.submit(
run_with_retry,
next_task,
retry_policy,
configurable={
CONFIG_KEY_SEND: partial(writer, next_task),
# CONFIG_KEY_CALL: partial(call, next_task),
},
__reraise_on_exit__=reraise,
)
] = next_task
fut = self.submit(
run_with_retry,
next_task,
retry_policy,
configurable={
CONFIG_KEY_SEND: partial(writer, next_task),
CONFIG_KEY_CALL: partial(call, next_task),
},
__reraise_on_exit__=reraise,
)
futures[fut] = next_task
rtn[idx - prev_length] = fut
return [rtn.get(i) for i in range(len(writes))]
def call(
task, func: str | Callable[[Any], Union[Awaitable[Any], Any]], input: Any
) -> concurrent.futures.Future[Any]:
(fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)])
assert fut is not None, "writer did not return a future for call"
return fut
tasks = tuple(tasks)
futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {}
@@ -113,7 +127,7 @@ class PregelRunner:
retry_policy,
configurable={
CONFIG_KEY_SEND: partial(writer, t),
# CONFIG_KEY_CALL: partial(call, t),
CONFIG_KEY_CALL: partial(call, t),
},
)
self.commit(t, None)
@@ -138,7 +152,7 @@ class PregelRunner:
retry_policy,
configurable={
CONFIG_KEY_SEND: partial(writer, t),
# CONFIG_KEY_CALL: partial(call, t),
CONFIG_KEY_CALL: partial(call, t),
},
__reraise_on_exit__=reraise,
)
@@ -208,7 +222,7 @@ class PregelRunner:
if next_task := self.schedule_task(
task,
idx,
calls[idx] if calls is not None else None,
calls[idx - prev_length] if calls is not None else None,
):
# if the parent task was retried,
# the next task might already be running
@@ -231,7 +245,7 @@ class PregelRunner:
__reraise_on_exit__=reraise,
)
futures[cast(asyncio.Future, fut)] = next_task
rtn[idx] = fut
rtn[idx - prev_length] = fut
return [rtn.get(i) for i in range(len(writes))]
def call(