sbatch.py

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"])),
        )