[aDAG] Support multi-read of the same shm channel

[ruisearch42]

Why are these changes needed?

If the same method of the same actor is bound to the same node (i.e., reads from the same shared memory channel), aDAG execution hangs. This PR adds support to this case by caching results read from the channel.

Related issue number

Closes #47041

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[rkooo567]

QQ: is this approach working for this kind of case? (I believe it will not work?)

with InputNode as inp:
    out = s1.fwd(inp)
    dag = s2.fwd(inp, out)

Also have some concerns that deepcopy on the first result can affect perf negatively.

I actually wonder if it is viable approach to just not allow the additional write until all downstream tasks finish reading (I think with buffering input PR, it may work well?)

[rkooo567]

discussed offline. deep/shallow copy overhead is pretty big, so we will make some assumptions that input won't be changed (which is not very correct). we can allow copy inputs with some sort of flags

[stephanie-wang]

Could you add a description to the PR?

[rkooo567]

QQ: is intraprocess channel still needed after this change?

[stephanie-wang]

Looks like a good start but let's try not to introduce the new use_cached flag if possible (it seems to make the code more complicated) and ideally clean up the cache once we know no more tasks will need it.

One suggestion might be to have an args cache per task, instead of a cache per channel. Then we can also do a list lookup instead of dict. Like:

cache: List[List[Any]]

The outer list index is the task idx on that actor, and the inner list is the resolved args for that task. The first time we deserialize an arg, we put it into the inner lists for all reader tasks.

[ruisearch42]

Looks like a good start but let's try not to introduce the new use_cached flag if possible (it seems to make the code more complicated) and ideally clean up the cache once we know no more tasks will need it.

Sounds good.

One suggestion might be to have an args cache per task, instead of a cache per channel. Then we can also do a list lookup instead of dict.

Why is a list lookup preferred than a dict? It is performance concerns (hash function overhead)?

The first time we deserialize an arg, we put it into the inner lists for all reader tasks.

How do we know where to put into the inner lists? We still need a dict to maintain that info? Otherwise we need to go through all items of all inner lists and replace?

[stephanie-wang]

One suggestion might be to have an args cache per task, instead of a cache per channel. Then we can also do a list lookup instead of dict.

Why is a list lookup preferred than a dict? It is performance concerns (hash function overhead)?

The first time we deserialize an arg, we put it into the inner lists for all reader tasks.

How do we know where to put into the inner lists? We still need a dict to maintain that info? Otherwise we need to go through all items of all inner lists and replace?

Yes, I thought list lookup would be better for performance and I think it also makes the garbage collection simpler. But the new approach here also looks OK.

[ruisearch42]

Still WIP after using a new approach, will clean up more!

[rkooo567]

@ruisearch42 plz remove author-action-required when it is ready!

[ruisearch42]

Looking into some GPU test CI failures, but PR is ready for another look.

[rkooo567]

Looks very clean! Can you run a microbenchmark and address last test comments? let's merge it after that

[ruisearch42]

Kicked off microbenchmark: https://buildkite.com/ray-project/release/builds/21924#0191a221-42bb-4bd7-8f01-25adf035540e

[rkooo567]

@ruisearch42 plz comment if there's any change here! if not, let's just merge it

[ruisearch42]

@rkooo567 microbenchmark aligns with past runs:
https://buildkite.com/ray-project/release/builds/21950#0191a40e-897c-44b1-b1f9-388ea3761ff6

The latest code just added a TODO comment compared to the last version which passed full CI.

[rkooo567]

@ruisearch42 auto merged enabled. so we should follow up with removing intra process channel right?

[rkooo567]

premerge failure. maybe consider to merge latest master?