Restructuring & added default template (#1)

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)
+  ```

pyproject.toml

@@ -1,8 +1,9 @@
 [tool.poetry]
 name = "slurmer"
-version = "0.1.0"
+version = "0.2.0"
 description = ""
 authors = ["Minjoon Hong <[email protected]>"]
+include = ["slurmer/"]
 
 [tool.poetry.dependencies]
 python = "^3.8"

slurmer/__init__.py

@@ -1,6 +1,6 @@
 from .config import TemplateManager
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 
 def list_templates():

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

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

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

tests/test_slurmer.py

@@ -2,4 +2,4 @@
 
 
 def test_version():
-    assert __version__ == '0.1.0'
+    assert __version__ == "0.2.0"