Options to pass additional parameters to slurm

README.md

@@ -82,10 +82,14 @@ These are all the possible arguments to `schedule`:
 - `job_name: str = "MiniZinc Benchmark"` - The SLURM job name.
 - `cpus_per_task: int = 1` - The number of CPU cores required for each task.
 - `memory: int = 4096` - The maximum memory used for each task.
+- `mem_per_cpu: bool = False` - If True, memory limit will be per CPU (--mem-per-cpu), 
+- else --mem will be used.
 - `debug: bool = False` - Directly capture the output of individual jobs
   and store them in a `./logs/` directory.
 - `wait: bool = False` - The scheduling process will wait for all jobs to
   finish.
+- Additional arguments `key_word=value` are directly passed to slurm unchecked, 
+  after transformation to `--key-word=value`.
 
 A `Configuration` object has the following attributes:
 

src/mzn_bench/mzn_slurm.py

@@ -112,9 +112,11 @@ def schedule(
     job_name: str = "MiniZinc Benchmark",
     cpus_per_task: int = 1,
     memory: int = 4096,
+    mem_per_cpu: bool = False,
     debug: bool = False,
     nice: Optional[int] = None,
     wait: bool = False,
+    **kwargs,
 ) -> NoReturn:
     # Count number of instances
     assert instances.exists()
@@ -151,9 +153,9 @@ def schedule(
     cmd = [
         "sbatch",
         f"--output={slurm_output}",
-        f'--job-name="{job_name}"',
+        f'--job-name={job_name}',
         f"--cpus-per-task={cpus_per_task}",
-        f"--mem={memory}",
+        f"--mem-per-cpu={memory}" if mem_per_cpu else f"--mem={memory}",
         f"--array=1-{n_tasks}",
         f"--time={timeout + timedelta(minutes=1)}",  # Set hard timeout as failsafe
     ]
@@ -169,6 +171,8 @@ def schedule(
         cmd.append(f"--nice={nice}")
     if wait:
         cmd.append("--wait")
+    for option, value in kwargs.items():
+        cmd.append(f"--{option.replace('_', '-')}={value}")
     cmd.extend(
         [
             str(this_script.resolve()),