Turn SharedData into generic for improved code intelligence / type checking

[nagyv]

It would be great if SharedData could be specified as a custom type (e.g. TypedDict). This would help greatly linters to catch type errors in prep and post calls.

[shenyimings]

I agree!

And shouldn't SharedData be an asyncio-compatible shared variable with locks to ensure concurrency safety?

[mouramax]

Here are my two cents on this discussion.

@shenyimings, when we talk about asyncio compatibility in the context of the shared state, we're basically concerned about race conditions, specifically write access, right? But there's a key design insight into PocketFlow: the framework appears to follow a sequential node execution model, even in its asynchronous versions (AsyncFlow). Look at the orchestrator methods (_orch and even _orch_async). They execute one node, wait for it to complete, and then determine the next node to run. The flow is a Directed Acyclic Graph (DAG), which was a very elegant design decision, and so at any given time, only one node is executing and has write access to SharedData. This means a large portion of real-world use cases are already naturally protected by this design choice.

A true parallelism edge case might occur within nodes like AsyncParallelBatchNode. In this specific case, the _exec method receives a list of items and is responsible for processing them. The crucial point is that the prep phase happens before this parallel execution, and the post phase happens after! This implies that SharedData modification is likely intended to happen only in the post method, which is called just once after all parallel tasks (where the real danger lies) are complete.

@nagyv, your concern is quite valid. SharedData as a simple Dict[str, Any] serves its purpose well. It's functional, simple, flexible, and, most importantly, concise (let's not forget the 100-line code limit). It's a solution that works out-of-the-box for most small projects. However, the permissiveness of a Dict starts to become a drawback, especially as a project begins to scale. A simple typo when reading (or writing) a dictionary key can lead to unexpected behavior, requiring an analysis of the entire flow to understand what SharedData was expected to look like at that point, right? So, in an effort to reconcile the 100-line code constraint while allowing each project to extend SharedData for its own needs, I think it's possible to leverage Generics. Something like this:

from typing import TypeVar, Dict, Any, Generic

_SharedDataT = TypeVar('_SharedDataT', bound=Dict[str, Any])

Now, we can just update the class definitions, like:

class BaseNode(Generic[_PrepResult, _ExecResult, _PostResult, _SharedDataT]):
    def prep(self, shared: _SharedDataT) -> _PrepResult: ...
    def post(self, shared: _SharedDataT, ...) -> _PostResult: ...
    def _run(self, shared: _SharedDataT) -> _PostResult: ...
    def run(self, shared: _SharedDataT) -> _PostResult: ...

class Flow(BaseNode[_PrepResult, Any, _PostResult, _SharedDataT]):
    def _orch(self, shared: _SharedDataT, ...) -> Any: ...

This is non-breaking. If a user doesn't specify a type, it defaults to Dict[str, Any], so existing code works perfectly, while at the same time allowing users to opt-in for stronger type safety on a per-flow basis.

[nagyv]

This sounds great @mouramax! -------- Eredeti üzenet --------2025. 08. 12. 17:35-kor, Max Moura ezt írta:mouramax left a comment (The-Pocket/PocketFlow#106) Here are my two cents on this discussion. @shenyimings, when we talk about asyncio compatibility in the context of the shared state, we're basically concerned about race conditions, specifically write access, right? But there's a key design insight into PocketFlow: the framework appears to follow a sequential node execution model, even in its asynchronous versions (AsyncFlow). Look at the orchestrator methods (_orch and even _orch_async). They execute one node, wait for it to complete, and then determine the next node to run. The flow is a Directed Acyclic Graph (DAG), which was a very elegant design decision, and so at any given time, only one node is executing and has write access to SharedData. This means a large portion of real-world use cases are already naturally protected by this design choice. A true parallelism edge case might occur within nodes like AsyncParallelBatchNode. In this specific case, the _exec method receives a list of items and is responsible for processing them. The crucial point is that the prep phase happens before this parallel execution, and the post phase happens after! This implies that SharedData modification is likely intended to happen only in the post method, which is called just once after all parallel tasks (where the real danger lies) are complete. @nagyv, your concern is quite valid. SharedData as a simple Dict[str, Any] serves its purpose well. It's functional, simple, flexible, and, most importantly, concise (let's not forget the 100-line code limit). It's a solution that works out-of-the-box for most small projects. However, the permissiveness of a Dict starts to become a drawback, especially as a project begins to scale. A simple typo when reading (or writing) a dictionary key can lead to unexpected behavior, requiring an analysis of the entire flow to understand what SharedData was expected to look like at that point, right? So, in an effort to reconcile the 100-line code constraint while allowing each project to extend SharedData for its own needs, I think it's possible to leverage Generics. Something like this: from typing import TypeVar, Dict, Any, Generic _SharedDataT = TypeVar('_SharedDataT', bound=Dict[str, Any]) Now, we can just update the class definitions, like: class BaseNode(Generic[_PrepResult, _ExecResult, _PostResult, _SharedDataT]): def prep(self, shared: _SharedDataT) -> _PrepResult: ... def post(self, shared: _SharedDataT, ...) -> _PostResult: ... def _run(self, shared: _SharedDataT) -> _PostResult: ... def run(self, shared: _SharedDataT) -> _PostResult: ... class Flow(BaseNode[_PrepResult, Any, _PostResult, _SharedDataT]): def _orch(self, shared: _SharedDataT, ...) -> Any: ... This is non-breaking. If a user doesn't specify a type, it defaults to Dict[str, Any], so existing code works perfectly, while at the same time allowing users to opt-in for stronger type safety on a per-flow basis. —Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you were mentioned.Message ID: ***@***.***> [ { ***@***.***": "http://schema.org", ***@***.***": "EmailMessage", "potentialAction": { ***@***.***": "ViewAction", "target": "#106 (comment)", "url": "#106 (comment)", "name": "View Issue" }, "description": "View this Issue on GitHub", "publisher": { ***@***.***": "Organization", "name": "GitHub", "url": "https://github.com" } } ]

[shenyimings]

@mouramax 👍

Thank you for your insightful reply! I now understand that my concern exists in specific scenarios.

Let me explain why I raised this question: I'm currently using AsyncParallelBatchFlow, and I noticed that line 99 in the source code directly discards the return results from each parallel sub-flow. This means that each flow can seemingly only pass the result by modifying the SharedData concurrently. That's why I had this concern about thread safety.

Additionally, I'm indeed reconsidering whether I really need to use AsyncParallelBatchFlow. I haven't yet deeply understood the core differences between BatchFlow and BatchNode.

Thank you again for taking the time to reply to my question.

[mouramax]

@shenyimings, you've got a point. Actually, you've hit on a huge point regarding how AsyncParallelBatchFlow works:

I noticed that line 99 in the source code directly discards the return results from each parallel sub-flow.

In the source code, we can see it doesn't assign the list of results from each task to a variable. The results are effectively discarded:

class AsyncParallelBatchFlow(AsyncFlow,BatchFlow):
    async def _run_async(self,shared): 
        pr=await self.prep_async(shared) or []
        # 👇 This next line runs all the flows in parallel
        await asyncio.gather(*(self._orch_async(shared,{**self.params,**bp}) for bp in pr))
        # 👇 The results from asyncio.gather go to Narnia!
        return await self.post_async(shared,pr,None)

This leads directly to the concern you raised: If the return values are ignored, the only way for a parallel sub-flow to communicate its result is by modifying the SharedData dictionary. And if multiple flows try to write to the same key in SharedData at the same time, you do, in fact, risk a race condition, where the final value is unpredictable.

⚠️ So, I want to go on record here and say that the edge case you've pointed out seems incredibly valid and definitely deserves a look from the project maintainers!

Now let's think about whether you actually need to use an AsyncParallelBatchFlow for your use case. Since you didn't provide a minimal reproducible example, I'll speak in general terms so you can decide for yourself if there's a better solution. Let's review some concepts (based on my own understanding, which might be limited):

• A Node is a worker with one specific task. It processes one item.

• A BatchNode is a worker doing one task on a list of items sequentially.

• An AsyncParallelBatchNode is a worker doing one task on a list of items in parallel.

• A Flow is a manager who directs the sequence of workers. It processes one item.

• A BatchFlow is a manager running the whole workflow multiple times (on a list of items) sequentially.

• An AsyncParallelBatchFlow is a manager running multiple workflows (multi-step workflow) concurrently.

So your race condition concern with SharedData is most dangerous with AsyncParallelBatchFlow, because you have entire, independent workflows running in parallel. Are you still with me? Right.

Now you need to look at your use case and think: are your parallel tasks truly independent? This is crucial to choosing the right pattern.

If each parallel flow can work without interfering with the others, you can just use AsyncParallelBatchFlow by safely partitioning your data, like partitioning a hard drive. Instead of all flows writing to the same key, give each one its own "slot" in the SharedData dictionary. Then, in the async def post_async() method of each AsyncNode, you can safely write shared["results"][task_id] = result_data.

The other scenario is that you need to collect results from a parallel step. This seems more likely to be what you're after, like wanting to run a bunch of I/O tasks in parallel and then do something with all the results after they've completed. In this case, a much safer and clearer pattern, in my opinion, is to use an AsyncParallelBatchNode within a regular AsyncFlow. This pattern contains the parallelism to a single, well-defined step in your flow, something like this:

1. Node 1 (preparation): Get the list of items to process.
2. Node 2 (an AsyncParallelBatchNode): Execute the parallel I/O calls. This node's _exec will return a list of all the results.
3. Node 3 (aggregation): This node runs only after the parallel node is completely finished. It receives the list of results from the previous step and can safely process them without the fear of any race conditions.

I hope this discussion has helped you as much as it has helped me better understand this very promising framework.