Fix lint issues

This commit is contained in:
Tat Dat Duong
2025-03-19 14:21:43 +01:00
parent 1e751a2256
commit 764929afd9
+20 -14
View File
@@ -1415,7 +1415,7 @@ class Pregel(PregelProtocol):
valid_updates.append((values, as_node))
tasks: list[tuple[str, PregelTaskWrites]] = []
run_tasks: list[tuple[str, PregelTaskWrites]] = []
for values, as_node in valid_updates:
# create task to run all writers of the chosen node
writers = self.nodes[as_node].flat_writers
@@ -1424,9 +1424,9 @@ class Pregel(PregelProtocol):
writes: deque[tuple[str, Any]] = deque()
task = PregelTaskWrites((), as_node, writes, [INTERRUPT])
task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT))
tasks.append((task_id, task))
run_tasks.append((task_id, task))
for _, task in tasks:
for _, task in run_tasks:
run = RunnableSequence(*writers) if len(writers) > 1 else writers[0]
# execute task
run.invoke(
@@ -1455,7 +1455,7 @@ class Pregel(PregelProtocol):
)
# save task writes
for _, task in tasks:
for _, task in run_tasks:
channel_writes = [w for w in task.writes if w[0] != PUSH]
# channel writes are saved to current checkpoint
@@ -1464,7 +1464,10 @@ class Pregel(PregelProtocol):
# apply to checkpoint and save
mv_writes = apply_writes(
checkpoint, channels, tasks, checkpointer.get_next_version
checkpoint,
channels,
[t for _, t in run_tasks],
checkpointer.get_next_version,
)
assert not mv_writes, "Can't write to SharedValues from update_state"
@@ -1485,7 +1488,7 @@ class Pregel(PregelProtocol):
),
)
for task_id, task in tasks:
for task_id, task in run_tasks:
if push_writes := [w for w in task.writes if w[0] == PUSH]:
checkpointer.put_writes(next_config, push_writes, task_id)
@@ -1747,7 +1750,7 @@ class Pregel(PregelProtocol):
valid_updates.append((values, as_node))
tasks: list[tuple[str, PregelTaskWrites]] = []
run_tasks: list[tuple[str, PregelTaskWrites]] = []
for values, as_node in valid_updates:
# create task to run all writers of the chosen node
@@ -1757,9 +1760,9 @@ class Pregel(PregelProtocol):
writes: deque[tuple[str, Any]] = deque()
task = PregelTaskWrites((), as_node, writes, [INTERRUPT])
task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT))
tasks.append((task_id, task))
run_tasks.append((task_id, task))
for _, task in tasks:
for _, task in run_tasks:
run = RunnableSequence(*writers) if len(writers) > 1 else writers[0]
# execute task
await run.ainvoke(
@@ -1788,7 +1791,7 @@ class Pregel(PregelProtocol):
)
# save task writes
for _, task in tasks:
for _, task in run_tasks:
# channel writes are saved to current checkpoint
channel_writes = [w for w in task.writes if w[0] != PUSH]
if saved and channel_writes:
@@ -1798,7 +1801,10 @@ class Pregel(PregelProtocol):
# apply to checkpoint and save
mv_writes = apply_writes(
checkpoint, channels, [task], checkpointer.get_next_version
checkpoint,
channels,
[t for _, t in run_tasks],
checkpointer.get_next_version,
)
assert not mv_writes, "Can't write to SharedValues from update_state"
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
@@ -1818,7 +1824,7 @@ class Pregel(PregelProtocol):
),
)
for task_id, task in tasks:
for task_id, task in run_tasks:
# save push writes
if push_writes := [w for w in task.writes if w[0] == PUSH]:
await checkpointer.aput_writes(next_config, push_writes, task_id)
@@ -1835,7 +1841,7 @@ class Pregel(PregelProtocol):
node `as_node`. If `as_node` is not provided, it will be set to the last node
that updated the state, if not ambiguous.
"""
return self.bulk_update_state(config, [(values, as_node)])
return self.bulk_update_state(config, [BulkUpdate(values, as_node)])
async def aupdate_state(
self,
@@ -1847,7 +1853,7 @@ class Pregel(PregelProtocol):
node `as_node`. If `as_node` is not provided, it will be set to the last node
that updated the state, if not ambiguous.
"""
return await self.bulk_update_state(config, [(values, as_node)])
return await self.abulk_update_state(config, [BulkUpdate(values, as_node)])
def _defaults(
self,