Add config_validator.py and refactor config (#487)

Author: chenyushuo

.github/workflows/docker/docker-compose.yaml

@@ -1,6 +1,8 @@
 services:
   trinity-node-1:
     image: trinity-rft-unittest:20260115
+    cap_add:
+      - SYS_PTRACE
     pull_policy: never
     command: bash -c "source /opt/venv/bin/activate && uv pip install -e .[dev] && ray start --head --dashboard-host 0.0.0.0 --include-dashboard true --block"
     environment:
@@ -31,6 +33,8 @@ services:
 
   trinity-node-2:
     image: trinity-rft-unittest:20260115
+    cap_add:
+      - SYS_PTRACE
     pull_policy: never
     command: bash -c "source /opt/venv/bin/activate && uv pip install -e .[dev] && ray start --address=trinity-node-1:6379 --block"
     environment:

pyproject.toml

@@ -41,8 +41,8 @@ dependencies = [
     "jsonlines",
     "sortedcontainers",
     "word2number",
-    "transformers",
-    "datasets",
+    "transformers>=4.51.0",
+    "datasets>=4.0.0",
 ]
 
 [project.scripts]

tests/common/config_test.py

@@ -91,7 +91,7 @@ def test_update_config_from_ray_cluster(self):
         config.cluster.node_num = None
         config.cluster.gpu_per_node = None
 
-        config._update_config_from_ray_cluster()
+        config.check_and_update()
         self.assertEqual(config.cluster.node_num, 2)
         self.assertEqual(config.cluster.gpu_per_node, 2)
 

tests/tools.py

@@ -1,3 +1,4 @@
+import gc
 import os
 import unittest
 from collections import defaultdict
@@ -243,21 +244,51 @@ def metric_list(self, metric_prefix: str) -> List[str]:
         return [name for name in self._metrics if name.startswith(metric_prefix)]
 
 
-class RayUnittestBase(unittest.TestCase):
+class RayCleanupPlugin:
+    @classmethod
+    def _cleanup_ray_data_state(cls):
+        """clean up the global states of Ray Data"""
+        try:
+            # reset execution context
+            if hasattr(ray.data._internal.execution.streaming_executor, "_execution_context"):
+                ray.data._internal.execution.streaming_executor._execution_context = None
+
+            # trigger gc.collect() on all workers in the cluster
+            ray._private.internal_api.global_gc()
+
+            # clean up stats manager
+            from ray.data._internal.stats import StatsManager
+
+            if hasattr(StatsManager, "_instance"):
+                StatsManager._instance = None
+
+        except Exception:
+            pass
+
+
+class RayUnittestBase(unittest.TestCase, RayCleanupPlugin):
     @classmethod
     def setUpClass(cls):
         ray.init(ignore_reinit_error=True, namespace="trinity_unittest")
 
+        # erase existing resources
+        cls._cleanup_ray_data_state()
+        gc.collect()
+
     @classmethod
     def tearDownClass(cls):
         ray.shutdown(_exiting_interpreter=True)
 
 
-class RayUnittestBaseAsync(unittest.IsolatedAsyncioTestCase):
+class RayUnittestBaseAsync(unittest.IsolatedAsyncioTestCase, RayCleanupPlugin):
     @classmethod
     def setUpClass(cls):
         ray.init(ignore_reinit_error=True, namespace="trinity_unittest")
 
+        # erase existing resources
+        cls._cleanup_ray_data_state()
+        gc.collect()
+
     @classmethod
     def tearDownClass(cls):
         ray.shutdown(_exiting_interpreter=True)

tests/trainer/trainer_test.py

@@ -350,6 +350,7 @@ def test_trainer(self, mock_load):
 
         with self.assertRaises(Exception):
             run(config_path="dummy.yaml")
+        ray.shutdown(_exiting_interpreter=True)
 
         stage_configs = [cfg.check_and_update() for cfg in deepcopy(self.config)]
 
@@ -372,6 +373,7 @@ def test_trainer(self, mock_load):
 
         self.config.stages[1].buffer.explorer_input.taskset.path = old_taskset_path
         mock_load.return_value = deepcopy(self.config)
+        ray.init(ignore_reinit_error=True, namespace=self.config.ray_namespace)
         run(config_path="dummy.yaml")
 
         # grpo stage