fix mode check

Author: chenyushuo

trinity/cli/launcher.py

@@ -37,6 +37,7 @@ def bench(config: Config) -> None:
 
 def explore(config: Config) -> None:
     """Run explorer."""
+    check_and_run_task_pipeline(config)
     try:
         explorer = Explorer.get_actor(config)
         ray.get(explorer.prepare.remote())
@@ -81,6 +82,7 @@ def both(config: Config) -> None:
     the latest step. The specific number of experiences may vary for different
     algorithms and tasks.
     """
+    check_and_run_task_pipeline(config)
     try:
         explorer = Explorer.get_actor(config)
         trainer = Trainer.get_actor(config)
@@ -151,7 +153,6 @@ def run_stage(config: Config) -> None:
     )
     pprint(config)
     try:
-        check_and_run_task_pipeline(config)
         MODE_MAP[config.mode](config)
     finally:
         if config.monitor.enable_ray_timeline:

trinity/common/config.py

@@ -853,8 +853,8 @@ def _check_interval(self) -> None:
                 )
 
     def _check_explorer_input(self) -> None:
-        if self.mode == "train":
-            # no need to check explorer_input in train mode
+        if self.mode in {"train", "bench", "serve"}:
+            # no need to check explorer_input in train/bench/serve mode
             return
 
         explorer_input = self.buffer.explorer_input
@@ -866,9 +866,8 @@ def _check_explorer_input(self) -> None:
             explorer_input.taskset = None
         elif len(explorer_input.tasksets) == 0:
             raise ValueError("At least one taskset should be provided in explorer_input!")
-        tasksets = explorer_input.tasksets
 
-        for i, taskset in enumerate(tasksets):
+        for i, taskset in enumerate(explorer_input.tasksets):
             if self.mode != "train" and not taskset.path:
                 raise ValueError(
                     "`buffer.explorer_input.taskset.path` is required, please set it to the path of the taskset."
@@ -914,6 +913,10 @@ def _check_explorer_input(self) -> None:
             set_if_none(dataset.rollout_args, "max_tokens", self.model.max_response_tokens)
 
     def _check_trainer_input(self) -> None:
+        if self.mode in {"explore", "bench", "serve"}:
+            # no need to check trainer_input in train/bench/serve mode
+            return
+
         trainer_input = self.buffer.trainer_input
         experience_buffer = trainer_input.experience_buffer
 
@@ -973,7 +976,7 @@ def _default_storage_path(self, storage_type: StorageType, name: str) -> str:
     def _check_data_processor(self) -> None:
         # check input/output buffers in pipelines
         experience_pipeline = self.data_processor.experience_pipeline
-        if experience_pipeline is not None:
+        if experience_pipeline is not None and self.mode in {"explore", "both", "serve"}:
             if experience_pipeline.save_input and experience_pipeline.input_save_path is None:
                 experience_pipeline.input_save_path = os.path.join(
                     self.buffer.cache_dir, "explorer_output.jsonl"  # type: ignore[arg-type]
@@ -983,7 +986,7 @@ def _check_data_processor(self) -> None:
                 )
 
         task_pipeline = self.data_processor.task_pipeline
-        if task_pipeline is not None:
+        if task_pipeline is not None and self.mode in {"explore", "both"}:
             if task_pipeline.output is None:
                 if self.mode != "train":
                     task_pipeline.output = self.buffer.explorer_input.tasksets[0]

trinity/explorer/explorer.py

@@ -52,7 +52,9 @@ def __init__(self, config: Config):
         self.models, self.auxiliary_models = create_inference_models(config)
         self.experience_pipeline = self._init_experience_pipeline()
         self.taskset = (
-            TasksetScheduler(explorer_state, config) if self.config.mode != "serve" else None
+            TasksetScheduler(explorer_state, config)
+            if self.config.mode not in {"bench", "serve"}
+            else None
         )
         self.scheduler = None
         self.monitor = MONITOR.get(self.config.monitor.monitor_type)(
@@ -406,6 +408,8 @@ async def is_alive(self) -> bool:
 
     def _init_experience_pipeline(self) -> ray.actor.ActorHandle:
         """Init experience pipeline for the explorer."""
+        if self.config.mode == "bench":
+            return None
         node_id = ray.get_runtime_context().get_node_id()
         return (
             ray.remote(ExperiencePipeline)