[dataset] Pipeline task submission during reduce stage in push-based shuffle

[stephanie-wang]

Why are these changes needed?

Reduce stage in push-based shuffle fails to complete at 100k output partitions or more. This is likely because of driver or raylet load from having too many tasks in flight at once.

We can fix this from ray core too, but for now, this PR adds pipelining for the reduce stage, to limit the total number of reduce tasks in flight at the same time. This is currently set to 2 * available parallelism in the cluster. We have to pick which reduce tasks to submit carefully since these are pinned to specific nodes. The PR does this by assigning tasks round-robin according to the corresponding merge task (which get spread throughout the cluster).

In addition, this PR refactors the map, merge, and reduce stages to use a common pipelined iterator pattern, since they all have a similar pattern of submitting a round of tasks at a time, then waiting for a previous round to finish before submitting more.

Related issue number

Closes #25412.

Checks

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

Adds a multinode test, but it would be nice to add a test that hooks into ray.remote and ray.get to check that we are actually submitting the right tasks at the right time.

[stephanie-wang]

(also includes changes from #25734)

[ericl]

The key config here is num_merge_tasks_per_round, is that right? Given this is a temporary workaround, I think we should make this feature flagged, and add a TODO to remove this once we fix the core issue. 100k tasks should fit comfortably within our scalability envelope, so this is a bit odd.

[ericl]

It would also be great to have a unit test that things work with pipelining enabled/disabled, so we can easily remove it in the future.

[stephanie-wang]

The key config here is num_merge_tasks_per_round, is that right? Given this is a temporary workaround, I think we should make this feature flagged, and add a TODO to remove this once we fix the core issue. 100k tasks should fit comfortably within our scalability envelope, so this is a bit odd.

Actually it's the reduce stage that is failing, so I think the numbers that matter are num rounds (= num map tasks / parallelism) * num reducers. Yeah, I also thought it was strange but I think it may have something to do with the fact that each reduce task also has a lot of plasma args, which we don't test in the scalability envelope. Fewer than in simple shuffle, but it's still on the order of 100s of args.

[ericl]

Feature flag lgtm... rubber stamping on the rest.