Merge branch 'main' into doc/add_exp_replay_example

yanxi-chen · yanxi-chen · commit cea6f0fb6401 · 2025-10-27T15:34:53.000+08:00
diff --git a/docs/sphinx_doc/source/tutorial/example_mix_algo.md b/docs/sphinx_doc/source/tutorial/example_mix_algo.md
@@ -105,7 +105,7 @@ class MixSampleStrategy(SampleStrategy):
 
     async def sample(self, step: int) -> Tuple[Experiences, Dict, List]:
         metrics = {}
-        with Timer(metrics, "read_time"):
+        with Timer(metrics, "time/read_experience"):
             usual_exp_list = await self.usual_exp_buffer.read_async()
             for exp in usual_exp_list:
                 if exp.info is None:
@@ -131,7 +131,7 @@ class MixSampleStrategy(SampleStrategy):
             exp_list = usual_exp_list + expert_exp_list
             repr_samples = representative_sample(exp_list)
 
-        with Timer(metrics, "gather_time"):
+        with Timer(metrics, "time/gather_experience"):
             exps = Experiences.gather_experiences(
                 experiences=exp_list,
                 pad_token_id=self.pad_token_id,  # type: ignore [arg-type]
diff --git a/docs/sphinx_doc/source_zh/tutorial/example_mix_algo.md b/docs/sphinx_doc/source_zh/tutorial/example_mix_algo.md
@@ -97,7 +97,7 @@ class MixSampleStrategy(SampleStrategy):
 
     async def sample(self, step: int) -> Tuple[Experiences, Dict, List]:
         metrics = {}
-        with Timer(metrics, "read_time"):
+        with Timer(metrics, "time/read_experience"):
             usual_exp_list = await self.usual_exp_buffer.read_async()
             for exp in usual_exp_list:
                 if exp.info is None:
@@ -123,7 +123,7 @@ class MixSampleStrategy(SampleStrategy):
             exp_list = usual_exp_list + expert_exp_list
             repr_samples = representative_sample(exp_list)
 
-        with Timer(metrics, "gather_time"):
+        with Timer(metrics, "time/gather_experience"):
             exps = Experiences.gather_experiences(
                 experiences=exp_list,
                 pad_token_id=self.pad_token_id,  # type: ignore [arg-type]
diff --git a/tests/algorithm/advantage_fn_test.py b/tests/algorithm/advantage_fn_test.py
@@ -107,8 +107,8 @@ def test_grpo_reward_std(self):
 
         exps, metrics = advantage_fn(exps)
         self.assertEqual(len(exps), 0)
-        self.assertIn("group_advantages/skipped_count/mean", metrics)
-        self.assertEqual(metrics["group_advantages/skipped_count/mean"], 5)
+        self.assertIn("filtered_count", metrics)
+        self.assertEqual(metrics["filtered_count"], 15)
 
     def test_grpo_correct_bias(self):
         advantage_fn_cls = ADVANTAGE_FN.get("grpo")
diff --git a/tests/buffer/experience_pipeline_test.py b/tests/buffer/experience_pipeline_test.py
@@ -73,7 +73,7 @@ async def test_experience_pipeline(self):
         experiences = get_experiences(task_num=task_num, repeat_times=repeat_times)
         metrics = await pipeline.process.remote(experiences)
         self.assertEqual(
-            metrics["pipeline/experience_count"], task_num * (repeat_times - 1)
+            metrics["experience_pipeline/experience_count"], task_num * (repeat_times - 1)
         )  # first experience of each task will be filtered out by the reward filter
 
         # tests
diff --git a/tests/buffer/queue_test.py b/tests/buffer/queue_test.py
@@ -326,7 +326,7 @@ async def test_priority_queue_reuse_count_control(self):
             path=BUFFER_FILE_PATH,
             replay_buffer=ReplayBufferConfig(
                 enable=True,
-                priority_fn="linear_decay_use_count_control_randomization",
+                priority_fn="decay_limit_randomization",
                 reuse_cooldown_time=0.5,
                 priority_fn_args={"decay": 1.2, "use_count_limit": 2, "sigma": 0.0},
             ),
diff --git a/tests/explorer/explorer_test.py b/tests/explorer/explorer_test.py
@@ -108,8 +108,8 @@ def test_explorer(self):
         eval_metrics = parser.metric_list("eval")
         self.assertTrue(len(eval_metrics) == 0)
         self.assertEqual(parser.metric_max_step(rollout_metrics[0]), 4)
-        self.assertTrue(parser.metric_exist("pipeline/experience_count"))
-        experience_counts = parser.metric_values("pipeline/experience_count")
+        self.assertTrue(parser.metric_exist("experience_pipeline/experience_count"))
+        experience_counts = parser.metric_values("experience_pipeline/experience_count")
         self.assertTrue(len(experience_counts) == 4)
         for count in experience_counts:
             self.assertTrue(count >= 0)
diff --git a/tests/trainer/trainer_test.py b/tests/trainer/trainer_test.py
@@ -819,7 +819,7 @@ def test_trainer(self):
         self.assertTrue(len(rollout_metrics) > 0)
         self.assertEqual(parser.metric_max_step(rollout_metrics[0]), 4)
         self.assertEqual(
-            parser.metric_values("pipeline/experience_count")[1], 16
+            parser.metric_values("experience_pipeline/experience_count")[1], 16
         )  # 16 rft experiences
         # test actor metrics
         actor_metrics = parser.metric_list("actor")
diff --git a/trinity/algorithm/advantage_fn/advantage_fn.py b/trinity/algorithm/advantage_fn/advantage_fn.py
@@ -76,10 +76,7 @@ def process(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]:
         for group_id, group_exps in exp_groups.items():
             group_exps, group_metrics = self.calculate_group_advantage(group_id, group_exps)
             metric_list.append(group_metrics)
-        try:
-            metrics = gather_metrics(metric_list, "group_advantages")
-        except ValueError:
-            metrics = {}  # empty metric list causes ValueError, ignore it
+        metrics = gather_metrics(metric_list, "group_advantages")
         exps = [exp for group in exp_groups.values() for exp in group]  # Flatten the list
         return exps, metrics
 
diff --git a/trinity/algorithm/advantage_fn/grpo_advantage.py b/trinity/algorithm/advantage_fn/grpo_advantage.py
@@ -214,11 +214,11 @@ def process(self, exps):
                 group_id, group_exps, precomputed_std=precomputed_std
             )
             metric_list.append(group_metrics)
-        try:
-            # TODO: sum skipped count
-            metrics = gather_metrics(metric_list, "group_advantages")
-        except ValueError:
-            metrics = {}  # empty metric list causes ValueError, ignore it
+
+        # Update the filtered_count metric
+        filtered_count = sum(metric.pop("skipped_count", 0) for metric in metric_list)
+        metrics = gather_metrics(metric_list, "group_advantages")
+        metrics["filtered_count"] = filtered_count
         if self.duplicate_experiences and self.std_threshold is not None:
             exps = self._duplicate_experiences(exp_groups)
         else:
diff --git a/trinity/algorithm/advantage_fn/multi_step_grpo_advantage.py b/trinity/algorithm/advantage_fn/multi_step_grpo_advantage.py
@@ -142,11 +142,8 @@ def process(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]:
                 cnt += len(exps)
                 result_exps.extend(exps)
 
-        try:
-            metrics = gather_metrics(metric_list, "group_advantages")
-            metrics["experience_count"] = cnt
-        except ValueError:
-            metrics = {}  # empty metric list causes ValueError, ignore it
+        metrics = gather_metrics(metric_list, "group_advantages")
+        metrics["experience_count"] = cnt
         return result_exps, metrics
 
     def __call__(self, exps, **kwargs):
diff --git a/trinity/algorithm/sample_strategy/mix_sample_strategy.py b/trinity/algorithm/sample_strategy/mix_sample_strategy.py
@@ -59,7 +59,7 @@ def __init__(self, buffer_config: BufferConfig, **kwargs):
 
     async def sample(self, step: int) -> Tuple[Experiences, Dict, List]:
         metrics = {}
-        with Timer(metrics, "read_time"):
+        with Timer(metrics, "time/read_experience"):
             usual_exp_list = await self.usual_exp_buffer.read_async()
             for exp in usual_exp_list:
                 if exp.info is None:
@@ -85,7 +85,7 @@ async def sample(self, step: int) -> Tuple[Experiences, Dict, List]:
             exp_list = usual_exp_list + expert_exp_list
             repr_samples = representative_sample(exp_list)
 
-        with Timer(metrics, "gather_time"):
+        with Timer(metrics, "time/gather_experience"):
             exps = Experiences.gather_experiences(
                 experiences=exp_list,
                 pad_token_id=self.pad_token_id,  # type: ignore [arg-type]
diff --git a/trinity/algorithm/utils.py b/trinity/algorithm/utils.py
@@ -102,4 +102,5 @@ def prefix_metrics(src_metrics: dict, prefix: str, dst_metrics: dict = None) ->
         dst_metrics = {}
     for k, v in src_metrics.items():
         dst_metrics[f"{prefix}/{k}"] = v
+
     return dst_metrics
diff --git a/trinity/buffer/pipelines/experience_pipeline.py b/trinity/buffer/pipelines/experience_pipeline.py
@@ -1,3 +1,4 @@
+import time
 import traceback
 from typing import Dict, List, Optional
 
@@ -14,6 +15,7 @@
 from trinity.common.experience import Experience
 from trinity.utils.log import get_logger
 from trinity.utils.plugin_loader import load_plugins
+from trinity.utils.timer import Timer
 
 
 def get_input_buffers(pipeline_config: ExperiencePipelineConfig) -> Dict:
@@ -104,26 +106,33 @@ async def process(self, exps: List[Experience]) -> Dict:
         Returns:
             Dict: A dictionary containing metrics collected during the processing of experiences.
         """
+        st = time.time()
         if self.input_store is not None:
             await self.input_store.write_async(exps)
 
         metrics = {}
 
         # Process experiences through operators
-        for operator in self.operators:
-            exps, metric = operator.process(exps)
-            metrics.update(metric)
-
+        for idx, operator in enumerate(self.operators):
+            with Timer(
+                metrics, f"time/experience_pipeline/operator/{idx}_{operator.__class__.__name__}"
+            ):
+                exps, metric = operator.process(exps)
+                metrics.update(metric)
         metrics["experience_count"] = len(exps)
 
         # Write processed experiences to output buffer
-        await self.output.write_async(exps)
+        with Timer(metrics, "time/experience_pipeline/write"):
+            await self.output.write_async(exps)
+        metrics["time/experience_pipeline/total"] = time.time() - st
 
         # prefix metrics keys with 'pipeline/'
         result_metrics = {}
         for key, value in metrics.items():
-            if isinstance(value, (int, float)):
-                result_metrics[f"pipeline/{key}"] = float(value)
+            if key.startswith("time/"):
+                result_metrics[key] = value
+            elif isinstance(value, (int, float)):
+                result_metrics[f"experience_pipeline/{key}"] = float(value)
         if SELECTOR_METRIC in metrics:
             result_metrics[SELECTOR_METRIC] = metrics[SELECTOR_METRIC]
 
diff --git a/trinity/buffer/storage/queue.py b/trinity/buffer/storage/queue.py
@@ -5,8 +5,7 @@
 from abc import ABC, abstractmethod
 from collections import deque
 from copy import deepcopy
-from functools import partial
-from typing import List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import ray
@@ -28,48 +27,82 @@ def is_json_file(path: str) -> bool:
 
 
 PRIORITY_FUNC = Registry("priority_fn")
-"""
-Each priority_fn,
-    Args:
-        item: List[Experience], assume that all experiences in it have the same model_version and use_count
-        kwargs: storage_config.replay_buffer_kwargs (except priority_fn)
-    Returns:
-        priority: float
-        put_into_queue: bool, decide whether to put item into queue
-Note that put_into_queue takes effect both for new item from the explorer and for item sampled from the buffer.
-"""
+
+
+class PriorityFunction(ABC):
+    """
+    Each priority_fn,
+        Args:
+            item: List[Experience], assume that all experiences in it have the same model_version and use_count
+            priority_fn_args: Dict, the arguments for priority_fn
+
+        Returns:
+            priority: float
+            put_into_queue: bool, decide whether to put item into queue
+
+    Note that put_into_queue takes effect both for new item from the explorer and for item sampled from the buffer.
+    """
+
+    @abstractmethod
+    def __call__(self, item: List[Experience]) -> Tuple[float, bool]:
+        """Calculate the priority of item."""
+
+    @classmethod
+    @abstractmethod
+    def default_config(cls) -> Dict:
+        """Return the default config."""
 
 
 @PRIORITY_FUNC.register_module("linear_decay")
-def linear_decay_priority(
-    item: List[Experience],
-    decay: float = 2.0,
-) -> Tuple[float, bool]:
+class LinearDecayPriority(PriorityFunction):
     """Calculate priority by linear decay.
 
     Priority is calculated as `model_version - decay * use_count. The item is always put back into the queue for reuse (as long as `reuse_cooldown_time` is not None).
     """
-    priority = float(item[0].info["model_version"] - decay * item[0].info["use_count"])
-    put_into_queue = True
-    return priority, put_into_queue
-
-
-@PRIORITY_FUNC.register_module("linear_decay_use_count_control_randomization")
-def linear_decay_use_count_control_priority(
-    item: List[Experience],
-    decay: float = 2.0,
-    use_count_limit: int = 3,
-    sigma: float = 0.0,
-) -> Tuple[float, bool]:
+
+    def __init__(self, decay: float = 2.0):
+        self.decay = decay
+
+    def __call__(self, item: List[Experience]) -> Tuple[float, bool]:
+        priority = float(item[0].info["model_version"] - self.decay * item[0].info["use_count"])
+        put_into_queue = True
+        return priority, put_into_queue
+
+    @classmethod
+    def default_config(cls) -> Dict:
+        return {
+            "decay": 2.0,
+        }
+
+
+@PRIORITY_FUNC.register_module("decay_limit_randomization")
+class LinearDecayUseCountControlPriority(PriorityFunction):
     """Calculate priority by linear decay, use count control, and randomization.
 
     Priority is calculated as `model_version - decay * use_count`; if `sigma` is non-zero, priority is further perturbed by random Gaussian noise with standard deviation `sigma`.  The item will be put back into the queue only if use count does not exceed `use_count_limit`.
     """
-    priority = float(item[0].info["model_version"] - decay * item[0].info["use_count"])
-    if sigma > 0.0:
-        priority += float(np.random.randn() * sigma)
-    put_into_queue = item[0].info["use_count"] < use_count_limit if use_count_limit > 0 else True
-    return priority, put_into_queue
+
+    def __init__(self, decay: float = 2.0, use_count_limit: int = 3, sigma: float = 0.0):
+        self.decay = decay
+        self.use_count_limit = use_count_limit
+        self.sigma = sigma
+
+    def __call__(self, item: List[Experience]) -> Tuple[float, bool]:
+        priority = float(item[0].info["model_version"] - self.decay * item[0].info["use_count"])
+        if self.sigma > 0.0:
+            priority += float(np.random.randn() * self.sigma)
+        put_into_queue = (
+            item[0].info["use_count"] < self.use_count_limit if self.use_count_limit > 0 else True
+        )
+        return priority, put_into_queue
+
+    @classmethod
+    def default_config(cls) -> Dict:
+        return {
+            "decay": 2.0,
+            "use_count_limit": 3,
+            "sigma": 0.0,
+        }
 
 
 class QueueBuffer(ABC):
@@ -168,7 +201,10 @@ def __init__(
         self.capacity = capacity
         self.item_count = 0
         self.priority_groups = SortedDict()  # Maps priority -> deque of items
-        self.priority_fn = partial(PRIORITY_FUNC.get(priority_fn), **(priority_fn_args or {}))
+        priority_fn_cls = PRIORITY_FUNC.get(priority_fn)
+        kwargs = priority_fn_cls.default_config()
+        kwargs.update(priority_fn_args or {})
+        self.priority_fn = priority_fn_cls(**kwargs)
         self.reuse_cooldown_time = reuse_cooldown_time
         self._condition = asyncio.Condition()  # For thread-safe operations
         self._closed = False
diff --git a/trinity/explorer/explorer.py b/trinity/explorer/explorer.py
diff --git a/trinity/manager/config_manager.py b/trinity/manager/config_manager.py
diff --git a/trinity/manager/config_registry/buffer_config_manager.py b/trinity/manager/config_registry/buffer_config_manager.py
diff --git a/trinity/service/data_juicer/client.py b/trinity/service/data_juicer/client.py
diff --git a/trinity/trainer/trainer.py b/trinity/trainer/trainer.py
diff --git a/trinity/trainer/verl_trainer.py b/trinity/trainer/verl_trainer.py
diff --git a/trinity/utils/monitor.py b/trinity/utils/monitor.py
