Spaces:
Sleeping
Sleeping
File size: 12,447 Bytes
be5548b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 |
import sys
from pathlib import Path
from datetime import date
import subprocess
import shutil
import os
import stat
import getpass
import re
import glob
def process_arg_string(expe_args): # function to extract flagged (with a *) arguments as details for experience name
details_string = ''
processed_arg_string = expe_args.replace('*', '') # keep a version of args cleaned from exp name related flags
# args = [arg_chunk.split(' -') for arg_chunk in expe_args.split(' --')]
arg_chunks = [arg_chunk for arg_chunk in expe_args.split(' --')]
args_list = []
for arg in arg_chunks:
if ' -' in arg and arg.split(' -')[1].isalpha():
args_list.extend(arg.split(' -'))
else:
args_list.append(arg)
# args_list = [item for sublist in args for item in sublist] # flatten
for arg in args_list:
if arg == '':
continue
if arg[0] == '*':
if arg[-1] == ' ':
arg = arg[:-1]
details_string += '_' + arg[1:].replace(' ', '_').replace('/', '-')
return details_string, processed_arg_string
slurm_confs = {'curta_extra_long': "#SBATCH -p inria\n"
"#SBATCH -t 119:00:00\n",
'curta_long': "#SBATCH -p inria\n"
"#SBATCH -t 72:00:00\n",
'curta_medium': "#SBATCH -p inria\n"
"#SBATCH -t 48:00:00\n",
'curta_short': "#SBATCH -p inria\n"
"#SBATCH -t 24:00:00\n",
'jz_super_short_gpu':
'#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:1\n'
"#SBATCH -t 9:59:00\n"
"#SBATCH --qos=qos_gpu-t3\n",
'jz_short_gpu': '#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:1\n'
"#SBATCH -t 19:59:00\n"
"#SBATCH --qos=qos_gpu-t3\n",
'jz_short_gpu_chained': '#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:1\n'
"#SBATCH -t 19:59:00\n"
"#SBATCH --qos=qos_gpu-t3\n",
'jz_short_2gpus_chained': '#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:2\n'
"#SBATCH -t 19:59:00\n"
"#SBATCH --qos=qos_gpu-t3\n",
'jz_medium_gpu': '#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:1\n'
"#SBATCH -t 48:00:00\n"
"#SBATCH --qos=qos_gpu-t4\n",
'jz_super_short_2gpus': '#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:2\n'
"#SBATCH -t 14:59:00\n"
"#SBATCH --qos=qos_gpu-t3\n",
'jz_short_2gpus': '#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:2\n'
"#SBATCH -t 19:59:00\n"
"#SBATCH --qos=qos_gpu-t3\n",
'jz_short_2gpus_32g': '#SBATCH -A imi@v100\n'
'#SBATCH -C v100-32g\n'
'#SBATCH --gres=gpu:2\n'
"#SBATCH -t 19:59:00\n"
"#SBATCH --qos=qos_gpu-t3\n",
'jz_medium_2gpus': '#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:2\n'
"#SBATCH -t 48:00:00\n"
"#SBATCH --qos=qos_gpu-t4\n",
'jz_medium_2gpus_32g': '#SBATCH -A imi@v100\n'
'#SBATCH -C v100-32g\n'
'#SBATCH --gres=gpu:2\n'
"#SBATCH -t 48:00:00\n"
"#SBATCH --qos=qos_gpu-t4\n",
'jz_long_gpu': '#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:1\n'
"#SBATCH -t 72:00:00\n"
"#SBATCH --qos=qos_gpu-t4\n",
'jz_long_2gpus': '#SBATCH -A imi@v100\n'
'#SBATCH --gres=gpu:2\n'
'#SBATCH -t 72:00:00\n'
'#SBATCH --qos=qos_gpu-t4\n',
'jz_long_2gpus_32g': '#SBATCH -A imi@v100\n'
'#SBATCH -C v100-32g\n'
'#SBATCH --gres=gpu:2\n'
"#SBATCH -t 72:00:00\n"
"#SBATCH --qos=qos_gpu-t4\n",
'jz_super_long_2gpus_32g': '#SBATCH -A imi@v100\n'
'#SBATCH -C v100-32g\n'
'#SBATCH --gres=gpu:2\n'
"#SBATCH -t 99:00:00\n"
"#SBATCH --qos=qos_gpu-t4\n",
'jz_short_cpu': '#SBATCH -A imi@cpu\n'
"#SBATCH -t 19:59:00\n"
"#SBATCH --qos=qos_cpu-t3\n",
'jz_medium_cpu': '#SBATCH -A imi@cpu\n'
"#SBATCH -t 48:00:00\n"
"#SBATCH --qos=qos_cpu-t4\n",
'jz_long_cpu': '#SBATCH -A imi@cpu\n'
"#SBATCH -t 72:00:00\n"
"#SBATCH --qos=qos_cpu-t4\n",
'plafrim_cpu_medium': "#SBATCH -t 48:00:00\n",
'plafrim_cpu_long': "#SBATCH -t 72:00:00\n",
'plafrim_gpu_medium': '#SBATCH -p long_sirocco\n'
"#SBATCH -t 48:00:00\n"
'#SBATCH --gres=gpu:1\n'
}
cur_path = str(Path.cwd())
date = date.today().strftime("%d-%m")
# create campain log dir if not already done
Path(cur_path + "/campain_logs/jobouts/").mkdir(parents=True, exist_ok=True)
Path(cur_path + "/campain_logs/scripts/").mkdir(parents=True, exist_ok=True)
# Load txt file containing experiments to run (give it as argument to this script)
filename = 'to_run.txt'
if len(sys.argv) >= 2:
filename = sys.argv[1]
launch = True
# Save a copy of txt file
shutil.copyfile(cur_path + "/" + filename, cur_path + '/campain_logs/scripts/' + date + '_' + filename)
# one_launch_per_n_seeds = 8
one_launch_per_n_seeds = 4
global_seed_offset = 0
incremental = False
if len(sys.argv) >= 3:
if sys.argv[2] == 'nolaunch':
launch = False
if sys.argv[2] == 'seed_offset':
global_seed_offset = int(sys.argv[3])
if sys.argv[2] == 'incremental_seed_offset':
global_seed_offset = int(sys.argv[3])
incremental = True
if launch:
print('Creating and Launching slurm scripts given arguments from {}'.format(filename))
# time.sleep(1.0)
expe_list = []
with open(filename, 'r') as f:
expe_list = [line.rstrip() for line in f]
exp_names = set()
for expe_args in expe_list:
seed_offset_to_use = global_seed_offset
if len(expe_args) == 0:
# empty line
continue
if expe_args[0] == '#':
# comment line
continue
exp_config = expe_args.split('--')[1:5]
if not [arg.split(' ')[0] for arg in exp_config] == ['slurm_conf', 'nb_seeds', 'frames', 'model']:
raise ValueError("Arguments must be in the following order {}".format(
['slurm_conf', 'nb_seeds', 'frames', 'model']))
slurm_conf_name, nb_seeds, frames, exp_name = [arg.split(' ')[1] for arg in exp_config]
user = getpass.getuser()
if 'curta' in slurm_conf_name:
gpu = ''
PYTHON_INTERP = "$HOME/anaconda3/envs/act_and_speak/bin/python"
n_cpus = 1
elif 'plafrim' in slurm_conf_name:
gpu = ''
PYTHON_INTERP = '/home/{}/USER/conda/envs/act_and_speak/bin/python'.format(user)
n_cpus = 1
elif 'jz' in slurm_conf_name:
if user == "utu57ed":
PYTHON_INTERP='/gpfsscratch/rech/imi/{}/miniconda3/envs/social_ai/bin/python'.format(user)
elif user == "uxo14qj":
PYTHON_INTERP='/gpfswork/rech/imi/{}/miniconda3/envs/act_and_speak/bin/python'.format(user)
else:
if user != "flowers":
raise ValueError("Who are you? User {} unknown.".format(user))
gpu = '' # '--gpu_id 0'
n_cpus = 2
n_cpus = 4
assert n_cpus*one_launch_per_n_seeds == 16 # cpus_per_task is 8 will result in 16 cpus
else:
raise Exception("Unrecognized conf name: {} ".format(slurm_conf_name))
# assert ((int(nb_seeds) % 8) == 0), 'number of seeds should be divisible by 8'
assert ((int(nb_seeds) % 4) == 0), 'number of seeds should be divisible by 8'
run_args = expe_args.split(exp_name, 1)[
1] # WARNING: assumes that exp_name comes after slurm_conf and nb_seeds and frames in txt
# prepare experiment name formatting (use --* or -* instead of -- or - to use argument in experiment name
# print(expe_args.split(exp_name))
exp_details, run_args = process_arg_string(run_args)
exp_name = date + '_' + exp_name + exp_details
# no two trains are to be put in the same dir
assert exp_names not in exp_names
exp_names.add(exp_name)
slurm_script_fullname = cur_path + "/campain_logs/scripts/{}".format(exp_name) + ".sh"
# create corresponding slurm script
# calculate how many chained jobs we need
chained_training = "chained" in slurm_conf_name
frames = int(frames)
if chained_training:
# assume 10M frames per 20h (fps 140 - very conservative)
timelimit = slurm_confs[slurm_conf_name].split("-t ")[-1].split("\n")[0]
assert timelimit == '19:59:00'
one_script_frames = 10000000
print(f"One script frames: {one_script_frames}")
num_chained_jobs = frames // one_script_frames + bool(frames % one_script_frames)
else:
one_script_frames = frames
num_chained_jobs = 1 # no chaining
assert "--frames " not in run_args
current_script_frames = min(one_script_frames, frames)
# launch scripts (1 launch per 4 seeds)
if launch:
for i in range(int(nb_seeds) // one_launch_per_n_seeds):
# continue jobs
cont_job_i = num_chained_jobs # last job
exp_name_no_date = exp_name[5:]
continue_slurm_script_fullname = cur_path + "/campain_logs/scripts/*{}_continue_{}".format(exp_name_no_date, "*")
matched_scripts = glob.glob(continue_slurm_script_fullname)
matched_scripts.sort(key=os.path.getctime)
for last_script in reversed(matched_scripts):
# start from the latest written script and start the first encountered that has a err file (that was ran)
p = re.compile("continue_(.*).sh")
last_job_id = int(p.search(last_script).group(1))
last_script_name = os.path.basename(last_script)[:-3].replace("_continue_", "_cont_")
if len(glob.glob(cur_path + "/campain_logs/jobouts/"+last_script_name+"*.sh.err")) == 1:
# error file found -> script was ran -> this is the script that crashed
break
print(f"Continuing job id: {last_job_id}")
# last_err_log = glob.glob(cur_path + "/campain_logs/jobouts/"+last_script_name+"*.sh.err")[0]
#
# print("Then ended with:\n")
# print('"""\n')
# for l in open(last_err_log).readlines():
# print("\t"+l, end='')
# print('"""\n')
# write continue script
cont_script_name = "{}_continue_{}.sh".format(exp_name, last_job_id)
continue_slurm_script_fullname = cur_path + "/campain_logs/scripts/"+cont_script_name
current_script_frames = min(one_script_frames*(2+cont_job_i), frames)
# run continue job
sbatch_pipe = subprocess.Popen(
['sbatch', 'campain_logs/scripts/{}'.format(os.path.basename(last_script)), str((i * one_launch_per_n_seeds) + seed_offset_to_use)], # 0 4 8 12
stdout=subprocess.PIPE
)
if incremental:
global_seed_offset += int(nb_seeds)
|