[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

@@ -138,6 +138,7 @@ def __init__(
                 self.keys = [key]
         else:
             self.keys = key
+        self._key_lock = Lock()
 
         # record invalid keys and skip them when requesting API
         # - keys have insufficient_quota
@@ -160,6 +161,23 @@ def __init__(
 
         self.path = path
 
+    def _next_valid_key(self):
+        with self._key_lock:
+            if len(self.invalid_keys) == len(self.keys):
+                raise RuntimeError('All keys have insufficient quota.')
+
+            # find the next valid key
+            while True:
+                self.key_ctr += 1
+                if self.key_ctr == len(self.keys):
+                    self.key_ctr = 0
+
+                if self.keys[self.key_ctr] not in self.invalid_keys:
+                    break
+
+            key = self.keys[self.key_ctr]
+        return key
+
     def generate(
         self,
         inputs: List[PromptType],
@@ -185,6 +203,10 @@ def generate(
         if self.temperature is not None:
             temperature = self.temperature
 
+        if len(inputs) == 1:
+            # Forget multi-thread for single inference.
+            return [self._generate(inputs[0], max_out_len, temperature)]
+
         with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
             results = list(
                 tqdm(
@@ -224,22 +246,7 @@ def _generate(self, input: PromptType, max_out_len: int,
 
         max_num_retries = 0
         while max_num_retries < self.retry:
-            self.wait()
-
-            with Lock():
-                if len(self.invalid_keys) == len(self.keys):
-                    raise RuntimeError('All keys have insufficient quota.')
-
-                # find the next valid key
-                while True:
-                    self.key_ctr += 1
-                    if self.key_ctr == len(self.keys):
-                        self.key_ctr = 0
-
-                    if self.keys[self.key_ctr] not in self.invalid_keys:
-                        break
-
-                key = self.keys[self.key_ctr]
+            key = self._next_valid_key()
 
             header = {
                 'Authorization': f'Bearer {key}',
@@ -254,6 +261,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):
@@ -314,23 +322,13 @@ def _generate(self, input: PromptType, max_out_len: int,
                         self.logger.debug(
                             f'Get response from {self.proxy_url}')
 
-            except requests.ConnectionError:
-                self.logger.error('Got connection error, retrying...')
-                continue
-            try:
                 if raw_response.status_code != 200:
                     self.logger.error(f'Request failed with status code '
                                       f'{raw_response.status_code}, response: '
                                       f'{raw_response.content.decode()}')
                     continue
                 response = raw_response.json()
-            except requests.JSONDecodeError:
-                self.logger.error(f'JsonDecode error, got status code '
-                                  f'{raw_response.status_code}, response: '
-                                  f'{raw_response.content.decode()}')
-                continue
-            self.logger.debug(str(response))
-            try:
+                self.logger.debug(str(response))
                 if self.logprobs:
                     return response['choices']
                 else:
@@ -356,6 +354,12 @@ def _generate(self, input: PromptType, max_out_len: int,
                             return reasoning_content
                     else:
                         return content.strip()
+            except requests.ConnectionError:
+                self.logger.error('Got connection error, retrying...')
+            except requests.JSONDecodeError:
+                self.logger.error(f'JsonDecode error, got status code '
+                                  f'{raw_response.status_code}, response: '
+                                  f'{raw_response.content.decode()}')
             except KeyError:
                 if 'error' in response:
                     if response['error']['code'] == 'rate_limit_exceeded':
@@ -377,6 +381,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 '
@@ -575,7 +581,7 @@ def __init__(
         query_per_second: int = 1,
         rpm_verbose: bool = False,
         retry: int = 2,
-        key: str | List[str] = 'ENV',
+        key: str = 'ENV',
         org: str | List[str] | None = None,
         meta_template: Dict | None = None,
         openai_api_base: str | List[str] = OPENAISDK_API_BASE,
@@ -671,7 +677,6 @@ def _generate(
 
         num_retries = 0
         while num_retries < self.retry:
-            self.wait()
             if any(model in self.path for model in OAI_REASONING_MODEL_LIST):
                 self.logger.warning(
                     f"'max_token' is unsupported for model {self.path}")
@@ -697,6 +702,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 +795,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 +933,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 +1061,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.')