565 lines
21 KiB
Python
565 lines
21 KiB
Python
import wandb
|
|
import yaml
|
|
import os
|
|
import math
|
|
import time
|
|
import random
|
|
import copy
|
|
import re
|
|
import itertools
|
|
from collections.abc import *
|
|
from functools import partial
|
|
from multiprocessing import Process
|
|
from threading import Thread
|
|
import git
|
|
import datetime
|
|
from pprint import pprint
|
|
|
|
import pdb
|
|
d = pdb.set_trace
|
|
|
|
REQUIRE_CONFIG_CONSUMED = False
|
|
DEFAULT_START_METHOD = 'fork'
|
|
DEFAULT_REINIT = True
|
|
|
|
Parallelization_Primitive = Process # Thread
|
|
|
|
try:
|
|
import pyslurm
|
|
except ImportError:
|
|
slurm_avaible = False
|
|
print('[!] Slurm not available.')
|
|
else:
|
|
slurm_avaible = True
|
|
|
|
class Slate():
|
|
def __init__(self, runners):
|
|
self.runners = {
|
|
'void': Void_Runner,
|
|
'printConfig': Print_Config_Runner,
|
|
'pdb': PDB_Runner,
|
|
}
|
|
self.runners.update(runners)
|
|
self._version = False
|
|
self.job_id = os.environ.get('SLURM_JOB_ID', False)
|
|
self.task_id = None
|
|
self.run_id = -1
|
|
self._tmp_path = os.path.expandvars('$TMP')
|
|
self.verify = False
|
|
|
|
def load_config(self, filename, name):
|
|
emptyStack = []
|
|
config, stack = self._load_config(filename, name, stack=emptyStack)
|
|
print('[i] Merged Configs: ', stack)
|
|
self._config = copy.deepcopy(config)
|
|
self.consume(config, 'vars', {})
|
|
return config
|
|
|
|
def _load_config(self, filename, name, stack):
|
|
stack.append(f'{filename}:{name}')
|
|
with open(filename, 'r') as f:
|
|
docs = yaml.safe_load_all(f)
|
|
for doc in docs:
|
|
if 'name' in doc:
|
|
if doc['name'] == name:
|
|
if 'import' in doc:
|
|
imports = doc['import'].split(',')
|
|
del doc['import']
|
|
for imp in imports:
|
|
if imp[0] == ' ':
|
|
imp = imp[1:]
|
|
if imp == "$":
|
|
imp = ':DEFAULT'
|
|
rel_path, *opt = imp.split(':')
|
|
if len(opt) == 0:
|
|
nested_name = 'DEFAULT'
|
|
elif len(opt) == 1:
|
|
nested_name = opt[0]
|
|
else:
|
|
raise Exception('Malformed import statement. Must be <import file:exp>, <import :exp>, <import file> for file:DEFAULT or <import $> for :DEFAULT.')
|
|
nested_path = os.path.normpath(os.path.join(os.path.dirname(filename), rel_path)) if len(rel_path) else filename
|
|
child, stack = self._load_config(nested_path, nested_name, stack=stack)
|
|
doc = self.deep_update(child, doc)
|
|
return doc, stack
|
|
raise Exception(f'Unable to find experiment <{name}> in <{filename}>')
|
|
|
|
def deep_update(self, d, u, traverse_dot_notation=True):
|
|
for kstr, v in u.items():
|
|
if traverse_dot_notation:
|
|
ks = kstr.split('.')
|
|
else:
|
|
ks = [kstr]
|
|
head = d
|
|
for k in ks:
|
|
if k in ['parameters']:
|
|
traverse_dot_notation = False
|
|
last_head = head
|
|
if k not in head:
|
|
head[k] = {}
|
|
head = head[k]
|
|
if isinstance(v, Mapping):
|
|
last_head[ks[-1]] = self.deep_update(d.get(k, {}), v, traverse_dot_notation=traverse_dot_notation)
|
|
else:
|
|
last_head[ks[-1]] = v
|
|
return d
|
|
|
|
def expand_vars(self, string, delta_desc='BASE', **kwargs):
|
|
if isinstance(string, str):
|
|
rand = int(random.random() * 99999999)
|
|
if string == '{rand}':
|
|
return rand
|
|
return string.format(delta_desc=delta_desc, **kwargs, rand=rand, tmp=self._tmp_path, job_id=(self.job_id or 'LOCAL'), task_id=(self.task_id or 0), run_id=self.run_id)
|
|
return string
|
|
|
|
def apply_nested(self, d, f):
|
|
for k, v in d.items():
|
|
if isinstance(v, dict):
|
|
self.apply_nested(v, f)
|
|
elif isinstance(v, list):
|
|
for i, e in enumerate(v):
|
|
ptr = {'PTR': d[k][i]}
|
|
self.apply_nested(ptr, f)
|
|
d[k][i] = ptr['PTR']
|
|
else:
|
|
d[k] = f(v)
|
|
|
|
def deep_expand_vars(self, dict, **kwargs):
|
|
self.apply_nested(dict, lambda x: self.expand_vars(x, **kwargs))
|
|
|
|
def consume(self, conf, key, default=None, expand=False, **kwargs):
|
|
if key == '':
|
|
if expand:
|
|
self.deep_expand_vars(conf, config=self._config, **kwargs)
|
|
elif type(conf) == str:
|
|
while conf.find('{') != -1:
|
|
conf = self.expand_vars(conf, config=self._config, **kwargs)
|
|
return conf
|
|
keys_arr = key.split('.')
|
|
if len(keys_arr) == 1:
|
|
k = keys_arr[0]
|
|
if default is not None:
|
|
if isinstance(conf, Mapping):
|
|
val = conf.get(k, default)
|
|
else:
|
|
if default is not None:
|
|
return default
|
|
raise Exception('')
|
|
else:
|
|
val = conf[k]
|
|
if k in conf:
|
|
del conf[k]
|
|
|
|
if expand:
|
|
self.deep_expand_vars(val, config=self._config, **kwargs)
|
|
elif type(val) == str:
|
|
while val.find('{') != -1:
|
|
val = self.expand_vars(val, config=self._config, **kwargs)
|
|
|
|
return val
|
|
child = conf.get(keys_arr[0], {})
|
|
child_keys = '.'.join(keys_arr[1:])
|
|
return self.consume(child, child_keys, default=default, expand=expand, **kwargs)
|
|
|
|
def get_version(self):
|
|
if not self._version:
|
|
repo = git.Repo(search_parent_directories=True)
|
|
sha = repo.head.object.hexsha
|
|
self._version = sha
|
|
return self._version
|
|
|
|
def _calc_num_jobs(self, schedC, num_conv_versions):
|
|
schedulerC = copy.deepcopy(schedC)
|
|
reps = self.consume(schedulerC, 'repetitions', self.consume(schedulerC, 'reps_per_version', 1) * num_conv_versions)
|
|
agents_per_job = self.consume(schedulerC, 'agents_per_job', 1)
|
|
reps_per_agent = self.consume(schedulerC, 'reps_per_agent', 1)
|
|
reps_per_job = reps_per_agent * agents_per_job
|
|
jobs_needed = math.ceil(reps / reps_per_job)
|
|
return jobs_needed
|
|
|
|
def _reps_for_job(self, schedC, task_id, num_conv_versions):
|
|
schedulerC = copy.deepcopy(schedC)
|
|
num_jobs = self._calc_num_jobs(schedulerC, num_conv_versions)
|
|
reps = self.consume(schedulerC, 'repetitions', self.consume(schedulerC, 'reps_per_version', 1) * num_conv_versions)
|
|
if task_id is None:
|
|
return list(range(0, reps))
|
|
reps_for_job = [[] for _ in range(num_jobs)]
|
|
for i in range(reps):
|
|
reps_for_job[i % num_jobs].append(i)
|
|
return reps_for_job[task_id]
|
|
|
|
def _make_configs_for_runs(self, config_exp_pairs):
|
|
"""
|
|
Expand configurations across all provided experiments, grid, and ablation variants.
|
|
|
|
Parameters:
|
|
config_exp_pairs (list): A list of tuples where each tuple contains (filename, experiment name).
|
|
|
|
Returns:
|
|
list: A list of expanded configurations ready for execution.
|
|
"""
|
|
all_expanded_configs = []
|
|
|
|
for config, exp in config_exp_pairs:
|
|
config_data = self.load_config(config, exp)
|
|
grid_versions = self._make_grid_versions(config_data)
|
|
exp_variants = self._make_ablative_versions(config_data, grid_versions)
|
|
all_expanded_configs.extend(exp_variants)
|
|
|
|
return all_expanded_configs
|
|
|
|
def run_local(self, config_exp_pairs, task_id):
|
|
"""
|
|
Run all expanded configurations locally, handling all variants and their repetitions concurrently.
|
|
|
|
Parameters:
|
|
config_exp_pairs (list): A list of tuples where each tuple contains (filename, experiment name).
|
|
task_id (int): The task ID for the experiments.
|
|
"""
|
|
self.task_id = task_id
|
|
all_configs = self._make_configs_for_runs(config_exp_pairs)
|
|
|
|
num_conv_versions = len(all_configs)
|
|
schedulerC = copy.deepcopy(all_configs[0].get('scheduler', {}))
|
|
rep_ids = self._reps_for_job(schedulerC, task_id, num_conv_versions)
|
|
self._fork_processes(all_configs, rep_ids)
|
|
|
|
def run_slurm(self, original_config_exp_string, config_exp_pairs):
|
|
"""
|
|
Schedule all expanded configurations on SLURM within a single job.
|
|
|
|
Parameters:
|
|
original_config_exp_string (str): The original string of config:experiment pairs provided by the user.
|
|
config_exp_pairs (list): A list of tuples where each tuple contains (filename, experiment name).
|
|
"""
|
|
all_configs = self._make_configs_for_runs(config_exp_pairs)
|
|
|
|
slurmC = self.consume(all_configs[0], 'slurm', expand=True)
|
|
s_name = self.consume(slurmC, 'name')
|
|
|
|
python_script = 'main.py'
|
|
sh_lines = ['#!/bin/bash']
|
|
sh_lines += self.consume(slurmC, 'sh_lines', [])
|
|
if venv := self.consume(slurmC, 'venv', False):
|
|
sh_lines += [f'source activate {venv}']
|
|
|
|
# Use the original config:experiment string to avoid verbosity
|
|
final_line = f'{python_script} {original_config_exp_string} -t $SLURM_ARRAY_TASK_ID'
|
|
|
|
if self.consume(slurmC, 'python_exec', False):
|
|
final_line = f'./omni_sif_python {final_line}'
|
|
else:
|
|
final_line = f'python3 {final_line}'
|
|
if self.consume(slurmC, 'xvfb', False):
|
|
final_line = f'xvfb-run {final_line}'
|
|
sh_lines.append(final_line)
|
|
|
|
script = "\n".join(sh_lines)
|
|
|
|
num_jobs = self._calc_num_jobs(all_configs[0].get('scheduler', {}), len(all_configs))
|
|
last_job_idx = num_jobs - 1
|
|
num_parallel_jobs = min(self.consume(slurmC, 'num_parallel_jobs', num_jobs), num_jobs)
|
|
array = f'0-{last_job_idx}%{num_parallel_jobs}'
|
|
job = pyslurm.JobSubmitDescription(name=s_name, script=script, array=array, **slurmC)
|
|
|
|
if self.verify:
|
|
input(f'[!] Press Enter to submit the job to SLURM.')
|
|
|
|
job_id = job.submit()
|
|
print(f'[>] Job submitted to SLURM with id {job_id}')
|
|
|
|
# Log file entry optimization
|
|
with open('job_hist.log', 'a') as f:
|
|
config_logs = {}
|
|
for config, exp in config_exp_pairs:
|
|
if config not in config_logs:
|
|
config_logs[config] = []
|
|
config_logs[config].append(exp)
|
|
|
|
for config, exps in config_logs.items():
|
|
exps_str = ",".join(exps)
|
|
f.write(f'{config}:{exps_str} submitted to SLURM with id {job_id}\n')
|
|
|
|
def _fork_processes(self, configs, rep_ids):
|
|
"""
|
|
Fork processes to run all expanded configurations concurrently.
|
|
|
|
Parameters:
|
|
configs (list): A list of expanded configurations.
|
|
rep_ids (list): A list of repetition identifiers for the configurations.
|
|
"""
|
|
schedC = self.consume(configs[0], 'scheduler', {})
|
|
agents_per_job = self.consume(schedC, 'agents_per_job', 1)
|
|
reps_per_agent = self.consume(schedC, 'reps_per_agent', 1)
|
|
|
|
node_reps = len(rep_ids)
|
|
num_p = min(agents_per_job, math.ceil(node_reps / reps_per_agent))
|
|
|
|
if num_p == 1:
|
|
print('[i] Running within main thread')
|
|
self._run_process(configs, rep_ids=rep_ids, p_ind=0)
|
|
return
|
|
|
|
procs = []
|
|
reps_done = 0
|
|
|
|
for p in range(num_p):
|
|
print(f'[i] Spawning separate thread/process ({p+1}/{num_p})')
|
|
num_reps = min(node_reps - reps_done, reps_per_agent)
|
|
proc_rep_ids = [rep_ids[i] for i in list(range(reps_done, reps_done + num_reps))]
|
|
proc_configs = [configs[i % len(configs)] for i in proc_rep_ids] # Distribute configs across processes
|
|
proc = Parallelization_Primitive(target=partial(self._run_process, proc_configs, proc_rep_ids, p_ind=p))
|
|
proc.start()
|
|
procs.append(proc)
|
|
reps_done += num_reps
|
|
|
|
for proc in procs:
|
|
proc.join()
|
|
print('[i] All threads/processes have terminated')
|
|
|
|
def _run_process(self, orig_configs, rep_ids, p_ind):
|
|
"""
|
|
Run a single process for a subset of configurations.
|
|
|
|
Parameters:
|
|
configs (list): A list of configurations to run.
|
|
rep_ids (list): A list of repetition identifiers for the configurations.
|
|
p_ind (int): Process index.
|
|
"""
|
|
for r in rep_ids:
|
|
self.run_id = r
|
|
config = orig_configs[r % len(orig_configs)]
|
|
self._run_single(config, [r], p_ind=p_ind)
|
|
|
|
def _run_single(self, orig_config, rep_ids, p_ind):
|
|
print(f'[P{p_ind}] I will work on reps {rep_ids}')
|
|
self._config = orig_config
|
|
runnerName = self.consume(orig_config, 'runner')
|
|
project = self.consume(orig_config, 'wandb.project', orig_config.get('project', orig_config.get('name')))
|
|
|
|
Runner = self.runners[runnerName]
|
|
|
|
if self.consume(orig_config, 'scheduler.bind_agent_to_core', False):
|
|
os.sched_setaffinity(0, [p_ind % os.cpu_count()])
|
|
|
|
for r in rep_ids:
|
|
self.run_id = r
|
|
runnerConf = copy.deepcopy(orig_config)
|
|
wandbC = self.consume(runnerConf, 'wandb', {}, expand=True, delta_desc=runnerConf.pop('delta_desc', 'BASE'))
|
|
if 'job_type' in wandbC and len(wandbC['job_type']) > 62:
|
|
wandbC['job_type'] = "..."+wandbC['job_type'][-50:]
|
|
|
|
retry = 5
|
|
while retry:
|
|
try:
|
|
with wandb.init(
|
|
project=project,
|
|
config=copy.deepcopy(runnerConf),
|
|
reinit=self.consume(wandbC, 'reinit', DEFAULT_REINIT),
|
|
settings=wandb.Settings(**self.consume(wandbC, 'settings', {})),
|
|
**wandbC
|
|
) as run:
|
|
runner = Runner(self, runnerConf)
|
|
runner.setup(wandbC['group']+wandbC['job_type'])
|
|
runner.run(run)
|
|
except wandb.errors.CommError as e:
|
|
retry -= 1
|
|
if retry:
|
|
print('Caught CommErr; retrying...')
|
|
time.sleep(int(60*random.random()))
|
|
else:
|
|
print('Caught CommErr; not retrying')
|
|
raise e
|
|
else:
|
|
retry = 0
|
|
|
|
if runnerConf != {}:
|
|
msg = ('Config was not completely consumed: ', runnerConf)
|
|
if REQUIRE_CONFIG_CONSUMED:
|
|
raise Exception(msg)
|
|
else:
|
|
print(msg)
|
|
orig_config = {}
|
|
|
|
def _get_num_conv_versions(self, config):
|
|
return len(self._make_configs_for_runs(config))
|
|
|
|
def _make_grid_versions(self, config):
|
|
if 'grid' in config:
|
|
return params_combine(config, 'grid', itertools.product)
|
|
return [config]
|
|
|
|
def _make_ablative_versions(self, config, grid_versions):
|
|
if 'ablative' in config:
|
|
return grid_versions + ablative_expand(grid_versions)
|
|
else:
|
|
return grid_versions
|
|
|
|
def from_args(self):
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("config_experiments", nargs='+', help="List of config:experiment pairs")
|
|
parser.add_argument("-s", "--slurm", action="store_true")
|
|
parser.add_argument("-w", "--worker", action="store_true")
|
|
parser.add_argument("-t", "--task_id", default=None, type=int)
|
|
parser.add_argument("--ask_verify", action="store_true")
|
|
|
|
args = parser.parse_args()
|
|
|
|
print(f'[i] I have task_id {args.task_id}')
|
|
print(f'[i] Running on version [git:{self.get_version()}]')
|
|
|
|
if args.worker:
|
|
raise Exception('Worker mode not yet implemented')
|
|
|
|
config_exp_pairs = []
|
|
for config_exp in args.config_experiments:
|
|
config, exps = config_exp.split(":")
|
|
exp_list = exps.split(",")
|
|
for exp in exp_list:
|
|
config_exp_pairs.append((config, exp))
|
|
|
|
if args.slurm:
|
|
if args.ask_verify:
|
|
self.verify = True
|
|
self.run_slurm(' '.join(args.config_experiments), config_exp_pairs)
|
|
else:
|
|
self.run_local(config_exp_pairs, args.task_id)
|
|
|
|
|
|
def params_combine(config: dict, key: str, iter_func):
|
|
if iter_func is None:
|
|
return [config]
|
|
|
|
combined_configs = []
|
|
tuple_dict = flatten_dict_to_tuple_keys(config[key])
|
|
_param_names = ['.'.join(t) for t in tuple_dict]
|
|
|
|
for values in iter_func(*tuple_dict.values()):
|
|
_config = copy.deepcopy(config)
|
|
del _config[key]
|
|
|
|
for i, t in enumerate(tuple_dict.keys()):
|
|
insert_deep_dictionary(d=_config, t=t, value=values[i])
|
|
|
|
_config = extend_config_name(_config, _param_names, values)
|
|
combined_configs.append(_config)
|
|
return combined_configs
|
|
|
|
|
|
def ablative_expand(conf_list):
|
|
combined_configs = []
|
|
for config in conf_list:
|
|
tuple_dict = flatten_dict_to_tuple_keys(config['ablative'])
|
|
_param_names = ['.'.join(t) for t in tuple_dict]
|
|
|
|
for i, key in enumerate(tuple_dict):
|
|
for val in tuple_dict[key]:
|
|
_config = copy.deepcopy(config)
|
|
insert_deep_dictionary(_config, key, val)
|
|
_config = extend_config_name(_config, [_param_names[i]], [val])
|
|
combined_configs.append(_config)
|
|
return combined_configs
|
|
|
|
|
|
def flatten_dict_to_tuple_keys(d: MutableMapping):
|
|
flat_dict = {}
|
|
for k, v in d.items():
|
|
if isinstance(v, MutableMapping):
|
|
sub_dict = flatten_dict_to_tuple_keys(v)
|
|
flat_dict.update({(k, *sk): sv for sk, sv in sub_dict.items()})
|
|
elif isinstance(v, MutableSequence):
|
|
flat_dict[(k,)] = v
|
|
return flat_dict
|
|
|
|
|
|
def insert_deep_dictionary(d: MutableMapping, t: tuple, value):
|
|
if type(t) is tuple:
|
|
if len(t) == 1:
|
|
d[t[0]] = value
|
|
else:
|
|
if t[0] not in d:
|
|
d[t[0]] = dict()
|
|
insert_deep_dictionary(d[t[0]], t[1:], value)
|
|
else:
|
|
d[t] = value
|
|
|
|
|
|
def append_deep_dictionary(d: MutableMapping, t: tuple, value):
|
|
if type(t) is tuple:
|
|
if len(t) == 1:
|
|
if t[0] not in d:
|
|
d[t[0]] = []
|
|
d[t[0]].append(value)
|
|
else:
|
|
if t[0] not in d:
|
|
d[t[0]] = dict()
|
|
append_deep_dictionary(d[t[0]], t[1:], value)
|
|
else:
|
|
d[t] = value
|
|
|
|
|
|
def extend_config_name(config: dict, param_names: list, values: list) -> dict:
|
|
_converted_name = convert_param_names(param_names, values)
|
|
config['delta_desc'] = config['delta_desc'] + '_' + _converted_name if 'delta_desc' in config else _converted_name
|
|
return config
|
|
|
|
|
|
def convert_param_names(_param_names: list, values: list) -> str:
|
|
_converted_name = '_'.join("{}{}".format(shorten_param(k), v) for k, v in zip(_param_names, values))
|
|
_converted_name = re.sub("[' ]", '', _converted_name)
|
|
_converted_name = re.sub('["]', '', _converted_name)
|
|
_converted_name = re.sub("[(\[]", '_', _converted_name)
|
|
_converted_name = re.sub("[)\]]", '', _converted_name)
|
|
_converted_name = re.sub("[,]", '_', _converted_name)
|
|
return _converted_name
|
|
|
|
|
|
def shorten_param(_param_name):
|
|
name_parts = _param_name.split('.')
|
|
shortened_parts = '.'.join(map(lambda s: s[:3], name_parts[:-1]))
|
|
shortened_leaf = ''.join(map(lambda s: s[0], name_parts[-1].split('_')))
|
|
if shortened_parts:
|
|
return shortened_parts + '.' + shortened_leaf
|
|
else:
|
|
return shortened_leaf
|
|
|
|
|
|
class Slate_Runner():
|
|
def __init__(self, slate, config):
|
|
self.slate = slate
|
|
self.config = config
|
|
|
|
def setup(self, name):
|
|
pass
|
|
|
|
def run(self, run):
|
|
pass
|
|
|
|
|
|
class Print_Config_Runner(Slate_Runner):
|
|
def run(self, run):
|
|
slate, config = self.slate, self.config
|
|
|
|
pprint(config)
|
|
print('---')
|
|
pprint(slate.consume(config, '', expand=True))
|
|
for k in list(config.keys()):
|
|
del config[k]
|
|
|
|
|
|
class Void_Runner(Slate_Runner):
|
|
def run(self, run):
|
|
slate, config = self.slate, self.config
|
|
for k in list(config.keys()):
|
|
del config[k]
|
|
|
|
|
|
class PDB_Runner(Slate_Runner):
|
|
def run(self, run):
|
|
d()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
raise Exception('You are using it wrong...')
|