Merge pull request #4333 from AI-Hypercomputer:aireen/grain_test_cpu

PiperOrigin-RevId: 943539777

Author: Google-ML-Automation

tests/unit/grain_data_processing_test.py

@@ -181,6 +181,7 @@ def _make_config(self, **overrides):
     return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs)
 
 
+@pytest.mark.cpu_only
 class GrainArrayRecordProcessingTest(
     _GrainArrayRecordSetup, GrainDeterminismMixin, GrainBaseProcessingTest, unittest.TestCase
 ):
@@ -201,6 +202,7 @@ def test_batch_determinism(self):
     super().test_batch_determinism()
 
 
+@pytest.mark.cpu_only
 class GrainArrayRecordProcessingWithMultiSourceBlendingTest(
     _GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase
 ):
@@ -211,6 +213,7 @@ def setUp(self):
     self.config = self._make_config(grain_train_files=train_files_weighted)
 
 
+@pytest.mark.cpu_only
 class GrainArrayRecordProcessingWithMixtureConfigTest(_GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase):
 
   def setUp(self):
@@ -270,6 +273,7 @@ def setUp(self):
     self.config = self._make_config(grain_ram_budget_mb=512, grain_worker_count=-1)  # Enable auto-tuning
 
 
+@pytest.mark.cpu_only
 class GrainArrayRecordTiktokenTest(_GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase):
   """Test grain data processing with tiktoken tokenizer."""
 
@@ -281,6 +285,7 @@ def setUp(self):
     )
 
 
+@pytest.mark.cpu_only
 class GrainArrayRecordHFTokenizerTest(_GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase):
   """Test grain data processing with HuggingFace tokenizer."""
 
@@ -292,6 +297,7 @@ def setUp(self):
     )
 
 
+@pytest.mark.cpu_only
 class GrainArrayRecordBestFitPackingTest(_GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase):
   """Test grain data processing with best_fit packing strategy."""
 
@@ -381,6 +387,7 @@ def _make_config(self, **overrides):
     return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs)
 
 
+@pytest.mark.cpu_only
 class GrainParquetProcessingTest(_GrainParquetSetup, GrainDeterminismMixin, GrainBaseProcessingTest, unittest.TestCase):
   """Test grain data processing with Parquet format.
 
@@ -519,6 +526,7 @@ def _make_config(self, **overrides):
     return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs)
 
 
+@pytest.mark.cpu_only
 class GrainTFRecordProcessingTest(_GrainTFRecordSetup, GrainDeterminismMixin, GrainBaseProcessingTest, unittest.TestCase):
   """Test grain data processing with TFRecord format.
 
@@ -533,6 +541,7 @@ def setUpClass(cls):
     super().setUpClass()
 
 
+@pytest.mark.cpu_only
 class GrainTFRecordPreTokenizedProcessingTest(_GrainTFRecordSetup, GrainBaseProcessingTest, unittest.TestCase):
   """Test grain data processing with a pre-tokenized TFRecord dataset (tokenize_train_data=False).
 

tests/unit/hf_data_processing_test.py

@@ -19,6 +19,7 @@
 import os.path
 
 import jax
+import pytest
 from jax.sharding import Mesh
 from jax.experimental import mesh_utils
 
@@ -30,6 +31,7 @@
 from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory
 
 
+@pytest.mark.cpu_only
 class HfDataProcessingTest(unittest.TestCase):
 
   def setUp(self):

tests/unit/tfds_data_processing_test.py

@@ -33,6 +33,7 @@
 from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory
 
 
+@pytest.mark.cpu_only
 class TfdsDataProcessingTest(unittest.TestCase):
 
   def setUp(self):