[Refactor] Support concurrent inference accorss tasks.

Author: mzr1996

opencompass/cli/main.py

@@ -5,19 +5,36 @@
 import getpass
 import os
 import os.path as osp
+import threading
 from datetime import datetime
 
 from mmengine.config import Config, DictAction
 
 from opencompass.registry import PARTITIONERS, RUNNERS, build_from_cfg
 from opencompass.runners import SlurmRunner
 from opencompass.summarizers import DefaultSummarizer
-from opencompass.utils import (LarkReporter, get_logger, pretty_print_config,
-                               read_from_station, save_to_station)
+from opencompass.utils import (HeartBeatManager, LarkReporter, get_logger,
+                               pretty_print_config, read_from_station,
+                               save_to_station)
 from opencompass.utils.run import (fill_eval_cfg, fill_infer_cfg,
                                    get_config_from_arg)
 
 
+def _run_eval_tasks(runner, tasks):
+    if isinstance(tasks, list) and len(tasks) != 0 and isinstance(tasks[0],
+                                                                  list):
+        for task_part in tasks:
+            runner(task_part)
+    else:
+        runner(tasks)
+
+
+def _is_eval_daemon(task_type) -> bool:
+    if isinstance(task_type, str):
+        return task_type.endswith('OpenICLEvalWatchTask')
+    return getattr(task_type, '__name__', '') == 'OpenICLEvalWatchTask'
+
+
 def parse_args():
     parser = argparse.ArgumentParser(description='Run an evaluation task')
     parser.add_argument('config', nargs='?', help='Train config file path')
@@ -318,7 +335,15 @@ def main():
     if args.config_verbose:
         pretty_print_config(cfg)
 
-    # infer
+    infer_tasks = None
+    infer_runner = None
+    eval_tasks = None
+    eval_runner = None
+    eval_daemon = False
+
+    # ========================
+    #  Setup Configuration
+    # ========================
     if args.mode in ['all', 'infer']:
         # When user have specified --slurm or --dlc, or have not set
         # "infer" in config, we will provide a default configuration
@@ -358,7 +383,8 @@ def main():
         if args.dump_res_length:
             for task in tasks:
                 task.dump_res_length = True
-        runner(tasks)
+        infer_tasks = tasks
+        infer_runner = runner
 
     # evaluate
     if args.mode in ['all', 'eval']:
@@ -397,14 +423,35 @@ def main():
         if args.dry_run:
             return
         runner = RUNNERS.build(cfg.eval.runner)
-
-        # For meta-review-judge in subjective evaluation
-        if isinstance(tasks, list) and len(tasks) != 0 and isinstance(
-                tasks[0], list):
-            for task_part in tasks:
-                runner(task_part)
-        else:
-            runner(tasks)
+        task_type = getattr(cfg.eval.runner, 'task', {}).get('type', '')
+        eval_daemon = _is_eval_daemon(task_type)
+
+        eval_tasks = tasks
+        eval_runner = runner
+
+    # =================
+    #  Startup Runner
+    # =================
+    if infer_runner and eval_runner and eval_daemon:
+        heartbeat = HeartBeatManager(cfg['work_dir'])
+        stop_event, hb_thread = heartbeat.start_heartbeat()
+
+        eval_thread = threading.Thread(target=_run_eval_tasks,
+                                       args=(eval_runner, eval_tasks),
+                                       daemon=True)
+        eval_thread.start()
+
+        infer_runner(infer_tasks)
+
+        stop_event.set()
+        hb_thread.join()
+        logger.info('All infer tasks finished, stop heartbeat.')
+        eval_thread.join()
+    else:
+        if infer_runner is not None:
+            infer_runner(infer_tasks)
+        if eval_runner is not None:
+            _run_eval_tasks(eval_runner, eval_tasks)
 
     # save to station
     if args.station_path is not None or cfg.get('station_path') is not None:

opencompass/models/openai_api.py

@@ -185,6 +185,10 @@ def generate(
         if self.temperature is not None:
             temperature = self.temperature
 
+        if len(inputs) == 1:
+            # Forget multi-thread for single infernece.
+            return [self._generate(inputs[0], max_out_len, temperature)]
+
         with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
             results = list(
                 tqdm(
@@ -254,6 +258,7 @@ def _generate(self, input: PromptType, max_out_len: int,
                         self.org_ctr = 0
                 header['OpenAI-Organization'] = self.orgs[self.org_ctr]
 
+            self.acquire()
             try:
                 if any(model in self.path
                        for model in OAI_REASONING_MODEL_LIST):
@@ -377,6 +382,8 @@ def _generate(self, input: PromptType, max_out_len: int,
                         'Find error message in response: ',
                         str(response['error']),
                     )
+            finally:
+                self.release()
             max_num_retries += 1
 
         raise RuntimeError('Calling OpenAI failed after retrying for '
@@ -697,6 +704,7 @@ def _generate(
             if self.openai_extra_kwargs:
                 query_data.update(self.openai_extra_kwargs)
 
+            self.acquire()
             try:
                 if self.verbose:
                     self.logger.info('Start calling OpenAI API')
@@ -789,6 +797,8 @@ def _generate(
             except Exception as e:
                 self.logger.error(f'error occurs at {self.openai_api_base}')
                 self.logger.error(e)
+            finally:
+                self.release()
             num_retries += 1
         raise RuntimeError('Calling OpenAI API failed after retrying for '
                            f'{self.retry} times. Check the logs for details.')
@@ -925,6 +935,7 @@ def _generate(
             if self.openai_extra_kwargs:
                 query_data.update(self.openai_extra_kwargs)
 
+            self.acquire()
             try:
                 if self.verbose:
                     self.logger.info('Start calling OpenAI API')
@@ -1052,6 +1063,8 @@ def _generate(
             except Exception as e:
                 self.logger.error(f'error occurs at {self.openai_api_base}')
                 self.logger.error(e)
+            finally:
+                self.release()
             num_retries += 1
         raise RuntimeError('Calling OpenAI API failed after retrying for '
                            f'{self.retry} times. Check the logs for details.')

opencompass/models/openai_streaming.py

@@ -86,6 +86,7 @@ def __init__(self,
         self.openai_extra_kwargs = openai_extra_kwargs
         self.timeout = timeout
         self.finish_reason_confirm = finish_reason_confirm
+        self.openai_client = self._create_fresh_client()
 
     def _create_fresh_client(self):
         """Create a fresh OpenAI client for each request to avoid
@@ -117,11 +118,15 @@ def _create_fresh_client(self):
                 'https://': self.proxy_url,
             }
 
+        limits = httpx.Limits(max_keepalive_connections=2048,
+                              max_connections=4096)
+
         return OpenAI(
             base_url=self.openai_api_base,
             api_key=current_key,
             http_client=httpx.Client(**http_client_cfg,
-                                     timeout=httpx.Timeout(self.timeout))
+                                     timeout=httpx.Timeout(self.timeout),
+                                     limits=limits)
             if http_client_cfg or True else None,
         )
 
@@ -185,6 +190,7 @@ def _generate(
             if self.openai_extra_kwargs:
                 query_data.update(self.openai_extra_kwargs)
 
+            self.acquire()
             try:
                 if self.verbose:
                     thread_id = threading.get_ident()
@@ -193,22 +199,13 @@ def _generate(
                         f'with streaming enabled')
 
                 if self.stream:
-                    # Create fresh client for each request to avoid
-                    # concurrency issues
-                    fresh_client = self._create_fresh_client()
-
                     # Handle streaming response with shorter timeout
-                    response_stream = fresh_client.chat.completions.create(
+                    response_stream = self.openai_client.chat.completions.create(
                         **query_data, timeout=self.timeout)
 
                     result = self._handle_stream_response(
                         response_stream, thread_id if self.verbose else None)
 
-                    # Clean up the client
-                    if (hasattr(fresh_client, '_client')
-                            and hasattr(fresh_client._client, 'close')):
-                        fresh_client._client.close()
-
                     return result
                 else:
                     # Fallback to non-streaming (use parent method)
@@ -237,6 +234,8 @@ def _generate(
                 import traceback
                 self.logger.error(f'[Thread {thread_id}] Traceback: '
                                   f'{traceback.format_exc()}')
+            finally:
+                self.release()
             num_retries += 1
 
         raise RuntimeError('Calling OpenAI API failed after retrying for '

opencompass/openicl/icl_inferencer/__init__.py

@@ -2,9 +2,11 @@
 from .icl_attack_inferencer import AttackInferencer  # noqa
 from .icl_base_inferencer import BaseInferencer  # noqa
 from .icl_chat_inferencer import ChatInferencer  # noqa
+from .icl_chat_inferencer_parallel import ParallelChatInferencer  # noqa
 from .icl_chatml_inferencer import ChatMLInferencer  # noqa
 from .icl_clp_inferencer import CLPInferencer  # noqa
 from .icl_gen_inferencer import GenInferencer  # noqa
+from .icl_gen_inferencer_parallel import ParallelGenInferencer  # noqa
 from .icl_inference_ppl_only_inferencer import \
     InferencePPLOnlyInferencer  # noqa
 from .icl_ll_inferencer import LLInferencer  # noqa

opencompass/openicl/icl_inferencer/icl_chat_inferencer_parallel.py

@@ -0,0 +1,137 @@
+"""Parallel Chat Inferencer."""
+import os
+import os.path as osp
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List, Optional
+
+import mmengine
+
+from opencompass.registry import ICL_INFERENCERS
+
+from ..icl_prompt_template import PromptTemplate
+from ..icl_retriever import BaseRetriever
+from ..utils.logging import get_logger
+from .icl_chat_inferencer import ChatInferencer
+
+logger = get_logger(__name__)
+
+
+@ICL_INFERENCERS.register_module()
+class ParallelChatInferencer(ChatInferencer):
+    """Parallel chat inferencer with thread pool over samples."""
+
+    def __init__(
+            self,
+            model,
+            output_json_filepath: Optional[str] = './icl_inference_output',
+            output_json_filename: Optional[str] = 'predictions',
+            save_every: Optional[int] = 1,
+            infer_mode: str = 'last',
+            max_out_len: int = 512,
+            max_infer_workers: Optional[int] = None,
+            **kwargs) -> None:
+        super().__init__(
+            model=model,
+            output_json_filename=output_json_filename,
+            output_json_filepath=output_json_filepath,
+            save_every=save_every,
+            infer_mode=infer_mode,
+            max_out_len=max_out_len,
+            **kwargs,
+        )
+        self.max_infer_workers = max_infer_workers
+        self.progress_tracker = None
+
+    def _resolve_max_workers(self) -> int:
+        if self.max_infer_workers is not None:
+            return self.max_infer_workers
+        max_workers = getattr(self.model, 'max_workers', None)
+        if max_workers is not None:
+            return max_workers
+        cpu_count = os.cpu_count() or 1
+        return min(32, cpu_count + 4)
+
+    def _progress_update(self, count: int = 1) -> None:
+        if self.progress_tracker is not None:
+            self.progress_tracker.incr(count)
+
+    def inference(self,
+                  retriever: BaseRetriever,
+                  ice_template: Optional[PromptTemplate] = None,
+                  prompt_template: Optional[PromptTemplate] = None,
+                  output_json_filepath: Optional[str] = None,
+                  output_json_filename: Optional[str] = None) -> dict:
+        output_handler = self.HandlerType()
+
+        if output_json_filepath is None:
+            output_json_filepath = self.output_json_filepath
+        if output_json_filename is None:
+            output_json_filename = self.output_json_filename
+
+        ice_idx_list = retriever.retrieve()
+
+        chat_list = self.get_chat_list(
+            ice_idx_list,
+            retriever,
+            prompt_template=prompt_template,
+        )
+
+        total_samples = len(chat_list)
+        if self.progress_tracker is not None:
+            self.progress_tracker.set_total(total_samples)
+
+        todo = list(range(total_samples))
+        tmp_json_filepath = os.path.join(output_json_filepath,
+                                         'tmp_' + output_json_filename)
+        if osp.exists(tmp_json_filepath):
+            try:
+                tmp_result_dict = mmengine.load(tmp_json_filepath)
+            except Exception:
+                pass
+            else:
+                output_handler.results_dict = tmp_result_dict
+                todo = [i for i in todo if str(i) not in tmp_result_dict.keys()]
+        if self.progress_tracker is not None:
+            self.progress_tracker.set_completed(total_samples - len(todo))
+
+        chats = [chat_list[i] for i in todo]
+
+        logger.info('Starting parallel chat inference process...')
+
+        def _infer_one(chat, idx):
+            local_handler = self.HandlerType()
+            if self.infer_mode == 'last':
+                self.infer_last(chat, idx, local_handler)
+            elif self.infer_mode == 'every':
+                self.infer_every(chat, idx, local_handler)
+            elif self.infer_mode == 'every_with_gt':
+                self.infer_every_with_gt(chat, idx, local_handler)
+            return local_handler.results_dict
+
+        max_workers = self._resolve_max_workers()
+        completed = total_samples - len(todo)
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = [
+                executor.submit(_infer_one, chat, idx)
+                for idx, chat in zip(todo, chats)
+            ]
+            for future in as_completed(futures):
+                result_dict = future.result()
+                output_handler.results_dict.update(result_dict)
+                delta = len(result_dict)
+                completed += delta
+                self._progress_update(delta)
+                if (self.save_every is not None
+                        and completed % self.save_every == 0
+                        and self.is_main_process):
+                    output_handler.write_to_json(output_json_filepath,
+                                                 'tmp_' + output_json_filename)
+
+        if self.is_main_process:
+            os.makedirs(output_json_filepath, exist_ok=True)
+            output_handler.write_to_json(output_json_filepath,
+                                         output_json_filename)
+            if osp.exists(tmp_json_filepath):
+                os.remove(tmp_json_filepath)
+
+        return output_handler.results_dict