From 0aff02e18048960005900e166b271183dd1f3715 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 12 Jun 2025 11:08:37 -0700 Subject: [PATCH] Remove code paths no longer needed - These were only used by the kafka scheduler --- libs/langgraph/langgraph/constants.py | 13 --- libs/langgraph/langgraph/errors.py | 13 --- libs/langgraph/langgraph/pregel/__init__.py | 13 ++- libs/langgraph/langgraph/pregel/loop.py | 88 ++------------------- libs/langgraph/langgraph/types.py | 1 - 5 files changed, 11 insertions(+), 117 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 3c89dab4c..ac53718e1 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -39,8 +39,6 @@ ERROR = sys.intern("__error__") # for errors raised by nodes NO_WRITES = sys.intern("__no_writes__") # marker to signal node didn't write anything -SCHEDULED = sys.intern("__scheduled__") -# marker to signal node was scheduled (in distributed mode) TASKS = sys.intern("__pregel_tasks") # for Send objects returned by nodes/edges, corresponds to PUSH below RETURN = sys.intern("__return__") @@ -71,13 +69,6 @@ CONFIG_KEY_RESUMING = sys.intern("__pregel_resuming") # holds a boolean indicating if subgraphs should resume from a previous checkpoint CONFIG_KEY_TASK_ID = sys.intern("__pregel_task_id") # holds the task ID for the current task -CONFIG_KEY_DEDUPE_TASKS = sys.intern("__pregel_dedupe_tasks") -# holds a boolean indicating if tasks should be deduplicated (for distributed mode) -CONFIG_KEY_ENSURE_LATEST = sys.intern("__pregel_ensure_latest") -# holds a boolean indicating whether to assert the requested checkpoint is the latest -# (for distributed mode) -CONFIG_KEY_DELEGATE = sys.intern("__pregel_delegate") -# holds a boolean indicating whether to delegate subgraphs (for distributed mode) CONFIG_KEY_THREAD_ID = sys.intern("thread_id") # holds the thread ID for the current invocation CONFIG_KEY_CHECKPOINT_MAP = sys.intern("checkpoint_map") @@ -121,7 +112,6 @@ RESERVED = { RESUME, ERROR, NO_WRITES, - SCHEDULED, # reserved config.configurable keys CONFIG_KEY_SEND, CONFIG_KEY_READ, @@ -132,9 +122,6 @@ RESERVED = { CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_RESUMING, CONFIG_KEY_TASK_ID, - CONFIG_KEY_DEDUPE_TASKS, - CONFIG_KEY_ENSURE_LATEST, - CONFIG_KEY_DELEGATE, CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 09e5052e2..6213ff68e 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -78,13 +78,6 @@ class NodeInterrupt(GraphInterrupt): super().__init__([Interrupt(value=value)]) -class GraphDelegate(GraphBubbleUp): - """Raised when a graph is delegated (for distributed mode).""" - - def __init__(self, *args: dict[str, Any]) -> None: - super().__init__(*args) - - class ParentCommand(GraphBubbleUp): args: tuple[Command] @@ -102,9 +95,3 @@ class TaskNotFound(Exception): """Raised when the executor is unable to find a task (for distributed mode).""" pass - - -class CheckpointNotLatest(Exception): - """Raised when the checkpoint is not the latest version (for distributed mode).""" - - pass diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index bd77d1487..2e8f9b01e 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -60,7 +60,6 @@ from langgraph.constants import ( NS_SEP, NULL_TASK_ID, PUSH, - SCHEDULED, TASKS, ) from langgraph.errors import ( @@ -1008,7 +1007,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ) if apply_pending_writes and saved.pending_writes: for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): + if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue @@ -1127,7 +1126,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ) if apply_pending_writes and saved.pending_writes: for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): + if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue @@ -1466,7 +1465,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ) # apply writes from tasks that already ran for tid, k, v in saved.pending_writes or []: - if k in (ERROR, INTERRUPT, SCHEDULED): + if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue @@ -1630,7 +1629,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ) # apply writes for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): + if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue @@ -1886,7 +1885,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ) # apply writes from tasks that already ran for tid, k, v in saved.pending_writes or []: - if k in (ERROR, INTERRUPT, SCHEDULED): + if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue @@ -2049,7 +2048,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou self.trigger_to_nodes, ) for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): + if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7d43ca445..df6c84dba 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio import binascii import concurrent.futures -import dataclasses from collections import defaultdict, deque from collections.abc import Iterator, Mapping, Sequence from contextlib import ( @@ -46,9 +45,6 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, - CONFIG_KEY_DEDUPE_TASKS, - CONFIG_KEY_DELEGATE, - CONFIG_KEY_ENSURE_LATEST, CONFIG_KEY_RESUME_MAP, CONFIG_KEY_RESUMING, CONFIG_KEY_SCRATCHPAD, @@ -65,13 +61,10 @@ from langgraph.constants import ( NULL_TASK_ID, PUSH, RESUME, - SCHEDULED, TAG_HIDDEN, ) from langgraph.errors import ( - CheckpointNotLatest, EmptyInputError, - GraphDelegate, GraphInterrupt, ) from langgraph.managed.base import ( @@ -252,10 +245,7 @@ class PregelLoop: self.interrupt_before = interrupt_before self.manager = manager self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {}) - self.skip_done_tasks = ( - CONFIG_KEY_CHECKPOINT_ID not in config[CONF] - or CONFIG_KEY_DEDUPE_TASKS in config[CONF] - ) + self.skip_done_tasks = CONFIG_KEY_CHECKPOINT_ID not in config[CONF] self._migrate_checkpoint = migrate_checkpoint self.trigger_to_nodes = trigger_to_nodes self.retry_policy = retry_policy @@ -265,9 +255,7 @@ class PregelLoop: if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) scratchpad: PregelScratchpad | None = config[CONF].get(CONFIG_KEY_SCRATCHPAD) - if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and isinstance( - scratchpad, PregelScratchpad - ): + if isinstance(scratchpad, PregelScratchpad): # if count is > 0, append to checkpoint_ns # if count is 0, leave as is if cnt := scratchpad.subgraph_counter(): @@ -589,18 +577,6 @@ class PregelLoop: self.status = "done" return False - # check if we should delegate (used by subgraphs in distributed mode) - if self.config[CONF].get(CONFIG_KEY_DELEGATE): - assert self.input is INPUT_RESUMING - raise GraphDelegate( - { - "config": patch_configurable( - self.config, {CONFIG_KEY_DELEGATE: False} - ), - "input": None, - } - ) - # if there are pending writes from a previous loop, apply them if self.skip_done_tasks and self.checkpoint_pending_writes: self._match_writes(self.tasks) @@ -643,14 +619,7 @@ class PregelLoop: if k in (ERROR, INTERRUPT, RESUME): continue if task := tasks.get(tid): - if k == SCHEDULED: - if v == max( - self.checkpoint["versions_seen"].get(INTERRUPT, {}).values(), - default=None, - ): - self.tasks[tid] = dataclasses.replace(task, scheduled=True) - else: - task.writes.append((k, v)) + task.writes.append((k, v)) def _first(self, *, input_keys: str | Sequence[str]) -> set[str] | None: # resuming from previous checkpoint requires @@ -720,17 +689,6 @@ class PregelLoop: self.input = INPUT_RESUMING # map inputs to channel updates elif input_writes := deque(map_input(input_keys, self.input)): - # TODO shouldn't these writes be passed to put_writes too? - # check if we should delegate (used by subgraphs in distributed mode) - if self.config[CONF].get(CONFIG_KEY_DELEGATE): - raise GraphDelegate( - { - "config": patch_configurable( - self.config, {CONFIG_KEY_DELEGATE: False} - ), - "input": self.input, - } - ) # discard any unfinished tasks from previous checkpoint discard_tasks = prepare_next_tasks( self.checkpoint, @@ -1105,25 +1063,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): # context manager def __enter__(self) -> Self: - if self.config.get(CONF, {}).get( - CONFIG_KEY_ENSURE_LATEST - ) and self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID): - if self.checkpointer is None: - raise RuntimeError( - "Cannot ensure latest checkpoint without checkpointer" - ) - saved = self.checkpointer.get_tuple( - patch_configurable( - self.checkpoint_config, {CONFIG_KEY_CHECKPOINT_ID: None} - ) - ) - if ( - saved is None - or saved.checkpoint["id"] - != self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID] - ): - raise CheckpointNotLatest - elif self.checkpointer: + if self.checkpointer: saved = self.checkpointer.get_tuple(self.checkpoint_config) else: saved = None @@ -1298,25 +1238,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): # context manager async def __aenter__(self) -> Self: - if self.config.get(CONF, {}).get( - CONFIG_KEY_ENSURE_LATEST - ) and self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID): - if self.checkpointer is None: - raise RuntimeError( - "Cannot ensure latest checkpoint without checkpointer" - ) - saved = await self.checkpointer.aget_tuple( - patch_configurable( - self.checkpoint_config, {CONFIG_KEY_CHECKPOINT_ID: None} - ) - ) - if ( - saved is None - or saved.checkpoint["id"] - != self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID] - ): - raise CheckpointNotLatest - elif self.checkpointer: + if self.checkpointer: saved = await self.checkpointer.aget_tuple(self.checkpoint_config) else: saved = None diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index ddcca374f..3e907e3eb 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -203,7 +203,6 @@ class PregelExecutableTask: cache_key: CacheKey | None id: str path: tuple[str | int | tuple, ...] - scheduled: bool = False writers: Sequence[Runnable] = () subgraphs: Sequence[PregelProtocol] = ()