Spaces:
Sleeping
Sleeping
| import os | |
| import argparse | |
| import copy | |
| from typing import List, Dict, Optional | |
| from omegaconf import OmegaConf, DictConfig | |
| from src.infra.utils import cfg_to_args_str | |
| #---------------------------------------------------------------------------- | |
| HYDRA_ARGS = "hydra.run.dir=. hydra.output_subdir=null hydra/job_logging=disabled hydra/hydra_logging=disabled" | |
| #---------------------------------------------------------------------------- | |
| def batch_launch(launcher: str, experiments_dir: os.PathLike, cfg: DictConfig, datasets: List[str], print_only: bool, time: str, use_qos: bool=False, other_args: Dict={}, num_gpus: int=4, *args, **kwargs): | |
| for dataset in datasets: | |
| for exp_args in construct_experiments_args(cfg, *args, **kwargs): | |
| exp_args['sbatch_args.time'] = time | |
| exp_args['experiments_dir'] = experiments_dir | |
| exp_args['dataset'] = dataset | |
| exp_args['env'] = 'ibex' | |
| exp_args['use_qos'] = use_qos | |
| exp_args = {**exp_args, **other_args} | |
| curr_exp_args_str = cfg_to_args_str(exp_args, use_dashes=False) | |
| launching_command = f"{launcher} num_gpus={num_gpus} {curr_exp_args_str}" | |
| if print_only: | |
| os.makedirs(exp_args['experiments_dir'], exist_ok=True) | |
| print(launching_command) | |
| else: | |
| os.system(launching_command) | |
| #---------------------------------------------------------------------------- | |
| def construct_experiments_args(cfg: DictConfig, experiments_list: Optional[List[str]]=None, suffix: str="") -> List[Dict]: | |
| args_dicts = [] | |
| common_cfg = cfg.get('common_args', {}) | |
| for exp_name, exp_cfg in to_dict(cfg.experiments).items(): | |
| if not experiments_list is None and not exp_name in experiments_list: | |
| continue | |
| curr_exp_cfg = {**copy.deepcopy(to_dict(common_cfg)), **to_dict(exp_cfg)} | |
| curr_exp_cfg['exp_suffix'] = f'{exp_name}{suffix}' | |
| args_dicts.append(curr_exp_cfg) | |
| return args_dicts | |
| #---------------------------------------------------------------------------- | |
| def to_dict(cfg) -> Dict: | |
| return OmegaConf.to_container(OmegaConf.create({**cfg})) | |
| #---------------------------------------------------------------------------- | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Experiments launcher") | |
| parser.add_argument('-e', '--series_name', type=str, required=True, help="Which experiments series to launch?") | |
| parser.add_argument('-d', '--datasets', required=True, type=str, help='Comma-separate list of datasets') | |
| parser.add_argument('-p', '--print_only', action='store_true', help='Just print commands and exit?') | |
| parser.add_argument('-t', '--time', type=str, default='1-0', help='Which time to specify for the sbatch command?') | |
| parser.add_argument('-q', '--use_qos', action='store_true', help='Should we use QoS to launch jobs?') | |
| parser.add_argument('--experiments_list', type=str, help='Should we run only some specific experiments from this experiments series?') | |
| parser.add_argument('--other_args', type=str, default="", help='Additional arguments for the experiments') | |
| parser.add_argument('--suffix', type=str, default="", help='Additional suffix for the experiments') | |
| parser.add_argument('--num_gpus', type=int, default=4, help='Number of GPUs to use per each experiment') | |
| parser.add_argument('--project_dir', type=str, default=os.getcwd(), help='Project directory path') | |
| parser.add_argument('--project_dir_for_exps_cfg', type=str, help="Overwrite the project directory to use for experiments.yaml. Useful for debugging the config.") | |
| args = parser.parse_args() | |
| os.chdir(args.project_dir) | |
| user = os.environ.get('USER', 'unknown') | |
| python_bin = os.path.join(args.project_dir, 'env/bin/python') | |
| launcher = f"{python_bin} src/infra/launch.py {HYDRA_ARGS} +quiet=true slurm=true" | |
| experiments_dir = f'experiments/{user}/{args.series_name}' | |
| exps_cfg_path = os.path.join(args.project_dir if args.project_dir_for_exps_cfg is None else args.project_dir_for_exps_cfg, 'src/infra/experiments.yaml') | |
| all_exp_series = OmegaConf.load(exps_cfg_path) | |
| assert args.series_name in all_exp_series, f"Experiments series not found: {args.series_name}" | |
| cfg = all_exp_series[args.series_name] | |
| datasets = args.datasets.split(',') | |
| experiments_list = None if args.experiments_list is None else args.experiments_list.split(',') | |
| other_args = {kv.split('=')[0]: kv.split('=')[1] for kv in args.other_args.split(',') if len(kv.split('=')) == 2} | |
| batch_launch( | |
| launcher=launcher, | |
| experiments_dir=experiments_dir, | |
| cfg=cfg, | |
| datasets=datasets, | |
| print_only=args.print_only, | |
| time=args.time, | |
| use_qos=args.use_qos, | |
| experiments_list=experiments_list, | |
| other_args=other_args, | |
| suffix=args.suffix, | |
| num_gpus=args.num_gpus, | |
| ) | |
| #---------------------------------------------------------------------------- | |