mirror of
https://github.com/Findus23/halo_comparison.git
synced 2024-09-19 16:03:50 +02:00
70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
import random
|
|
from pathlib import Path
|
|
from subprocess import run
|
|
from time import sleep
|
|
from typing import List, Union
|
|
|
|
import yaml
|
|
|
|
config_file = Path("~/.config/jobrun.yaml").expanduser()
|
|
|
|
|
|
def jobrun(
|
|
args: List[Union[str, Path]],
|
|
pwd: Path = None,
|
|
time: str = "05:10:00",
|
|
tasks: int = None,
|
|
nodes: int = None,
|
|
mem: int = None,
|
|
qos: str = None,
|
|
account: str = None,
|
|
source: Path = None,
|
|
sbatch: bool = True
|
|
):
|
|
with config_file.open() as f:
|
|
config = yaml.safe_load(f)
|
|
jobscript_dir = Path(config["jobscript_dir"]).expanduser().absolute()
|
|
name = random.randint(0, 100000)
|
|
jobscript_file = jobscript_dir / f"{name}.sh"
|
|
jobscript_lines = ["#!/bin/bash"]
|
|
sbatch_options = {
|
|
"mail-type": "ALL",
|
|
"mail-user": config["mail"],
|
|
"time": time,
|
|
"nodes": nodes,
|
|
"tasks": tasks,
|
|
"mem": f"{mem}G",
|
|
"qos": qos,
|
|
"account": account,
|
|
"job-name": name,
|
|
"output": jobscript_dir / f"out-{name}-%a_%a.txt"
|
|
}
|
|
for key, value in sbatch_options.items():
|
|
if value is not None:
|
|
jobscript_lines.append(f"#SBATCH --{key}={value}")
|
|
|
|
if source is not None:
|
|
jobscript_lines.append(f"source {source.absolute()}")
|
|
|
|
if pwd is not None:
|
|
assert pwd.exists()
|
|
jobscript_lines.append(f"cd {pwd.absolute()}")
|
|
|
|
strargs: List[str] = []
|
|
for a in args:
|
|
if isinstance(a, Path):
|
|
strargs.append(str(a.absolute()))
|
|
else:
|
|
strargs.append(a)
|
|
jobscript_lines.append(" ".join(strargs))
|
|
jobscript_lines.append("")
|
|
jobscript_file.write_text("\n".join(jobscript_lines))
|
|
|
|
print(f"created {jobscript_file.name}")
|
|
|
|
if sbatch:
|
|
run([
|
|
"sbatch",
|
|
str(jobscript_file)
|
|
], check=True)
|
|
sleep(2)
|