diff --git a/pocketflow/__init__.py b/pocketflow/__init__.py index 0b71858b..0098b1f2 100644 --- a/pocketflow/__init__.py +++ b/pocketflow/__init__.py @@ -1,4 +1,4 @@ -import asyncio, warnings, copy, time +import asyncio, warnings, copy, time, contextvars class BaseNode: def __init__(self): self.params,self.successors={},{} @@ -57,6 +57,12 @@ def _run(self,shared): return self.post(shared,pr,None) class AsyncNode(Node): + def __init__(self,max_retries=1,wait=0): + super().__init__(max_retries,wait); self._cur_retry=contextvars.ContextVar("cur_retry",default=0) + @property + def cur_retry(self): return self._cur_retry.get() + @cur_retry.setter + def cur_retry(self,value): self._cur_retry.set(value) async def prep_async(self,shared): pass async def exec_async(self,prep_res): pass async def exec_fallback_async(self,prep_res,exc): raise exc @@ -97,4 +103,4 @@ class AsyncParallelBatchFlow(AsyncFlow,BatchFlow): async def _run_async(self,shared): pr=await self.prep_async(shared) or [] await asyncio.gather(*(self._orch_async(shared,{**self.params,**bp}) for bp in pr)) - return await self.post_async(shared,pr,None) \ No newline at end of file + return await self.post_async(shared,pr,None) diff --git a/tests/test_async_parallel_batch_node.py b/tests/test_async_parallel_batch_node.py index 34d9aced..9befa708 100644 --- a/tests/test_async_parallel_batch_node.py +++ b/tests/test_async_parallel_batch_node.py @@ -139,5 +139,37 @@ async def exec_async(self, item): self.assertLess(execution_order.index(1), execution_order.index(0)) self.assertLess(execution_order.index(3), execution_order.index(2)) + def test_retry_state_is_isolated_per_item(self): + """Each parallel item should observe and control its own retry attempt.""" + class RetryProcessor(AsyncParallelBatchNode): + def __init__(self): + super().__init__(max_retries=2) + self.second_attempt_started = asyncio.Event() + self.attempts = {'a': [], 'b': []} + + async def prep_async(self, shared_storage): + return ['a', 'b'] + + async def exec_async(self, item): + self.attempts[item].append(self.cur_retry) + if self.cur_retry == 0: + if item == 'a': + await self.second_attempt_started.wait() + raise RuntimeError(f"retry {item}") + + if item == 'b': + self.second_attempt_started.set() + return f"{item}-ok" + + async def post_async(self, shared_storage, prep_result, exec_result): + shared_storage['results'] = exec_result + + shared_storage = {} + processor = RetryProcessor() + self.loop.run_until_complete(processor.run_async(shared_storage)) + + self.assertEqual(shared_storage['results'], ['a-ok', 'b-ok']) + self.assertEqual(processor.attempts, {'a': [0, 1], 'b': [0, 1]}) + if __name__ == '__main__': - unittest.main() \ No newline at end of file + unittest.main()