perf: cache UntrackedValue isinstance scan in put_writes

The `put_writes` method was scanning all channels with
`any(isinstance(ch, UntrackedValue) ...)` on every call. For the
sequential_1000 benchmark this produced 1M+ isinstance calls through
the ABC machinery, consuming 40% of total runtime.

Cache the result as `_has_untracked_channels` once in __enter__ and
__aenter__, replacing both scan sites (put_writes and checkpoint
sanitization). This yields a 1.8-2.3x speedup on sequential_1000.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
John Kennedy
2026-07-27 17:08:11 +00:00
committed by John Kennedy
co-authored by Claude Opus 4.6
parent 30c4d58db8
commit f61e9b45b6
+8 -6
View File
@@ -437,9 +437,7 @@ class PregelLoop:
writes_to_save = writes
# check if any writes are to an UntrackedValue channel
if any(
isinstance(channel, UntrackedValue) for channel in self.channels.values()
):
if self._has_untracked_channels:
# we do not persist untracked values in checkpoints
writes_to_save = [
# sanitize UntrackedValues that are nested within Send packets
@@ -1161,9 +1159,7 @@ class PregelLoop:
elif "counters_since_delta_snapshot" in self.checkpoint_metadata:
del self.checkpoint_metadata["counters_since_delta_snapshot"]
# sanitize TASK channel in the checkpoint before saving (durability=="exit")
if TASKS in self.checkpoint["channel_values"] and any(
isinstance(channel, UntrackedValue) for channel in self.channels.values()
):
if TASKS in self.checkpoint["channel_values"] and self._has_untracked_channels:
sanitized_tasks = [
sanitize_untracked_values_in_send(value, self.channels)
if isinstance(value, Send)
@@ -1695,6 +1691,9 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
saver=self.checkpointer,
config=self.checkpoint_config,
)
self._has_untracked_channels = any(
isinstance(ch, UntrackedValue) for ch in self.channels.values()
)
self.stack.push(self._suppress_interrupt)
self.status = "input"
self.step = self.checkpoint_metadata["step"] + 1
@@ -1955,6 +1954,9 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
saver=self.checkpointer,
config=self.checkpoint_config,
)
self._has_untracked_channels = any(
isinstance(ch, UntrackedValue) for ch in self.channels.values()
)
self.stack.push(self._suppress_interrupt)
self.status = "input"
self.step = self.checkpoint_metadata["step"] + 1