handle CP2K out-of-memory using ulimit

Author: svandenhaute

configs/hortense.yaml

@@ -30,7 +30,7 @@ ModelTraining:
 CP2K:
   cores_per_worker: 64
   max_evaluation_time: 30
-  launch_command: 'apptainer exec --memory 50000000000 -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -np 32 -bind-to core cp2k.psmp -i cp2k.inp'
+  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -np 32 -bind-to core cp2k.psmp'
   slurm:
     partition: "cpu_rome"
     account: "2024_079"

configs/lumi.yaml

@@ -6,7 +6,7 @@ default_threads: 8
 CP2K:
   cores_per_worker: 32
   max_evaluation_time: 20
-  launch_command: 'singularity exec --memory 50000000000 -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -np 32 cp2k.psmp -i cp2k.inp'
+  launch_command: 'singularity exec -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -np 32 cp2k.psmp'
   slurm:
     partition: "standard"
     account: "project_465001125"

configs/threadpool.yaml

@@ -13,17 +13,18 @@ ModelTraining:
 CP2K:
   cores_per_worker: 2
   max_evaluation_time: 0.3
-  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -bind-to core -np 2 -env OMP_NUM_THREADS 1 cp2k.psmp -i cp2k.inp'
+  memory_limit: 2GB
+  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -bind-to core -np 2 -env OMP_NUM_THREADS 1 cp2k.psmp'
 CP2K_container:
   cores_per_worker: 2
   max_evaluation_time: 0.3
-  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -bind-to core -np 2 -env OMP_NUM_THREADS 1 cp2k.psmp -i cp2k.inp'
+  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -bind-to core -np 2 -env OMP_NUM_THREADS 1 cp2k.psmp'
 GPAW:
   cores_per_worker: 2
   max_evaluation_time: 0.3
-  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/gpaw:24.1 /opt/entry.sh mpirun -np 2 gpaw python /opt/run_gpaw.py input.json'
+  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/gpaw:24.1 /opt/entry.sh mpirun -np 2 gpaw python /opt/run_gpaw.py'
 GPAW_container:
   cores_per_worker: 2
   max_evaluation_time: 0.3
-  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/gpaw:24.1 /opt/entry.sh mpirun -np 2 gpaw python /opt/run_gpaw.py input.json'
+  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/gpaw:24.1 /opt/entry.sh mpirun -np 2 gpaw python /opt/run_gpaw.py'
 ...

configs/wq.yaml

@@ -13,5 +13,5 @@ ModelTraining:
 CP2K:
   cores_per_worker: 2
   max_evaluation_time: 0.3
-  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2023.2 /opt/entry.sh mpirun -np 2 -x OMP_NUM_THREADS=1 cp2k.psmp -i cp2k.inp'
+  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2023.2 /opt/entry.sh mpirun -np 2 -x OMP_NUM_THREADS=1 cp2k.psmp'
 ...

psiflow/execution.py

@@ -1,5 +1,7 @@
 from __future__ import annotations  # necessary for type-guarding class methods
 
+import re
+
 import logging
 import math
 import shutil
@@ -301,18 +303,22 @@ def __init__(
         name: str,
         launch_command: Optional[str] = None,
         max_evaluation_time: Optional[float] = None,
+        memory_limit: Optional[str] = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
         self.name = name  # override name
+
+        if launch_command is None:
+            launch_command = self.default_launch_command
+        self.launch_command = launch_command
+
         if max_evaluation_time is None:
             max_evaluation_time = self.max_runtime / 60
         assert max_evaluation_time * 60 <= self.max_runtime
         self.max_evaluation_time = max_evaluation_time
 
-        if launch_command is None:
-            launch_command = self.default_launch_command
-        self.launch_command = launch_command
+        self.memory_limit = memory_limit
 
     @property
     def default_launch_command(self):
@@ -334,14 +340,32 @@ def default_launch_command(self):
             raise ValueError('provide path to ORCA executable via "launch_command"')
 
     def command(self):
-        max_time = 0.9 * (60 * self.max_evaluation_time)
-        command = " ".join(
-            [
-                "timeout -s 9 {}s".format(max_time),
-                self.launch_command,
-                "|| true",
-            ]
-        )
+        extra_commands = []
+
+        launch_command = self.launch_command
+        if self.name.startswith("CP2K"):
+            launch_command += " -i cp2k.inp"  # add input file
+        elif self.name.startswith("GPAW"):
+            launch_command += " input.json"
+        if self.max_evaluation_time is not None:
+            max_time = 0.9 * (60 * self.max_evaluation_time)
+            launch_command = "timeout -s 9 {}s {}".format(max_time, launch_command)
+        if self.memory_limit is not None:
+            # based on https://stackoverflow.com/a/42865957/2002471
+            units = {"B": 1, "KB": 2**10, "MB": 2**20, "GB": 2**30, "TB": 2**40}
+
+            def parse_size(size):
+                size = size.upper()
+                if not re.match(r' ', size):
+                    size = re.sub(r'([KMGT]?B)', r' \1', size)
+                number, unit = [string.strip() for string in size.split()]
+                return int(float(number) * units[unit])
+
+            actual = parse_size(self.memory_limit)
+            extra_commands.append("ulimit -v {}".format(actual))
+
+        body = "; ".join(extra_commands + [launch_command, "exit 0"]) + "; "
+        command = " { " + body + " } "
         return command
 
     def wq_resources(self):

psiflow/reference/_cp2k.py

@@ -260,7 +260,8 @@ def get_single_atom_references(self, element):
                     "ot": {"minimizer": "CG"}
                 }
             # necessary for oxygen calculation, at least in 2024.1
-            cp2k_input_dict['force_eval']['dft']['scf']['ignore_convergence_failure'] = True
+            key = 'ignore_convergence_failure'
+            cp2k_input_dict['force_eval']['dft']['scf'][key] = "TRUE"
 
             reference = CP2K(
                 dict_to_str(cp2k_input_dict),

pyproject.toml

@@ -16,7 +16,7 @@ dependencies = [
     "parsl @ git+https://github.com/parsl/parsl.git@benc-status-refactor",
     "prettytable",
     "psutil",
-    "cp2k-input-tools",
+    "cp2k-input-tools @ git+https://github.com/cp2k/cp2k-input-tools",  # need 2024.1
     "pytimeparse",
     ]