From 0e11a96987f6562f06f18baad2a94a8d1b58bf26 Mon Sep 17 00:00:00 2001
From: Kenny Workman <kennyworkman@sbcglobal.net>
Date: Wed, 4 Oct 2023 19:12:08 -0700
Subject: [PATCH] retry job on conda create exception

---
 latch_cli/snakemake/serialize.py |  3 ++-
 latch_cli/snakemake/workflow.py  | 27 +++++++++++++++++++--------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/latch_cli/snakemake/serialize.py b/latch_cli/snakemake/serialize.py
index 382bddcd..35a5261b 100644
--- 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
index 300419a3..17014be0 100644
--- 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