Harden retry loops in run_mmseqs2

colabfold/colabfold.py

@@ -89,9 +89,6 @@ def submit(seqs, mode, N=101):
         # https://requests.readthedocs.io/en/latest/user/advanced/#advanced
         # "good practice to set connect timeouts to slightly larger than a multiple of 3"
         res = requests.post(f'{host_url}/{submission_endpoint}', data={ 'q': query, 'mode': mode }, timeout=6.02, headers=headers)
-      except requests.exceptions.Timeout:
-        logger.warning("Timeout while submitting to MSA server. Retrying...")
-        continue
       except Exception as e:
         error_count += 1
         logger.warning(f"Error while fetching result from MSA server. Retrying... ({error_count}/5)")
@@ -110,19 +107,16 @@ def submit(seqs, mode, N=101):
     return out
 
   def status(ID):
+    error_count = 0
     while True:
-      error_count = 0
       try:
         res = requests.get(f'{host_url}/ticket/{ID}', timeout=6.02, headers=headers)
-      except requests.exceptions.Timeout:
-        logger.warning("Timeout while fetching status from MSA server. Retrying...")
-        continue
       except Exception as e:
         error_count += 1
         logger.warning(f"Error while fetching result from MSA server. Retrying... ({error_count}/5)")
         logger.warning(f"Error: {e}")
         time.sleep(5)
-        if error_count > 5:
+        if error_count >= 5:
           raise
         continue
       break
@@ -138,15 +132,12 @@ def download(ID, path):
     while True:
       try:
         res = requests.get(f'{host_url}/result/download/{ID}', timeout=6.02, headers=headers)
-      except requests.exceptions.Timeout:
-        logger.warning("Timeout while fetching result from MSA server. Retrying...")
-        continue
       except Exception as e:
         error_count += 1
         logger.warning(f"Error while fetching result from MSA server. Retrying... ({error_count}/5)")
         logger.warning(f"Error: {e}")
         time.sleep(5)
-        if error_count > 5:
+        if error_count >= 5:
           raise
         continue
       break
@@ -272,21 +263,18 @@ def download(ID, path):
         os.mkdir(TMPL_PATH)
         TMPL_LINE = ",".join(TMPL[:20])
         response = None
+        error_count = 0
         while True:
-          error_count = 0
           try:
             # https://requests.readthedocs.io/en/latest/user/advanced/#advanced
             # "good practice to set connect timeouts to slightly larger than a multiple of 3"
             response = requests.get(f"{host_url}/template/{TMPL_LINE}", stream=True, timeout=6.02, headers=headers)
-          except requests.exceptions.Timeout:
-            logger.warning("Timeout while submitting to template server. Retrying...")
-            continue
           except Exception as e:
             error_count += 1
             logger.warning(f"Error while fetching result from template server. Retrying... ({error_count}/5)")
             logger.warning(f"Error: {e}")
             time.sleep(5)
-            if error_count > 5:
+            if error_count >= 5:
               raise
             continue
           break

tests/test_mmseqs_api.py

@@ -0,0 +1,155 @@
+"""Regression tests for the run_mmseqs2 retry loops.
+
+Each of the four retry loops in run_mmseqs2 (submit, status, download, templates)
+caps at a fixed number of attempts. status() and the templates loop used to
+reset their counter inside the loop, which silently turned them into
+unbounded retries; the Timeout branch used to retry without sleeping or
+counting, also unbounded. These tests pin the bounded behavior across both
+failure modes.
+"""
+
+import pytest
+import requests
+
+from colabfold.colabfold import run_mmseqs2
+
+# Safety bound so a regression to unbounded retries fails fast instead of
+# spinning forever (time.sleep is monkeypatched to a no-op below).
+MAX_RETRY_CALLS = 50
+
+FAILURE_MODES = [
+    pytest.param(requests.exceptions.ConnectionError, id="connection-error"),
+    pytest.param(requests.exceptions.Timeout, id="timeout"),
+]
+
+
+class JsonResponse:
+    def __init__(self, payload):
+        self.payload = payload
+        self.text = str(payload)
+        self.content = b""
+
+    def json(self):
+        return self.payload
+
+
+def _failing_get(calls, exc_cls, message):
+    def get(*args, **kwargs):
+        calls["get"] += 1
+        if calls["get"] > MAX_RETRY_CALLS:
+            raise AssertionError(
+                f"retry loop exceeded {MAX_RETRY_CALLS} calls — "
+                "the bound on retries has regressed"
+            )
+        raise exc_cls(message)
+
+    return get
+
+
+@pytest.mark.parametrize("exc_cls", FAILURE_MODES)
+def test_submit_retries_are_bounded(monkeypatch, tmp_path, exc_cls):
+    calls = {"post": 0}
+
+    def post(*args, **kwargs):
+        calls["post"] += 1
+        if calls["post"] > MAX_RETRY_CALLS:
+            raise AssertionError(
+                f"retry loop exceeded {MAX_RETRY_CALLS} calls — "
+                "the bound on retries has regressed"
+            )
+        raise exc_cls("submit failed")
+
+    monkeypatch.setattr("colabfold.colabfold.requests.post", post)
+    monkeypatch.setattr("colabfold.colabfold.time.sleep", lambda _s: None)
+
+    with pytest.raises(exc_cls, match="submit failed"):
+        run_mmseqs2(
+            "ACDE",
+            str(tmp_path / "msa"),
+            use_env=False,
+            user_agent="colabfold/test",
+        )
+
+    assert calls["post"] == 5
+
+
+@pytest.mark.parametrize("exc_cls", FAILURE_MODES)
+def test_status_retries_are_bounded(monkeypatch, tmp_path, exc_cls):
+    calls = {"get": 0}
+
+    monkeypatch.setattr(
+        "colabfold.colabfold.requests.post",
+        lambda *a, **kw: JsonResponse({"status": "RUNNING", "id": "job-id"}),
+    )
+    monkeypatch.setattr(
+        "colabfold.colabfold.requests.get",
+        _failing_get(calls, exc_cls, "status failed"),
+    )
+    monkeypatch.setattr("colabfold.colabfold.time.sleep", lambda _s: None)
+
+    with pytest.raises(exc_cls, match="status failed"):
+        run_mmseqs2(
+            "ACDE",
+            str(tmp_path / "msa"),
+            use_env=False,
+            user_agent="colabfold/test",
+        )
+
+    assert calls["get"] == 5
+
+
+@pytest.mark.parametrize("exc_cls", FAILURE_MODES)
+def test_download_retries_are_bounded(monkeypatch, tmp_path, exc_cls):
+    calls = {"get": 0}
+
+    monkeypatch.setattr(
+        "colabfold.colabfold.requests.post",
+        lambda *a, **kw: JsonResponse({"status": "COMPLETE", "id": "job-id"}),
+    )
+    monkeypatch.setattr(
+        "colabfold.colabfold.requests.get",
+        _failing_get(calls, exc_cls, "download failed"),
+    )
+    monkeypatch.setattr("colabfold.colabfold.time.sleep", lambda _s: None)
+
+    with pytest.raises(exc_cls, match="download failed"):
+        run_mmseqs2(
+            "ACDE",
+            str(tmp_path / "msa"),
+            use_env=False,
+            user_agent="colabfold/test",
+        )
+
+    assert calls["get"] == 5
+
+
+@pytest.mark.parametrize("exc_cls", FAILURE_MODES)
+def test_templates_retries_are_bounded(monkeypatch, tmp_path, exc_cls):
+    # Pre-create the MSA-stage artifacts so run_mmseqs2 skips submit/status/
+    # download (tar_gz_file exists) and a3m extraction (uniref.a3m exists),
+    # going straight to the templates fetch we want to exercise.
+    msa_dir = tmp_path / "msa_all"
+    msa_dir.mkdir()
+    (msa_dir / "out.tar.gz").write_bytes(b"")
+    (msa_dir / "uniref.a3m").write_text(">seq1\nACDE\n")
+    (msa_dir / "pdb70.m8").write_text(
+        "101\t1abc_A\t50.0\t100\t10\t1\t1\t100\t1\t100\t1e-10\t100.0\n"
+    )
+
+    calls = {"get": 0}
+    monkeypatch.setattr(
+        "colabfold.colabfold.requests.get",
+        _failing_get(calls, exc_cls, "template failed"),
+    )
+    monkeypatch.setattr("colabfold.colabfold.time.sleep", lambda _s: None)
+
+    with pytest.raises(exc_cls, match="template failed"):
+        run_mmseqs2(
+            "ACDE",
+            str(tmp_path / "msa"),
+            use_env=False,
+            use_templates=True,
+            user_agent="colabfold/test",
+        )
+
+    assert calls["get"] == 5