| import os |
| import os.path as osp |
| import re |
| import subprocess |
| import time |
| import traceback |
| from functools import partial |
| from multiprocessing import Pipe, Pool |
| from typing import Any, Dict, List, Optional, Tuple |
|
|
| import mmengine |
| from mmengine.config import ConfigDict |
| from tqdm import tqdm |
|
|
| from opencompass.registry import RUNNERS, TASKS |
| from opencompass.utils import batched, get_logger |
|
|
| from .base import BaseRunner |
|
|
|
|
| @RUNNERS.register_module() |
| class SlurmSequentialRunner(BaseRunner): |
| """Distributed runner based on Slurm. It will launch tasks in parallel |
| using `srun` command. |
| |
| This runner launches tasks one by one for execution. A new task will only |
| be launched when and only when max_num_workers is not met, and the previous |
| task has been successfully allocated to a machine. Therefore, unlike the |
| `SlurmRunner`, at most only one task will be in the PENDING status at the |
| same time during a run, making the random_sleep strategy no longer |
| necessary. In addition, this runner also includes a feature to |
| automatically kill all jobs by the job_id on exit. |
| |
| The runner will obtain the job_id by reading the srun output similar to |
| `srun: Job 123456 scheduled successfully!`. If the output of srun does not |
| match this pattern, the runner will not work properly. |
| |
| Args: |
| task (ConfigDict): Task type config. |
| max_num_workers (int): Max number of workers to run in parallel. |
| Defaults to 32. |
| retry (int): Number of retries if the job failed. Defaults to 2. |
| partition (str): Slurm partition name. Defaults to None. |
| quotatype (str): Slurm quota type. Defaults to None. |
| qos (str): Slurm quality of service. Defaults to None. |
| debug (bool): Whether to run in debug mode. Defaults to False. |
| lark_bot_url (str): Lark bot url. Defaults to None. |
| extra_command (List, optional): Extra slurm command. |
| For example ['-c 12', '-w node1']. Defaults to None. |
| """ |
|
|
| def __init__(self, |
| task: ConfigDict, |
| task_prefix: str = '', |
| max_num_workers: int = 32, |
| retry: int = 2, |
| partition: str = None, |
| quotatype: str = None, |
| qos: str = None, |
| debug: bool = False, |
| lark_bot_url: str = None, |
| extra_command: Optional[List[str]] = None): |
| super().__init__(task=task, debug=debug, lark_bot_url=lark_bot_url) |
| self.max_num_workers = max_num_workers |
| self.retry = retry |
| self.partition = partition |
| self.quotatype = quotatype |
| self.qos = qos |
| self.task_prefix = task_prefix |
| if not extra_command: |
| extra_command = [] |
| assert isinstance(extra_command, list) |
| self.extra_command = extra_command |
|
|
| logger = get_logger() |
| if self.quotatype in ['spot', 'auto']: |
| logger.warning( |
| 'Quotatype spot or auto may cause stability issues, ' |
| 'reserved is recommended.') |
|
|
| def launch(self, tasks: List[Dict[str, Any]]) -> List[Tuple[str, int]]: |
| if not self.debug: |
| return self._launch_wo_debug(tasks) |
| else: |
| return [self._launch(task) for task in tasks] |
|
|
| def _launch_wo_debug(self, |
| tasks: List[Dict[str, Any]]) -> List[Tuple[str, int]]: |
| launched_bar = tqdm(total=len(tasks), desc='Launched') |
| finished_bar = tqdm(total=len(tasks), desc='Finished') |
| job_ids = [] |
| status = [] |
|
|
| def _update(result): |
| finished_bar.update() |
| status.append(result) |
| return result |
|
|
| def _err_update(err): |
| finished_bar.update() |
| traceback.print_exc() |
| status.append(('', -1)) |
|
|
| try: |
| parent_conns = [] |
| num_workers = max(min(self.max_num_workers, len(tasks)), 1) |
| with Pool(processes=num_workers) as pool: |
| for task in tasks: |
| parent_conn, child_conn = Pipe() |
| _ = pool.apply_async(self._launch, |
| kwds={ |
| 'cfg': task, |
| 'child_conn': child_conn |
| }, |
| callback=_update, |
| error_callback=_err_update) |
| time.sleep(0.5) |
|
|
| job_id = parent_conn.recv() |
| launched_bar.update() |
| parent_conns.append(parent_conn) |
| job_ids.append(job_id) |
|
|
| pool.close() |
| pool.join() |
| return status |
| except KeyboardInterrupt: |
| raise |
| finally: |
| launched_bar.close() |
| finished_bar.close() |
| for parent_conn in parent_conns: |
| while parent_conn.poll(): |
| try: |
| job_id = parent_conn.recv() |
| job_ids.append(job_id) |
| except EOFError: |
| break |
| parent_conn.close() |
|
|
| tbar = tqdm(total=len(job_ids), desc='clear sruns') |
| for batched_job_ids in batched(job_ids, 4): |
| while True: |
| ps = [] |
| try: |
| for job_id in batched_job_ids: |
| tbar.update() |
| if job_id is None: |
| continue |
| cmd = f'scancel {job_id}' |
| p = subprocess.Popen(cmd, |
| shell=True, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT) |
| ps.append(p) |
| break |
| except KeyboardInterrupt: |
| logger = get_logger() |
| logger.error('Ignoring KeyboardInterrupt...') |
| for p in ps: |
| p.wait() |
| tbar.close() |
|
|
| def _launch(self, cfg: ConfigDict, child_conn: Pipe = None): |
| logger = get_logger() |
|
|
| task = TASKS.build(dict(cfg=cfg, type=self.task_cfg['type'])) |
| num_gpus = task.num_gpus |
| task_name = task.name |
| task_name = self.task_prefix + task_name |
|
|
| |
| mmengine.mkdir_or_exist('tmp/') |
| param_file = f'tmp/{os.getpid()}_params.py' |
| process = None |
| try: |
| cfg.dump(param_file) |
|
|
| |
| tmpl = 'srun' |
| if self.partition: |
| tmpl += f' -p {self.partition}' |
| if self.quotatype: |
| tmpl += f' --quotatype={self.quotatype}' |
| if self.qos: |
| tmpl += f' --qos={self.qos}' |
| if num_gpus > 0: |
| tmpl += f' --gres=gpu:{num_gpus}' |
| for extra_cmd in self.extra_command: |
| tmpl += f' {extra_cmd}' |
| tmpl += f" -N1 -u -J '{task_name[:512]}'" + ' {task_cmd}' |
| get_cmd = partial(task.get_command, |
| cfg_path=param_file, |
| template=tmpl) |
| cmd = get_cmd() |
|
|
| logger.debug(f'Running command: {cmd}') |
|
|
| retry = self.retry |
| output_paths = task.get_output_paths() |
|
|
| if self.debug: |
| while True: |
| process = subprocess.Popen(cmd, shell=True, text=True) |
| process.communicate() |
| process.wait() |
| if self._job_failed(process.returncode, output_paths): |
| if retry > 0: |
| logger.warning( |
| f'task {task_name} failed, retrying...') |
| retry -= 1 |
| cmd = get_cmd() |
| else: |
| break |
| else: |
| break |
| else: |
| out_path = task.get_log_path(file_extension='out') |
| mmengine.mkdir_or_exist(osp.split(out_path)[0]) |
| stdout = open(out_path, 'w', encoding='utf-8') |
| stderr = subprocess.PIPE |
| while True: |
| process = subprocess.Popen(cmd, |
| shell=True, |
| text=True, |
| stdout=stdout, |
| stderr=stderr) |
| job_id = None |
| while True: |
| line = process.stderr.readline() |
| if not line: |
| break |
| match = re.search( |
| r'srun: Job (\d+) scheduled successfully!', line) |
| if match and job_id is None: |
| job_id = match.group(1) |
| child_conn.send(job_id) |
| stdout.write(line) |
| process.wait() |
| if self._job_failed(process.returncode, output_paths): |
| if retry > 0: |
| retry -= 1 |
| cmd = get_cmd() |
| else: |
| logger.warning( |
| f'task {task_name} fail, see\n{out_path}') |
| break |
| else: |
| break |
| except KeyboardInterrupt: |
| raise |
| finally: |
| |
| if child_conn is not None: |
| child_conn.send(None) |
| child_conn.close() |
| if process is not None: |
| process.kill() |
| os.remove(param_file) |
| return task_name, process.returncode |
|
|
| def _job_failed(self, return_code: int, output_paths: List[str]) -> bool: |
| return return_code != 0 or not all( |
| osp.exists(output_path) for output_path in output_paths) |
|
|
