Raising error for reserved keywords under function parameter options in get_allocation

smartsim/error/__init__.py

@@ -35,4 +35,6 @@
     SSInternalError,
     SSUnsupportedError,
     UserStrategyError,
+    SSReservedKeywordError,
+
 )

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
 
 

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
 
 

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 ----------------------------
 
 

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