revert metric modification

Author: chenyushuo

tests/explorer/explorer_test.py

@@ -24,10 +24,9 @@
 from trinity.buffer import get_buffer_reader
 from trinity.cli.launcher import explore, run_stage
 from trinity.common.config import ExperienceBufferConfig, InferenceModelConfig
-from trinity.common.constants import PLUGIN_DIRS_ENV_VAR, StorageType
+from trinity.common.constants import StorageType
 from trinity.explorer.explorer import Explorer
 from trinity.manager.state_manager import StateManager
-from trinity.utils.plugin_loader import load_plugins
 
 
 class BaseExplorerCase(RayUnittestBase):
@@ -46,22 +45,6 @@ def setUp(self):
         self.config.explorer.eval_interval = 4
 
 
-class TestExplorerCountdownMaxRepeatTimes(BaseExplorerCase):
-    def test_explorer(self):
-        self.config.buffer.explorer_input.taskset = get_unittest_dataset_config("countdown")
-        self.config.buffer.explorer_input.taskset.default_workflow_type = "custom_workflow"
-        self.config.algorithm.repeat_times = 4
-        self.config.explorer.max_repeat_times_per_runner = 3
-        self.config.name = f"explore-eval-{datetime.now().strftime('%Y%m%d%H%M%S')}"
-        self.config.check_and_update()
-        os.environ[PLUGIN_DIRS_ENV_VAR] = os.path.join("tests", "utils", "plugins")
-        load_plugins()
-        explore(self.config)
-        parser = TensorBoardParser(os.path.join(self.config.monitor.cache_dir, "tensorboard"))
-        custom_metric_mean = parser.metric_values("rollout/custom_metric/mean")
-        self.assertEqual(custom_metric_mean, [0.75] * 8)
-
-
 class TestExplorerCountdownEval(BaseExplorerCase):
     def test_explorer(self):
         self.config.buffer.explorer_input.taskset = get_unittest_dataset_config("countdown")

tests/utils/plugins/my_workflow.py

@@ -1,7 +1,6 @@
 from typing import List
 
 from trinity.common.workflows import WORKFLOWS, Workflow
-from trinity.common.workflows.workflow import MathWorkflow
 
 
 @WORKFLOWS.register_module("my_workflow")
@@ -18,12 +17,3 @@ def set_repeat_times(self, repeat_times, run_id_base):
 
     def run(self) -> List:
         return ["Hello world", "Hi"]
-
-
-@WORKFLOWS.register_module("custom_workflow")
-class CustomWorkflow(MathWorkflow):
-    def run(self):
-        responses = super().run()
-        for i, response in enumerate(responses):
-            response.metrics["custom_metric"] = i
-        return responses

trinity/explorer/explorer.py

@@ -26,7 +26,6 @@
 from trinity.common.models import create_inference_models
 from trinity.common.models.utils import get_checkpoint_dir_with_step_num
 from trinity.explorer.scheduler import Scheduler
-from trinity.explorer.workflow_runner import group_metrics
 from trinity.manager.state_manager import StateManager
 from trinity.manager.synchronizer import Synchronizer
 from trinity.utils.annotations import Experimental
@@ -363,7 +362,7 @@ async def _finish_explore_step(self, step: int, model_version: int) -> None:
         self.taskset.update(pipeline_metrics)
         metric.update(pipeline_metrics)
         if statuses:
-            metric.update(gather_metrics(group_metrics(statuses), "rollout"))
+            metric.update(gather_metrics([status.metric for status in statuses], "rollout"))
             self.monitor.log(metric, step=step)
 
     async def _finish_eval_step(self, step: Optional[int] = None, prefix: str = "eval") -> None:
@@ -377,7 +376,11 @@ async def _finish_eval_step(self, step: Optional[int] = None, prefix: str = "eva
                 return
             self.pending_eval_tasks.popleft()
             eval_results, _ = await self.scheduler.get_results(f"{step}/{eval_task_name}")
-            metric.update(gather_metrics(group_metrics(eval_results), f"{prefix}/{eval_task_name}"))
+            metric.update(
+                gather_metrics(
+                    [status.metric for status in eval_results], f"{prefix}/{eval_task_name}"
+                )
+            )
         if self.eval_start_time is not None:
             metric.update({"time/eval": time.time() - self.eval_start_time})
             self.eval_start_time = None

trinity/explorer/scheduler.py

@@ -381,21 +381,23 @@ async def get_results(
         statuses = []
         experiences = []
         completed_queue = self.completed_tasks.get(batch_id, deque())
-        completed_count = len(completed_queue)
+        for _ in range(min_num):
+            if completed_queue:
+                status, exps = completed_queue.pop()
+                statuses.append(status)
+                if isinstance(exps, list):
+                    experiences.extend(exps)
+                else:
+                    experiences.append(exps)
+
+        if batch_id in self.completed_tasks and not self.completed_tasks[batch_id]:
+            del self.completed_tasks[batch_id]
+
+        completed_count = len(statuses)
         if completed_count < min_num:
             self.logger.warning(
                 f"Timeout reached, only {completed_count}/{min_num} tasks completed"
             )
-        while completed_queue:
-            status, exps = completed_queue.pop()
-            statuses.append(status)
-            if isinstance(exps, list):
-                experiences.extend(exps)
-            else:
-                experiences.append(exps)
-
-        if batch_id in self.completed_tasks and not self.completed_tasks[batch_id]:
-            del self.completed_tasks[batch_id]
 
         return statuses, experiences
 

trinity/explorer/workflow_runner.py

@@ -5,7 +5,7 @@
 import traceback
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple
 
 from trinity.buffer import get_buffer_reader
 from trinity.common.config import Config
@@ -21,32 +21,8 @@ class Status:
     """Status of the task running result."""
 
     ok: bool
-    metric: dict[str, Union[float, List[float]]]
+    metric: dict[str, float]
     message: Optional[str] = None
-    task_id: Union[int, str] = ""
-
-
-def group_metrics(statuses: List[Status]):
-    task2metrics = {}
-    for status in statuses:
-        task_id = status.task_id
-        metric = status.metric
-        if task_id not in task2metrics:
-            task2metrics[task_id] = metric
-        else:
-            for k, v in metric.items():
-                task2metrics[task_id][k] += v  # type: ignore
-
-    metric_list = []
-    for metrics in task2metrics.values():
-        agg_metrics = {}
-        for k, v in metrics.items():
-            if isinstance(v, list):
-                agg_metrics[k] = sum(v) / len(v)
-            else:
-                agg_metrics[k] = v
-        metric_list.append(agg_metrics)
-    return metric_list
 
 
 class WorkflowRunner:
@@ -167,28 +143,23 @@ async def run_task(
                     exp.metrics = {}
                 for k, v in exp.metrics.items():
                     metrics[k].append(v)
-
-            metric: dict[str, Union[float, List[float]]] = {"time_per_task": time.time() - st}
-            metric.update(metrics)
+            # We get the average of metrics into the state
+            metric = {}
+            metric["time_per_task"] = time.time() - st
+            if metrics:
+                for k, v in metrics.items():
+                    metric[k] = sum(v) / len(v)  # type: ignore
 
             if task.is_eval:
                 # If the task is an evaluation task, we do not record the experiences to the buffer
-                return Status(True, metric=metric, task_id=task.task_id), []
+                return Status(True, metric=metric), []
             else:
-                return Status(True, metric=metric, task_id=task.task_id), exps
+                return Status(True, metric=metric), exps
 
         except Exception as e:
             error_trace_back = traceback.format_exc()
             self.logger.error(f"WorkflowRunner run task error: {e}\nTraceback:\n{error_trace_back}")
-            return (
-                Status(
-                    False,
-                    metric={"time_per_task": time.time() - st},
-                    message=str(e),
-                    task_id=task.task_id,
-                ),
-                [],
-            )
+            return Status(False, metric={"time_per_task": time.time() - st}, message=str(e)), []
 
 
 class DebugWorkflowRunner(WorkflowRunner):