[Data] Account internal inqueue to previous operator's memory usage

[bveeramani]

Stacked on

• [data] sanitize streaming executor resource management #42817

Why are these changes needed?

Related issue number

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 :(

[raulchen]

These metrics currently only work for map ops. I'm wondering if we can also extend them to also support other ops. So 1) we can report them to the dashboard for all ops; 2) we don't have 2 places reporting memory metrics (here and internal_inqueue_memory_usage).

The on_input_received function is already called on all ops. I think we'll also need to add other callback APIs like on_input_freed, on_output_generated, etc.

[bveeramani]

Yeah, I thought about that too.

OpRuntimeMetrics is tightly coupled with MapOperator, and this PR introduces a layer of indirection with PhysicalOperator.*_memory_usage() and OpRuntimeMetrics.obj_store_mem_*.

One concern I have is that if we rely on metrics directly, we need to ensure that they're updated correctly for all operators (for example, we need to make sure all operators call on_output_generated, or else obj_store_mem_internal_outqueue will be incorrect). But, I'm not sure how we can enforce that with the PhysicalOperator interface.

While I agree multiple abstractions for memory reporting isn't ideal, I think it might be the easiest way to implement this? Wdyt?

[bveeramani]

Documenting tradeoffs from offline discussion:

Disadvantages of adding methods to PhysicalOperator:

• OpRuntimeMetrics is tightly coupled with MapOperator
• Layer of indirection for MapOperator
• Multiple abstractions for memory reporting

Disadvantages of using OpRuntimeMetrics only:

• Updating metrics is implicitly required (not part of interface)

Given these tradeoffs, I've opened a new PR with the second approach https://github.com/ray-project/ray/pull/42930/files

[bveeramani]

Closing in favor of #42930