Skip to content

Commit 660c800

Browse files
authored
fix: quoting parameters (#355)
This PR should fix issue #354 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - Bug Fixes - Safer SBATCH parameter handling: flags are emitted only when appropriate, quoting and empty/None handling improved; workdir flag now only emitted when present. - Refactor - Internal parameter handling streamlined for more consistent and robust command construction; resource and CPU/memory selection logic clarified. - Tests - Test expectations updated to reflect the new command formatting and resource-flag behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent e8378ec commit 660c800

File tree

2 files changed

+40
-22
lines changed

2 files changed

+40
-22
lines changed

snakemake_executor_plugin_slurm/submit_string.py

Lines changed: 35 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,18 @@
11
from snakemake_executor_plugin_slurm_jobstep import get_cpu_setting
22
from types import SimpleNamespace
3+
import shlex
4+
5+
6+
def safe_quote(value):
7+
"""
8+
Safely quote a parameter value using shlex.quote.
9+
Handles None values and converts to string if needed.
10+
Returns empty quotes for empty strings.
11+
"""
12+
str_value = str(value)
13+
if str_value == "":
14+
return "''"
15+
return shlex.quote(str_value)
316

417

518
def get_submit_command(job, params):
@@ -10,37 +23,41 @@ def get_submit_command(job, params):
1023
params = SimpleNamespace(**params)
1124

1225
call = (
13-
f"sbatch "
14-
f"--parsable "
15-
f"--job-name {params.run_uuid} "
16-
f'--output "{params.slurm_logfile}" '
17-
f"--export=ALL "
18-
f'--comment "{params.comment_str}"'
26+
"sbatch "
27+
"--parsable "
28+
f"--job-name {safe_quote(params.run_uuid)} "
29+
f"--output {safe_quote(params.slurm_logfile)} "
30+
"--export=ALL "
31+
f"--comment {safe_quote(params.comment_str)}"
1932
)
2033

2134
# No accout or partition checking is required, here.
2235
# Checking is done in the submit function.
2336

2437
# here, only the string is used, as it already contains
25-
# '-A {account_name}'
38+
# "-A '{account_name}'"
2639
call += f" {params.account}"
2740
# here, only the string is used, as it already contains
28-
# '- p {partition_name}'
41+
# "- p '{partition_name}'"
2942
call += f" {params.partition}"
3043

3144
if job.resources.get("clusters"):
32-
call += f" --clusters {job.resources.clusters}"
45+
call += f" --clusters {safe_quote(job.resources.clusters)}"
3346

3447
if job.resources.get("runtime"):
35-
call += f" -t {job.resources.runtime}"
48+
call += f" -t {safe_quote(job.resources.runtime)}"
3649

37-
if job.resources.get("constraint") or isinstance(
38-
job.resources.get("constraint"), str
39-
):
40-
call += f" -C '{job.resources.get('constraint')}'"
50+
# Both, constraint and qos are optional.
51+
# If not set, they will not be added to the sbatch call.
52+
# If explicitly set to an empty string,
53+
# `--constraint ''` or `--qos ''` will be added.
54+
constraint = job.resources.get("constraint")
55+
if constraint is not None:
56+
call += f" -C {safe_quote(constraint)}"
4157

42-
if job.resources.get("qos") or isinstance(job.resources.get("qos"), str):
43-
call += f" --qos='{job.resources.qos}'"
58+
qos = job.resources.get("qos")
59+
if qos is not None:
60+
call += f" --qos={safe_quote(qos)}"
4461

4562
if job.resources.get("mem_mb_per_cpu"):
4663
call += f" --mem-per-cpu {job.resources.mem_mb_per_cpu}"
@@ -77,6 +94,7 @@ def get_submit_command(job, params):
7794
# ensure that workdir is set correctly
7895
# use short argument as this is the same in all slurm versions
7996
# (see https://github.com/snakemake/snakemake/issues/2014)
80-
call += f" -D '{params.workdir}'"
97+
if params.workdir:
98+
call += f" -D {safe_quote(params.workdir)}"
8199

82100
return call

tests/tests.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -388,7 +388,7 @@ def test_constraint_resource(self, mock_job):
388388
process_mock.returncode = 0
389389
mock_popen.return_value = process_mock
390390

391-
assert " -C 'haswell'" in get_submit_command(job, params)
391+
assert " -C haswell" in get_submit_command(job, params)
392392

393393
def test_qos_resource(self, mock_job):
394394
"""Test that the qos resource is correctly added to the sbatch command."""
@@ -412,7 +412,7 @@ def test_qos_resource(self, mock_job):
412412
process_mock.returncode = 0
413413
mock_popen.return_value = process_mock
414414

415-
assert " --qos='normal'" in get_submit_command(job, params)
415+
assert " --qos=normal" in get_submit_command(job, params)
416416

417417
def test_both_constraint_and_qos(self, mock_job):
418418
"""Test that both constraint and qos resources can be used together."""
@@ -439,8 +439,8 @@ def test_both_constraint_and_qos(self, mock_job):
439439

440440
# Assert both resources are correctly included
441441
sbatch_command = get_submit_command(job, params)
442-
assert " --qos='high'" in sbatch_command
443-
assert " -C 'haswell'" in sbatch_command
442+
assert " --qos=high" in sbatch_command
443+
assert " -C haswell" in sbatch_command
444444

445445
def test_no_resources(self, mock_job):
446446
"""
@@ -517,7 +517,7 @@ def test_empty_qos(self, mock_job):
517517
process_mock.communicate.return_value = ("123", "")
518518
process_mock.returncode = 0
519519
mock_popen.return_value = process_mock
520-
# Assert the qoes is included (even if empty)
520+
# Assert the qos is included (even if empty)
521521
assert "--qos=''" in get_submit_command(job, params)
522522

523523
def test_taks(self, mock_job):

0 commit comments

Comments
 (0)