Merge branch 'main' into feat/exp_pipeline

Author: HYLcool

docs/sphinx_doc/source/tutorial/trinity_configs.md

@@ -107,11 +107,13 @@ Used to log training metrics during execution.
 ```yaml
 monitor:
   monitor_type: wandb
+  enable_ray_timeline: False
 ```
 
 - `monitor_type`: Type of monitoring system. Options:
   - `wandb`: Logs to [Weights & Biases](https://docs.wandb.ai/quickstart/). Requires logging in and setting `WANDB_API_KEY`. Project and run names match the `project` and `name` fields in global configs.
   - `tensorboard`: Logs to [TensorBoard](https://www.tensorflow.org/tensorboard). Files are saved under `<checkpoint_root_dir>/<project>/<name>/monitor/tensorboard`.
+- `enable_ray_timeline`: Whether to export the ray timeline. If set to `True`, a `timeline.json` file will be exported to `<checkpoint_root_dir>/<project>/<name>/monitor`. You can view the timeline file in Chrome at [chrome://tracing](chrome://tracing).
 
 ---
 

tests/tools.py

@@ -158,7 +158,7 @@ def metric_list(self, metric_prefix: str) -> List[str]:
 class RayUnittestBase(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        ray.init(ignore_reinit_error=True)
+        ray.init(ignore_reinit_error=True, namespace="trinity_unittest")
 
     @classmethod
     def tearDownClass(cls):

tests/trainer/trainer_test.py

@@ -27,7 +27,6 @@ def setUp(self):
         self.config.model.model_path = get_model_path()
         self.config.explorer.rollout_model.engine_type = "vllm_async"
         self.config.algorithm.repeat_times = 3
-        self.config.explorer.rollout_model.use_v1 = False
         self.config.project = "Trainer-unittest"
         self.config.name = f"trainer-{datetime.now().strftime('%Y%m%d%H%M%S')}"
         self.config.monitor.monitor_type = "tensorboard"
@@ -45,6 +44,7 @@ class TestTrainerCountdown(BaseTrainerCase):
     def test_trainer(self):
         """Test the both and bench mode."""
         # test both mode
+        self.config.explorer.rollout_model.use_v1 = False
         self.config.buffer.explorer_input.taskset = get_unittest_dataset_config("countdown")
         self.config.buffer.explorer_input.eval_tasksets.append(
             get_unittest_dataset_config("countdown", "test")
@@ -149,6 +149,7 @@ def test_trainer(self):
         response_metrics = parser.metric_list("response_length")
         self.assertTrue(len(response_metrics) > 0)
         self.assertEqual(parser.metric_max_step(response_metrics[0]), 4)
+        ray.timeline(filename="timeline.json")
         ray.shutdown(_exiting_interpreter=True)
         # check checkpoint
         from trinity.common.models.utils import get_checkpoint_dir_with_step_num

trinity/buffer/queue.py

@@ -3,8 +3,6 @@
 from copy import deepcopy
 from typing import List
 
-import ray
-
 from trinity.buffer.writer.file_writer import JSONWriter
 from trinity.buffer.writer.sql_writer import SQLWriter
 from trinity.common.config import BufferConfig, StorageConfig
@@ -20,7 +18,6 @@ def is_json_file(path: str) -> bool:
     return path.endswith(".json") or path.endswith(".jsonl")
 
 
-@ray.remote
 class QueueActor:
     """An asyncio.Queue based queue actor."""
 

trinity/buffer/ray_wrapper.py

@@ -55,6 +55,7 @@ def get_wrapper(cls, storage_config: StorageConfig, config: BufferConfig):
                 ray.remote(cls)
                 .options(
                     name=f"sql-{storage_config.name}",
+                    namespace=ray.get_runtime_context().namespace,
                     get_if_exists=True,
                 )
                 .remote(storage_config, config)
@@ -154,6 +155,7 @@ def get_wrapper(cls, storage_config: StorageConfig, config: BufferConfig):
                 ray.remote(cls)
                 .options(
                     name=f"json-{storage_config.name}",
+                    namespace=ray.get_runtime_context().namespace,
                     get_if_exists=True,
                 )
                 .remote(storage_config, config)

trinity/buffer/reader/queue_reader.py

@@ -19,10 +19,15 @@ class QueueReader(BufferReader):
     def __init__(self, storage_config: StorageConfig, config: BufferConfig):
         assert storage_config.storage_type == StorageType.QUEUE
         self.read_batch_size = config.read_batch_size
-        self.queue = QueueActor.options(
-            name=f"queue-{storage_config.name}",
-            get_if_exists=True,
-        ).remote(storage_config, config)
+        self.queue = (
+            ray.remote(QueueActor)
+            .options(
+                name=f"queue-{storage_config.name}",
+                namespace=ray.get_runtime_context().namespace,
+                get_if_exists=True,
+            )
+            .remote(storage_config, config)
+        )
 
     def read(
         self, batch_size: Optional[int] = None, strategy: Optional[ReadStrategy] = None

trinity/buffer/writer/queue_writer.py

@@ -18,10 +18,15 @@ class QueueWriter(BufferWriter):
     def __init__(self, meta: StorageConfig, config: BufferConfig):
         assert meta.storage_type == StorageType.QUEUE
         self.config = config
-        self.queue = QueueActor.options(
-            name=f"queue-{meta.name}",
-            get_if_exists=True,
-        ).remote(meta, config)
+        self.queue = (
+            ray.remote(QueueActor)
+            .options(
+                name=f"queue-{meta.name}",
+                namespace=ray.get_runtime_context().namespace,
+                get_if_exists=True,
+            )
+            .remote(meta, config)
+        )
 
     def write(self, data: List) -> None:
         ray.get(self.queue.put_batch.remote(data))

trinity/cli/launcher.py

@@ -20,7 +20,14 @@
 
 def bench(config: Config) -> None:
     """Evaluate model."""
-    explorer = ray.remote(Explorer).options(name=EXPLORER_NAME).remote(config)
+    explorer = (
+        ray.remote(Explorer)
+        .options(
+            name=EXPLORER_NAME,
+            namespace=ray.get_runtime_context().namespace,
+        )
+        .remote(config)
+    )
     try:
         ray.get(explorer.prepare.remote())
         ray.get(explorer.benchmark.remote())
@@ -34,7 +41,14 @@ def bench(config: Config) -> None:
 def explore(config: Config) -> None:
     """Run explorer."""
     try:
-        explorer = ray.remote(Explorer).options(name=EXPLORER_NAME).remote(config)
+        explorer = (
+            ray.remote(Explorer)
+            .options(
+                name=EXPLORER_NAME,
+                namespace=ray.get_runtime_context().namespace,
+            )
+            .remote(config)
+        )
         ray.get(explorer.prepare.remote())
         ray.get(explorer.sync_weight.remote())
         ray.get(explorer.explore.remote())
@@ -47,7 +61,14 @@ def explore(config: Config) -> None:
 def train(config: Config) -> None:
     """Run trainer."""
     try:
-        trainer = ray.remote(Trainer).options(name=TRAINER_NAME).remote(config)
+        trainer = (
+            ray.remote(Trainer)
+            .options(
+                name=TRAINER_NAME,
+                namespace=ray.get_runtime_context().namespace,
+            )
+            .remote(config)
+        )
         ray.get(trainer.prepare.remote())
         ray.get(trainer.sync_weight.remote())
         ray.get(trainer.train.remote())
@@ -67,8 +88,23 @@ def both(config: Config) -> None:
     the latest step. The specific number of experiences may vary for different
     algorithms and tasks.
     """
-    explorer = ray.remote(Explorer).options(name=EXPLORER_NAME).remote(config)
-    trainer = ray.remote(Trainer).options(name=TRAINER_NAME).remote(config)
+    namespace = ray.get_runtime_context().namespace
+    explorer = (
+        ray.remote(Explorer)
+        .options(
+            name=EXPLORER_NAME,
+            namespace=namespace,
+        )
+        .remote(config)
+    )
+    trainer = (
+        ray.remote(Trainer)
+        .options(
+            name=TRAINER_NAME,
+            namespace=namespace,
+        )
+        .remote(config)
+    )
     ray.get([explorer.__ray_ready__.remote(), trainer.__ray_ready__.remote()])
     ray.get(
         [
@@ -192,30 +228,36 @@ def run(config_path: str, dlc: bool = False, plugin_dir: str = None):
         activate_data_module(
             f"{data_processor_config.data_processor_url}/{DataProcessorPipelineType.EXPERIENCE.value}", config_path
         )
-    ray_namespace = config.ray_namespace
     if dlc:
         from trinity.utils.dlc_utils import setup_ray_cluster
 
-        setup_ray_cluster(namespace=ray_namespace)
+        setup_ray_cluster(namespace=config.ray_namespace)
     else:
         from trinity.utils.dlc_utils import is_running
 
         if not is_running:
             raise RuntimeError("Ray is not running, please start it by `ray start --head`.")
-        ray.init(namespace=ray_namespace, ignore_reinit_error=True)
-    if config.mode == "explore":
-        explore(config)
-    elif config.mode == "train":
-        train(config)
-    elif config.mode == "both":
-        both(config)
-    elif config.mode == "bench":
-        bench(config)
-
-    if dlc:
-        from trinity.utils.dlc_utils import stop_ray_cluster
-
-        stop_ray_cluster()
+        ray.init(namespace=config.ray_namespace, ignore_reinit_error=True)
+    try:
+        if config.mode == "explore":
+            explore(config)
+        elif config.mode == "train":
+            train(config)
+        elif config.mode == "both":
+            both(config)
+        elif config.mode == "bench":
+            bench(config)
+    finally:
+        if config.monitor.enable_ray_timeline:
+            timeline_file = os.path.join(config.monitor.cache_dir, "timeline.json")
+            logger.info(f"Exporting Ray timeline to {timeline_file}...")
+            ray.timeline(filename=timeline_file)
+            logger.info("Done. You can open the timeline file in `chrome://tracing`")
+
+        if dlc:
+            from trinity.utils.dlc_utils import stop_ray_cluster
+
+            stop_ray_cluster(namespace=config.ray_namespace)
 
 
 def studio(port: int = 8501):

trinity/common/config.py

@@ -192,7 +192,6 @@ class InferenceModelConfig:
 
     # ! DO NOT SET
     bundle_indices: str = ""
-    ray_namespace: str = ""
 
 
 @dataclass
@@ -331,6 +330,9 @@ class MonitorConfig:
     monitor_type: str = "tensorboard"
     # the default args for monitor
     monitor_args: Dict = field(default_factory=dict)
+    # whether to enable ray timeline profile
+    # the output file will be saved to `cache_dir/timeline.json`
+    enable_ray_timeline: bool = False
     # ! DO NOT SET, automatically generated as checkpoint_job_dir/monitor
     cache_dir: str = ""
 
@@ -365,7 +367,7 @@ class Config:
     checkpoint_root_dir: str = ""
     # ! DO NOT SET, automatically generated as `checkpoint_root_dir/project/name`
     checkpoint_job_dir: str = ""
-    # ! DO NOT SET, automatically generated as f"{config.project}-{config.name}"
+    # If not set, automatically generated as f"{config.project}-{config.name}"
     ray_namespace: str = ""
 
     algorithm: AlgorithmConfig = field(default_factory=AlgorithmConfig)
@@ -590,7 +592,8 @@ def check_and_update(self) -> None:  # noqa: C901
         self._check_deprecated()
 
         # set namespace
-        self.ray_namespace = f"{self.project}-{self.name}"
+        if self.ray_namespace is None or len(self.ray_namespace) == 0:
+            self.ray_namespace = f"{self.project}-{self.name}"
 
         # check algorithm
         self._check_algorithm()
@@ -622,9 +625,6 @@ def check_and_update(self) -> None:  # noqa: C901
             self.explorer.rollout_model.max_prompt_tokens = self.model.max_prompt_tokens
         if self.explorer.rollout_model.max_response_tokens is None:
             self.explorer.rollout_model.max_response_tokens = self.model.max_response_tokens
-        self.explorer.rollout_model.ray_namespace = self.ray_namespace
-        for model in self.explorer.auxiliary_models:
-            model.ray_namespace = self.ray_namespace
 
         # check synchronizer
         self.synchronizer.explorer_world_size = (

trinity/common/models/__init__.py

@@ -89,7 +89,7 @@ def create_inference_models(
     for bundle_id, node_id in bundle_node_map.items():
         node_bundle_map[node_id].append(bundle_id)
     allocator = _BundleAllocator(node_bundle_map)
-
+    namespace = ray.get_runtime_context().namespace
     # create rollout models
     for _ in range(config.explorer.rollout_model.engine_num):
         bundles_for_engine = allocator.allocate(config.explorer.rollout_model.tensor_parallel_size)
@@ -101,6 +101,7 @@ def create_inference_models(
             .options(
                 num_cpus=0,
                 num_gpus=0 if config.explorer.rollout_model.tensor_parallel_size > 1 else 1,
+                namespace=namespace,
                 scheduling_strategy=PlacementGroupSchedulingStrategy(
                     placement_group=pg,
                     placement_group_capture_child_tasks=True,
@@ -128,6 +129,7 @@ def create_inference_models(
                 .options(
                     num_cpus=0,
                     num_gpus=0 if model_config.tensor_parallel_size > 1 else 1,
+                    namespace=namespace,
                     scheduling_strategy=PlacementGroupSchedulingStrategy(
                         placement_group=pg,
                         placement_group_capture_child_tasks=True,