Merge pull request #2355 from langchain-ai/nc/6nov/loop-match-writes

lib: Split out _match_writes util in PregelLoop
This commit is contained in:
Nuno Campos
2024-11-06 08:55:10 -08:00
committed by GitHub
+15 -14
View File
@@ -381,20 +381,7 @@ class PregelLoop(LoopProtocol):
# if there are pending writes from a previous loop, apply them
if self.skip_done_tasks and self.checkpoint_pending_writes:
for tid, k, v in self.checkpoint_pending_writes:
if k in (ERROR, INTERRUPT):
continue
if task := self.tasks.get(tid):
if k == SCHEDULED:
if v == max(
self.checkpoint["versions_seen"]
.get(INTERRUPT, {})
.values(),
default=None,
):
self.tasks[tid] = task._replace(scheduled=True)
else:
task.writes.append((k, v))
self._match_writes(self.tasks)
elif not self.skip_done_tasks:
# "not skip_done_tasks" only applies to first tick after resuming
self.skip_done_tasks = True
@@ -429,6 +416,20 @@ class PregelLoop(LoopProtocol):
# private
def _match_writes(self, tasks: Mapping[str, PregelExecutableTask]) -> None:
for tid, k, v in self.checkpoint_pending_writes:
if k in (ERROR, INTERRUPT):
continue
if task := tasks.get(tid):
if k == SCHEDULED:
if v == max(
self.checkpoint["versions_seen"].get(INTERRUPT, {}).values(),
default=None,
):
self.tasks[tid] = task._replace(scheduled=True)
else:
task.writes.append((k, v))
def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None:
# resuming from previous checkpoint requires
# - finding a previous checkpoint