Skip to content

Commit

Permalink
feat: first half of auto selection between cpus-per-task and cpus-per…
Browse files Browse the repository at this point in the history
…-gpu - needs adjustment in jobstep executor code, too.
  • Loading branch information
cmeesters committed Jan 10, 2025
1 parent cb94eb9 commit 3f36f86
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions snakemake_executor_plugin_slurm/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,8 @@ def run_job(self, job: JobExecutorInterface):

# fixes #40 - set ntasks regardless of mpi, because
# SLURM v22.05 will require it for all jobs
if job.resources.get("gpus") or "gpu" in job.resources.get("gres", ""):
gpu_job = job.resources.get("gpus") or "gpu" in job.resources.get("gres", "")
if gpu_job:
call += f" --ntasks-per-gpu={job.resources.get('tasks', 1)}"
else
call += f" --ntasks={job.resources.get('tasks', 1)}"
Expand All @@ -305,8 +306,9 @@ def run_job(self, job: JobExecutorInterface):
"specified. Assuming 'tasks_per_node=1'."
"Probably not what you want."
)

call += f" --cpus-per-task={get_cpus_per_task(job)}"
# we need to set cpus-per-task OR cpus-per-gpu, the function
# will return a string with the corresponding value
call += f" {get_cpus_per_task(job, gpu_job)}"

if job.resources.get("slurm_extra"):
self.check_slurm_extra(job)
Expand Down

0 comments on commit 3f36f86

Please sign in to comment.