From 8acd22b473667ce3a4bb5aaf01bc0383770c0ef2 Mon Sep 17 00:00:00 2001
From: aireenmei <aireenmei@gmail.com>
Date: Thu, 2 Jul 2026 00:22:01 +0000
Subject: [PATCH] Move data_processing_test to cpu

---
 tests/unit/grain_data_processing_test.py | 9 +++++++++
 tests/unit/hf_data_processing_test.py    | 2 ++
 tests/unit/tfds_data_processing_test.py  | 1 +
 3 files changed, 12 insertions(+)

diff --git a/tests/unit/grain_data_processing_test.py b/tests/unit/grain_data_processing_test.py
index dc129b5ebc..ae0507ed64 100644
--- a/tests/unit/grain_data_processing_test.py
+++ b/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).
 
diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py
index 262c56ff9b..93564514a3 100644
--- a/tests/unit/hf_data_processing_test.py
+++ b/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):
diff --git a/tests/unit/tfds_data_processing_test.py b/tests/unit/tfds_data_processing_test.py
index fe1bb75f24..c9700d701f 100644
--- a/tests/unit/tfds_data_processing_test.py
+++ b/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):