refactor: extract _track helper for available_channels sync in apply_writes

Replace 5 repeated inline blocks that sync available_channels with
a local _track() helper that checks is_available() and updates the
set, returning the availability bool for callers that also need to
update updated_channels.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
John Kennedy
2026-07-27 17:01:37 +00:00
committed by John Kennedy
co-authored by Claude Opus 4.6
parent 3921da01cf
commit b916a2e183
2 changed files with 20 additions and 29 deletions
+20 -28
View File
@@ -282,6 +282,18 @@ def apply_writes(
None,
)
# Sync available_channels with channel's actual availability state.
# Returns True if the channel is available (for callers that also need
# to update updated_channels).
def _track(chan: str) -> bool:
avail = channels[chan].is_available()
if available_channels is not None:
if avail:
available_channels.add(chan)
else:
available_channels.discard(chan)
return avail
# Consume all channels that were read
for chan in {
chan
@@ -291,11 +303,7 @@ def apply_writes(
}:
if channels[chan].consume() and next_version is not None:
checkpoint["channel_versions"][chan] = next_version
if available_channels is not None:
if channels[chan].is_available():
available_channels.add(chan)
else:
available_channels.discard(chan)
_track(chan)
# Group writes by channel
pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
@@ -325,17 +333,10 @@ def apply_writes(
if channels[chan].update(vals) and next_version is not None:
checkpoint["channel_versions"][chan] = next_version
# unavailable channels can't trigger tasks, so don't add them
if channels[chan].is_available():
if _track(chan):
updated_channels.add(chan)
if available_channels is not None:
available_channels.add(chan)
elif available_channels is not None:
available_channels.discard(chan)
elif available_channels is not None:
if channels[chan].is_available():
available_channels.add(chan)
else:
available_channels.discard(chan)
else:
_track(chan)
# Channels that weren't updated in this step are notified of a new step
if bump_step:
@@ -352,10 +353,8 @@ def apply_writes(
if channels[chan].update(EMPTY_SEQ) and next_version is not None:
checkpoint["channel_versions"][chan] = next_version
# unavailable channels can't trigger tasks, so don't add them
if channels[chan].is_available():
if _track(chan):
updated_channels.add(chan)
elif available_channels is not None:
available_channels.discard(chan)
# If this is (tentatively) the last superstep, notify all channels of finish
if bump_step and updated_channels.isdisjoint(trigger_to_nodes):
@@ -363,17 +362,10 @@ def apply_writes(
if channels[chan].finish() and next_version is not None:
checkpoint["channel_versions"][chan] = next_version
# unavailable channels can't trigger tasks, so don't add them
if channels[chan].is_available():
if _track(chan):
updated_channels.add(chan)
if available_channels is not None:
available_channels.add(chan)
elif available_channels is not None:
available_channels.discard(chan)
elif available_channels is not None:
if channels[chan].is_available():
available_channels.add(chan)
else:
available_channels.discard(chan)
else:
_track(chan)
# Return managed values writes to be applied externally
return updated_channels
-1
View File
@@ -198,7 +198,6 @@ class PregelLoop:
_migrate_checkpoint: Callable[[Checkpoint], None] | None
submit: Submit
channels: Mapping[str, BaseChannel]
_has_untracked_channels: bool
_available_channels: set[str]
# Futures from `checkpointer.put_writes` calls that produced delta-channel
# writes. `_checkpointer_put_after_previous` drains this list (swap to a