retry job on conda create exception

latchbio · Oct 5, 2023 · 0e11a96 · 0e11a96
1 parent 4d97a9a
commit 0e11a96
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 9 deletions.
diff --git a/latch_cli/snakemake/serialize.py b/latch_cli/snakemake/serialize.py
@@ -321,12 +321,13 @@ def generate_snakemake_entrypoint(
         import shutil
         import subprocess
         from subprocess import CalledProcessError
+        import traceback
         from typing import NamedTuple
         import stat
         import sys
 
         from flytekit.extras.persistence import LatchPersistence
-        import traceback
+        from snakemake.exceptions import CreateCondaEnvironmentException
 
         from latch.resources.tasks import custom_task
         from latch.types.directory import LatchDir

diff --git a/latch_cli/snakemake/workflow.py b/latch_cli/snakemake/workflow.py
@@ -1269,14 +1269,25 @@ def get_fn_code(
 
                     print("\n\n\n")
                     try:
-                        subprocess.run(
-                            [sys.executable,{','.join(repr(x) for x in snakemake_args)}],
-                            check=True,
-                            env={{
-                                **os.environ,
-                                "LATCH_SNAKEMAKE_DATA": {repr(json.dumps(snakemake_data))}
-                            }}
-                        )
+                        conda_retries = 0
+                        while conda_retries < 3:
+                            try:
+                                subprocess.run(
+                                    [sys.executable,{','.join(repr(x) for x in snakemake_args)}],
+                                    check=True,
+                                    env={{
+                                        **os.environ,
+                                        "LATCH_SNAKEMAKE_DATA": {repr(json.dumps(snakemake_data))}
+                                    }}
+                                )
+                                break
+                            except CreateCondaEnvironmentException as e:
+                                if conda_retries < 3:
+                                    print("Retrying snakemake\n\n")
+                                    conda_retries +=1
+                                    continue
+                                else:
+                                    raise e
                     finally:
                         if tail is not None:
                             import signal