climategan / sbatch.py
NimaBoscarino's picture
copy the climategan repo in here
6e601ed
import datetime
import itertools
import os
import re
import subprocess
import sys
from collections import defaultdict
from pathlib import Path
import numpy as np
import yaml
def flatten_conf(conf, to={}, parents=[]):
"""
Flattens a configuration dict: nested dictionaries are flattened
as key1.key2.key3 = value
conf.yaml:
```yaml
a: 1
b:
c: 2
d:
e: 3
g:
sample: sequential
from: [4, 5]
```
Is flattened to
{
"a": 1,
"b.c": 2,
"b.d.e": 3,
"b.g": {
"sample": "sequential",
"from": [4, 5]
}
}
Does not affect sampling dicts.
Args:
conf (dict): the configuration to flatten
new (dict, optional): the target flatenned dict. Defaults to {}.
parents (list, optional): a final value's list of parents. Defaults to [].
"""
for k, v in conf.items():
if isinstance(v, dict) and "sample" not in v:
flatten_conf(v, to, parents + [k])
else:
new_k = ".".join([str(p) for p in parents + [k]])
to[new_k] = v
def env_to_path(path):
"""Transorms an environment variable mention in a json
into its actual value. E.g. $HOME/clouds -> /home/vsch/clouds
Args:
path (str): path potentially containing the env variable
"""
path_elements = path.split("/")
new_path = []
for el in path_elements:
if "$" in el:
new_path.append(os.environ[el.replace("$", "")])
else:
new_path.append(el)
return "/".join(new_path)
class C:
HEADER = "\033[95m"
OKBLUE = "\033[94m"
OKGREEN = "\033[92m"
WARNING = "\033[93m"
FAIL = "\033[91m"
ENDC = "\033[0m"
BOLD = "\033[1m"
UNDERLINE = "\033[4m"
ITALIC = "\33[3m"
BEIGE = "\33[36m"
def escape_path(path):
p = str(path)
return p.replace(" ", "\ ").replace("(", "\(").replace(")", "\)") # noqa: W605
def warn(*args, **kwargs):
print("{}{}{}".format(C.WARNING, " ".join(args), C.ENDC), **kwargs)
def parse_jobID(command_output):
"""
get job id from successful sbatch command output like
`Submitted batch job 599583`
Args:
command_output (str): sbatch command's output
Returns:
int: the slurm job's ID
"""
command_output = command_output.strip()
if isinstance(command_output, str):
if "Submitted batch job" in command_output:
return int(command_output.split()[-1])
return -1
def now():
return str(datetime.datetime.now()).replace(" ", "_")
def cols():
try:
col = os.get_terminal_size().columns
except Exception:
col = 50
return col
def print_box(txt):
if not txt:
txt = "{}{}ERROR ⇪{}".format(C.BOLD, C.FAIL, C.ENDC)
lt = 7
else:
lt = len(txt)
nlt = lt + 12
txt = "|" + " " * 5 + txt + " " * 5 + "|"
line = "-" * nlt
empty = "|" + " " * (nlt - 2) + "|"
print(line)
print(empty)
print(txt)
print(empty)
print(line)
def print_header(idx):
b = C.BOLD
bl = C.OKBLUE
e = C.ENDC
char = "≡"
c = cols()
txt = " " * 20
txt += f"{b}{bl}Run {idx}{e}"
txt += " " * 20
ln = len(txt) - len(b) - len(bl) - len(e)
t = int(np.floor((c - ln) / 2))
tt = int(np.ceil((c - ln) / 2))
print(char * c)
print(char * t + " " * ln + char * tt)
print(char * t + txt + char * tt)
print(char * t + " " * ln + char * tt)
print(char * c)
def print_footer():
c = cols()
char = "﹎"
print()
print(char * (c // len(char)))
print()
print(" " * (c // 2) + "•" + " " * (c - c // 2 - 1))
print()
def extend_summary(summary, tmp_train_args_dict, tmp_template_dict, exclude=[]):
exclude = set(exclude)
if summary is None:
summary = defaultdict(list)
for k, v in tmp_template_dict.items():
if k not in exclude:
summary[k].append(v)
for k, v in tmp_train_args_dict.items():
if k not in exclude:
if isinstance(v, list):
v = str(v)
summary[k].append(v)
return summary
def search_summary_table(summary, summary_dir=None):
# filter out constant values
summary = {k: v for k, v in summary.items() if len(set(v)) > 1}
# if everything is constant: no summary
if not summary:
return None, None
# find number of searches
n_searches = len(list(summary.values())[0])
# print section title
print(
"{}{}{}Varying values across {} experiments:{}\n".format(
C.OKBLUE,
C.BOLD,
C.UNDERLINE,
n_searches,
C.ENDC,
)
)
# first column holds the Exp. number
first_col = {
"len": 8, # length of a column, to split columns according to terminal width
"str": ["| Exp. |", "|:----:|"]
+ [
"| {0:^{1}} |".format(i, 4) for i in range(n_searches)
], # list of values to print
}
print_columns = [[first_col]]
file_columns = [first_col]
for k in sorted(summary.keys()):
v = summary[k]
col_title = f" {k} |"
col_blank_line = f":{'-' * len(k)}-|"
col_values = [
" {0:{1}} |".format(
crop_string(
str(crop_float(v[idx], min([5, len(k) - 2]))), len(k)
), # crop floats and long strings
len(k),
)
for idx in range(len(v))
]
# create column object
col = {"len": len(k) + 3, "str": [col_title, col_blank_line] + col_values}
# if adding a new column would overflow the terminal and mess up printing, start
# new set of columns
if sum(c["len"] for c in print_columns[-1]) + col["len"] >= cols():
print_columns.append([first_col])
# store current column to latest group of columns
print_columns[-1].append(col)
file_columns.append(col)
print_table = ""
# print each column group individually
for colgroup in print_columns:
# print columns line by line
for i in range(n_searches + 2):
# get value of column for current line i
for col in colgroup:
print_table += col["str"][i]
# next line for current columns
print_table += "\n"
# new lines for new column group
print_table += "\n"
file_table = ""
for i in range(n_searches + 2):
# get value of column for current line i
for col in file_columns:
file_table += col["str"][i]
# next line for current columns
file_table += "\n"
summary_path = None
if summary_dir is not None:
summary_path = summary_dir / (now() + ".md")
with summary_path.open("w") as f:
f.write(file_table.strip())
return print_table, summary_path
def clean_arg(v):
"""
chain cleaning function
Args:
v (any): arg to pass to train.py
Returns:
str: parsed value to string
"""
return stringify_list(crop_float(quote_string(resolve_env(v))))
def resolve_env(v):
"""
resolve env variables in paths
Args:
v (any): arg to pass to train.py
Returns:
str: try and resolve an env variable
"""
if isinstance(v, str):
try:
if "$" in v:
if "/" in v:
v = env_to_path(v)
else:
_v = os.environ.get(v)
if _v is not None:
v = _v
except Exception:
pass
return v
def stringify_list(v):
"""
Stringify list (with double quotes) so that it can be passed a an argument
to train.py's hydra command-line parsing
Args:
v (any): value to clean
Returns:
any: type of v, str if v was a list
"""
if isinstance(v, list):
return '"{}"'.format(str(v).replace('"', "'"))
if isinstance(v, str):
if v.startswith("[") and v.endswith("]"):
return f'"{v}"'
return v
def quote_string(v):
"""
Add double quotes around string if it contains a " " or an =
Args:
v (any): value to clean
Returns:
any: type of v, quoted if v is a string with " " or =
"""
if isinstance(v, str):
if " " in v or "=" in v:
return f'"{v}"'
return v
def crop_float(v, k=5):
"""
If v is a float, crop precision to 5 digits and return v as a str
Args:
v (any): value to crop if float
Returns:
any: cropped float as str if v is a float, original v otherwise
"""
if isinstance(v, float):
return "{0:.{1}g}".format(v, k)
return v
def compute_n_search(conf):
"""
Compute the number of searchs to do if using -1 as n_search and using
cartesian or sequential search
Args:
conf (dict): experimental configuration
Returns:
int: size of the cartesian product or length of longest sequential field
"""
samples = defaultdict(list)
for k, v in conf.items():
if not isinstance(v, dict) or "sample" not in v:
continue
samples[v["sample"]].append(v)
totals = []
if "cartesian" in samples:
total = 1
for s in samples["cartesian"]:
total *= len(s["from"])
totals.append(total)
if "sequential" in samples:
total = max(map(len, [s["from"] for s in samples["sequential"]]))
totals.append(total)
if totals:
return max(totals)
raise ValueError(
"Used n_search=-1 without any field being 'cartesian' or 'sequential'"
)
def crop_string(s, k=10):
if len(s) <= k:
return s
else:
return s[: k - 2] + ".."
def sample_param(sample_dict):
"""sample a value (hyperparameter) from the instruction in the
sample dict:
{
"sample": "range | list",
"from": [min, max, step] | [v0, v1, v2 etc.]
}
if range, as np.arange is used, "from" MUST be a list, but may contain
only 1 (=min) or 2 (min and max) values, not necessarily 3
Args:
sample_dict (dict): instructions to sample a value
Returns:
scalar: sampled value
"""
if not isinstance(sample_dict, dict) or "sample" not in sample_dict:
return sample_dict
if sample_dict["sample"] == "cartesian":
assert isinstance(
sample_dict["from"], list
), "{}'s `from` field MUST be a list, found {}".format(
sample_dict["sample"], sample_dict["from"]
)
return "__cartesian__"
if sample_dict["sample"] == "sequential":
assert isinstance(
sample_dict["from"], list
), "{}'s `from` field MUST be a list, found {}".format(
sample_dict["sample"], sample_dict["from"]
)
return "__sequential__"
if sample_dict["sample"] == "range":
return np.random.choice(np.arange(*sample_dict["from"]))
if sample_dict["sample"] == "list":
return np.random.choice(sample_dict["from"])
if sample_dict["sample"] == "uniform":
return np.random.uniform(*sample_dict["from"])
raise ValueError("Unknown sample type in dict " + str(sample_dict))
def sample_sequentials(sequential_keys, exp, idx):
"""
Samples sequentially from the "from" values specified in each key of the
experimental configuration which have sample == "sequential"
Unlike `cartesian` sampling, `sequential` sampling iterates *independently*
over each keys
Args:
sequential_keys (list): keys to be sampled sequentially
exp (dict): experimental config
idx (int): index of the current sample
Returns:
conf: sampled dict
"""
conf = {}
for k in sequential_keys:
v = exp[k]["from"]
conf[k] = v[idx % len(v)]
return conf
def sample_cartesians(cartesian_keys, exp, idx):
"""
Returns the `idx`th item in the cartesian product of all cartesian keys to
be sampled.
Args:
cartesian_keys (list): keys in the experimental configuration that are to
be used in the full cartesian product
exp (dict): experimental configuration
idx (int): index of the current sample
Returns:
dict: sampled point in the cartesian space (with keys = cartesian_keys)
"""
conf = {}
cartesian_values = [exp[key]["from"] for key in cartesian_keys]
product = list(itertools.product(*cartesian_values))
for k, v in zip(cartesian_keys, product[idx % len(product)]):
conf[k] = v
return conf
def resolve(hp_conf, nb):
"""
Samples parameters parametrized in `exp`: should be a dict with
values which fit `sample_params(dic)`'s API
Args:
exp (dict): experiment's parametrization
nb (int): number of experiments to sample
Returns:
dict: sampled configuration
"""
if nb == -1:
nb = compute_n_search(hp_conf)
confs = []
for idx in range(nb):
conf = {}
cartesians = []
sequentials = []
for k, v in hp_conf.items():
candidate = sample_param(v)
if candidate == "__cartesian__":
cartesians.append(k)
elif candidate == "__sequential__":
sequentials.append(k)
else:
conf[k] = candidate
if sequentials:
conf.update(sample_sequentials(sequentials, hp_conf, idx))
if cartesians:
conf.update(sample_cartesians(cartesians, hp_conf, idx))
confs.append(conf)
return confs
def get_template_params(template):
"""
extract args in template str as {arg}
Args:
template (str): sbatch template string
Returns:
list(str): Args required to format the template string
"""
return map(
lambda s: s.replace("{", "").replace("}", ""),
re.findall("\{.*?\}", template), # noqa: W605
)
def read_exp_conf(name):
"""
Read hp search configuration from shared/experiment/
specified with or without the .yaml extension
Args:
name (str): name of the template to find in shared/experiment/
Returns:
Tuple(Path, dict): file path and loaded dict
"""
if ".yaml" not in name:
name += ".yaml"
paths = []
dirs = ["shared", "config"]
for d in dirs:
path = Path(__file__).parent / d / "experiment" / name
if path.exists():
paths.append(path.resolve())
if len(paths) == 0:
failed = [Path(__file__).parent / d / "experiment" for d in dirs]
s = "Could not find search config {} in :\n".format(name)
for fd in failed:
s += str(fd) + "\nAvailable:\n"
for ym in fd.glob("*.yaml"):
s += " " + ym.name + "\n"
raise ValueError(s)
if len(paths) == 2:
print(
"Warning: found 2 relevant files for search config:\n{}".format(
"\n".join(paths)
)
)
print("Using {}".format(paths[-1]))
with paths[-1].open("r") as f:
conf = yaml.safe_load(f)
flat_conf = {}
flatten_conf(conf, to=flat_conf)
return (paths[-1], flat_conf)
def read_template(name):
"""
Read template from shared/template/ specified with or without the .sh extension
Args:
name (str): name of the template to find in shared/template/
Returns:
str: file's content as 1 string
"""
if ".sh" not in name:
name += ".sh"
paths = []
dirs = ["shared", "config"]
for d in dirs:
path = Path(__file__).parent / d / "template" / name
if path.exists():
paths.append(path)
if len(paths) == 0:
failed = [Path(__file__).parent / d / "template" for d in dirs]
s = "Could not find template {} in :\n".format(name)
for fd in failed:
s += str(fd) + "\nAvailable:\n"
for ym in fd.glob("*.sh"):
s += " " + ym.name + "\n"
raise ValueError(s)
if len(paths) == 2:
print("Warning: found 2 relevant template files:\n{}".format("\n".join(paths)))
print("Using {}".format(paths[-1]))
with paths[-1].open("r") as f:
return f.read()
def is_sampled(key, conf):
"""
Is a key sampled or constant? Returns true if conf is empty
Args:
key (str): key to check
conf (dict): hyper parameter search configuration dict
Returns:
bool: key is sampled?
"""
return not conf or (
key in conf and isinstance(conf[key], dict) and "sample" in conf[key]
)
if __name__ == "__main__":
"""
Notes:
* Must provide template name as template=name
* `name`.sh should be in shared/template/
"""
# -------------------------------
# ----- Default Variables -----
# -------------------------------
args = sys.argv[1:]
command_output = ""
user = os.environ.get("USER")
home = os.environ.get("HOME")
exp_conf = {}
dev = False
escape = False
verbose = False
template_name = None
hp_exp_name = None
hp_search_nb = None
exp_path = None
resume = None
force_sbatchs = False
sbatch_base = Path(home) / "climategan_sbatchs"
summary_dir = Path(home) / "climategan_exp_summaries"
hp_search_private = set(["n_search", "template", "search", "summary_dir"])
sbatch_path = "hash"
# --------------------------
# ----- Sanity Check -----
# --------------------------
for arg in args:
if "=" not in arg or " = " in arg:
raise ValueError(
"Args should be passed as `key=value`. Received `{}`".format(arg)
)
# --------------------------------
# ----- Parse Command Line -----
# --------------------------------
args_dict = {arg.split("=")[0]: arg.split("=")[1] for arg in args}
assert "template" in args_dict, "Please specify template=xxx"
template = read_template(args_dict["template"])
template_dict = {k: None for k in get_template_params(template)}
train_args = []
for k, v in args_dict.items():
if k == "verbose":
if v != "0":
verbose = True
elif k == "sbatch_path":
sbatch_path = v
elif k == "sbatch_base":
sbatch_base = Path(v).resolve()
elif k == "force_sbatchs":
force_sbatchs = v.lower() == "true"
elif k == "dev":
if v.lower() != "false":
dev = True
elif k == "escape":
if v.lower() != "false":
escape = True
elif k == "template":
template_name = v
elif k == "exp":
hp_exp_name = v
elif k == "n_search":
hp_search_nb = int(v)
elif k == "resume":
resume = f'"{v}"'
template_dict[k] = f'"{v}"'
elif k == "summary_dir":
if v.lower() == "none":
summary_dir = None
else:
summary_dir = Path(v)
elif k in template_dict:
template_dict[k] = v
else:
train_args.append(f"{k}={v}")
# ------------------------------------
# ----- Load Experiment Config -----
# ------------------------------------
if hp_exp_name is not None:
exp_path, exp_conf = read_exp_conf(hp_exp_name)
if "n_search" in exp_conf and hp_search_nb is None:
hp_search_nb = exp_conf["n_search"]
assert (
hp_search_nb is not None
), "n_search should be specified in a yaml file or from the command line"
hps = resolve(exp_conf, hp_search_nb)
else:
hps = [None]
# ---------------------------------
# ----- Run All Experiments -----
# ---------------------------------
if summary_dir is not None:
summary_dir.mkdir(exist_ok=True, parents=True)
summary = None
for hp_idx, hp in enumerate(hps):
# copy shared values
tmp_template_dict = template_dict.copy()
tmp_train_args = train_args.copy()
tmp_train_args_dict = {
arg.split("=")[0]: arg.split("=")[1] for arg in tmp_train_args
}
print_header(hp_idx)
# override shared values with run-specific values for run hp_idx/n_search
if hp is not None:
for k, v in hp.items():
if k == "resume" and resume is None:
resume = f'"{v}"'
# hp-search params to ignore
if k in hp_search_private:
continue
if k == "codeloc":
v = escape_path(v)
if k == "output":
Path(v).parent.mkdir(parents=True, exist_ok=True)
# override template params depending on exp config
if k in tmp_template_dict:
if template_dict[k] is None or is_sampled(k, exp_conf):
tmp_template_dict[k] = v
# store sampled / specified params in current tmp_train_args_dict
else:
if k in tmp_train_args_dict:
if is_sampled(k, exp_conf):
# warn if key was specified from the command line
tv = tmp_train_args_dict[k]
warn(
"\nWarning: overriding sampled config-file arg",
"{} to command-line value {}\n".format(k, tv),
)
else:
tmp_train_args_dict[k] = v
# create sbatch file where required
tmp_sbatch_path = None
if sbatch_path == "hash":
tmp_sbatch_name = "" if hp_exp_name is None else hp_exp_name[:14] + "_"
tmp_sbatch_name += now() + ".sh"
tmp_sbatch_path = sbatch_base / tmp_sbatch_name
tmp_sbatch_path.parent.mkdir(parents=True, exist_ok=True)
tmp_train_args_dict["sbatch_file"] = str(tmp_sbatch_path)
tmp_train_args_dict["exp_file"] = str(exp_path)
else:
tmp_sbatch_path = Path(sbatch_path).resolve()
summary = extend_summary(
summary, tmp_train_args_dict, tmp_template_dict, exclude=["sbatch_file"]
)
# format train.py's args and crop floats' precision to 5 digits
tmp_template_dict["train_args"] = " ".join(
sorted(
[
"{}={}".format(k, clean_arg(v))
for k, v in tmp_train_args_dict.items()
]
)
)
if "resume.py" in template and resume is None:
raise ValueError("No `resume` value but using a resume.py template")
# format template with clean dict (replace None with "")
sbatch = template.format(
**{
k: v if v is not None else ""
for k, v in tmp_template_dict.items()
if k in template_dict
}
)
# --------------------------------------
# ----- Execute `sbatch` Command -----
# --------------------------------------
if not dev or force_sbatchs:
if tmp_sbatch_path.exists():
print(f"Warning: overwriting {sbatch_path}")
# write sbatch file
with open(tmp_sbatch_path, "w") as f:
f.write(sbatch)
if not dev:
# escape special characters such as " " from sbatch_path's parent dir
parent = str(tmp_sbatch_path.parent)
if escape:
parent = escape_path(parent)
# create command to execute in a subprocess
command = "sbatch {}".format(tmp_sbatch_path.name)
# execute sbatch command & store output
command_output = subprocess.run(
command.split(), stdout=subprocess.PIPE, cwd=parent
)
command_output = "\n" + command_output.stdout.decode("utf-8") + "\n"
print(f"Running from {parent}:")
print(f"$ {command}")
# ---------------------------------
# ----- Summarize Execution -----
# ---------------------------------
if verbose:
print(C.BEIGE + C.ITALIC, "\n" + sbatch + C.ENDC)
if not dev:
print_box(command_output.strip())
jobID = parse_jobID(command_output.strip())
summary["Slurm JOBID"].append(jobID)
summary["Comet Link"].append(f"[{hp_idx}][{hp_idx}]")
print(
"{}{}Summary{} {}:".format(
C.UNDERLINE,
C.OKGREEN,
C.ENDC,
f"{C.WARNING}(DEV){C.ENDC}" if dev else "",
)
)
print(
" "
+ "\n ".join(
"{:10}: {}".format(k, v) for k, v in tmp_template_dict.items()
)
)
print_footer()
print(f"\nRan a total of {len(hps)} jobs{' in dev mode.' if dev else '.'}\n")
table, sum_path = search_summary_table(summary, summary_dir if not dev else None)
if table is not None:
print(table)
print(
"Add `[i]: https://...` at the end of a markdown document",
"to fill in the comet links.\n",
)
if summary_dir is None:
print("Add summary_dir=path to store the printed markdown table ⇪")
else:
print("Saved table in", str(sum_path))
if not dev:
print(
"Cancel entire experiment? \n$ scancel",
" ".join(map(str, summary["Slurm JOBID"])),
)