From df8becd5cfc74bb09ab2d7cbdf17c5559114b8d8 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Fri, 14 Nov 2025 16:13:03 -0800 Subject: [PATCH] refactor: separate prepare_push_* functions (#6450) Extract two common cases from the big switch statement of `prepare_single_task` since it's a tad more composable. All this does is shift/extract code to separate functions --- libs/langgraph/langgraph/pregel/_algo.py | 602 +++++++++++++---------- 1 file changed, 352 insertions(+), 250 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index d1d7f32da..9f265f80f 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -258,9 +258,11 @@ def apply_writes( next_version = None else: next_version = get_next_version( - max(checkpoint["channel_versions"].values()) - if checkpoint["channel_versions"] - else None, + ( + max(checkpoint["channel_versions"].values()) + if checkpoint["channel_versions"] + else None + ), None, ) @@ -491,6 +493,11 @@ def prepare_next_tasks( PUSH_TRIGGER = (PUSH,) +class _TaskIDFn(Protocol): + def __call__(self, namespace: bytes, *parts: str | bytes) -> str: + pass + + def prepare_single_task( task_path: tuple[Any, ...], task_id_checksum: str | None, @@ -520,250 +527,50 @@ def prepare_single_task( task_id_func = _xxhash_str if checkpoint["v"] > 1 else _uuid5_str 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, Call], task_path) - call = task_path_t[-1] - proc_ = get_runnable_for_task(call.func) - name = proc_.name - if name is None: - raise ValueError("`call` functions must have a `__name__` attribute") - # create task id - triggers: Sequence[str] = PUSH_TRIGGER - checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name - task_id = task_id_func( - checkpoint_id_bytes, - checkpoint_ns, - str(step), - name, - PUSH, - task_path_str(task_path[1]), - str(task_path[2]), + return prepare_push_task_functional( + cast(tuple[str, tuple, int, str, Call], task_path), + task_id_checksum, + checkpoint=checkpoint, + checkpoint_id_bytes=checkpoint_id_bytes, + pending_writes=pending_writes, + channels=channels, + managed=managed, + config=config, + step=step, + stop=stop, + for_execution=for_execution, + store=store, + checkpointer=checkpointer, + manager=manager, + cache_policy=cache_policy, + retry_policy=retry_policy, + parent_ns=parent_ns, + task_id_func=task_id_func, + ) + + elif task_path[0] == PUSH: + return prepare_push_task_send( + cast(tuple[str, tuple], task_path), + task_id_checksum, + checkpoint=checkpoint, + checkpoint_id_bytes=checkpoint_id_bytes, + pending_writes=pending_writes, + channels=channels, + managed=managed, + config=config, + step=step, + processes=processes, + stop=stop, + for_execution=for_execution, + store=store, + checkpointer=checkpointer, + manager=manager, + cache_policy=cache_policy, + retry_policy=retry_policy, + parent_ns=parent_ns, + task_id_func=task_id_func, ) - task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" - # we append True to the task path to indicate that a call is being - # made, so we should not return interrupts from this task (responsibility lies with the parent) - task_path = (*task_path[:3], True) - metadata = { - "langgraph_step": step, - "langgraph_node": name, - "langgraph_triggers": triggers, - "langgraph_path": task_path, - "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() - cache_policy = call.cache_policy or cache_policy - if cache_policy: - args_key = cache_policy.key_func(*call.input[0], **call.input[1]) - cache_key: CacheKey | None = CacheKey( - ( - CACHE_NS_WRITES, - (identifier(call.func) or "__dynamic__"), - ), - xxh3_128_hexdigest( - args_key.encode() if isinstance(args_key, str) else args_key, - ), - cache_policy.ttl, - ) - else: - cache_key = None - scratchpad = _scratchpad( - config[CONF].get(CONFIG_KEY_SCRATCHPAD), - pending_writes, - task_id, - xxh3_128_hexdigest(task_checkpoint_ns.encode()), - config[CONF].get(CONFIG_KEY_RESUME_MAP), - step, - stop, - ) - runtime = cast( - Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME) - ) - runtime = runtime.override(store=store) - return PregelExecutableTask( - name, - call.input, - proc_, - writes, - patch_config( - merge_configs(config, {"metadata": metadata}), - run_name=name, - callbacks=call.callbacks - or (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: writes.extend, - CONFIG_KEY_READ: partial( - local_read, - scratchpad, - channels, - managed, - PregelTaskWrites(task_path, name, writes, triggers), - ), - 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, - CONFIG_KEY_SCRATCHPAD: scratchpad, - CONFIG_KEY_RUNTIME: runtime, - }, - ), - triggers, - call.retry_policy or retry_policy, - cache_key, - task_id, - task_path, - ) - else: - return PregelTask(task_id, name, task_path) - elif task_path[0] == PUSH: - if len(task_path) == 2: - # SEND tasks, executed in superstep n+1 - # (PUSH, idx of pending send) - idx = cast(int, task_path[1]) - if not channels[TASKS].is_available(): - return - sends: Sequence[Send] = channels[TASKS].get() - if idx < 0 or idx >= len(sends): - return - packet = sends[idx] - if not isinstance(packet, Send): - logger.warning( - f"Ignoring invalid packet type {type(packet)} in pending sends" - ) - return - if packet.node not in processes: - logger.warning( - f"Ignoring unknown node name {packet.node} in pending sends" - ) - return - # find process - proc = processes[packet.node] - proc_node = proc.node - if proc_node is None: - return - # create task id - triggers = PUSH_TRIGGER - checkpoint_ns = ( - f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node - ) - task_id = task_id_func( - checkpoint_id_bytes, - checkpoint_ns, - str(step), - packet.node, - PUSH, - str(idx), - ) - else: - logger.warning(f"Ignoring invalid PUSH task path {task_path}") - return - task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" - # we append False to the task path to indicate that a call is not being made - # so we should return interrupts from this task - task_path = (*task_path[:3], False) - metadata = { - "langgraph_step": step, - "langgraph_node": packet.node, - "langgraph_triggers": triggers, - "langgraph_path": task_path, - "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: - if proc.metadata: - metadata.update(proc.metadata) - writes = deque() - cache_policy = proc.cache_policy or cache_policy - if cache_policy: - args_key = cache_policy.key_func(packet.arg) - cache_key = CacheKey( - ( - CACHE_NS_WRITES, - (identifier(proc) or "__dynamic__"), - packet.node, - ), - xxh3_128_hexdigest( - args_key.encode() if isinstance(args_key, str) else args_key, - ), - cache_policy.ttl, - ) - else: - cache_key = None - scratchpad = _scratchpad( - config[CONF].get(CONFIG_KEY_SCRATCHPAD), - pending_writes, - task_id, - xxh3_128_hexdigest(task_checkpoint_ns.encode()), - config[CONF].get(CONFIG_KEY_RESUME_MAP), - step, - stop, - ) - runtime = cast( - Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME) - ) - runtime = runtime.override( - store=store, previous=checkpoint["channel_values"].get(PREVIOUS, None) - ) - additional_config: RunnableConfig = { - "metadata": metadata, - "tags": proc.tags, - } - return PregelExecutableTask( - packet.node, - packet.arg, - proc_node, - writes, - patch_config( - merge_configs(config, additional_config), - run_name=packet.node, - 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: writes.extend, - CONFIG_KEY_READ: partial( - local_read, - scratchpad, - channels, - managed, - PregelTaskWrites(task_path, packet.node, writes, triggers), - ), - 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, - CONFIG_KEY_SCRATCHPAD: scratchpad, - CONFIG_KEY_RUNTIME: runtime, - }, - ), - triggers, - proc.retry_policy or retry_policy, - cache_key, - task_id, - task_path, - writers=proc.flat_writers, - subgraphs=proc.subgraphs, - ) - else: - return PregelTask(task_id, packet.node, task_path) elif task_path[0] == PULL: # (PULL, node name) name = cast(str, task_path[1]) @@ -834,7 +641,7 @@ def prepare_single_task( if node := proc.node: if proc.metadata: metadata.update(proc.metadata) - writes = deque() + writes: deque[tuple[str, Any]] = deque() cache_policy = proc.cache_policy or cache_policy if cache_policy: args_key = cache_policy.key_func(val) @@ -845,9 +652,11 @@ def prepare_single_task( name, ), xxh3_128_hexdigest( - args_key.encode() - if isinstance(args_key, str) - else args_key, + ( + args_key.encode() + if isinstance(args_key, str) + else args_key + ), ), cache_policy.ttl, ) @@ -870,7 +679,9 @@ def prepare_single_task( node, writes, patch_config( - merge_configs(config, additional_config), + merge_configs( + config, cast(RunnableConfig, additional_config) + ), run_name=name, callbacks=( manager.get_child(f"graph:step:{step}") @@ -919,6 +730,297 @@ def prepare_single_task( return PregelTask(task_id, name, task_path[:3]) +def prepare_push_task_functional( + task_path: tuple[str, tuple, int, str, Call], + # (PUSH, parent task path, idx of PUSH write, id of parent task, Call) + task_id_checksum: str | None, + *, + checkpoint: Checkpoint, + checkpoint_id_bytes: bytes, + pending_writes: list[PendingWrite], + channels: Mapping[str, BaseChannel], + managed: ManagedValueMapping, + config: RunnableConfig, + step: int, + stop: int, + for_execution: bool, + store: BaseStore | None = None, + checkpointer: BaseCheckpointSaver | None = None, + manager: None | ParentRunManager | AsyncParentRunManager = None, + cache_policy: CachePolicy | None = None, + retry_policy: Sequence[RetryPolicy] = (), + parent_ns: str, + # namespace: bytes, *parts: str | bytes + task_id_func: _TaskIDFn, +) -> PregelTask | PregelExecutableTask: + """Prepare a push task with an attached caller. Used for the functional API.""" + configurable = config.get(CONF, {}) + + call = task_path[-1] + proc_ = get_runnable_for_task(call.func) + name = proc_.name + if name is None: + raise ValueError("`call` functions must have a `__name__` attribute") + # create task id + triggers: Sequence[str] = PUSH_TRIGGER + checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name + task_id = task_id_func( + checkpoint_id_bytes, + checkpoint_ns, + str(step), + name, + PUSH, + task_path_str(task_path[1]), + str(task_path[2]), + ) + task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" + # we append True to the task path to indicate that a call is being + # made, so we should not return interrupts from this task (responsibility lies with the parent) + in_progress_task_path = (*task_path[:3], True) + metadata = { + "langgraph_step": step, + "langgraph_node": name, + "langgraph_triggers": triggers, + "langgraph_path": in_progress_task_path, + "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() + cache_policy = call.cache_policy or cache_policy + if cache_policy: + args_key = cache_policy.key_func(*call.input[0], **call.input[1]) + cache_key: CacheKey | None = CacheKey( + ( + CACHE_NS_WRITES, + (identifier(call.func) or "__dynamic__"), + ), + xxh3_128_hexdigest( + args_key.encode() if isinstance(args_key, str) else args_key, + ), + cache_policy.ttl, + ) + else: + cache_key = None + scratchpad = _scratchpad( + configurable.get(CONFIG_KEY_SCRATCHPAD), + pending_writes, + task_id, + xxh3_128_hexdigest(task_checkpoint_ns.encode()), + configurable.get(CONFIG_KEY_RESUME_MAP), + step, + stop, + ) + runtime = cast(Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)) + runtime = runtime.override(store=store) + return PregelExecutableTask( + name, + call.input, + proc_, + writes, + patch_config( + merge_configs(config, {"metadata": metadata}), + run_name=name, + callbacks=call.callbacks + or (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: writes.extend, + CONFIG_KEY_READ: partial( + local_read, + scratchpad, + channels, + managed, + PregelTaskWrites(in_progress_task_path, name, writes, triggers), + ), + 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, + CONFIG_KEY_SCRATCHPAD: scratchpad, + CONFIG_KEY_RUNTIME: runtime, + }, + ), + triggers, + call.retry_policy or retry_policy, + cache_key, + task_id, + in_progress_task_path, + ) + else: + return PregelTask(task_id, name, in_progress_task_path) + + +def prepare_push_task_send( + task_path: tuple[str, tuple], + # (PUSH, parent task path) + task_id_checksum: str | None, + *, + checkpoint: Checkpoint, + checkpoint_id_bytes: bytes, + pending_writes: list[PendingWrite], + channels: Mapping[str, BaseChannel], + managed: ManagedValueMapping, + config: RunnableConfig, + step: int, + stop: int, + for_execution: bool, + store: BaseStore | None = None, + checkpointer: BaseCheckpointSaver | None = None, + manager: None | ParentRunManager | AsyncParentRunManager = None, + cache_policy: CachePolicy | None = None, + retry_policy: Sequence[RetryPolicy] = (), + parent_ns: str, + task_id_func: _TaskIDFn, + processes: Mapping[str, PregelNode], +) -> PregelTask | PregelExecutableTask | None: + if len(task_path) == 2: + # SEND tasks, executed in superstep n+1 + # (PUSH, idx of pending send) + idx = cast(int, task_path[1]) + if not channels[TASKS].is_available(): + return + sends: Sequence[Send] = channels[TASKS].get() + if idx < 0 or idx >= len(sends): + return + packet = sends[idx] + if not isinstance(packet, Send): + logger.warning( + f"Ignoring invalid packet type {type(packet)} in pending sends" + ) + return + + if packet.node not in processes: + logger.warning(f"Ignoring unknown node name {packet.node} in pending sends") + return + # find process + proc = processes[packet.node] + proc_node = proc.node + if proc_node is None: + return + # create task id + triggers = PUSH_TRIGGER + checkpoint_ns = ( + f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node + ) + task_id = task_id_func( + checkpoint_id_bytes, + checkpoint_ns, + str(step), + packet.node, + PUSH, + str(idx), + ) + else: + logger.warning(f"Ignoring invalid PUSH task path {task_path}") + return + configurable = config.get(CONF, {}) + task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" + # we append False to the task path to indicate that a call is not being made + # so we should return interrupts from this task + translated_task_path = (*task_path[:3], False) + metadata = { + "langgraph_step": step, + "langgraph_node": packet.node, + "langgraph_triggers": triggers, + "langgraph_path": translated_task_path, + "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: + if proc.metadata: + metadata.update(proc.metadata) + writes: deque[tuple[str, Any]] = deque() + cache_policy = proc.cache_policy or cache_policy + if cache_policy: + args_key = cache_policy.key_func(packet.arg) + cache_key = CacheKey( + ( + CACHE_NS_WRITES, + (identifier(proc) or "__dynamic__"), + packet.node, + ), + xxh3_128_hexdigest( + args_key.encode() if isinstance(args_key, str) else args_key, + ), + cache_policy.ttl, + ) + else: + cache_key = None + scratchpad = _scratchpad( + config[CONF].get(CONFIG_KEY_SCRATCHPAD), + pending_writes, + task_id, + xxh3_128_hexdigest(task_checkpoint_ns.encode()), + config[CONF].get(CONFIG_KEY_RESUME_MAP), + step, + stop, + ) + runtime = cast(Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)) + runtime = runtime.override( + store=store, previous=checkpoint["channel_values"].get(PREVIOUS, None) + ) + additional_config: RunnableConfig = { + "metadata": metadata, + "tags": proc.tags, + } + return PregelExecutableTask( + packet.node, + packet.arg, + proc_node, + writes, + patch_config( + merge_configs(config, additional_config), + run_name=packet.node, + 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: writes.extend, + CONFIG_KEY_READ: partial( + local_read, + scratchpad, + channels, + managed, + PregelTaskWrites( + translated_task_path, packet.node, writes, triggers + ), + ), + 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, + CONFIG_KEY_SCRATCHPAD: scratchpad, + CONFIG_KEY_RUNTIME: runtime, + }, + ), + triggers, + proc.retry_policy or retry_policy, + cache_key, + task_id, + translated_task_path, + writers=proc.flat_writers, + subgraphs=proc.subgraphs, + ) + else: + return PregelTask(task_id, packet.node, translated_task_path) + + def checkpoint_null_version( checkpoint: Checkpoint, ) -> V | None: