Allow variable memory for motioncorr via slurm

Dockerfiles/class_selection

@@ -1,4 +1,6 @@
 # This Dockerfile is used for Relion class selection
+# The second stage needs an image built with the common Relion Dockerfile
+
 # First stage creates a conda environment
 FROM rockylinux:8 AS conda-build
 
@@ -28,43 +30,21 @@ RUN tar -xzf /tmp/env.tar.gz -C /install/venv
 RUN eval "$(conda shell.bash hook)" && conda activate /install/venv && /install/venv/bin/conda-unpack
 
 
-# Second stage builds relion
-FROM rockylinux:8 AS relion-build
-
-# Copy python environment
-COPY --from=conda-build --chown="${userid}":"${groupid}" /install/venv /install/venv
-ENV PATH="/install/venv/bin:${PATH}"
-
-# Get required build packages and libraries
-RUN yum install gcc gcc-c++ cmake openmpi -y
-RUN yum install fftw-devel libtiff-devel libpng-devel libjpeg-devel zlib-devel -y
-
-# Build Relion - need to be on the ver5.0-mc-devolve tag
-RUN mkdir -p /install/relion5.0
-RUN curl -L -o relion.tar.gz https://github.com/d-j-hatton/relion/archive/ver5.0-mc-devolve.tar.gz
-RUN tar -xf relion.tar.gz -C /install
-RUN mkdir /install/relion-ver5.0-mc-devolve/build
-
-RUN mkdir /torch_home
-
-RUN cmake -DCMAKE_INSTALL_PREFIX=/install/relion5.0 -DCMAKE_C_COMPILER=/usr/bin/gcc -DCMAKE_CXX_COMPILER=/usr/bin/g++ -DMPI_C_COMPILER=/usr/lib64/openmpi/bin/mpicc -DMPI_CXX_COMPILER=/usr/lib64/openmpi/bin/mpicxx -DMPI_C_LIBRARIES=/usr/lib64/openmpi/lib/libmpi.so -DMPI_CXX_LIBRARIES=/usr/lib64/openmpi/lib/libmpi.so -DGUI=OFF -DALTCPU=ON -DPYTHON_EXE_PATH=/install/venv/bin/python -DTORCH_HOME_PATH=/torch_home -B/install/relion-ver5.0-mc-devolve/build -S/install/relion-ver5.0-mc-devolve
-RUN make --directory=/install/relion-ver5.0-mc-devolve/build/ install
-
-
-# Third stage combines the Relion and conda builds
+# Second stage combines the Relion and conda builds
 FROM rockylinux:8
 
 # Create EM user
 ARG groupid
 ARG userid
 ARG groupname
+ARG relionbuild
 RUN groupadd -r -g "${groupid}" "${groupname}" && useradd -r -M "${groupname}" -u "${userid}" -g "${groupname}"
 
 # Install libraries for Relion
 RUN yum install openmpi fftw-devel libtiff-devel libpng-devel libjpeg-devel zlib-devel -y
 
 # Copy Relion
-COPY --from=relion-build --chown="${userid}":"${groupid}" /install/relion5.0 /install/relion5.0
+COPY --from="${relionbuild}" --chown="${userid}":"${groupid}" /install/relion5.0 /install/relion5.0
 ENV PATH="/install/relion5.0/bin:${PATH}"
 
 # Copy python environment (from Relion build)

Dockerfiles/motioncor_relion

@@ -1,35 +1,19 @@
 # This Dockerfile is used for Relion's own motion correction and Relion postprocessing
-# First stage builds relion
-FROM rockylinux:8 AS relion-build
-
-# Get required build packages and libraries
-RUN yum install gcc gcc-c++ cmake openmpi -y
-RUN yum install fftw-devel libtiff-devel libpng-devel libjpeg-devel zlib-devel -y
-
-# Build Relion - need to be on the ver5.0-mc-devolve branch
-RUN mkdir -p /install/relion5.0
-RUN curl -L -o relion.tar.gz https://github.com/d-j-hatton/relion/archive/ver5.0-mc-devolve.tar.gz
-RUN tar -C /install -xf relion.tar.gz
-RUN mkdir /install/relion-ver5.0-mc-devolve/build
-
-RUN cmake -DCMAKE_INSTALL_PREFIX=/install/relion5.0 -DCMAKE_C_COMPILER=/usr/bin/gcc -DCMAKE_CXX_COMPILER=/usr/bin/g++ -DMPI_C_COMPILER=/usr/lib64/openmpi/bin/mpicc -DMPI_CXX_COMPILER=/usr/lib64/openmpi/bin/mpicxx -DMPI_C_LIBRARIES=/usr/lib64/openmpi/lib/libmpi.so -DMPI_CXX_LIBRARIES=/usr/lib64/openmpi/lib/libmpi.so -DGUI=OFF -DALTCPU=ON -DDoublePrec_CPU=OFF -DFORCE_OWN_FFTW=ON -DAMDFFTW=ON -B/install/relion-ver5.0-mc-devolve/build -S/install/relion-ver5.0-mc-devolve
-RUN make --directory=/install/relion-ver5.0-mc-devolve/build/ install
-
-
-# Second stage combines the Relion build with a python environment
+# Needs to be built on an image built with the common Relion Dockerfile
 FROM rockylinux:8
 
 # Create EM user
 ARG groupid
 ARG userid
 ARG groupname
+ARG relionbuild
 RUN groupadd -r -g "${groupid}" "${groupname}" && useradd -r -M "${groupname}" -u "${userid}" -g "${groupname}"
 
 # Install libraries for Relion
 RUN yum install openmpi fftw-devel libtiff-devel libpng-devel libjpeg-devel python3.11 zlib-devel -y
 
 # Copy Relion
-COPY --from=relion-build --chown="${userid}":"${groupid}" /install/relion5.0 /install/relion5.0
+COPY --from="${relionbuild}" --chown="${userid}":"${groupid}" /install/relion5.0 /install/relion5.0
 ENV PATH="/install/relion5.0/bin:${PATH}"
 
 # Make the python virtual environment and install cryoem-services

Dockerfiles/relion

@@ -0,0 +1,25 @@
+# Common Dockerfile to build relion
+# Used for motion correction and class selection
+FROM rockylinux:8
+
+# Get required build packages and libraries
+RUN yum install gcc gcc-c++ cmake openmpi -y
+RUN yum install fftw-devel libtiff-devel libpng-devel libjpeg-devel zlib-devel -y
+
+# Build Relion - need to be on the ver5.0-mc-devolve branch
+RUN mkdir -p /install/relion5.0
+RUN curl -L -o relion.tar.gz https://github.com/d-j-hatton/relion/archive/ver5.0-mc-devolve.tar.gz
+RUN tar -C /install -xf relion.tar.gz
+RUN mkdir /install/relion-ver5.0-mc-devolve/build
+
+RUN mkdir /torch_home
+
+RUN cmake -DCMAKE_INSTALL_PREFIX=/install/relion5.0 -DCMAKE_C_COMPILER=/usr/bin/gcc \
+    -DCMAKE_CXX_COMPILER=/usr/bin/g++ -DMPI_C_COMPILER=/usr/lib64/openmpi/bin/mpicc \
+    -DMPI_CXX_COMPILER=/usr/lib64/openmpi/bin/mpicxx \
+    -DMPI_C_LIBRARIES=/usr/lib64/openmpi/lib/libmpi.so \
+    -DMPI_CXX_LIBRARIES=/usr/lib64/openmpi/lib/libmpi.so \
+    -DGUI=OFF -DALTCPU=ON -DDoublePrec_CPU=OFF -DFORCE_OWN_FFTW=ON -DAMDFFTW=ON \
+    -DPYTHON_EXE_PATH=/install/venv/bin/python -DTORCH_HOME_PATH=/torch_home \
+    -B/install/relion-ver5.0-mc-devolve/build -S/install/relion-ver5.0-mc-devolve
+RUN make --directory=/install/relion-ver5.0-mc-devolve/build/ install

src/cryoemservices/services/motioncorr.py

@@ -62,6 +62,7 @@ class MotionCorrParameters(BaseModel):
     split_sum: Optional[int] = None
     dose_motionstats_cutoff: float = 4.0
     do_icebreaker_jobs: bool = True
+    slurm_memory: int = 20000
     mc_uuid: int
     app_id: Optional[int] = None
     picker_uuid: int
@@ -173,7 +174,7 @@ def motioncor2(self, command: List[str], mrc_out: Path):
         self.parse_mc2_stdout(result.stdout.decode("utf8", "replace"))
         return result
 
-    def motioncor2_slurm(self, command: List[str], mrc_out: Path):
+    def motioncor2_slurm(self, command: List[str], mrc_out: Path, memory: int):
         """Submit MotionCor2 jobs to a slurm cluster via the RestAPI"""
         slurm_outcome = slurm_submission_for_services(
             log=self.log,
@@ -187,6 +188,7 @@ def motioncor2_slurm(self, command: List[str], mrc_out: Path):
             use_gpu=True,
             use_singularity=True,
             cif_name=os.environ["MOTIONCOR2_SIF"],
+            memory_request=memory,
             extra_singularity_directories=["/lib64"],
         )
 
@@ -218,7 +220,7 @@ def relion_motioncorr(self, command: List[str], mrc_out: Path):
             result.returncode = 1
         return result
 
-    def relion_motioncorr_slurm(self, command: List[str], mrc_out: Path):
+    def relion_motioncorr_slurm(self, command: List[str], mrc_out: Path, memory: int):
         """Submit Relion's own motion correction to a slurm cluster via the RestAPI"""
         result = slurm_submission_for_services(
             log=self.log,
@@ -231,6 +233,7 @@ def relion_motioncorr_slurm(self, command: List[str], mrc_out: Path):
             cpus=4,
             use_gpu=False,
             use_singularity=False,
+            memory_request=memory,
             script_extras="module load EM/relion/motioncorr",
         )
         if Path(mrc_out).with_suffix(".star").exists():
@@ -400,7 +403,9 @@ def motion_correction(self, rw, header: dict, message: dict):
                         command.extend((mc2_flags[k], str(v)))
             # Run MotionCor2
             if mc_params.submit_to_slurm:
-                result = self.motioncor2_slurm(command, Path(mc_params.mrc_out))
+                result = self.motioncor2_slurm(
+                    command, Path(mc_params.mrc_out), mc_params.slurm_memory
+                )
             else:
                 result = self.motioncor2(command, Path(mc_params.mrc_out))
 
@@ -451,7 +456,9 @@ def motion_correction(self, rw, header: dict, message: dict):
             command.extend(("--dose_weighting", "--i", "dummy"))
             # Run Relion motion correction
             if mc_params.submit_to_slurm:
-                result = self.relion_motioncorr_slurm(command, Path(mc_params.mrc_out))
+                result = self.relion_motioncorr_slurm(
+                    command, Path(mc_params.mrc_out), mc_params.slurm_memory
+                )
             else:
                 result = self.relion_motioncorr(command, Path(mc_params.mrc_out))
 

tests/services/test_motioncorr_service.py

@@ -997,6 +997,7 @@ def test_motioncor2_slurm_service_spa(mock_requests, offline_transport, tmp_path
         "patch_sizes": {"x": 5, "y": 5},
         "mc_uuid": 0,
         "picker_uuid": 0,
+        "slurm_memory": 12000,
         "relion_options": {},
     }
     output_relion_options = dict(RelionServiceOptions())
@@ -1235,7 +1236,7 @@ def test_motioncor_superres_does_slurm(mock_requests, offline_transport, tmp_pat
                 "partition": "partition",
                 "prefer": "preference",
                 "tasks": 1,
-                "memory_per_node": {"number": 12000, "set": True, "infinite": False},
+                "memory_per_node": {"number": 20000, "set": True, "infinite": False},
                 "time_limit": {"number": 60, "set": True, "infinite": False},
                 "tres_per_job": "gres/gpu:1",
             },
@@ -1353,6 +1354,7 @@ def test_motioncor2_slurm_parameters(mock_slurm, offline_transport, tmp_path):
         use_gpu=True,
         use_singularity=True,
         cif_name="MotionCor2_SIF",
+        memory_request=20000,
         extra_singularity_directories=["/lib64"],
     )
 
@@ -1389,6 +1391,7 @@ def test_motioncor_relion_slurm_parameters(mock_slurm, offline_transport, tmp_pa
         "patch_sizes": {"x": 5, "y": 5},
         "mc_uuid": 0,
         "picker_uuid": 0,
+        "slurm_memory": 12000,
         "relion_options": {},
     }
 
@@ -1456,6 +1459,7 @@ def test_motioncor_relion_slurm_parameters(mock_slurm, offline_transport, tmp_pa
         cpus=4,
         use_gpu=False,
         use_singularity=False,
+        memory_request=12000,
         script_extras="module load EM/relion/motioncorr",
     )