File size: 1,186 Bytes
27ca8b3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
"""
utils for submitting to clusters, such as slurm
"""
import os
from omegaconf import DictConfig, OmegaConf
from datetime import datetime
from pathlib import Path
from utils.print_utils import cyan
# This is set below.
REPO_DIR = None
def submit_slurm_job(
cfg: DictConfig,
python_args: str,
project_root: Path,
):
log_dir = project_root / "slurm_logs" / f"{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}-{cfg.name}"
log_dir.mkdir(exist_ok=True, parents=True)
(project_root / "slurm_logs" / "latest").unlink(missing_ok=True)
(project_root / "slurm_logs" / "latest").symlink_to(log_dir, target_is_directory=True)
params = dict(name=cfg.name, log_dir=log_dir, project_root=project_root, python_args=python_args)
params.update(cfg.cluster.params)
slurm_script = cfg.cluster.launch_template.format(**params)
slurm_script_path = log_dir / "job.slurm"
with slurm_script_path.open("w") as f:
f.write(slurm_script)
os.system(f"chmod +x {slurm_script_path}")
os.system(f"sbatch {slurm_script_path}")
print(f"\n{cyan('script:')} {slurm_script_path}\n{cyan('slurm errors and logs:')} {log_dir}\n")
return log_dir
|