[Data] Add GCE variant for shuffle_data_loader (ray-project#34632)

This PR configures BuildKite to run the shuffle_data_loader Data release tests on GCE. This test was previously excluded in ray-project#34105. --------- Signed-off-by: Balaji Veeramani <[email protected]> Signed-off-by: Jack He <[email protected]>
ProjectsByJackHe · May 4, 2023 · ea74a49 · ea74a49
1 parent 9e1603a
commit ea74a49
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 20 deletions.
diff --git a/release/nightly_tests/dataset/dataset_shuffle_data_loader.py b/release/nightly_tests/dataset/dataset_shuffle_data_loader.py
@@ -5,13 +5,20 @@
 
 import ray
 
+from pyarrow import fs
 import numpy as np
 import torch
 
-PATH = [
-    f"s3://shuffling-data-loader-benchmarks/data/input_data_{i}.parquet.snappy"
-    for i in range(0, 25)
-]
+PATHS = {
+    "aws": [
+        f"s3://shuffling-data-loader-benchmarks/data/input_data_{i}.parquet.snappy"
+        for i in range(0, 25)
+    ],
+    "gcp": [
+        f"gcs://shuffling-data-loader-benchmarks/data/input_data_{i}.parquet.snappy"
+        for i in range(0, 25)
+    ],
+}
 
 
 def create_parser():
@@ -26,6 +33,7 @@ def create_parser():
     )
     parser.add_argument("--num-workers", type=int, default=4)
     parser.add_argument("--repeat-times", type=int, default=16)
+    parser.add_argument("--cloud", type=str, choices=["aws", "gcp"])
     return parser
 
 
@@ -83,9 +91,14 @@ def create_torch_iterator(split, batch_size, rank=None):
     return torch_iterator
 
 
-def create_dataset(filenames, repeat_times):
+def create_dataset(filenames, repeat_times, cloud):
+    if cloud == "gcp":
+        filesystem = fs.GcsFileSystem()
+    else:
+        filesystem = None
+
     pipeline = (
-        ray.data.read_parquet(list(filenames))
+        ray.data.read_parquet(list(filenames), filesystem=filesystem)
         .repeat(times=repeat_times)
         .random_shuffle_each_window()
     )
@@ -100,7 +113,7 @@ def create_dataset(filenames, repeat_times):
 
     start = time.time()
 
-    pipeline = create_dataset(PATH, args.repeat_times)
+    pipeline = create_dataset(PATHS[args.cloud], args.repeat_times, args.cloud)
     splits = pipeline.split(args.num_workers)
 
     @ray.remote(num_gpus=1)

diff --git a/release/nightly_tests/dataset/shuffle_app_config.yaml b/release/nightly_tests/dataset/shuffle_app_config.yaml
@@ -3,7 +3,6 @@ base_image: {{ env["RAY_IMAGE_ML_NIGHTLY_GPU"] | default("anyscale/ray-ml:nightl
 python:
   pip_packages:
     - boto3
-    - pyarrow<7.0.0
   conda_packages: []
 
 post_build_cmds:

diff --git a/release/release_tests.yaml b/release/release_tests.yaml
@@ -3371,14 +3371,14 @@
   cluster:
     cluster_env: app_config.yaml
     cluster_compute: multi_node_checkpointing_compute_config.yaml
-  
+
   run:
     timeout: 3600
     script: pytest checkpointing_tests/test_learner_group_checkpointing.py
 
     wait_for_nodes:
       num_nodes: 3
-  
+
   alert: default
 
 
@@ -4982,17 +4982,17 @@
 
   run:
     timeout: 1800
-    script: python dataset_shuffle_data_loader.py
+    script: python dataset_shuffle_data_loader.py --cloud aws
 
-  # TODO: Port s3://shuffling-data-loader-benchmarks/ to GCS.
-  # variations:
-  #   - __suffix__: aws
-  #   - __suffix__: gce
-  #     env: gce
-  #     frequency: manual
-  #     cluster:
-  #       cluster_env: shuffle_app_config.yaml
-  #       cluster_compute: shuffle_compute_gce.yaml
+  variations:
+    - __suffix__: aws
+    - __suffix__: gce
+      env: gce
+      frequency: manual
+      cluster:
+        cluster_compute: shuffle_compute_gce.yaml
+      run:
+        script: python dataset_shuffle_data_loader.py --cloud gcp
 
 - name: parquet_metadata_resolution
   group: data-tests