From e67f52174360338020ee2aa2c3e92ec8aa203444 Mon Sep 17 00:00:00 2001 From: mjhong0708 <61532201+mjhong0708@users.noreply.github.com> Date: Thu, 20 Jan 2022 18:56:23 +0900 Subject: [PATCH] Restructuring & added default template (#1) --- README.md | 83 +++++++++-------------- pyproject.toml | 3 +- slurmer/__init__.py | 2 +- slurmer/config.py | 9 +-- slurmer/job.py | 127 +++++++++++++++++++++++++++++++++++ slurmer/script.py | 63 ----------------- slurmer/templates/default.sh | 16 +++++ tests/test_slurmer.py | 2 +- 8 files changed, 183 insertions(+), 122 deletions(-) create mode 100644 slurmer/job.py delete mode 100644 slurmer/script.py create mode 100644 slurmer/templates/default.sh diff --git a/README.md b/README.md index 6853ff5..87841b2 100644 --- a/README.md +++ b/README.md @@ -5,61 +5,40 @@ Automatic generation of `SLURM` job script based on `jinja` template engine ## Install -Clone repository, install `poetry` and run `poetry install`. +Install `master` branch: `pip install git+https://github.com/mjhong0708/slurmer` Or, download wheel from `release` tab and `pip install` it. ## Usage -Prepare a template file in `~/.config/slurmer/templates`. - -Example: `run_slurm.sh` - -```bash -#!/bin/bash -#SBATCH -J {{ job_name }} -#SBATCH -o myMPI.o%j # output and error file name (%j expands to jobID) -#SBATCH -p {{ node_partition }} -#SBATCH -N {{ num_nodes }} -#SBATCH -n {{ num_tasks }} -{% if node_list == 'none' %} -##SBATCH -w, --nodelist= -{% else %} -#SBATCH -w, --nodelist={{ node_list }} -{% endif %} -{% if node_exclude_list != 'none' %} -#SBATCH -x, --exclude={{ node_exclude_list }} -{% endif %} - -mpiexec.hydra -np $SLURM_NTASKS path/to/vasp > stdout.log -``` - -Then, call `get_script` to render template as job script. -```python -from slurmer.script import get_script -script = get_script( - template_file="run_slurm.sh", - job_name="myjob", - node_partition="g1", - num_nodes=2, - num_tasks=32, - node_list=[10, 12, 15], # optional -) -print(script) -``` - -Now you can see output: - -```bash -#!/bin/bash -#SBATCH -J myjob -#SBATCH -o myMPI.o%j # output and error file name (%j expands to jobID) -#SBATCH -p g1 -#SBATCH -N 2 -#SBATCH -n 32 - -#SBATCH -w, --nodelist=n010,n012,n015 - -mpiexec.hydra -np $SLURM_NTASKS path/to/vasp > stdout.log -``` +By default, [`default.sh`](https://github.com/mjhong0708/slurmer/blob/master/slurmer/templates/default.sh) is available as template. + +- Show list of available templates + + ```python + import slurmer + + slurmer.list_templates() + ``` +- Add directory to seek template files + + ```python + from slurmer.config import TemplateManager + + m = TemplateManager() + m.add_path('path/to/template') + ``` +- Submit job + + ```python + import os + from slurmer.job import SlurmJob + + job_dir = "my_job" + os.mkdir(job_dir) + + job = SlurmJob(job_dir, "default.sh", "myjob", "g1", 2, 32, exec_command="echo 'Hello'") + + job.submit(write_job_script=True) + ``` diff --git a/pyproject.toml b/pyproject.toml index ddfa234..0f32021 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,8 +1,9 @@ [tool.poetry] name = "slurmer" -version = "0.1.0" +version = "0.2.0" description = "" authors = ["Minjoon Hong "] +include = ["slurmer/"] [tool.poetry.dependencies] python = "^3.8" diff --git a/slurmer/__init__.py b/slurmer/__init__.py index 1b19e5b..e62c76e 100644 --- a/slurmer/__init__.py +++ b/slurmer/__init__.py @@ -1,6 +1,6 @@ from .config import TemplateManager -__version__ = "0.1.0" +__version__ = "0.2.0" def list_templates(): diff --git a/slurmer/config.py b/slurmer/config.py index c85f3fe..fcdb0e1 100644 --- a/slurmer/config.py +++ b/slurmer/config.py @@ -4,12 +4,13 @@ from tinydb import Query, TinyDB config_dir = Path.home() / ".config" / "slurmer" -default_template_dir = config_dir / "templates" +user_template_dir = config_dir / "templates" +default_template_dir = Path(__file__).parent / "templates" if not config_dir.is_dir(): config_dir.mkdir(parents=True) -if not default_template_dir.is_dir(): - default_template_dir.mkdir(parents=True) +if not user_template_dir.is_dir(): + user_template_dir.mkdir(parents=True) class TemplateManager: @@ -23,7 +24,7 @@ def __init__(self): def _update_dirs(self): """Updates template dirs with config db""" - self.template_dirs = [default_template_dir.absolute()] + self.template_dirs = [default_template_dir.absolute(), user_template_dir.absolute()] for d in self.db.all(): self.template_dirs.append(Path(d["path"])) diff --git a/slurmer/job.py b/slurmer/job.py new file mode 100644 index 0000000..f494dac --- /dev/null +++ b/slurmer/job.py @@ -0,0 +1,127 @@ +import os +import subprocess +from typing import Optional, Sequence, Union +from pathlib import Path +from shutil import copyfile +from tempfile import NamedTemporaryFile + +from jinja2 import Environment, FileSystemLoader, select_autoescape + +from .config import TemplateManager + +m = TemplateManager() +loader = FileSystemLoader(m.template_dirs) +env = Environment(loader=loader, autoescape=select_autoescape()) + + +class SlurmJob: + def __init__( + self, + workdir: Union[str, os.PathLike], + template_file: str, + job_name: str, + node_partition: str, + num_nodes: int, + num_tasks: int, + node_list: Optional[Sequence[int]] = None, + node_exclude_list: Optional[str] = None, + **kwargs, + ): + if not isinstance(workdir, Path): + self.workdir = Path(workdir) + else: + self.workdir = workdir + self.template_file = template_file + self.job_name = job_name + self.node_partition = node_partition + self.num_nodes = num_nodes + self.num_tasks = num_tasks + self.node_list = node_list + self.node_exclude_list = node_exclude_list + self.extra_args = kwargs + self.__job_script: Optional[str] = None + + @property + def job_script(self) -> str: + if self.__job_script is not None: + return self.__job_script + else: + job_script = generate_jobscript( + template_file=self.template_file, + job_name=self.job_name, + node_partition=self.node_partition, + num_nodes=self.num_nodes, + num_tasks=self.num_tasks, + node_list=self.node_list, + node_exclude_list=self.node_exclude_list, + **self.extra_args, + ) + return job_script + + def submit(self, write_job_script: bool = False, job_script_name: str = "job_script.sh"): + cwd = Path.cwd() + if not self.workdir.is_dir(): + raise RuntimeError("Workdir does not exists.") + + os.chdir(self.workdir) + with NamedTemporaryFile("w") as f: + f.write(self.job_script) + f.file.close() + + subprocess.check_call(["sbatch", f.name]) + if write_job_script: + copyfile(f.name, job_script_name) + os.chdir(cwd) + + +def generate_jobscript( + template_file: str, + job_name: str, + node_partition: str, + num_nodes: int, + num_tasks: int, + node_list: Optional[Sequence[int]] = None, + node_exclude_list: Optional[str] = None, + **kwargs, +) -> str: + """Get slurm job script. + + Args: + template_file (str): Path of template file. + job_name (str): Name of job. + node_partition (str): Partition of nodes. + num_nodes (int): The number of nodes. + num_tasks (int): The number of tasks(threads). + node_list (Optional[Sequence[int]], optional): List of nodes to use. Defaults to None. + node_exclude_list (Optional[str], optional): List of nodes to exclude. Defaults to None. + + Raises: + ValueError: Raised when both node_list and node_exclude_list are specified. + + Returns: + str: Rendered job script. + """ + if node_list is not None and node_exclude_list is not None: + raise ValueError("node_list and node_exclude list cannot be specified simultaneously.") + + if node_list is None: + _node_list = "none" + else: + _node_list = ",".join([f"n{node:0>3}" for node in node_list]) + + if node_exclude_list is None: + _node_exclude_list = "none" + else: + _node_exclude_list = node_exclude_list + + template = env.get_template(template_file.__str__()) + rendered = template.render( + job_name=job_name, + node_partition=node_partition, + num_nodes=num_nodes, + num_tasks=num_tasks, + node_list=_node_list, + node_exclude_list=_node_exclude_list, + **kwargs, + ) + return rendered diff --git a/slurmer/script.py b/slurmer/script.py deleted file mode 100644 index b8b5c22..0000000 --- a/slurmer/script.py +++ /dev/null @@ -1,63 +0,0 @@ -import os -from typing import Sequence, Optional - -from jinja2 import Environment, FileSystemLoader, select_autoescape - -from .config import TemplateManager - -m = TemplateManager() -loader = FileSystemLoader(m.template_dirs) -env = Environment(loader=loader, autoescape=select_autoescape()) - - -def get_script( - template_file: os.PathLike, - job_name: str, - node_partition: str, - num_nodes: int, - num_tasks: int, - node_list: Optional[Sequence[int]] = None, - node_exclude_list: Optional[str] = None, - **kwargs, -) -> str: - """Get slurm job script. - - Args: - template_file (os.PathLike): Path of template file. - job_name (str): Name of job. - node_partition (str): Partition of nodes. - num_nodes (int): The number of nodes. - num_tasks (int): The number of tasks(threads). - node_list (Optional[Sequence[int]], optional): List of nodes to use. Defaults to None. - node_exclude_list (Optional[str], optional): List of nodes to exclude. Defaults to None. - - Raises: - ValueError: Raised when both node_list and node_exclude_list are specified. - - Returns: - str: Rendered job script. - """ - if node_list is not None and node_exclude_list is not None: - raise ValueError("node_list and node_exclude list cannot be specified simultaneously.") - - if node_list is None: - _node_list = "none" - else: - _node_list = ",".join([f"n{node:0>3}" for node in node_list]) - - if node_exclude_list is None: - _node_exclude_list = "none" - else: - _node_exclude_list = node_exclude_list - - template = env.get_template(template_file.__str__()) - rendered = template.render( - job_name=job_name, - node_partition=node_partition, - num_nodes=num_nodes, - num_tasks=num_tasks, - node_list=_node_list, - node_exclude_list=_node_exclude_list, - **kwargs, - ) - return rendered diff --git a/slurmer/templates/default.sh b/slurmer/templates/default.sh new file mode 100644 index 0000000..ceb577d --- /dev/null +++ b/slurmer/templates/default.sh @@ -0,0 +1,16 @@ +#!/bin/bash +#SBATCH -J {{ job_name }} +#SBATCH -o myMPI.o%j # output and error file name (%j expands to jobID) +#SBATCH -p {{ node_partition }} +#SBATCH -N {{ num_nodes }} +#SBATCH -n {{ num_tasks }} +{% if node_list == 'none' %} +##SBATCH -w, --nodelist= +{% else %} +#SBATCH -w, --nodelist={{ node_list }} +{% endif %} +{% if node_exclude_list != 'none' %} +#SBATCH -x, --exclude={{ node_exclude_list }} +{% endif %} + +{{ exec_command }} > stdout.log \ No newline at end of file diff --git a/tests/test_slurmer.py b/tests/test_slurmer.py index 72f5e2c..fe0bc30 100644 --- a/tests/test_slurmer.py +++ b/tests/test_slurmer.py @@ -2,4 +2,4 @@ def test_version(): - assert __version__ == '0.1.0' + assert __version__ == "0.2.0"