Write SLURM with Python

Write SLURM with Python#

Helper Python function to write a SLURM array script from a CSV

Here is a helper function in Python that automatically creates a SLURM job array script based on a CSV parameter file. It scans the CSV, figures out how many lines there are (skipping the header), and then writes out a .slurm script with the right --array range.

NOTE This function was generated by AI and serves only as a guide – needs to be tested and modified to your needs.

def create_slurm_array_script(csv_file, job_name="ParamSweep", output_prefix="output", 
                              nodes=1, tasks=4, time="00:20:00", partition="normal",
                              python_script="run_analysis.py"):
    """
    Creates a SLURM job array script that will loop over a CSV file with parameters.
    It assumes the CSV has a header line and data starts on line 2.
    """

    # Count lines in the CSV
    with open(csv_file) as f:
        lines = f.readlines()
    num_jobs = len(lines) - 1  # minus header

    # Build the script content
    script_content = f"""#!/bin/bash
#SBATCH -J {job_name}
#SBATCH -o {output_prefix}.%A_%a.out
#SBATCH -e {output_prefix}.%A_%a.err
#SBATCH -p {partition}
#SBATCH -N {nodes}
#SBATCH -n {tasks}
#SBATCH -t {time}
#SBATCH --array=2-{num_jobs + 1}

module load python3

cd $SCRATCH/my_project
cp ~/my_local_dir/{csv_file} .

python {python_script} ${{SLURM_ARRAY_TASK_ID}}
"""

    # Write to file
    slurm_filename = f"{job_name}.slurm"
    with open(slurm_filename, "w") as f:
        f.write(script_content)

    print(f"Generated SLURM script '{slurm_filename}' for {num_jobs} jobs.")

# Example usage:
create_slurm_array_script("params.csv")

What it does#

  • Automatically reads your CSV to count how many runs you’ll have (minus the header line).

  • Writes a .slurm script with:

  • the right --array range (starts at 2 because 1 is typically the header line),

  • your preferred number of nodes, tasks, wall time, etc.,

  • your specified Python script.

Customize easily#

Want more cores or a different time? Just change:

create_slurm_array_script("params.csv", tasks=8, time="01:00:00", partition="development")

That gives you a turnkey system:

  • Edit your params.csv,

  • Auto-generate the SLURM job array script with this helper,

  • Submit with sbatch.

Contents