tweak cp2k setup

Dockerfile.cp2k

@@ -18,21 +18,20 @@ RUN git clone --recursive -b support/v2024.1 https://github.com/cp2k/cp2k.git /o
 # Build CP2K toolchain for target CPU generic
 WORKDIR /opt/cp2k/tools/toolchain
 RUN /bin/bash -c -o pipefail \
-    "./install_cp2k_toolchain.sh -j 12 \
+    "./install_cp2k_toolchain.sh -j 8 \
      --install-all \
      --enable-cuda=no \
      --target-cpu=generic \
      --with-cusolvermp=no \
      --with-gcc=system \
-     --with-mpich=system \
-     --with-libtorch=no"
+     --with-mpich=system"
 
 # Build CP2K for target CPU generic
 WORKDIR /opt/cp2k
 RUN /bin/bash -c -o pipefail \
     "cp ./tools/toolchain/install/arch/local.psmp ./arch/; \
      source ./tools/toolchain/install/setup; \
-     make -j 12 ARCH=local VERSION=psmp"
+     make -j 8 ARCH=local VERSION=psmp"
 
 # Collect components for installation and remove symbolic links
 RUN /bin/bash -c -o pipefail \
@@ -84,6 +83,7 @@ RUN printf "#!/bin/bash\n\
 ulimit -c 0 -s unlimited\n\
 \
 export OMP_STACKSIZE=16M\n\
+export OMP_NUM_THREADS=1\n\
 export PATH=/opt/cp2k/exe/local:\${PATH}\n\
 source /opt/cp2k/tools/toolchain/install/setup\n\
 \"\$@\"" \

build_containers.sh

@@ -12,6 +12,7 @@ psiflow=false
 gpaw=false
 cp2k=false
 build_sif=false
+#mpi=mpich
 
 # Parse command line options
 while [[ $# -gt 0 ]]; do

configs/hortense.yaml

@@ -28,14 +28,14 @@ ModelTraining:
     walltime: "12:00:00"
     scheduler_options: "#SBATCH --clusters=dodrio\n#SBATCH --gpus=1\n"
 CP2K:
-  cores_per_worker: 128
+  cores_per_worker: 64
   max_evaluation_time: 30
-  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2023.2 /opt/entry.sh mpirun -np 32 --bind-to core --map-by core -x 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 -np 32 -bind-to core cp2k.psmp -i cp2k.inp'
   slurm:
     partition: "cpu_rome"
     account: "2024_079"
     nodes_per_block: 1
-    cores_per_node: 32
+    cores_per_node: 64
     max_blocks: 25
     walltime: "06:00:00"
     scheduler_options: "#SBATCH --clusters=dodrio\n"

configs/lumi.yaml

@@ -6,7 +6,7 @@ default_threads: 8
 CP2K:
   cores_per_worker: 32
   max_evaluation_time: 20
-  launch_command: 'singularity exec -e --no-init oras://ghcr.io/molmod/cp2k:2023.2 /opt/entry.sh mpirun -np 32 -x OMP_NUM_THREADS=1 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 -i cp2k.inp'
   slurm:
     partition: "standard"
     account: "project_465001125"

examples/hortense.yaml

@@ -30,7 +30,7 @@ ModelTraining:
 CP2K:
   cores_per_worker: 64
   max_evaluation_time: 30
-  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2023.2 /opt/entry.sh mpirun -np 32 --bind-to core --map-by core -x 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 -np 32 -bind-to core cp2k.psmp -i cp2k.inp'
   slurm:
     partition: "cpu_rome"
     account: "2024_079"

examples/lumi.yaml

@@ -6,7 +6,7 @@ default_threads: 8
 CP2K:
   cores_per_worker: 32
   max_evaluation_time: 20
-  launch_command: 'singularity exec -e --no-init oras://ghcr.io/molmod/cp2k:2023.2 /opt/entry.sh mpirun -np 32 -x OMP_NUM_THREADS=1 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 -i cp2k.inp'
   slurm:
     partition: "standard"
     account: "project_465001125"

psiflow/models/_mace.py

@@ -208,6 +208,8 @@ def __init__(self, **config) -> None:
         config = MACEConfig(**config)  # validate input
         config.save_cpu = True  # assert model is saved to CPU after training
         config.device = "cpu"
+        if not config.swa:
+            config.start_swa = int(1e10)  # otherwise he fails to read
         self._config = asdict(config)
 
         self.model_future = None