| | import os |
| | import shutil |
| | import subprocess |
| | import time |
| | from collections import deque |
| | from copy import deepcopy |
| | from dataclasses import asdict, dataclass, field |
| | from typing import Any, Dict, List |
| |
|
| | import json |
| | import torch |
| |
|
| | from swift.llm import ExportArguments |
| | from swift.utils import find_free_port, get_device_count, get_logger |
| |
|
| | logger = get_logger() |
| |
|
| |
|
| | @dataclass |
| | class Experiment: |
| |
|
| | name: str |
| |
|
| | cmd: str |
| |
|
| | group: str |
| |
|
| | requirements: Dict = field(default_factory=dict) |
| |
|
| | eval_requirements: Dict = field(default_factory=dict) |
| |
|
| | eval_dataset: List = field(default_factory=list) |
| |
|
| | args: Dict = field(default_factory=dict) |
| |
|
| | env: Dict = field(default_factory=dict) |
| |
|
| | record: Dict = field(default_factory=dict) |
| |
|
| | create_time: float = None |
| |
|
| | runtime: Dict = field(default_factory=dict) |
| |
|
| | input_args: Any = None |
| |
|
| | do_eval = False |
| |
|
| | def __init__(self, |
| | name, |
| | cmd, |
| | group, |
| | requirements=None, |
| | eval_requirements=None, |
| | eval_dataset=None, |
| | args=None, |
| | input_args=None, |
| | **kwargs): |
| | self.name = name |
| | self.cmd = cmd |
| | self.group = group |
| | self.requirements = requirements or {} |
| | self.args = args or {} |
| | self.record = {} |
| | self.env = {} |
| | self.runtime = {} |
| | self.input_args = input_args |
| | self.eval_requirements = eval_requirements or {} |
| | self.eval_dataset = eval_dataset or [] |
| | if self.cmd == 'eval': |
| | self.do_eval = True |
| |
|
| | def load(self, _json): |
| | self.name = _json['name'] |
| | self.cmd = _json['cmd'] |
| | self.requirements = _json['requirements'] |
| | self.args = _json['args'] |
| | self.record = _json['record'] |
| | self.env = _json['env'] |
| | self.create_time = _json['create_time'] |
| |
|
| | @property |
| | def priority(self): |
| | return self.requirements.get('gpu', 0) |
| |
|
| | def to_dict(self): |
| | _dict = asdict(self) |
| | _dict.pop('runtime') |
| | _dict.pop('input_args') |
| | return _dict |
| |
|
| |
|
| | class ExpManager: |
| |
|
| | RESULT_FILE = 'result.jsonl' |
| |
|
| | def __init__(self): |
| | self.exps = [] |
| |
|
| | def assert_gpu_not_overlap(self): |
| | all_gpus = set() |
| | for exp in self.exps: |
| | gpus = exp.runtime['env']['CUDA_VISIBLE_DEVICES'].split(',') |
| | if all_gpus & set(gpus): |
| | raise ValueError(f'GPU overlap: {self.exps}!') |
| | all_gpus.update(gpus) |
| |
|
| | def run(self, exp: Experiment): |
| | if os.path.exists(os.path.join(exp.input_args.save_dir, exp.name + '.json')): |
| | with open(os.path.join(exp.input_args.save_dir, exp.name + '.json'), 'r', encoding='utf-8') as f: |
| | _json = json.load(f) |
| | if exp.eval_dataset and 'eval_result' not in _json['record']: |
| | if not exp.do_eval: |
| | logger.info(f'Experiment {exp.name} need eval, load from file.') |
| | exp.load(_json) |
| | exp.do_eval = True |
| | else: |
| | logger.warn(f'Experiment {exp.name} already done, skip') |
| | return |
| |
|
| | if exp.do_eval: |
| | runtime = self._build_eval_cmd(exp) |
| | exp.runtime = runtime |
| | envs = deepcopy(runtime.get('env', {})) |
| | envs.update(os.environ) |
| | logger.info(f'Running cmd: {runtime["running_cmd"]}, env: {runtime.get("env", {})}') |
| | os.makedirs('exp', exist_ok=True) |
| | log_file = os.path.join('exp', f'{exp.name}.eval.log') |
| | exp.handler = subprocess.Popen(runtime['running_cmd'] + f' > {log_file} 2>&1', env=envs, shell=True) |
| | self.exps.append(exp) |
| | self.assert_gpu_not_overlap() |
| | return |
| |
|
| | if any([exp.name == e.name for e in self.exps]): |
| | raise ValueError(f'Why exp name duplicate? {exp.name}') |
| | elif exp.cmd == 'export' and any([exp.cmd == 'export' for exp in self.exps]): |
| | raise AssertionError('Cannot run parallel export task.') |
| | else: |
| | exp.create_time = time.time() |
| | runtime = self._build_cmd(exp) |
| | exp.runtime = runtime |
| | envs = deepcopy(runtime.get('env', {})) |
| | envs.update(os.environ) |
| | logger.info(f'Running cmd: {runtime["running_cmd"]}, env: {runtime.get("env", {})}') |
| | os.makedirs('exp', exist_ok=True) |
| | log_file = os.path.join('exp', f'{exp.name}.{exp.cmd}.log') |
| | exp.handler = subprocess.Popen(runtime['running_cmd'] + f' > {log_file} 2>&1', env=envs, shell=True) |
| | self.exps.append(exp) |
| | self.assert_gpu_not_overlap() |
| |
|
| | def _build_eval_cmd(self, exp: Experiment): |
| | gpu = exp.eval_requirements.get('gpu', None) |
| | env = {} |
| | allocated = [] |
| | if gpu: |
| | allocated = self._find_free_gpu(int(gpu)) |
| | assert allocated, 'No free gpu for now!' |
| | allocated = [str(gpu) for gpu in allocated] |
| | env['CUDA_VISIBLE_DEVICES'] = ','.join(allocated) |
| |
|
| | best_model_checkpoint = exp.record.get('best_model_checkpoint') |
| | eval_dataset = exp.eval_dataset |
| | if best_model_checkpoint is not None: |
| | if not os.path.exists(os.path.join(best_model_checkpoint, 'args.json')): |
| | cmd = f'swift eval --ckpt_dir {best_model_checkpoint} ' \ |
| | + f'--infer_backend pt --train_type full --eval_dataset {" ".join(eval_dataset)}' |
| | else: |
| | cmd = f'swift eval --model {exp.args.get("model")} --infer_backend pt ' \ |
| | f'--eval_dataset {" ".join(eval_dataset)}' |
| |
|
| | return { |
| | 'running_cmd': cmd, |
| | 'gpu': allocated, |
| | 'env': env, |
| | } |
| |
|
| | def _build_cmd(self, exp: Experiment): |
| | gpu = exp.requirements.get('gpu', None) |
| | env = {} |
| | allocated = [] |
| | if gpu: |
| | allocated = self._find_free_gpu(int(gpu)) |
| | assert allocated, 'No free gpu for now!' |
| | allocated = [str(gpu) for gpu in allocated] |
| | env['CUDA_VISIBLE_DEVICES'] = ','.join(allocated) |
| | if int(exp.requirements.get('ddp', 1)) > 1: |
| | env['NPROC_PER_NODE'] = exp.requirements.get('ddp') |
| | env['MASTER_PORT'] = str(find_free_port()) |
| |
|
| | if exp.cmd == 'sft': |
| | from swift.llm import TrainArguments |
| | args = exp.args |
| | sft_args = TrainArguments(**args) |
| | args['output_dir'] = sft_args.output_dir |
| | args['logging_dir'] = sft_args.logging_dir |
| | args['add_version'] = False |
| | os.makedirs(sft_args.output_dir, exist_ok=True) |
| | os.makedirs(sft_args.logging_dir, exist_ok=True) |
| | cmd = 'swift sft ' |
| | for key, value in args.items(): |
| | cmd += f' --{key} {value}' |
| | elif exp.cmd == 'rlhf': |
| | from swift.llm import RLHFArguments |
| | args = exp.args |
| | rlhf_args = RLHFArguments(**args) |
| | args['output_dir'] = rlhf_args.output_dir |
| | args['logging_dir'] = rlhf_args.logging_dir |
| | args['add_version'] = False |
| | os.makedirs(rlhf_args.output_dir, exist_ok=True) |
| | os.makedirs(rlhf_args.logging_dir, exist_ok=True) |
| | cmd = 'swift rlhf ' |
| | for key, value in args.items(): |
| | cmd += f' --{key} {value}' |
| | elif exp.cmd == 'export': |
| | args = exp.args |
| | cmd = 'swift export ' |
| | for key, value in args.items(): |
| | cmd += f' --{key} {value}' |
| | else: |
| | raise ValueError(f'Unsupported cmd type: {exp.cmd}') |
| | return { |
| | 'running_cmd': cmd, |
| | 'gpu': allocated, |
| | 'env': env, |
| | 'logging_dir': args.get('logging_dir'), |
| | 'output_dir': args.get('output_dir', args.get('ckpt_dir')) |
| | } |
| |
|
| | def _find_free_gpu(self, n): |
| | all_gpus = set() |
| | for exp in self.exps: |
| | all_gpus.update(exp.runtime.get('gpu', set())) |
| | all_gpus = {int(g) for g in all_gpus} |
| | free_gpu = set(range(get_device_count())) - all_gpus |
| | if len(free_gpu) < n: |
| | return None |
| | return list(free_gpu)[:n] |
| |
|
| | def prepare_experiments(self, args: Any): |
| | experiments = [] |
| | for config_file in args.config: |
| | with open(config_file, 'r', encoding='utf-8') as f: |
| | group = os.path.basename(config_file) |
| | group = group[:-5] |
| | content = json.load(f) |
| | exps = content['experiment'] |
| | for exp in exps: |
| | main_cfg = deepcopy(content) |
| | name = exp['name'] |
| | cmd = main_cfg['cmd'] |
| | run_args = main_cfg['args'] |
| | env = main_cfg.get('env', {}) |
| | requirements = main_cfg.get('requirements', {}) |
| | eval_requirements = main_cfg.get('eval_requirements', {}) |
| | eval_dataset = main_cfg.get('eval_dataset', {}) |
| | if 'args' in exp: |
| | run_args.update(exp['args']) |
| | if 'requirements' in exp: |
| | requirements.update(exp['requirements']) |
| | if 'env' in exp: |
| | env.update(exp['env']) |
| | experiments.append( |
| | Experiment( |
| | group=group, |
| | name=name, |
| | cmd=cmd, |
| | args=run_args, |
| | env=env, |
| | requirements=requirements, |
| | eval_requirements=eval_requirements, |
| | eval_dataset=eval_dataset, |
| | input_args=args)) |
| | return experiments |
| |
|
| | @staticmethod |
| | def _get_metric(exp: Experiment): |
| | if exp.do_eval: |
| | if os.path.isfile(os.path.join('exp', f'{exp.name}.eval.log')): |
| | with open(os.path.join('exp', f'{exp.name}.eval.log'), 'r', encoding='utf-8') as f: |
| | for line in f.readlines(): |
| | if 'Final report:' in line: |
| | return json.loads(line.split('Final report:')[1].replace('\'', '"')) |
| | elif exp.cmd == 'export': |
| | exp_args = ExportArguments(**exp.args) |
| | if exp_args.quant_bits > 0: |
| | if exp_args.ckpt_dir is None: |
| | path = f'{exp_args.model_type}-{exp_args.quant_method}-int{exp_args.quant_bits}' |
| | else: |
| | ckpt_dir, ckpt_name = os.path.split(exp_args.ckpt_dir) |
| | path = os.path.join(ckpt_dir, f'{ckpt_name}-{exp_args.quant_method}-int{exp_args.quant_bits}') |
| | else: |
| | ckpt_dir, ckpt_name = os.path.split(exp_args.ckpt_dir) |
| | path = os.path.join(ckpt_dir, f'{ckpt_name}-merged') |
| | if os.path.exists(path): |
| | shutil.rmtree(exp.name, ignore_errors=True) |
| | os.makedirs(exp.name, exist_ok=True) |
| | shutil.move(path, os.path.join(exp.name, path)) |
| | return { |
| | 'best_model_checkpoint': os.path.join(exp.name, path), |
| | } |
| | else: |
| | logging_dir = exp.runtime.get('logging_dir') |
| | logging_file = os.path.join(logging_dir, '..', 'logging.jsonl') |
| | if os.path.isfile(logging_file): |
| | with open(logging_file, 'r', encoding='utf-8') as f: |
| | for line in f.readlines(): |
| | if 'model_info' in line: |
| | return json.loads(line) |
| | return None |
| |
|
| | @staticmethod |
| | def write_record(exp: Experiment): |
| | target_dir = exp.input_args.save_dir |
| | file = os.path.join(target_dir, exp.name + '.json') |
| | with open(file, 'w', encoding='utf-8') as f: |
| | f.write(json.dumps(exp.to_dict()) + '\n') |
| |
|
| | def _poll(self): |
| | while True: |
| | time.sleep(5) |
| |
|
| | has_finished = False |
| | for exp in self.exps: |
| | rt = exp.handler.poll() |
| | if rt is None: |
| | continue |
| |
|
| | has_finished = True |
| | if rt == 0: |
| | if not exp.do_eval: |
| | all_metric = self._get_metric(exp) |
| | if all_metric: |
| | exp.record.update(all_metric) |
| | if exp.eval_dataset: |
| | exp.do_eval = True |
| | self.exp_queue.appendleft(exp) |
| | self.write_record(exp) |
| | else: |
| | logger.error(f'Running {exp.name} task, but no result found') |
| | else: |
| | all_metric = self._get_metric(exp) |
| | exp.record['eval_result'] = all_metric |
| | if all_metric: |
| | self.write_record(exp) |
| | else: |
| | logger.error(f'Running {exp.name} eval task, but no eval result found') |
| | logger.info(f'Running {exp.name} finished with return code: {rt}') |
| |
|
| | if has_finished: |
| | self.exps = [exp for exp in self.exps if exp.handler.poll() is None] |
| | break |
| |
|
| | def begin(self, args: Any): |
| | exps = self.prepare_experiments(args) |
| | logger.info(f'all exps: {exps}') |
| | exps.sort(key=lambda e: e.priority) |
| | self.exp_queue = deque() |
| | for exp in exps: |
| | self.exp_queue.append(exp) |
| |
|
| | while len(self.exp_queue) or len(self.exps) > 0: |
| | while len(self.exp_queue): |
| | try: |
| | logger.info(f'Running exp: {self.exp_queue[0].name}') |
| | self.run(self.exp_queue[0]) |
| | except Exception as e: |
| | if not isinstance(e, AssertionError): |
| | logger.error(f'Adding exp {self.exp_queue[0].name} error because of:') |
| | logger.error(e) |
| | self.exp_queue.popleft() |
| | else: |
| | logger.info(f'Adding exp {self.exp_queue[0].name} error because of:', str(e)) |
| | if 'no free gpu' in str(e).lower(): |
| | break |
| | else: |
| | continue |
| | else: |
| | self.exp_queue.popleft() |
| | self._poll() |
| | logger.info(f'Run task finished because of exp queue: {self.exp_queue} and exps: {self.exps}') |
| |
|
| |
|
| | def find_all_config(dir_or_file: str): |
| | if os.path.isfile(dir_or_file): |
| | return [dir_or_file] |
| | else: |
| | configs = [] |
| | for dirpath, dirnames, filenames in os.walk(dir_or_file): |
| | for name in filenames: |
| | if name.endswith('.json') and 'ipynb' not in dirpath: |
| | configs.append(os.path.join(dirpath, name)) |
| | return configs |
| |
|
