fix: set slurm walltime & memory usage (#116)

* feat: set slurm walltime & memory usage

* docs: walltime & memory

* test: update test_use_template() with new walltime & memory keys

Author: kelly-sovacool

CHANGELOG.md

@@ -3,6 +3,7 @@
 - Fixes for `ccbr_tools install`: (#98, @kelly-sovacool)
   - Correct installation method for `spacesavers2`.
   - Change symlink permissions so group members can edit them.
+- Allow nextflow pipelines to specify walltime & memory usage for the slurm submission template. (#116, @kelly-sovacool)
 
 ## Tools 0.4.3
 

src/ccbr_tools/pipeline/nextflow.py

@@ -36,6 +36,8 @@ def run(
     debug=False,
     hpc=get_hpc(),
     hpc_modules="nextflow",
+    hpc_walltime="1-00:00:00",
+    hpc_memory="1G",
 ):
     """
     Runs a Nextflow workflow with support for local or SLURM (HPC) execution modes.
@@ -49,6 +51,8 @@ def run(
         debug (bool, optional): If True, prints commands without executing them. Defaults to False.
         hpc (object, optional): HPC environment object, used for SLURM execution and module loading. Defaults to result of `~ccbr_tools.pipeline.hpc.get_hpc()`.
         hpc_modules (str, optional): Name(s) of modules to load for Nextflow execution on HPC. Defaults to "nextflow".
+        hpc_walltime (str, optional): Walltime for SLURM job submission. Defaults to "1-00:00:00".
+        hpc_memory (str, optional): Memory allocation for SLURM job submission. Defaults to "1G".
 
     Behavior:
         - Constructs the Nextflow command with appropriate arguments and profiles.
@@ -116,7 +120,10 @@ def run(
     if mode == "slurm":
         slurm_filename = "submit_slurm.sh"
         use_template(
-            "submit_slurm.sh",
+            slurm_filename,
+            output_filepath=slurm_filename,
+            MEMORY=hpc_memory,
+            WALLTIME=hpc_walltime,
             PIPELINE=pipeline_name,
             MODULES=hpc_modules,
             ENV_VARS="\n".join(

src/ccbr_tools/templates/__init__.py

@@ -68,7 +68,7 @@ def use_template(template_name, output_filepath=None, **kwargs: str):
         IOError: If there is an error writing the output file.
 
     Examples:
-       >>> use_template("submit_slurm.sh", output_filepath="./submit_slurm.sh", PIPELINE="CCBR_nxf", MODULES="ccbrpipeliner nextflow", ENV_VARS="", RUN_COMMAND="nextflow run main.nf -stub")
+       >>> use_template("submit_slurm.sh", output_filepath="./submit_slurm.sh", MEMORY='1G', WALLTIME="1-00:00:00", PIPELINE="CCBR_nxf", MODULES="ccbrpipeliner nextflow", ENV_VARS="", RUN_COMMAND="nextflow run main.nf -stub")
     """
     template_str = read_template(template_name)
     if not output_filepath:

src/ccbr_tools/templates/submit_slurm.sh

@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 #SBATCH --cpus-per-task=1
-#SBATCH --mem=1g
-#SBATCH --time=1-00:00:00
+#SBATCH --mem={MEMORY}
+#SBATCH --time={WALLTIME}
 #SBATCH --parsable
 #SBATCH -J "{PIPELINE}"
 #SBATCH --mail-type=BEGIN,END,FAIL

tests/test_templates.py

@@ -23,6 +23,8 @@ def test_use_template():
         use_template(
             "submit_slurm.sh",
             output_filepath=out_filepath,
+            WALLTIME="4-00:00:00",
+            MEMORY="1G",
             PIPELINE="CCBR_nxf",
             MODULES="ccbrpipeliner nextflow",
             ENV_VARS="export HELLO=WORLD",
@@ -32,6 +34,7 @@ def test_use_template():
             template_str = outfile.read()
         assertions = [
             '#SBATCH -J "CCBR_nxf"' in template_str,
+            "#SBATCH --time=4-00:00:00" in template_str,
             "module load ccbrpipeliner nextflow" in template_str,
             "export HELLO=WORLD" in template_str,
             "nextflow run main.nf -stub" in template_str,
@@ -44,6 +47,8 @@ def generate_slurm_template():
     use_template(
         "submit_slurm.sh",
         output_filepath="tests/data/templates/submit_slurm.sh",
+        WALLTIME="1-00:00:00",
+        MEMORY="1G",
         PIPELINE="CCBR_nxf",
         MODULES=hpc.modules["nxf"],
         ENV_VARS=hpc.env_vars,