diff --git a/smartsim/error/__init__.py b/smartsim/error/__init__.py index 9e91f078b..383f82bb4 100644 --- a/smartsim/error/__init__.py +++ b/smartsim/error/__init__.py @@ -35,4 +35,6 @@ SSInternalError, SSUnsupportedError, UserStrategyError, + SSReservedKeywordError, + ) diff --git a/smartsim/error/errors.py b/smartsim/error/errors.py index 0acafc1e8..95bc92420 100644 --- a/smartsim/error/errors.py +++ b/smartsim/error/errors.py @@ -75,6 +75,10 @@ def create_message(file_path: str, read: bool) -> str: return msg +class SSReservedKeywordError(SmartSimError): + """Raised when a Reserved Keyword is used incorrectly""" + + # Internal Exceptions diff --git a/smartsim/wlm/slurm.py b/smartsim/wlm/slurm.py index 6ed37a2da..8fe12b3f9 100644 --- a/smartsim/wlm/slurm.py +++ b/smartsim/wlm/slurm.py @@ -32,7 +32,12 @@ from .._core.launcher.slurm.slurmParser import parse_salloc, parse_salloc_error from .._core.launcher.util.launcherUtil import ComputeNode, Partition from .._core.utils.helpers import init_default -from ..error import AllocationError, LauncherError, SmartSimError +from ..error import ( + AllocationError, + LauncherError, + SmartSimError, + SSReservedKeywordError, +) from ..log import get_logger logger = get_logger(__name__) @@ -241,24 +246,29 @@ def _get_alloc_cmd( "-J", "SmartSim", ] - # TODO check format here if time: salloc_args.extend(["-t", time]) if account: salloc_args.extend(["-A", str(account)]) + arguments = set(options.keys() if options is not None else {}) + invalid = {"t", "time", "N", "nodes", "A", "account"} + + if valid := arguments.intersection(invalid): + raise SSReservedKeywordError( + f"Expecting time, nodes, account as an argument. Also received: {valid}" + ) + for opt, val in (options or {}).items(): - if opt not in ["t", "time", "N", "nodes", "A", "account"]: - short_arg = bool(len(str(opt)) == 1) - prefix = "-" if short_arg else "--" - if not val: - salloc_args += [prefix + opt] + short_arg = bool(len(str(opt)) == 1) + prefix = "-" if short_arg else "--" + if not val: + salloc_args += [prefix + opt] + else: + if short_arg: + salloc_args += [prefix + opt, str(val)] else: - if short_arg: - salloc_args += [prefix + opt, str(val)] - else: - salloc_args += ["=".join((prefix + opt, str(val)))] - + salloc_args += ["=".join((prefix + opt, str(val)))] return salloc_args diff --git a/tests/full_wlm/test_slurm_allocation.py b/tests/full_wlm/test_slurm_allocation.py index 69554c1e4..01d40bf2f 100644 --- a/tests/full_wlm/test_slurm_allocation.py +++ b/tests/full_wlm/test_slurm_allocation.py @@ -28,7 +28,7 @@ import pytest -from smartsim.error import AllocationError +from smartsim.error import AllocationError, SSReservedKeywordError from smartsim.wlm import slurm # retrieved from pytest fixtures @@ -48,11 +48,31 @@ def test_get_release_allocation_w_options(wlmutils): """test slurm interface for obtaining allocations""" options = {"ntasks-per-node": 1} account = wlmutils.get_test_account() - alloc = slurm.get_allocation(nodes=1, time="00:05:00", options=options, account=account) + alloc = slurm.get_allocation( + nodes=1, time="00:05:00", options=options, account=account + ) time.sleep(5) # give slurm a rest slurm.release_allocation(alloc) +@pytest.mark.parametrize( + "func_parameters,test_parameters", + [ + ("t", "T"), + ("time", "TIME"), + ("nodes", "NODES"), + ("N", "N"), + ("account", "ACCOUNT"), + ("A", "A"), + ], +) +def test_get_allocation_bad_params(func_parameters, test_parameters): + """test get_allocation with reserved keywords as option""" + with pytest.raises(SSReservedKeywordError) as e: + alloc = slurm.get_allocation(options={func_parameters: test_parameters}) + assert "Expecting time, nodes, account as an argument" in e.value.args[0] + + # --------- Error handling ---------------------------- diff --git a/tests/test_slurm_get_alloc.py b/tests/test_slurm_get_alloc.py index 667ea2e03..de488b25f 100644 --- a/tests/test_slurm_get_alloc.py +++ b/tests/test_slurm_get_alloc.py @@ -47,24 +47,3 @@ def test_get_alloc_format(): "--ntasks-per-node=5", ] assert alloc_cmd == result - - -def test_get_alloc_format_overlap(): - """Test get alloc with collision between arguments and options""" - time = "10:00:00" - nodes = 5 - account = "A35311" - options = {"N": 10, "time": "15", "account": "S1242"} - alloc_cmd = _get_alloc_cmd(nodes, time, account, options) - result = [ - "--no-shell", - "-N", - "5", - "-J", - "SmartSim", - "-t", - "10:00:00", - "-A", - "A35311", - ] - assert result == alloc_cmd