[Data] External shuffle repartition release benchmark (manual)

[ShockYoungCHN]

Summary

• Add bench_shuffle.py: keyed TPC-H lineitem repartition benchmark comparing SHUFFLE_V2 in-memory vs external file-transport.
• Add manual release tests repartition_lineitem_in-memory and repartition_lineitem_external (512GB, 500 partitions, 32-worker all-to-all cluster).

Purpose

Benchmark-only PR, separate from ext-shuffle-3-planner (#65499). Trigger via buildkite/release Rebuild with filter:

name:repartition_lineitem

Test plan

• buildkite/release → Rebuild → filter name:repartition_lineitem_in-memory
• buildkite/release → Rebuild → filter name:repartition_lineitem_external
• Compare wall time / throughput in Buildkite artifacts (TEST_OUTPUT_JSON)

Made with Cursor

[gemini-code-assist]

Code Review

This pull request introduces support for the external (file-transport) hash-shuffle variant in Ray Data's SHUFFLE_V2 strategy. It updates the physical planner and operator fusion rules to handle the external map and reduce operators, adds comprehensive correctness and smoke tests, and introduces a new benchmark script (bench_shuffle.py) along with release test configurations. The review feedback highlights several critical improvements for the benchmark script, such as ensuring distributed cleanup of local parquet files on worker nodes, raising an error on fleet scaling timeouts, preventing blocking when object store memory is untracked, and ensuring the parent directory for the results JSON is created to avoid late-stage crashes. Additionally, a minor docstring mismatch in a test helper was identified.

[gemini-code-assist]

In a distributed Ray cluster, calling write_parquet with a local path (like /tmp/shuffle_output) will write files to the local filesystems of the worker nodes where the write tasks run. Calling shutil.rmtree(output_path) on the driver node only deletes the directory on the driver, leaving the written parquet files on the worker nodes' local disks. This can lead to disk space leaks on the workers.\n\nTo fix this, you should clean up the directory on all active nodes in the cluster using a Ray remote function. The same applies to lines 128-129.

    if write_parquet:\n        try:\n            @ray.remote\n            def _cleanup(p):\n                import shutil\n                shutil.rmtree(p, ignore_errors=True)\n            ray.get([_cleanup.options(resources={f"node:{node['NodeManagerAddress']}": 0.01}).remote(output_path) for node in ray.nodes() if node.get("Alive")])\n        except Exception:\n            shutil.rmtree(output_path, ignore_errors=True)

[gemini-code-assist]

If the fleet does not reach the target_cpu count within timeout_s, the function currently prints the final CPU count and exits silently. This can lead to running the benchmark with fewer resources than expected, which will produce skewed or incorrect benchmark results. It is safer to raise a RuntimeError if the target CPU count is not met.

def _wait_for_fleet(target_cpu: int, timeout_s: int = 1200) -> None:\n    print(f"Waiting for {target_cpu} CPU to come online ...", flush=True)\n    deadline = time.perf_counter() + timeout_s\n    while time.perf_counter() < deadline:\n        cur = ray.cluster_resources().get("CPU", 0)\n        if cur >= target_cpu:\n            break\n        time.sleep(10)\n    final = ray.cluster_resources().get("CPU", 0)\n    if final < target_cpu:\n        raise RuntimeError(\n            f"Timed out waiting for {target_cpu} CPUs. Only {final:.0f} CPUs came online."\n        )\n    print(f"Fleet ready: {final:.0f} CPU", flush=True)

[gemini-code-assist]

If "object_store_memory" is not present in ray.cluster_resources() (which can happen in certain test or local environments), used_pct will default to 100% and the function will block/time out for the full timeout_s (180 seconds). We should check if "object_store_memory" is tracked in the cluster resources before waiting.

def wait_for_object_store_to_drain(threshold_pct=20, timeout_s=180, poll_s=5):\n    if "object_store_memory" not in ray.cluster_resources():\n        return\n    deadline = time.perf_counter() + timeout_s\n    while time.perf_counter() < deadline:\n        mem = ray.cluster_resources().get("object_store_memory", 1)\n        avail = ray.available_resources().get("object_store_memory", 0)\n        used_pct = (1 - avail / mem) * 100 if mem > 0 else 0\n        if used_pct < threshold_pct:\n            return\n        print(f"    draining object store ({used_pct:.0f}% used)...", flush=True)\n        time.sleep(poll_s)\n    print(f"    object store drain timed out after {timeout_s}s", flush=True)

[gemini-code-assist]

If args.result_json is specified with a path whose parent directory does not exist, writing the JSON will fail with a FileNotFoundError. Since this benchmark can run for a long time, crashing at the very end due to a missing directory is highly frustrating. We should ensure the parent directory is created before writing.

        os.makedirs(os.path.dirname(os.path.abspath(args.result_json)), exist_ok=True)\n        with open(args.result_json, "w") as f:\n            json.dump(slim, f, indent=2)

[gemini-code-assist]

The docstring for _drive_external_shuffle has a mismatch with the actual return values. It states that it returns (map_out, reduce_out, reduce_op) (3 elements) in the summary, and (map_output, reduce_output, reduce_op, map_op, upstream) (5 elements) in the Returns section. However, the actual implementation returns 4 elements: map_output, reduce_op, map_op, upstream. This should be corrected to avoid confusion.

    """Wire map→reduce, drive to completion, return (map_output, reduce_op, map_op, upstream).\n\n    Caller must shut down ``reduce_op`` (and its upstream map / InputDataBuffer)\n    via the returned reduce op's input chain, or use the try/finally in tests.\n    Returns ``(map_output, reduce_op, map_op, upstream)``.\n    """

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 5930aa4. Configure here.}

[cursor]

Orphaned m5.2xlarge compute config

Low Severity

Switching repartition_lineitem_* to fixed_size_all_to_all_compute.yaml leaves fixed_size_all_to_all_compute_m5_2xlarge.yaml with no remaining references, while its header still claims it is used by this release benchmark.

Additional Locations (1)

• release/release_data_tests.yaml#L447-L448

 

^{Reviewed by Cursor Bugbot for commit 5930aa4. Configure here.}