let's run ray tests separately

Author: dlwh

.github/workflows/run_ray_tests.yaml

@@ -0,0 +1,28 @@
+name: Run tests that use ray
+
+on: [push]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10"]
+        jax-version: ["0.4.23"]
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install flake8 pytest
+          pip install . "jax[cpu]==${{ matrix.jax-version }}" "jaxlib==${{ matrix.jax-version }}"
+          pip install soundfile librosa
+      - name: Run entry tests with pytest
+        run: |
+          XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=$(pwd)/tests:$(pwd)/src:$(pwd):. pytest tests -m ray

.github/workflows/run_tests.yaml

@@ -25,4 +25,4 @@ jobs:
           pip install soundfile librosa
       - name: Test with pytest
         run: |
-          XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. pytest tests -m "not entry and not slow"
+          XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. pytest tests -m "not entry and not slow and not ray"

tests/test_shard_cache.py

@@ -63,6 +63,7 @@ def simple_process(processor, source):
     return result
 
 
+@pytest.mark.ray
 def test_cache_simple():
     td = tempfile.TemporaryDirectory()
     with td as tmpdir:
@@ -73,6 +74,7 @@ def test_cache_simple():
         assert list(ray_ds) == list(simple_processed)
 
 
+@pytest.mark.ray
 def test_cache_remembers_its_cached():
     directory = tempfile.TemporaryDirectory()
     with directory as tmpdir:
@@ -101,6 +103,7 @@ class _CustomException(Exception):
     pass
 
 
+@pytest.mark.ray
 def test_cache_recover_from_crash():
     class CrashingShardSource(ShardedDataset[List[int]]):
         def __init__(self, crash_point: int):
@@ -144,6 +147,7 @@ def open_shard_at_row(self, shard_name: str, row: int) -> Iterator[List[int]]:
         assert len(list(reader1)) == 40
 
 
+@pytest.mark.ray
 def test_no_hang_if_empty_shard_source():
     class EmptyShardSource(ShardedDataset[List[int]]):
         @property
@@ -158,6 +162,7 @@ def open_shard_at_row(self, shard_name: str, row: int) -> Iterator[List[int]]:
         assert list(reader) == []
 
 
+@pytest.mark.ray
 def test_chunk_ordering_is_correct_with_slow_shards():
     class SlowShardSource(ShardedDataset[List[int]]):
         @property
@@ -245,6 +250,7 @@ def back_to_py(batch: pa.RecordBatch):
         cache.await_finished(timeout=10)
 
 
+@pytest.mark.ray
 def test_shard_cache_crashes_if_processor_throws():
     class ThrowingProcessor(BatchProcessor[Sequence[int]]):
         def __call__(self, batch: Sequence[Sequence[int]]) -> pa.RecordBatch:
@@ -263,6 +269,7 @@ def num_cpus(self) -> int:
             build_cache(tmpdir, SimpleShardSource(), ThrowingProcessor(), await_finished=True)
 
 
+@pytest.mark.ray
 def test_map_batches_and_map_shard_cache():
     td = tempfile.TemporaryDirectory()
     with td as tmpdir:
@@ -289,6 +296,7 @@ def composite_fn(list):
         assert ray_entries == list(simple_processed)
 
 
+@pytest.mark.ray
 def test_serial_cache_writer():
     with tempfile.TemporaryDirectory() as tmpdir1, tempfile.TemporaryDirectory() as tmpdir2:
         source = SimpleShardSource(num_shards=4)

tests/test_tokenized_document_cache.py

@@ -26,6 +26,7 @@ def teardown_module(module):
     ray.shutdown()
 
 
+@pytest.mark.ray
 def test_index_empty_file():
     with tempfile.TemporaryDirectory() as tmpdir:
         empty_dataset = [""]
@@ -43,6 +44,7 @@ def test_index_empty_file():
             assert chunk["input_ids"].size == 0
 
 
+@pytest.mark.ray
 def test_index_no_files():
     with tempfile.TemporaryDirectory() as tmpdir:
         empty_dataset = []
@@ -60,6 +62,7 @@ def test_index_no_files():
             pytest.fail("Should not have any chunks")
 
 
+@pytest.mark.ray
 def test_doc_cache_reproduces_data_one_batch_per_shard():
     def doc_i(i: int):
         return BatchEncoding(data=dict(input_ids=[list(range(10 * i, 10 * (i + 1)))]))
@@ -96,6 +99,7 @@ def open_shard_at_row(self, shard_name: str, row: int):
             assert as_listed == docs[i]
 
 
+@pytest.mark.ray
 @pytest.mark.parametrize("batch_size", list([1, 2, 3, 8]))
 def test_doc_cache_reproduces_data_multi_docs_per_batch_sharded(batch_size):
     def batch_docs(doc_ids):
@@ -130,6 +134,7 @@ def list_in_list(a, b):
             assert found
 
 
+@pytest.mark.ray
 def test_doc_cache_sharding():
     def doc_i(i: int):
         return BatchEncoding(data=dict(input_ids=[list(range(10 * i, 10 * (i + 1)))]))