Merge branch 'main' into feat/slurm_signals

Author: cmeesters

snakemake_executor_plugin_slurm/__init__.py

@@ -1320,7 +1320,7 @@ async def check_active_jobs(
         # slurmdbd/accounting hiccups where job states are briefly incomplete.
         for i in range(status_attempts):
             async with self.status_rate_limiter:
-                (status_of_jobs, sacct_query_duration) = await query_job_status(
+                status_of_jobs, sacct_query_duration = await query_job_status(
                     status_command, self.logger
                 )
                 if status_of_jobs is None and sacct_query_duration is None:
@@ -1616,8 +1616,10 @@ def get_account_arg(self, job: JobExecutorInterface):
         if job.resources.get("slurm_account"):
             # split the account upon ',' and whitespace, to allow
             # multiple accounts being given
+            # Note: YAML parsing may convert numeric strings (e.g. "123456") to int.
+            # Ensure we always work with strings for re.split and shlex.quote.
             accounts = [
-                a for a in re.split(r"[,\s]+", job.resources.slurm_account) if a
+                a for a in re.split(r"[,\s]+", str(job.resources.slurm_account)) if a
             ]
             for account in accounts:
                 # here, we check whether the given or guessed account is valid

snakemake_executor_plugin_slurm/utils.py

@@ -474,59 +474,49 @@ def set_gres_string(job: JobExecutorInterface) -> str:
     Function to set the gres string for the SLURM job
     based on the resources requested in the job.
     """
-    # generic resources (GRES) arguments can be of type
+    # generic resources (GRES) arguments can be of type "string",
     # "string:int" or "string:string:int" with optional postfix 'T' or 'G' or 'M'
-    gres_re = re.compile(r"^[a-zA-Z0-9_]+(:[a-zA-Z0-9_\.]+)?:\d+[TGM]?$")
+    gres_re = re.compile(r"^[a-zA-Z0-9_]+(:[a-zA-Z0-9_\.]+)?(:\d+[TGM]?)?$")
 
     # gpu model arguments can be of type "string"
     # The model string may contain a dot for variants, see
     # https://github.com/snakemake/snakemake-executor-plugin-slurm/issues/387
     gpu_model_re = re.compile(r"^[a-zA-Z0-9_\.]+$")
+
     # any arguments should not start and end with ticks or
     # quotation marks:
     string_check = re.compile(r"^[^'\"].*[^'\"]$")
+
     # The Snakemake resources can be only be of type "int",
     # hence no further regex is needed.
 
-    gpu_string = None
-    if job.resources.get("gpu"):
-        gpu_string = str(job.resources.get("gpu"))
-
-    gpu_model = None
-    if job.resources.get("gpu_model"):
-        gpu_model = job.resources.get("gpu_model")
-
-    # ensure that gres is not set, if gpu and gpu_model are set
-    if job.resources.get("gres") and gpu_string:
-        raise WorkflowError(
-            "GRES and GPU are set. Please only set one of them.", rule=job.rule
-        )
-    elif not job.resources.get("gres") and not gpu_model and not gpu_string:
-        return ""
-
+    # GRES
+    gres = ""
     if job.resources.get("gres"):
         # Validate GRES format (e.g., "gpu:1", "gpu:tesla:2")
-        gres = job.resources.gres
-        if not gres_re.match(gres):
-            if not string_check.match(gres):
+        if not gres_re.match(job.resources.gres):
+            if not string_check.match(job.resources.gres):
                 raise WorkflowError(
                     "GRES format should not be a nested string (start "
                     "and end with ticks or quotation marks). "
-                    "Expected format: "
+                    "Expected format: '<name>', "
                     "'<name>:<number>' or '<name>:<type>:<number>' with an optional "
                     "'T' 'M' or 'G' postfix "
-                    "(e.g., 'gpu:1' or 'gpu:tesla:2')"
+                    "(e.g., 'gpu', 'gpu:2' or 'gpu:tesla:1')"
                 )
             else:
                 raise WorkflowError(
-                    f"Invalid GRES format: {gres}. Expected format: "
-                    "'<name>:<number>' or '<name>:<type>:<number>' with an optional "
-                    "'T' 'M' or 'G' postfix "
-                    "(e.g., 'gpu:1' or 'gpu:tesla:2') "
+                    f"Invalid GRES format: {job.resources.gres}. Expected format: "
+                    "'<name>', '<name>:<number>' or '<name>:<type>:<number>' "
+                    "with an optional 'T' 'M' or 'G' postfix "
+                    "(e.g., 'gpu', 'gpu:1' or 'gpu:tesla:2') "
                 )
-        return f" --gres={job.resources.gres}"
+        gres += f" --gres={job.resources.gres}"
 
-    if gpu_model and gpu_string:
+    # GPU
+    gpu_string = job.resources.get("gpu")
+    gpu_model = job.resources.get("gpu_model")
+    if gpu_model:
         # validate GPU model format
         if not gpu_model_re.match(gpu_model):
             if not string_check.match(gpu_model):
@@ -540,10 +530,12 @@ def set_gres_string(job: JobExecutorInterface) -> str:
                     f"Invalid GPU model format: {gpu_model}."
                     " Expected format: '<name>' (e.g., 'tesla')"
                 )
-        return f" --gpus={gpu_model}:{gpu_string}"
-    elif gpu_model and not gpu_string:
-        raise WorkflowError("GPU model is set, but no GPU number is given")
+        # Default GPU to 1
+        gpu_string = job.resources.get("gpu", "1")
+        gres += f" --gpus={gpu_model}:{gpu_string}"
     elif gpu_string:
         # we assume here, that the validator ensures that the 'gpu_string'
         # is an integer
-        return f" --gpus={gpu_string}"
+        gres += f" --gpus={gpu_string}"
+
+    return gres

tests/test_account.py

@@ -0,0 +1,81 @@
+"""Unit tests for get_account_arg() method."""
+
+from unittest.mock import MagicMock, patch
+
+from snakemake_executor_plugin_slurm import Executor
+
+
+class _Resources(dict):
+    """Dict-like resources with attribute access for known keys only."""
+
+    def __getattr__(self, name):
+        try:
+            return self[name]
+        except KeyError as e:
+            raise AttributeError(name) from e
+
+
+def _make_mock_job(
+    rule_name="myrule", name=None, wildcards=None, jobid=1, is_group=False, **resources
+):
+    """Return a minimal mock job compatible with get_account_arg."""
+    mock_resources = _Resources(resources)
+
+    mock_rule = MagicMock()
+    mock_rule.name = rule_name
+
+    job = MagicMock()
+    job.resources = mock_resources
+    job.rule = mock_rule
+    job.name = name if name is not None else rule_name
+    job.wildcards = wildcards if wildcards is not None else {}
+    job.is_group.return_value = is_group
+    job.threads = resources.get("threads", 1)
+    job.jobid = jobid
+    return job
+
+
+class TestGetAccountArg:
+    """Tests for get_account_arg() method handling string and integer accounts."""
+
+    def _make_executor_stub(self):
+        """Return a minimal Executor stub."""
+        executor = Executor.__new__(Executor)
+        executor.logger = MagicMock()
+        executor._fallback_account_arg = None
+        return executor
+
+    # Patches prevent actual SLURM command execution (sacct/sacctmgr).
+    # - get_account() would call 'sacct' to guess account from recent jobs.
+    # - test_account() would call 'sacctmgr' or 'sshare' to validate account.
+    # Mocks allow testing without a live SLURM cluster.
+    @patch("snakemake_executor_plugin_slurm.get_account")
+    @patch("snakemake_executor_plugin_slurm.test_account")
+    def test_string_account(self, mock_test_account, mock_get_account, tmp_path):
+        """String account values work correctly."""
+        executor = self._make_executor_stub()
+        job = _make_mock_job(slurm_account="123456789")
+        result = next(executor.get_account_arg(job))
+        assert result == " -A 123456789"
+        mock_test_account.assert_called_with("123456789", executor.logger)
+
+    @patch("snakemake_executor_plugin_slurm.get_account")
+    @patch("snakemake_executor_plugin_slurm.test_account")
+    def test_integer_account(self, mock_test_account, mock_get_account, tmp_path):
+        """Integer account values (from YAML parsing) work correctly."""
+        executor = self._make_executor_stub()
+        job = _make_mock_job(slurm_account=123456789)
+        result = next(executor.get_account_arg(job))
+        assert result == " -A 123456789"
+        mock_test_account.assert_called_with("123456789", executor.logger)
+
+    @patch("snakemake_executor_plugin_slurm.get_account")
+    @patch("snakemake_executor_plugin_slurm.test_account")
+    def test_multiple_accounts_string(
+        self, mock_test_account, mock_get_account, tmp_path
+    ):
+        """Multiple string accounts work correctly."""
+        executor = self._make_executor_stub()
+        job = _make_mock_job(slurm_account="acc1,acc2 acc3")
+        results = list(executor.get_account_arg(job))
+        assert results == [" -A acc1", " -A acc2", " -A acc3"]

tests/tests.py

@@ -327,7 +327,7 @@ def test_valid_gres_with_model(self, mock_job):
 
         assert set_gres_string(job) == " --gres=gpu:tesla:2"
 
-    def test_invalid_gres_format(self, mock_job):
+    def test_valid_gres_format_gpu(self, mock_job):
         """Test with invalid GRES format."""
         job = mock_job(gres="gpu")
 
@@ -338,12 +338,12 @@ def test_invalid_gres_format(self, mock_job):
             process_mock.communicate.return_value = ("123", "")
             process_mock.returncode = 0
             mock_popen.return_value = process_mock
-        with pytest.raises(WorkflowError, match="Invalid GRES format"):
-            set_gres_string(job)
 
-    def test_invalid_gres_format_missing_count(self, mock_job):
+        assert set_gres_string(job) == " --gres=gpu"
+
+    def test_valid_gres_format_gpu_model(self, mock_job):
         """Test with invalid GRES format (missing count)."""
-        job = mock_job(gres="gpu:tesla:")
+        job = mock_job(gres="gpu:tesla")
 
         # Patch subprocess.Popen to capture the sbatch command
         with patch("subprocess.Popen") as mock_popen:
@@ -353,8 +353,7 @@ def test_invalid_gres_format_missing_count(self, mock_job):
             process_mock.returncode = 0
             mock_popen.return_value = process_mock
 
-        with pytest.raises(WorkflowError, match="Invalid GRES format"):
-            set_gres_string(job)
+        assert set_gres_string(job) == " --gres=gpu:tesla"
 
     def test_valid_gpu_number(self, mock_job):
         """Test with valid GPU number."""
@@ -424,11 +423,7 @@ def test_gpu_model_without_gpu(self, mock_job):
             process_mock.returncode = 0
             mock_popen.return_value = process_mock
 
-            # test whether the resource setting raises the correct error
-            with pytest.raises(
-                WorkflowError, match="GPU model is set, but no GPU number is given"
-            ):
-                set_gres_string(job)
+        assert set_gres_string(job) == " --gpus=tesla:1"
 
     def test_tmpspace_gres_10G(self, mock_job):
         """Test with valid GRES format (simple)."""
@@ -456,11 +451,7 @@ def test_both_gres_and_gpu_set(self, mock_job):
             process_mock.returncode = 0
             mock_popen.return_value = process_mock
 
-            # Ensure the error is raised when both GRES and GPU are set
-            with pytest.raises(
-                WorkflowError, match="GRES and GPU are set. Please only set one"
-            ):
-                set_gres_string(job)
+        assert set_gres_string(job) == " --gres=gpu:1 --gpus=2"
 
     def test_nested_string_raise(self, mock_job):
         """Test error case when gres is a nested string."""