bug fix

Author: chenyushuo

tests/explorer/runner_pool_test.py

@@ -37,7 +37,7 @@ def run(self) -> List[Experience]:
         elif self.error_type == "exit":
             exit(1)
         elif self.error_type == "auxiliary_models":
-            assert self.auxiliary_models is not None and len(self.auxiliary_models) > 0
+            assert self.auxiliary_models is not None and len(self.auxiliary_models) == 2
         return [Experience(tokens=torch.zeros(5), prompt_length=2, prompt_text=self.error_type)]
 
 
@@ -221,12 +221,15 @@ def test_runner_pool_with_auxiliary_models(self):
         config.explorer.auxiliary_models = [
             InferenceModelConfig(
                 engine_num=1,
-            )
+            ),
+            InferenceModelConfig(
+                engine_num=1,
+            ),
         ]
         pool = RunnerPool(
             config,
             [DummyModel.remote(), DummyModel.remote()],
-            [DummyAuxiliaryModel.remote(), DummyAuxiliaryModel.remote()],
+            [[DummyAuxiliaryModel.remote()], [DummyAuxiliaryModel.remote()]],
         )
         taskset_config = get_unittest_dataset_config("countdown")
         tasks = [

trinity/common/models/__init__.py

@@ -34,7 +34,7 @@ def allocate(self, num: int) -> list:
 
 def create_inference_models(
     config: Config,
-) -> Tuple[List[InferenceModel], List[InferenceModel]]:
+) -> Tuple[List[InferenceModel], List[List[InferenceModel]]]:
     """Create `engine_num` rollout models.
 
     Each model has `tensor_parallel_size` workers.
@@ -116,11 +116,12 @@ def create_inference_models(
 
     # create auxiliary models
     for model_config in config.explorer.auxiliary_models:
+        engines = []
         for _ in range(model_config.engine_num):
             bundles_for_engine = allocator.allocate(model_config.tensor_parallel_size)
             model_config.enable_openai_api = True
             model_config.engine_type = "vllm_async"
-            auxiliary_engines.append(
+            engines.append(
                 ray.remote(vLLMAysncRolloutModel)
                 .options(
                     num_cpus=0,
@@ -132,8 +133,10 @@ def create_inference_models(
                 )
                 .remote(config=model_config)
             )
+        auxiliary_engines.append(engines)
     # all auxiliary engines run api server
-    for engine in auxiliary_engines:
-        engine.run_api_server.remote()
+    for engines in auxiliary_engines:
+        for engine in engines:
+            engine.run_api_server.remote()
 
     return rollout_engines, auxiliary_engines

trinity/explorer/runner_pool.py

@@ -5,6 +5,7 @@
 import ray
 
 from trinity.common.config import Config
+from trinity.common.models.model import InferenceModel
 from trinity.common.workflows import Task
 from trinity.explorer.workflow_runner import Status, WorkflowRunner
 from trinity.utils.log import get_logger
@@ -19,13 +20,17 @@ class RunnerPool:
     `config.explorer.max_timeout`.
     """
 
-    def __init__(self, config: Config, models: List, auxiliary_models: Optional[List] = None):
+    def __init__(
+        self,
+        config: Config,
+        models: List[InferenceModel],
+        auxiliary_models: Optional[List[List[InferenceModel]]] = None,
+    ):
         # actors to be used
         self.logger = get_logger(__name__)
         self.config = config
         self.models = models
         self.auxiliary_models = auxiliary_models or []
-        self.auxiliary_models = [self.auxiliary_models]  # TODO: support multiple auxiliary models
         self.timeout = config.explorer.max_timeout
         self.max_retry_times = config.explorer.max_retry_times
 
@@ -47,10 +52,7 @@ def __init__(self, config: Config, models: List, auxiliary_models: Optional[List
         # create new actors
         self.engine_status = [0] * config.explorer.rollout_model.engine_num
         self.auxiliary_engine_status_list = [
-            [0]
-            * len(self.auxiliary_models[0])
-            # TODO: support multiple auxiliary models
-            # [0] * cfg.engine_num for cfg in config.explorer.auxiliary_models
+            [0] * cfg.engine_num for cfg in config.explorer.auxiliary_models
         ]
         self._idle_actors = list()
         self.actor_to_engine_index = {}