Add optional fallback parameter to AsyncSlurmFoamCase.run

foamlib/_cases/_run.py

@@ -349,6 +349,7 @@ def _run_calls(
         cpus: Optional[int],
         check: bool,
         log: bool,
+        **kwargs: Any,
     ) -> Generator[Any, None, None]:
         if cmd is not None:
             if parallel:
@@ -381,6 +382,7 @@ def _run_calls(
                     env=self.__env(shell=isinstance(cmd, str)),
                     stdout=stdout,
                     stderr=stderr,
+                    **kwargs,
                 )
 
         else:
@@ -399,7 +401,7 @@ def _run_calls(
                     if cpus is None:
                         cpus = 1
 
-                yield self.run([script_path], parallel=False, cpus=cpus, check=check)
+                yield self.run([script_path], parallel=False, cpus=cpus, check=check, **kwargs)
 
             else:
                 if not self and (self.path / "0.orig").is_dir():
@@ -429,5 +431,5 @@ def _run_calls(
                         cpus = 1
 
                 yield self.run(
-                    [self.application], parallel=parallel, cpus=cpus, check=check
+                    [self.application], parallel=parallel, cpus=cpus, check=check, **kwargs
                 )

foamlib/_cases/_slurm.py

@@ -1,6 +1,7 @@
 import os
+import shutil
 import sys
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 if sys.version_info >= (3, 9):
     from collections.abc import Sequence
@@ -19,8 +20,13 @@ async def _run(
         cmd: Union[Sequence[Union[str, "os.PathLike[str]"]], str],
         *,
         cpus: int,
+        fallback: bool = False,
         **kwargs: Any,
     ) -> None:
+        if fallback and shutil.which("salloc") is None:
+            await AsyncFoamCase._run(cmd, cpus=cpus, **kwargs)
+            return
+
         if cpus >= 1:
             if isinstance(cmd, str):
                 cmd = ["/bin/sh", "-c", cmd]
@@ -31,3 +37,28 @@ async def _run(
             cmd = ["salloc", "-n", str(cpus), "--job-name", "foamlib", *cmd]
 
         await run_async(cmd, **kwargs)
+
+    async def run(
+        self,
+        cmd: Optional[Union[Sequence[Union[str, "os.PathLike[str]"]], str]] = None,
+        *,
+        parallel: Optional[bool] = None,
+        cpus: Optional[int] = None,
+        check: bool = True,
+        log: bool = True,
+        fallback: bool = False,
+    ) -> None:
+        """
+        Run this case, or a specified command in the context of this case.
+
+        :param cmd: The command to run. If None, run the case. If a sequence, the first element is the command and the rest are arguments. If a string, `cmd` is executed in a shell.
+        :param parallel: If True, run in parallel using MPI. If None, autodetect whether to run in parallel.
+        :param cpus: The number of CPUs to use. If None, autodetect according to the case. If 0, run locally.
+        :param check: If True, raise a CalledProcessError if any command returns a non-zero exit code.
+        :param log: If True, log the command output to a file.
+        :param fallback: If True, fall back to running the command locally if Slurm is not available.
+        """
+        for coro in self._run_calls(
+            cmd=cmd, parallel=parallel, cpus=cpus, check=check, log=log, fallback=fallback
+        ):
+            await coro

tests/test_cases/test_flange_async.py

@@ -22,9 +22,6 @@ async def flange(request: pytest.FixtureRequest) -> "AsyncGenerator[AsyncFoamCas
     case = request.param(path if path.exists() else of11_path)
     assert isinstance(case, AsyncFoamCase)
 
-    if isinstance(case, AsyncSlurmFoamCase) and shutil.which("salloc") is None:
-        pytest.skip("Slurm not available")
-
     async with case.clone() as clone:
         yield clone
 
@@ -39,11 +36,20 @@ async def test_run(flange: AsyncFoamCase, parallel: bool) -> None:
             assert d["method"] == "scotch"
             d["numberOfSubdomains"] = 2
 
-    await flange.run(parallel=parallel)
+    if isinstance(flange, AsyncSlurmFoamCase):
+        await flange.run(parallel=parallel, fallback=True)
+    else:
+        await flange.run(parallel=parallel)
+
     if parallel:
         await flange.reconstruct_par()
+
     await flange.clean()
-    await flange.run(parallel=parallel)
+
+    if isinstance(flange, AsyncSlurmFoamCase):
+        await flange.run(parallel=parallel, fallback=True)
+    else:
+        await flange.run(parallel=parallel)
 
 
 @pytest.mark.asyncio
@@ -64,9 +70,13 @@ async def test_run_cmd(flange: AsyncFoamCase) -> None:
             "-scale",
             "0.001",
         ],
+        cpus=0,
     )
-    await flange.run([flange.application])
-
+
+    if isinstance(flange, AsyncSlurmFoamCase):
+        await flange.run([flange.application], fallback=True)
+    else:
+        await flange.run([flange.application])
 
 @pytest.mark.asyncio
 async def test_run_cmd_shell(flange: AsyncFoamCase) -> None:
@@ -75,11 +85,16 @@ async def test_run_cmd_shell(flange: AsyncFoamCase) -> None:
 
     try:
         await flange.run(
-            'ansysToFoam "$FOAM_TUTORIALS/resources/geometry/flange.ans" -scale 0.001'
+            'ansysToFoam "$FOAM_TUTORIALS/resources/geometry/flange.ans" -scale 0.001',
+            cpus=0,
         )
     except CalledProcessError:
-        await flange.run('ansysToFoam "flange.ans" -scale 0.001')
-    await flange.run(flange.application)
+        await flange.run('ansysToFoam "flange.ans" -scale 0.001', cpus=0)
+
+    if isinstance(flange, AsyncSlurmFoamCase):
+        await flange.run(flange.application, fallback=True)
+    else: 
+        await flange.run(flange.application)
 
 
 def test_path(flange: AsyncFoamCase) -> None: