Fix some slurm issues

engibench/problems/airfoil/v0.py

@@ -396,9 +396,6 @@ def simulate(self, design: DesignType, config: dict[str, Any] | None = None, mpi
         if isinstance(design["angle_of_attack"], np.ndarray):
             design["angle_of_attack"] = design["angle_of_attack"][0]
 
-        # docker pull image if not already pulled
-        if container.RUNTIME is not None and self.container_id is not None:
-            container.pull(self.container_id)
         # pre-process the design and run the simulation
 
         # Prepares the airfoil_analysis.py script with the simulation configuration
@@ -465,9 +462,6 @@ def optimize(
         if isinstance(starting_point["angle_of_attack"], np.ndarray):
             starting_point["angle_of_attack"] = starting_point["angle_of_attack"][0]
 
-        # docker pull image if not already pulled
-        if container.RUNTIME is not None and self.container_id is not None:
-            container.pull(self.container_id)
         # pre-process the design and run the simulation
         filename = "candidate_design"
 

engibench/problems/heatconduction2d/v0.py

@@ -255,8 +255,8 @@ def random_design(self, dataset_split: str = "train", design_key: str = "optimal
                 np.ndarray: The valid random design.
                 int: The random index selected.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split][design_key]))
-        return np.array(self.dataset[dataset_split][design_key][rnd]), int(rnd)
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split][design_key]), dtype=int)
+        return np.array(self.dataset[dataset_split][design_key][rnd]), rnd
 
     def render(self, design: npt.NDArray, *, open_window: bool = False) -> Any:
         """Renders the design in a human-readable format.

engibench/utils/container.py

@@ -2,6 +2,7 @@
 
 from collections.abc import Sequence
 import os
+import shutil
 import subprocess
 import tempfile
 
@@ -46,7 +47,10 @@ def run(  # noqa: PLR0913
         result = RUNTIME.run(command, image, mounts, env, name, sync_uid=sync_uid)
         result.check_returncode()
     except subprocess.CalledProcessError as e:
-        msg = f"Container command failed with exit code {e.returncode}:\nCommand: {' '.join(command)}\nOutput: {e.output.decode() if e.output else 'No output'}"
+        msg = f"""Container command failed with exit code {e.returncode}:
+Command: {" ".join(command)}
+stdout: {result.stdout.decode() if result.stdout else "No output"}
+stderr: {result.stderr.decode() if result.stderr else "No output"}"""
         raise RuntimeError(msg) from e
 
 
@@ -197,12 +201,18 @@ def run(  # noqa: PLR0913
                 *command,
             ],
             check=False,
+            capture_output=True,
+            env=cls._env(),
         )
 
     @classmethod
     def _user_args(cls) -> tuple[str, ...]:
         return ("--user", str(os.getuid()))
 
+    @classmethod
+    def _env(cls) -> dict[str, str] | None:
+        return None
+
 
 class Podman(Docker):
     """Podman 🦭 runtime."""
@@ -234,6 +244,16 @@ def is_available(cls) -> bool:
     def _user_args(cls) -> tuple[str, ...]:
         return ("--userns=keep-id", "--user", str(os.getuid()))
 
+    @classmethod
+    def _env(cls) -> dict[str, str] | None:
+        # podman needs to have pasta in the PATH variable to configure
+        # network namespaces
+        pasta_executable = shutil.which("pasta")
+        if pasta_executable is None:
+            msg = "pasta executable not available. This is needed for podman to work properly"
+            raise RuntimeError(msg)
+        return {"PATH": os.path.dirname(pasta_executable)}
+
 
 DOCKER_PREFIX = "docker://"
 
@@ -317,6 +337,7 @@ def run(  # noqa: PLR0913
 
         # Get sif filename
         sif_image = cls.sif_filename(image)
+        cls.pull(image)
 
         return subprocess.run(
             [