Merge pull request #4 from znamlab/dev

Dev

CHANGELOG.md

@@ -1,6 +1,11 @@
 
 ## Changelog
 
+### [v0.5] - 2024-03-27
+
+- Option to print job ID in log when starting.
+
+
 ### [v0.4] - 2024-01-17
 
 - Job dependencies that are list or tuples are automatically formatted as `afterok:jobid1:jobid2:...`

ReadMe.md

@@ -86,6 +86,10 @@ The decorated function will have 4 new keyword arguments:
 
 ```
 use_slurm (bool): whether to use slurm or not
+dependency_type (str, optional): Type of dependence on previous jobs.
+    Defaults to "afterok" which only runs the next job if all previous
+    jobs have finished successfully. Other options are "after", "afterany",
+    "aftercorr" and "afternotok". See sbatch documentation for more details.
 job_dependency (str): job id to depend on
 slurm_folder (str): where to write the slurm script and logs
 scripts_name (str): name of the slurm script and python file

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="znamutils",
-    version="v0.4",
+    version="v0.5",
     packages=find_packages(exclude=["tests"]),
     url="https://github.com/znamlab/znamutils",
     license="MIT",

tests/test_slurm_helpers.py

@@ -9,3 +9,16 @@
 
 def test_create_slurm_sbatch():
     p = Path(flz.PARAMETERS["data_root"]["processed"]) / "test"
+    slurm_helper.create_slurm_sbatch(
+        p,
+        print_job_id=True,
+        conda_env="cottage_analysis",
+        python_script="test.py",
+        script_name="test.sh",
+    )
+    with open(p / "test.sh") as f:
+        txt = f.read()
+
+
+if __name__ == "__main__":
+    test_create_slurm_sbatch()

znamutils/decorators.py

@@ -27,6 +27,10 @@ def slurm_it(
 
     The decorated function will have 4 new keyword arguments:
         use_slurm (bool): whether to use slurm or not
+        dependency_type (str, optional): Type of dependence on previous jobs.
+            Defaults to "afterok" which only runs the next job if all previous
+            jobs have finished successfully. Other options are "after", "afterany",
+            "aftercorr" and "afternotok". See sbatch documentation for more details.
         job_dependency (str): job id to depend on
         slurm_folder (str): where to write the slurm script and logs
         scripts_name (str): name of the slurm script and python file
@@ -65,6 +69,9 @@ def slurm_it(
     use_slurm = Parameter(
         "use_slurm", kind=Parameter.POSITIONAL_OR_KEYWORD, default=False
     )
+    dependency_type = Parameter(
+        "dependency_type", kind=Parameter.POSITIONAL_OR_KEYWORD, default=None
+    )
     job_dependency = Parameter(
         "job_dependency", kind=Parameter.POSITIONAL_OR_KEYWORD, default=None
     )
@@ -80,7 +87,14 @@ def slurm_it(
 
     new_sig = add_signature_parameters(
         func_sig,
-        last=(use_slurm, job_dependency, slurm_folder, script_name, slurm_options),
+        last=(
+            use_slurm,
+            dependency_type,
+            job_dependency,
+            slurm_folder,
+            script_name,
+            slurm_options,
+        ),
     )
     from_imports = from_imports or {func.__module__: func.__name__}
     # create the new function with modified signature
@@ -89,6 +103,7 @@ def new_func(*args, **kwargs):
 
         # pop the slurm only arguments
         use_slurm = kwargs.pop("use_slurm")
+        dependency_type = kwargs.pop("dependency_type")
         job_dependency = kwargs.pop("job_dependency")
         slurm_folder = kwargs.pop("slurm_folder")
         scripts_name = kwargs.pop("scripts_name")
@@ -136,7 +151,13 @@ def new_func(*args, **kwargs):
             from_imports=from_imports,
         )
 
-        return slurm_helper.run_slurm_batch(sbatch_file, job_dependency=job_dependency)
+        return slurm_helper.run_slurm_batch(
+            sbatch_file,
+            dependency_type=dependency_type
+            if dependency_type is not None
+            else "afterok",
+            job_dependency=job_dependency,
+        )
 
     # return the new function
     return new_func

znamutils/slurm_helper.py

@@ -5,19 +5,23 @@
 import inspect
 
 
-def run_slurm_batch(script_path, job_dependency=None):
+def run_slurm_batch(script_path, dependency_type="afterok", job_dependency=None):
     """Run a slurm script
 
     Args:
         script_path (str): Full path to the script
+        dependency_type (str, optional): Type of dependence on previous jobs.
+            Defaults to "afterok" which only runs the next job if all previous
+            jobs have finished successfully. Other options are "after", "afterany",
+            "aftercorr" and "afternotok". See sbatch documentation for more details.
         job_dependency (str, optional): Job ID that needs to finish before running
             sbtach. Defaults to None.
 
     Returns:
         str: Job ID of the sbatch job
     """
     if job_dependency is not None:
-        dep = f"--dependency=afterok:{job_dependency} "
+        dep = f"--dependency={dependency_type}:{job_dependency} "
     else:
         dep = ""
     command = f"sbatch {dep}{script_path}"
@@ -36,6 +40,7 @@ def create_slurm_sbatch(
     slurm_options=None,
     module_list=None,
     split_err_out=False,
+    print_job_id=False,
 ):
     """Create a slurm sh script that will call a python script
 
@@ -79,6 +84,10 @@ def create_slurm_sbatch(
             boiler = "\n" + "\n".join([f"ml {module}" for module in module_list]) + "\n"
         else:
             boiler = "\n"
+
+        if print_job_id:
+            boiler += f'echo "Job ID: $SLURM_JOB_ID"\n'
+
         boiler += "\n".join(
             [
                 "ml Anaconda3",