From 0461d45d76d4af9f1bd214a8b77792d88e59fa2b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 8 Nov 2024 16:33:07 -0800 Subject: [PATCH] Finish impl --- libs/langgraph/langgraph/pregel/algo.py | 107 +++++++++++++++++--- libs/langgraph/langgraph/pregel/call.py | 113 ++++++++++++++++++++++ libs/langgraph/langgraph/pregel/retry.py | 10 +- libs/langgraph/langgraph/pregel/runner.py | 52 ++++++---- 4 files changed, 241 insertions(+), 41 deletions(-) create mode 100644 libs/langgraph/langgraph/pregel/call.py diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index a8d39674c..c95c699c8 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -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: diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py new file mode 100644 index 000000000..8932697b1 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/call.py @@ -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 == "": + 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] = {} diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 093859579..29faaab21 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -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] diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 067394bdb..6155ce74e 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -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(