diff --git a/backend/app/jobs/boltz_jobs.py b/backend/app/jobs/boltz_jobs.py index 7d3bc279..ce15eaa8 100644 --- a/backend/app/jobs/boltz_jobs.py +++ b/backend/app/jobs/boltz_jobs.py @@ -1,7 +1,9 @@ import io import logging +import re import string import subprocess +import time from pathlib import Path from tempfile import TemporaryDirectory @@ -21,6 +23,12 @@ ) from app.models import Fold, Invokation +# Maximum time to wait for MSA server jobs stuck in PENDING state before failing. +# The ColabFold MSA server can drop or deprioritize jobs, causing them to stay +# PENDING forever. Failing fast allows RQ retry to resubmit with a fresh job ID. +# See: https://github.com/sokrypton/ColabFold/issues/606 +MSA_PENDING_TIMEOUT_SECONDS = 10 * 60 # 10 minutes + def try_check_smiles_string_validity(smiles_string): """Try to check if a smiles string is valid.""" @@ -120,9 +128,36 @@ def run_boltz(fold_id, invokation_id): text=True, ) + # Track MSA PENDING state to detect stuck jobs + pending_start_time = None + pending_pattern = re.compile(r"PENDING.*\d+/\d+") + for line in iter(process.stdout.readline, ""): logging.info(line.strip()) + # Check if this line indicates MSA jobs stuck in PENDING + if pending_pattern.search(line): + if pending_start_time is None: + pending_start_time = time.time() + logging.info( + f"MSA server jobs entered PENDING state. " + f"Will timeout after {MSA_PENDING_TIMEOUT_SECONDS // 60} minutes." + ) + elif time.time() - pending_start_time > MSA_PENDING_TIMEOUT_SECONDS: + process.kill() + process.wait() + raise TimeoutError( + f"MSA server jobs stuck in PENDING state for over " + f"{MSA_PENDING_TIMEOUT_SECONDS // 60} minutes. " + f"The ColabFold server may have dropped this job. " + f"See: https://github.com/sokrypton/ColabFold/issues/606" + ) + elif "RUNNING" in line or "COMPLETE" in line: + # Reset timer if jobs start making progress + if pending_start_time is not None: + logging.info("MSA server jobs progressing, resetting PENDING timer.") + pending_start_time = None + process.stdout.close() process.wait() diff --git a/backend/app/util.py b/backend/app/util.py index c6468081..05a2c051 100644 --- a/backend/app/util.py +++ b/backend/app/util.py @@ -119,6 +119,7 @@ def start_stage(fold_id: int, stage: str, email_on_completion: bool) -> None: get_job_type_replacement(fold, "boltz"), job_timeout="12h", result_ttl=48 * 60 * 60, # 2 days + retry=Retry(max=3, interval=[60, 120, 300]), # Retry on MSA PENDING timeout **email_args, ) add_meta_to_job(boltz_job, fold, "boltz", fold.id) diff --git a/worker/Dockerfile.boltz b/worker/Dockerfile.boltz index 0546006f..3dd36b9f 100644 --- a/worker/Dockerfile.boltz +++ b/worker/Dockerfile.boltz @@ -44,8 +44,7 @@ RUN /scripts/conda_env_setup.sh RUN /opt/conda/bin/conda create -y -n boltzenv python=3.11 # RUN /opt/conda/envs/boltzenv/bin/python -m pip install boltz==1.0.0 -RUN /root/.local/bin/uv pip install --python /opt/conda/envs/boltzenv/bin/python boltz==2.1.1 -# RUN /opt/conda/envs/boltzenv/bin/python -m pip install boltz==2.2.0 +RUN /root/.local/bin/uv pip install --python /opt/conda/envs/boltzenv/bin/python boltz==2.2.1 # # Chat suggested this, for when on a weird network. # RUN echo "nameserver 8.8.8.8" > /etc/resolv.conf && \