OpenCompass/opencompass/runners/dlc.py

import datetime
import json
import os
import os.path as osp
import random
import re
import subprocess
import sys
import time
from functools import partial
from typing import Any, Dict, List, Optional, Tuple

import mmengine
from mmengine.config import ConfigDict
from mmengine.utils import track_parallel_progress

from opencompass.registry import RUNNERS, TASKS
from opencompass.utils import get_logger

from .base import BaseRunner


@RUNNERS.register_module()
class DLCRunner(BaseRunner):
    """Distributed runner based on Alibaba Cloud Deep Learning Cluster (DLC).
    It will launch multiple tasks in parallel with 'dlc' command. Please
    install and configure DLC first before using this runner.

    Args:
        task (ConfigDict): Task type config.
        aliyun_cfg (ConfigDict): Alibaba Cloud config.
        max_num_workers (int): Max number of workers. Default: 32.
        retry (int): Number of retries when job failed. Default: 2.
        debug (bool): Whether to run in debug mode. Default: False.
        lark_bot_url (str): Lark bot url. Default: None.
    """

    def __init__(self,
                 task: ConfigDict,
                 aliyun_cfg: ConfigDict,
                 max_num_workers: int = 32,
                 eval_with_gpu: list = ['plugin_eval'],
                 retry: int = 2,
                 debug: bool = False,
                 lark_bot_url: str = None):
        super().__init__(task=task, debug=debug, lark_bot_url=lark_bot_url)
        self.aliyun_cfg = aliyun_cfg
        self.max_num_workers = max_num_workers
        self.retry = retry

        self.eval_with_gpu = eval_with_gpu

        logger = get_logger()
        logger.warning(
            'To ensure the integrity of the log results, the log displayed '
            f'by {self.__class__.__name__} has a 10-second delay.')

    def launch(self, tasks: List[Dict[str, Any]]) -> List[Tuple[str, int]]:
        """Launch multiple tasks.

        Args:
            tasks (list[dict]): A list of task configs, usually generated by
                Partitioner.

        Returns:
            list[tuple[str, int]]: A list of (task name, exit code).
        """

        if not self.debug:
            status = track_parallel_progress(self._launch,
                                             tasks,
                                             nproc=self.max_num_workers,
                                             keep_order=False)
        else:
            status = [self._launch(task, random_sleep=False) for task in tasks]
        return status

    def _launch(self, cfg: ConfigDict, random_sleep: Optional[bool] = None):
        """Launch a single task.

        Args:
            cfg (ConfigDict): Task config.
            random_sleep (bool): Whether to sleep for a random time before
                running the command. When Aliyun has many tasks to schedule,
                its stability decreases. Therefore, when we need to submit a
                large number of tasks at once, we adopt the "random_sleep"
                strategy. Tasks that would have been submitted all at once are
                now evenly spread out over a 10-second period. Default: None.

        Returns:
            tuple[str, int]: Task name and exit code.
        """
        if random_sleep is None:
            random_sleep = (self.max_num_workers > 32)

        task = TASKS.build(dict(cfg=cfg, type=self.task_cfg['type']))
        num_gpus = task.num_gpus
        task_name = task.name

        is_eval_task = 'OpenICLEval' in task_name
        if is_eval_task and num_gpus == 0:
            for check_name in self.eval_with_gpu:
                if check_name in task_name:
                    num_gpus = 1
                    break

        # Dump task config to file
        mmengine.mkdir_or_exist('tmp/')
        param_file = f'tmp/{os.getpid()}_params.py'
        pwd = os.getcwd()
        try:
            cfg.dump(param_file)
            if self.aliyun_cfg.get('bashrc_path') is not None:
                # using user's conda env
                bashrc_path = self.aliyun_cfg['bashrc_path']
                assert osp.exists(bashrc_path)
                assert self.aliyun_cfg.get('conda_env_name') is not None
                conda_env_name = self.aliyun_cfg['conda_env_name']
                shell_cmd = (f'source {bashrc_path}; '
                             f'conda activate {conda_env_name}; ')
            else:
                # using public conda env
                # users can also set `python_env_path` to their
                # own env python path
                assert self.aliyun_cfg.get('python_env_path') is not None
                shell_cmd = (
                    f'export PATH={self.aliyun_cfg["python_env_path"]}/bin:$PATH; '  # noqa: E501
                    f'export PYTHONPATH={pwd}:$PYTHONPATH; ')

            huggingface_cache = self.aliyun_cfg.get('huggingface_cache')
            if huggingface_cache is not None:
                # HUGGINGFACE_HUB_CACHE is a Legacy env variable, here we set
                # `HF_HUB_CACHE` and `HUGGINGFACE_HUB_CACHE` for bc
                shell_cmd += f'export HF_HUB_CACHE={huggingface_cache}; '
                shell_cmd += f'export HUGGINGFACE_HUB_CACHE={huggingface_cache}; '  # noqa: E501

            torch_cache = self.aliyun_cfg.get('torch_cache')
            if torch_cache is not None:
                shell_cmd += f'export TORCH_HOME={torch_cache}; '

            hf_offline = self.aliyun_cfg.get('hf_offline', True)
            if hf_offline:
                shell_cmd += 'export HF_DATASETS_OFFLINE=1; export TRANSFORMERS_OFFLINE=1; export HF_EVALUATE_OFFLINE=1; '  # noqa: E501

            http_proxy = self.aliyun_cfg.get('http_proxy')
            if http_proxy is not None:
                shell_cmd += f'export http_proxy={http_proxy}; export https_proxy={http_proxy}; '  # noqa: E501
                shell_cmd += f'export HTTP_PROXY={http_proxy}; export HTTPS_PROXY={http_proxy}; '  # noqa: E501

            hf_endpoint = self.aliyun_cfg.get('hf_endpoint')
            if hf_endpoint is not None:
                shell_cmd += f'export HF_ENDPOINT={hf_endpoint}; '

            shell_cmd += f'cd {pwd}; '
            shell_cmd += '{task_cmd}'

            tmpl = ('dlc create job'
                    f" --command '{shell_cmd}'"
                    f' --name {task_name[:512]}'
                    ' --kind BatchJob'
                    f" -c {self.aliyun_cfg['dlc_config_path']}"
                    f" --workspace_id {self.aliyun_cfg['workspace_id']}"
                    ' --worker_count 1'
                    f' --worker_cpu {max(num_gpus * 8, 32)}'
                    f' --worker_gpu {num_gpus}'
                    f' --worker_memory {max(num_gpus * 128, 256)}'
                    f" --worker_image {self.aliyun_cfg['worker_image']}")
            get_cmd = partial(task.get_command,
                              cfg_path=param_file,
                              template=tmpl)
            cmd = get_cmd()

            logger = get_logger()
            logger.debug(f'Running command: {cmd}')

            # Run command with retry
            if self.debug:
                stdout = sys.stdout
            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')

            if random_sleep:
                time.sleep(random.randint(0, 10))

            def _run_within_retry():
                output = subprocess.getoutput(cmd)
                match = re.search(r'\|\s+(dlc[0-9a-z]+)\s+\|', output)
                if match is None:
                    raise RuntimeError(
                        f'Failed to launch dlc job for {output}')
                else:
                    job_id = match.group(1)
                stdout.write(output)

                pod_create_time = None
                pri_time = None
                initial_time = datetime.datetime.now()
                while True:
                    # 1. Avoid to request dlc too frequently.
                    # 2. DLC job may not be ready immediately after creation.
                    for _ in range(5):
                        time.sleep(2)
                        try:
                            job_info = json.loads(
                                subprocess.getoutput(f'dlc get job {job_id}'))
                            break
                        except:  # noqa: E722
                            pass
                    else:
                        raise RuntimeError(
                            f'Failed to get job info for {job_id}')

                    status = job_info['Status']
                    if status == 'Failed':
                        return -1
                    elif status == 'Succeeded':
                        return 0
                    elif status != 'Running':
                        continue

                    # The pod time could be different from the real time.
                    # Therefore we need to extract the pod start time from
                    # the `job_info` and calculate the `start_time` and
                    # `end_time` in pod.
                    if pod_create_time is None:
                        pod_create_time = job_info['GmtCreateTime']
                        pri_time = pod_create_time
                        pod_create_time = datetime.datetime.strptime(
                            pod_create_time, '%Y-%m-%dT%H:%M:%SZ')
                    elasped_time = datetime.datetime.now() - initial_time
                    cur_time = (pod_create_time +
                                elasped_time).strftime('%Y-%m-%dT%H:%M:%SZ')
                    logs_cmd = ('dlc logs'
                                f' {job_id} {job_id}-worker-0'
                                f" -c {self.aliyun_cfg['dlc_config_path']}"
                                f' --start_time {pri_time}'
                                f' --end_time {cur_time}')
                    log_output = subprocess.getoutput(logs_cmd)

                    if '[WARN] No logs found for the pod' not in log_output:
                        pri_time = cur_time
                        stdout.write(log_output)
                        stdout.flush()

            return_code = _run_within_retry()
            retry = self.retry
            output_paths = task.get_output_paths()
            while self._job_failed(return_code, output_paths) and retry > 0:
                retry -= 1
                cmd = get_cmd()
                return_code = _run_within_retry()
        finally:
            # Clean up
            os.remove(param_file)

        return task_name, return_code

    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)