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

from mmengine.config import Config

from opencompass.partitioners import NaivePartitioner, SizePartitioner
Tong Gao's avatar
Tong Gao committed
10
from opencompass.registry import PARTITIONERS, RUNNERS
gaotongxiao's avatar
gaotongxiao committed
11
12
13
14
15
16
17
from opencompass.runners import DLCRunner, LocalRunner, SlurmRunner
from opencompass.utils import LarkReporter, Summarizer, get_logger


def parse_args():
    parser = argparse.ArgumentParser(description='Run an evaluation task')
    parser.add_argument('config', help='Train config file path')
Tong Gao's avatar
Tong Gao committed
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
    # 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')
gaotongxiao's avatar
gaotongxiao committed
33
34
35
36
37
38
39
    # 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
40
41
42
43
44
45
    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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
    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
67
68
69
70
71
                        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
72
73
74
75
76
77
78
79
                        default=None,
                        type=str)
    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
80
81
                        help='The maximum size of an infer task. Only '
                        'effective when "infer" is missing from the config.',
gaotongxiao's avatar
gaotongxiao committed
82
83
84
85
                        type=int,
                        default=2000),
    parser.add_argument(
        '--gen-task-coef',
Tong Gao's avatar
Tong Gao committed
86
87
        help='The dataset cost measurement coefficient for generation tasks, '
        'Only effective when "infer" is missing from the config.',
gaotongxiao's avatar
gaotongxiao committed
88
89
90
        type=int,
        default=20)
    parser.add_argument('--max-num-workers',
Tong Gao's avatar
Tong Gao committed
91
92
93
                        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
94
95
96
97
                        type=int,
                        default=32)
    parser.add_argument(
        '--retry',
Tong Gao's avatar
Tong Gao committed
98
99
        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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
        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)
    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
114
            'When launching tasks using dlc, it needs to be configured '
gaotongxiao's avatar
gaotongxiao committed
115
116
117
118
119
120
            '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
121
    """These args are all for slurm launch."""
gaotongxiao's avatar
gaotongxiao committed
122
123
124
125
126
127
128
129
    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
130
                              default=None,
gaotongxiao's avatar
gaotongxiao committed
131
132
133
134
                              type=str)


def parse_dlc_args(dlc_parser):
Tong Gao's avatar
Tong Gao committed
135
    """These args are all for dlc launch."""
gaotongxiao's avatar
gaotongxiao committed
136
137
138
139
140
141
142
143
    dlc_parser.add_argument('--aliyun-cfg',
                            help='The config path for aliyun config',
                            default='~/.aliyun.cfg',
                            type=str)


def main():
    args = parse_args()
Leymore's avatar
Leymore committed
144
145
    if args.dry_run:
        args.debug = True
gaotongxiao's avatar
gaotongxiao committed
146
147
148
    # initialize logger
    logger = get_logger(log_level='DEBUG' if args.debug else 'INFO')

149
    cfg = Config.fromfile(args.config, format_python_code=False)
gaotongxiao's avatar
gaotongxiao committed
150
151
152
153
154
155
156
157
158
    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
159
160
161
162
163
164
            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
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
        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
182
    cfg = Config.fromfile(output_config_path, format_python_code=False)
gaotongxiao's avatar
gaotongxiao committed
183
184
185
186
187
188
189
190
191

    # 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
192
193
194
        # 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
195
196
197
198
199
200
201
202
203
204
205
206
        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.')
        if args.dlc or args.slurm or cfg.get('infer', None) is None:
            # Use SizePartitioner to split into subtasks
            partitioner = SizePartitioner(
                osp.join(cfg['work_dir'], 'predictions/'),
                max_task_size=args.max_partition_size,
                gen_task_coef=args.gen_task_coef)
            tasks = partitioner(cfg)
Leymore's avatar
Leymore committed
207
208
            if args.dry_run:
                return
Tong Gao's avatar
Tong Gao committed
209
210
            # execute the infer subtasks
            exec_infer_runner(tasks, args, cfg)
Tong Gao's avatar
Tong Gao committed
211
212
        # If they have specified "infer" in config and haven't used --slurm
        # or --dlc, just follow the config
Tong Gao's avatar
Tong Gao committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
        else:
            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
            else:
                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)
Leymore's avatar
Leymore committed
229
230
            if args.dry_run:
                return
Tong Gao's avatar
Tong Gao committed
231
232
            runner = RUNNERS.build(cfg.infer.runner)
            runner(tasks)
gaotongxiao's avatar
gaotongxiao committed
233
234
235

    # evaluate
    if args.mode in ['all', 'eval']:
Tong Gao's avatar
Tong Gao committed
236
237
238
        # 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
239
240
241
242
243
244
245
246
247
248
        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.')
        if args.dlc or args.slurm or cfg.get('eval', None) is None:
            # Use NaivePartitioner,not split
            partitioner = NaivePartitioner(
                osp.join(cfg['work_dir'], 'results/'))
            tasks = partitioner(cfg)
Leymore's avatar
Leymore committed
249
250
            if args.dry_run:
                return
Tong Gao's avatar
Tong Gao committed
251
252
            # execute the eval tasks
            exec_eval_runner(tasks, args, cfg)
Tong Gao's avatar
Tong Gao committed
253
254
        # If they have specified "eval" in config and haven't used --slurm
        # or --dlc, just follow the config
Tong Gao's avatar
Tong Gao committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
        else:
            if args.partition is not None:
                if RUNNERS.get(cfg.infer.runner.type) == SlurmRunner:
                    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)
Leymore's avatar
Leymore committed
271
272
            if args.dry_run:
                return
Tong Gao's avatar
Tong Gao committed
273
274
            runner = RUNNERS.build(cfg.eval.runner)
            runner(tasks)
gaotongxiao's avatar
gaotongxiao committed
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299

    # visualize
    if args.mode in ['all', 'eval', 'viz']:
        summarizer = Summarizer(cfg)
        summarizer.summarize(time_str=cfg_time_str)


def exec_infer_runner(tasks, args, cfg):
    """execute infer runner according to args."""
    if args.slurm:
        runner = SlurmRunner(dict(type='OpenICLInferTask'),
                             max_num_workers=args.max_num_workers,
                             partition=args.partition,
                             quotatype=args.quotatype,
                             retry=args.retry,
                             debug=args.debug,
                             lark_bot_url=cfg['lark_bot_url'])
    elif args.dlc:
        runner = DLCRunner(dict(type='OpenICLInferTask'),
                           max_num_workers=args.max_num_workers,
                           aliyun_cfg=Config.fromfile(args.aliyun_cfg),
                           retry=args.retry,
                           debug=args.debug,
                           lark_bot_url=cfg['lark_bot_url'])
    else:
Tong Gao's avatar
Tong Gao committed
300
301
302
303
        runner = LocalRunner(task=dict(type='OpenICLInferTask'),
                             max_num_workers=args.max_num_workers,
                             debug=args.debug,
                             lark_bot_url=cfg['lark_bot_url'])
gaotongxiao's avatar
gaotongxiao committed
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
    runner(tasks)


def exec_eval_runner(tasks, args, cfg):
    """execute infer runner according to args."""
    if args.slurm:
        runner = SlurmRunner(dict(type='OpenICLEvalTask'),
                             max_num_workers=args.max_num_workers,
                             partition=args.partition,
                             quotatype=args.quotatype,
                             retry=args.retry,
                             debug=args.debug,
                             lark_bot_url=cfg['lark_bot_url'])
    elif args.dlc:
        runner = DLCRunner(dict(type='OpenICLEvalTask'),
                           max_num_workers=args.max_num_workers,
                           aliyun_cfg=Config.fromfile(args.aliyun_cfg),
                           retry=args.retry,
                           debug=args.debug,
                           lark_bot_url=cfg['lark_bot_url'])
    else:
Tong Gao's avatar
Tong Gao committed
325
326
327
328
        runner = LocalRunner(task=dict(type='OpenICLEvalTask'),
                             max_num_workers=args.max_num_workers,
                             debug=args.debug,
                             lark_bot_url=cfg['lark_bot_url'])
gaotongxiao's avatar
gaotongxiao committed
329
330
331
332
333
    runner(tasks)


if __name__ == '__main__':
    main()