[nightly-test] add 4-nodes shuffle-data-loader test
diff --git a/release/nightly_tests/nightly_tests.yaml b/release/nightly_tests/nightly_tests.yaml
index d2be830a56845..56cafc554f948 100644
--- a/release/nightly_tests/nightly_tests.yaml
+++ b/release/nightly_tests/nightly_tests.yaml
@@ -283,6 +283,29 @@
       --data-dir s3://core-nightly-test/shuffle-data/
       --no-stats
 
+# Stress test shuffle_data_loader.
+- name: shuffle_data_loader_4_nodes
+  cluster:
+    app_config: shuffle_data_loader/shuffle_data_loader_app_config.yaml
+    compute_template: shuffle_data_loader/shuffle_data_loader_compute_4_nodes.yaml
+
+  run:
+    timeout: 7200
+    prepare: python wait_cluster.py 4 600
+    script: >
+      python shuffle_data_loader/benchmark.py
+      --num-rows 400_000_000
+      --num-files 50
+      --num-row-groups-per-file 5
+      --num-reducers 32 --num-trainers 16
+      --batch-size 250000
+      --cluster
+      --num-trials 1
+      --num-epochs 10
+      --max-concurrent-epochs 2
+      --data-dir s3://core-nightly-test/shuffle-data/
+      --no-stats
+
 - name: dask_on_ray_1tb_sort
   cluster:
     app_config: dask_on_ray/dask_on_ray_app_config.yaml
diff --git a/release/nightly_tests/shuffle_data_loader/benchmark.py b/release/nightly_tests/shuffle_data_loader/benchmark.py
index e20452649fa93..0627053272f7d 100644
--- a/release/nightly_tests/shuffle_data_loader/benchmark.py
+++ b/release/nightly_tests/shuffle_data_loader/benchmark.py
@@ -125,6 +125,11 @@ def run_trials(num_epochs,
     """
     print("Using from-memory shuffler.")
     all_stats = []
+    pg = ray.util.placement_group(
+        [{
+            "CPU": 0.1
+        } for _ in range(num_trainers)], strategy="SPREAD")
+    ray.get(pg.ready())
     if collect_stats:
         stats_collector = TrialStatsCollector.remote(
             num_epochs, len(filenames), num_reducers, num_trainers)
@@ -137,7 +142,7 @@ def run_trials(num_epochs,
         except AttributeError:
             # Python 3.6 doesn't support nullcontext().
             object_store_stats_collector = contextlib.suppress()
-    batch_consumer = BatchConsumer(num_trainers, num_epochs, None,
+    batch_consumer = BatchConsumer(num_trainers, num_epochs, pg,
                                    max_concurrent_epochs, stats_collector)
     # Wait until batch consumer actors have been created.
     batch_consumer.actors_ready()
diff --git a/release/nightly_tests/shuffle_data_loader/shuffle_data_loader_compute_4_nodes.yaml b/release/nightly_tests/shuffle_data_loader/shuffle_data_loader_compute_4_nodes.yaml
new file mode 100644
index 0000000000000..4c99d3f1d7056
--- /dev/null
+++ b/release/nightly_tests/shuffle_data_loader/shuffle_data_loader_compute_4_nodes.yaml
@@ -0,0 +1,25 @@
+cloud_id: cld_17WvYIBBkdgLwEUNcLeRAE
+region: us-west-2
+
+max_workers: 4
+
+aws:
+    IamInstanceProfile: {"Name": "ray-autoscaler-v1"}
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+            VolumeSize: 500
+
+head_node_type:
+    name: head_node
+    instance_type: i3.4xlarge
+    resources: {"object_store_memory": 53687091200}
+
+worker_node_types:
+    - name: worker_node
+      instance_type: m5.4xlarge
+      min_workers: 4
+      max_workers: 4
+      use_spot: false
+      resources:
+        cpu: 16

add shuffle_data_loader_4_nodes that run shuffle across 4 nodes with 400million rows.
Test plan:
-[x] schedule a run https://beta.anyscale.com/o/anyscale-internal/projects/prj_SVFGM5yBqK6DHCfLtRMryXHM/clusters/ses_Fie17wKrUbaLTRFKuDtQa6ee