run.py 14.2 KB
Newer Older
gaotongxiao's avatar
gaotongxiao committed
1
2
3
4
5
6
import argparse
import getpass
import os
import os.path as osp
from datetime import datetime

7
from mmengine.config import Config, DictAction
gaotongxiao's avatar
gaotongxiao committed
8

9
from opencompass.partitioners import MultimodalNaivePartitioner
Leymore's avatar
Leymore committed
10
from opencompass.registry import PARTITIONERS, RUNNERS, build_from_cfg
11
from opencompass.runners import SlurmRunner
Leymore's avatar
Leymore committed
12
13
from opencompass.summarizers import DefaultSummarizer
from opencompass.utils import LarkReporter, get_logger
14
15
from opencompass.utils.run import (exec_mm_infer_runner, fill_eval_cfg,
                                   fill_infer_cfg, get_config_from_arg)
gaotongxiao's avatar
gaotongxiao committed
16
17
18
19


def parse_args():
    parser = argparse.ArgumentParser(description='Run an evaluation task')
20
21
    parser.add_argument('config', nargs='?', help='Train config file path')

Tong Gao's avatar
Tong Gao committed
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
    # add mutually exclusive args `--slurm` `--dlc`, defaults to local runner
    # if "infer" or "eval" not specified
    launch_method = parser.add_mutually_exclusive_group()
    launch_method.add_argument('--slurm',
                               action='store_true',
                               default=False,
                               help='Whether to force tasks to run with srun. '
                               'If True, `--partition(-p)` must be set. '
                               'Defaults to False')
    launch_method.add_argument('--dlc',
                               action='store_true',
                               default=False,
                               help='Whether to force tasks to run on dlc. If '
                               'True, `--aliyun-cfg` must be set. Defaults'
                               ' to False')
37
38
39
40
41
    # multi-modal support
    parser.add_argument('--mm-eval',
                        help='Whether or not enable multimodal evaluation',
                        action='store_true',
                        default=False)
42
    # Add shortcut parameters (models, datasets and summarizer)
43
44
    parser.add_argument('--models', nargs='+', help='', default=None)
    parser.add_argument('--datasets', nargs='+', help='', default=None)
45
    parser.add_argument('--summarizer', help='', default=None)
gaotongxiao's avatar
gaotongxiao committed
46
47
48
49
50
51
52
    # add general args
    parser.add_argument('--debug',
                        help='Debug mode, in which scheduler will run tasks '
                        'in the single process, and output will not be '
                        'redirected to files',
                        action='store_true',
                        default=False)
Leymore's avatar
Leymore committed
53
54
55
56
57
58
    parser.add_argument('--dry-run',
                        help='Dry run mode, in which the scheduler will not '
                        'actually run the tasks, but only print the commands '
                        'to run',
                        action='store_true',
                        default=False)
gaotongxiao's avatar
gaotongxiao committed
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
    parser.add_argument('-m',
                        '--mode',
                        help='Running mode. You can choose "infer" if you '
                        'only want the inference results, or "eval" if you '
                        'already have the results and want to evaluate them, '
                        'or "viz" if you want to visualize the results.',
                        choices=['all', 'infer', 'eval', 'viz'],
                        default='all',
                        type=str)
    parser.add_argument('-r',
                        '--reuse',
                        nargs='?',
                        type=str,
                        const='latest',
                        help='Reuse previous outputs & results, and run any '
                        'missing jobs presented in the config. If its '
                        'argument is not specified, the latest results in '
                        'the work_dir will be reused. The argument should '
                        'also be a specific timestamp, e.g. 20230516_144254'),
    parser.add_argument('-w',
                        '--work-dir',
Tong Gao's avatar
Tong Gao committed
80
81
82
83
84
                        help='Work path, all the outputs will be '
                        'saved in this path, including the slurm logs, '
                        'the evaluation results, the summary results, etc.'
                        'If not specified, the work_dir will be set to '
                        './outputs/default.',
gaotongxiao's avatar
gaotongxiao committed
85
86
                        default=None,
                        type=str)
87
88
89
90
91
92
    parser.add_argument(
        '--config-dir',
        default='configs',
        help='Use the custom config directory instead of config/ to '
        'search the configs for datasets, models and summarizers',
        type=str)
gaotongxiao's avatar
gaotongxiao committed
93
94
95
96
97
98
    parser.add_argument('-l',
                        '--lark',
                        help='Report the running status to lark bot',
                        action='store_true',
                        default=False)
    parser.add_argument('--max-partition-size',
Tong Gao's avatar
Tong Gao committed
99
100
                        help='The maximum size of an infer task. Only '
                        'effective when "infer" is missing from the config.',
gaotongxiao's avatar
gaotongxiao committed
101
                        type=int,
102
                        default=40000),
gaotongxiao's avatar
gaotongxiao committed
103
104
    parser.add_argument(
        '--gen-task-coef',
Tong Gao's avatar
Tong Gao committed
105
106
        help='The dataset cost measurement coefficient for generation tasks, '
        'Only effective when "infer" is missing from the config.',
gaotongxiao's avatar
gaotongxiao committed
107
108
109
        type=int,
        default=20)
    parser.add_argument('--max-num-workers',
Tong Gao's avatar
Tong Gao committed
110
111
112
                        help='Max number of workers to run in parallel. '
                        'Will be overrideen by the "max_num_workers" argument '
                        'in the config.',
gaotongxiao's avatar
gaotongxiao committed
113
114
                        type=int,
                        default=32)
115
116
117
118
    parser.add_argument('--max-workers-per-gpu',
                        help='Max task to run in parallel on one GPU. '
                        'It will only be used in the local runner.',
                        type=int,
liushz's avatar
liushz committed
119
                        default=1)
gaotongxiao's avatar
gaotongxiao committed
120
121
    parser.add_argument(
        '--retry',
Tong Gao's avatar
Tong Gao committed
122
123
        help='Number of retries if the job failed when using slurm or dlc. '
        'Will be overrideen by the "retry" argument in the config.',
gaotongxiao's avatar
gaotongxiao committed
124
125
126
127
128
129
130
131
        type=int,
        default=2)
    # set srun args
    slurm_parser = parser.add_argument_group('slurm_args')
    parse_slurm_args(slurm_parser)
    # set dlc args
    dlc_parser = parser.add_argument_group('dlc_args')
    parse_dlc_args(dlc_parser)
132
133
134
    # set hf args
    hf_parser = parser.add_argument_group('hf_args')
    parse_hf_args(hf_parser)
gaotongxiao's avatar
gaotongxiao committed
135
136
137
138
139
140
    args = parser.parse_args()
    if args.slurm:
        assert args.partition is not None, (
            '--partition(-p) must be set if you want to use slurm')
    if args.dlc:
        assert os.path.exists(args.aliyun_cfg), (
Tong Gao's avatar
Tong Gao committed
141
            'When launching tasks using dlc, it needs to be configured '
gaotongxiao's avatar
gaotongxiao committed
142
143
144
145
146
147
            'in "~/.aliyun.cfg", or use "--aliyun-cfg $ALiYun-CFG_Path"'
            ' to specify a new path.')
    return args


def parse_slurm_args(slurm_parser):
Tong Gao's avatar
Tong Gao committed
148
    """These args are all for slurm launch."""
gaotongxiao's avatar
gaotongxiao committed
149
150
151
152
153
154
155
156
    slurm_parser.add_argument('-p',
                              '--partition',
                              help='Slurm partition name',
                              default=None,
                              type=str)
    slurm_parser.add_argument('-q',
                              '--quotatype',
                              help='Slurm quota type',
Tong Gao's avatar
Tong Gao committed
157
                              default=None,
gaotongxiao's avatar
gaotongxiao committed
158
                              type=str)
Haonan Li's avatar
Haonan Li committed
159
160
161
162
    slurm_parser.add_argument('--qos',
                              help='Slurm quality of service',
                              default=None,
                              type=str)
gaotongxiao's avatar
gaotongxiao committed
163
164
165


def parse_dlc_args(dlc_parser):
Tong Gao's avatar
Tong Gao committed
166
    """These args are all for dlc launch."""
gaotongxiao's avatar
gaotongxiao committed
167
168
169
170
171
172
    dlc_parser.add_argument('--aliyun-cfg',
                            help='The config path for aliyun config',
                            default='~/.aliyun.cfg',
                            type=str)


173
174
175
176
177
def parse_hf_args(hf_parser):
    """These args are all for the quick construction of HuggingFace models."""
    hf_parser.add_argument('--hf-path', type=str)
    hf_parser.add_argument('--peft-path', type=str)
    hf_parser.add_argument('--tokenizer-path', type=str)
178
179
180
181
182
183
184
185
    hf_parser.add_argument('--model-kwargs',
                           nargs='+',
                           action=DictAction,
                           default={})
    hf_parser.add_argument('--tokenizer-kwargs',
                           nargs='+',
                           action=DictAction,
                           default={})
186
187
188
189
190
191
192
    hf_parser.add_argument('--max-out-len', type=int)
    hf_parser.add_argument('--max-seq-len', type=int)
    hf_parser.add_argument('--no-batch-padding',
                           action='store_true',
                           default=False)
    hf_parser.add_argument('--batch-size', type=int)
    hf_parser.add_argument('--num-gpus', type=int)
193
    hf_parser.add_argument('--pad-token-id', type=int)
194
195


gaotongxiao's avatar
gaotongxiao committed
196
197
def main():
    args = parse_args()
Leymore's avatar
Leymore committed
198
199
    if args.dry_run:
        args.debug = True
gaotongxiao's avatar
gaotongxiao committed
200
201
202
    # initialize logger
    logger = get_logger(log_level='DEBUG' if args.debug else 'INFO')

203
    cfg = get_config_from_arg(args)
gaotongxiao's avatar
gaotongxiao committed
204
205
206
207
208
209
210
211
212
    if args.work_dir is not None:
        cfg['work_dir'] = args.work_dir
    else:
        cfg.setdefault('work_dir', './outputs/default/')

    # cfg_time_str defaults to the current time
    cfg_time_str = dir_time_str = datetime.now().strftime('%Y%m%d_%H%M%S')
    if args.reuse:
        if args.reuse == 'latest':
Leymore's avatar
Leymore committed
213
214
215
216
217
218
            if not os.path.exists(cfg.work_dir) or not os.listdir(
                    cfg.work_dir):
                logger.warning('No previous results to reuse!')
            else:
                dirs = os.listdir(cfg.work_dir)
                dir_time_str = sorted(dirs)[-1]
gaotongxiao's avatar
gaotongxiao committed
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
        else:
            dir_time_str = args.reuse
        logger.info(f'Reusing experiements from {dir_time_str}')
    elif args.mode in ['eval', 'viz']:
        raise ValueError('You must specify -r or --reuse when running in eval '
                         'or viz mode!')

    # update "actual" work_dir
    cfg['work_dir'] = osp.join(cfg.work_dir, dir_time_str)
    os.makedirs(osp.join(cfg.work_dir, 'configs'), exist_ok=True)

    # dump config
    output_config_path = osp.join(cfg.work_dir, 'configs',
                                  f'{cfg_time_str}.py')
    cfg.dump(output_config_path)
    # Config is intentally reloaded here to avoid initialized
    # types cannot be serialized
236
    cfg = Config.fromfile(output_config_path, format_python_code=False)
gaotongxiao's avatar
gaotongxiao committed
237
238
239
240
241
242
243
244
245

    # report to lark bot if specify --lark
    if not args.lark:
        cfg['lark_bot_url'] = None
    elif cfg.get('lark_bot_url', None):
        content = f'{getpass.getuser()}\'s task has been launched!'
        LarkReporter(cfg['lark_bot_url']).post(content)

    if args.mode in ['all', 'infer']:
Tong Gao's avatar
Tong Gao committed
246
247
248
        # When user have specified --slurm or --dlc, or have not set
        # "infer" in config, we will provide a default configuration
        # for infer
Tong Gao's avatar
Tong Gao committed
249
250
251
252
253
        if (args.dlc or args.slurm) and cfg.get('infer', None):
            logger.warning('You have set "infer" in the config, but '
                           'also specified --slurm or --dlc. '
                           'The "infer" configuration will be overridden by '
                           'your runtime arguments.')
Yuan Liu's avatar
Yuan Liu committed
254
255
256
257
258
259
260
        # Check whether run multimodal evaluation
        if args.mm_eval:
            partitioner = MultimodalNaivePartitioner(
                osp.join(cfg['work_dir'], 'predictions/'))
            tasks = partitioner(cfg)
            exec_mm_infer_runner(tasks, args, cfg)
            return
261
262
263
264
265
266
267
268

        if args.dlc or args.slurm or cfg.get('infer', None) is None:
            fill_infer_cfg(cfg, args)

        if args.partition is not None:
            if RUNNERS.get(cfg.infer.runner.type) == SlurmRunner:
                cfg.infer.runner.partition = args.partition
                cfg.infer.runner.quotatype = args.quotatype
Tong Gao's avatar
Tong Gao committed
269
        else:
270
271
272
273
274
275
276
277
278
279
280
281
282
            logger.warning('SlurmRunner is not used, so the partition '
                           'argument is ignored.')
        if args.debug:
            cfg.infer.runner.debug = True
        if args.lark:
            cfg.infer.runner.lark_bot_url = cfg['lark_bot_url']
        cfg.infer.partitioner['out_dir'] = osp.join(cfg['work_dir'],
                                                    'predictions/')
        partitioner = PARTITIONERS.build(cfg.infer.partitioner)
        tasks = partitioner(cfg)
        if args.dry_run:
            return
        runner = RUNNERS.build(cfg.infer.runner)
283
284
285
286
287
        # Add extra attack config if exists
        if hasattr(cfg, 'attack'):
            for task in tasks:
                cfg.attack.dataset = task.datasets[0][0].abbr
                task.attack = cfg.attack
288
        runner(tasks)
gaotongxiao's avatar
gaotongxiao committed
289
290
291

    # evaluate
    if args.mode in ['all', 'eval']:
Tong Gao's avatar
Tong Gao committed
292
293
294
        # When user have specified --slurm or --dlc, or have not set
        # "eval" in config, we will provide a default configuration
        # for eval
Tong Gao's avatar
Tong Gao committed
295
296
297
298
299
        if (args.dlc or args.slurm) and cfg.get('eval', None):
            logger.warning('You have set "eval" in the config, but '
                           'also specified --slurm or --dlc. '
                           'The "eval" configuration will be overridden by '
                           'your runtime arguments.')
300

Tong Gao's avatar
Tong Gao committed
301
        if args.dlc or args.slurm or cfg.get('eval', None) is None:
302
303
304
            fill_eval_cfg(cfg, args)

        if args.partition is not None:
Leymore's avatar
Leymore committed
305
            if RUNNERS.get(cfg.eval.runner.type) == SlurmRunner:
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
                cfg.eval.runner.partition = args.partition
                cfg.eval.runner.quotatype = args.quotatype
            else:
                logger.warning('SlurmRunner is not used, so the partition '
                               'argument is ignored.')
        if args.debug:
            cfg.eval.runner.debug = True
        if args.lark:
            cfg.eval.runner.lark_bot_url = cfg['lark_bot_url']
        cfg.eval.partitioner['out_dir'] = osp.join(cfg['work_dir'], 'results/')
        partitioner = PARTITIONERS.build(cfg.eval.partitioner)
        tasks = partitioner(cfg)
        if args.dry_run:
            return
        runner = RUNNERS.build(cfg.eval.runner)
        runner(tasks)
gaotongxiao's avatar
gaotongxiao committed
322
323
324

    # visualize
    if args.mode in ['all', 'eval', 'viz']:
Leymore's avatar
Leymore committed
325
326
327
328
329
        summarizer_cfg = cfg.get('summarizer', {})
        if not summarizer_cfg or summarizer_cfg.get('type', None) is None:
            summarizer_cfg['type'] = DefaultSummarizer
        summarizer_cfg['config'] = cfg
        summarizer = build_from_cfg(summarizer_cfg)
gaotongxiao's avatar
gaotongxiao committed
330
331
332
333
334
        summarizer.summarize(time_str=cfg_time_str)


if __name__ == '__main__':
    main()