add more slurm scripts with logging

.gitignore

@@ -167,3 +167,4 @@ experiments/samples/*.hdf5
 **/bash/**/*.sh
 **/logs/**
 **log**.txt
+**/jobs_out/**

scripts/slurm/slurm_job.py

@@ -0,0 +1,45 @@
+"""
+Submit sbatch job with default settings. To be used in other scripts specifying `cmd`.
+"""
+
+import secrets
+from pathlib import Path
+
+JOB_DIR = Path(__file__).parent.joinpath("jobs_out")
+assert JOB_DIR.exists(), f"Need to create '{JOB_DIR}' dir"
+JOB_SEED = secrets.randbelow(1_000_000_000)
+
+
+def setup_sbatch_job_gpu(
+    jobname: str,
+    time: str = "00:10",
+    nodes: int = 1,
+    n_tasks_per_node: int = 1,
+    qos: str = "debug",
+) -> Path:
+    """Returns path to sbatch ready job file, commands need to be appended."""
+
+    # time formating
+    assert len(time) == 5 and time.count(":") == 1
+
+    # prepare files and directories
+    jobfile_name = f"{jobname}_{JOB_SEED}.sbatch"
+    job_dir = Path(JOB_DIR)
+    jobfile = job_dir.joinpath(jobfile_name)
+
+    with open(jobfile, "w") as f:
+        f.writelines(
+            "#!/bin/bash\n"
+            f"#SBATCH --job-name={jobname}\n"
+            f"#SBATCH --account=m1727\n"
+            f"#SBATCH --qos={qos}\n"
+            f"#SBATCH --mail-type=begin,end,fail\n"
+            f"#SBATCH [email protected]\n"
+            f"#SBATCH -C gpu\n"
+            f"#SBATCH --output={JOB_DIR}/{jobname}_%j.out\n"
+            f"#SBATCH --time={time}:00\n"
+            f"#SBATCH --nodes={nodes}\n"
+            f"#SBATCH --ntasks-per-node={n_tasks_per_node}\n"
+        )
+
+    return jobfile

scripts/slurm/slurm_toy_shear_gpu.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+import subprocess
+
+import click
+from slurm_job import setup_sbatch_job_gpu
+
+from bpd import DATA_DIR
+
+
+@click.command()
+@click.option("--seed", required=True, type=int)
+@click.option("--jobname", required=True, type=str)
+@click.option("--time", default="00:20", type=str, show_default=True, help="HH:MM")
+@click.option("--mem-per-gpu", default="20G", type=str, show_default=True)
+@click.option("--n-vec", default=50, type=int, show_default=True)
+@click.option("--k", default=1000, type=int, show_default=True)
+@click.option("--trim", default=10, type=int, show_default=True)
+@click.option("--n-seeds-per-task", default=250, type=int, show_default=True)
+@click.option("--n-gals", default=1000, type=int, show_default=True)
+@click.option("--n-samples-shear", default=3000, type=int, show_default=True)
+@click.option(
+    "--add-extra",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Adding additional runs to previously existing experiment.",
+)
+@click.option("--qos", default="debug", type=str, show_default=True)  # debug, regular
+def main(
+    seed: int,
+    jobname: str,
+    time: str,
+    mem_per_gpu: str,
+    n_vec: int,
+    k: int,
+    trim: int,
+    n_seeds_per_task: int,
+    n_gals: int,
+    n_samples_shear: int,
+    add_extra: bool,
+    qos: str,
+):
+
+    tagpath = DATA_DIR / "cache_chains" / jobname
+    if not add_extra:
+        assert not tagpath.exists()
+
+    jobfile = setup_sbatch_job_gpu(jobname, time, nodes=1, n_tasks_per_node=4, qos=qos)
+
+    template_cmd = "python /global/u2/i/imendoza/BPD/scripts/vect_toy_shear_gpu.py --n-samples-gals {n_gals} --n-samples-shear {n_samples_shear} --n-vec {n_vec} --seed {{seed}}  --n-seeds {n_seeds_per_task} --tag {jobname} --k {k} --trim {trim}"
+
+    base_cmd = template_cmd.format(
+        n_gals=n_gals,
+        n_samples_shear=n_samples_shear,
+        n_vec=n_vec,
+        n_seeds_per_task=n_seeds_per_task,
+        jobname=jobname,
+        k=k,
+        trim=trim,
+    )
+
+    # append to jobfile the  commands.
+    with open(jobfile, "a") as f:
+        f.write("\n")
+
+    for ii in range(4):
+        cmd_seed = int(f"{seed}{ii}")
+        cmd = base_cmd.format(seed=cmd_seed)
+        srun_cmd = f"srun --exact -u -n 1 -c 1 --gpus-per-task 1 --mem-per-gpu={mem_per_gpu} {cmd}  &\n"
+
+        with open(jobfile, "a") as f:
+            f.write(srun_cmd)
+
+    with open(jobfile, "a") as f:
+        f.write("\nwait")
+
+    subprocess.run(f"sbatch {jobfile.as_posix()}", shell=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/vect_toy_shear_gpu.py

@@ -12,9 +12,7 @@
 
 
 @click.command()
-@click.option(
-    "--n-vec", type=int, default=100, help="# shear chains to run in parallel"
-)
+@click.option("--n-vec", type=int, default=50, help="# shear chains in parallel")
 @click.option("--tag", type=str, required=True)
 @click.option("--seed", type=int, required=True)
 @click.option("--n-seeds", type=int, required=True)