From 208aa4b39e3b2cc734d77da3116907f542585581 Mon Sep 17 00:00:00 2001
From: Rahul Desai <rahul@latch.bio>
Date: Fri, 19 Jan 2024 10:26:38 -0800
Subject: [PATCH] add toggle for using conda and container envs

---
 latch/types/metadata.py         | 40 +++++++++++++++++++++++++++++++++
 latch_cli/snakemake/workflow.py | 14 ++++++++----
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/latch/types/metadata.py b/latch/types/metadata.py
index d49bbbf2..372fe47d 100644
--- a/latch/types/metadata.py
+++ b/latch/types/metadata.py
@@ -523,16 +523,56 @@ def _parameter_str(t: Tuple[str, LatchParameter]):
 
 @dataclass
 class DockerMetadata:
+    """Class describing credentials for private docker repositories"""
+
     username: str
+    """
+    The account username for the private repository
+    """
     secret_name: str
+    """
+    The name of the Latch Secret that contains the password for the private repository
+    """
+
+
+@dataclass
+class EnvironmentConfig:
+    """Class describing environment for spawning Snakemake tasks"""
+
+    use_conda: bool = True
+    """
+    Use Snakemake `conda` directive to spawn tasks in conda environments
+    """
+    use_containers: bool = True
+    """
+    Use Snakemake `container` directive to spawn tasks in Docker containers
+    """
 
 
 @dataclass
 class SnakemakeMetadata(LatchMetadata):
+    """Class for organizing Snakemake workflow metadata"""
+
     output_dir: Optional[LatchDir] = None
+    """
+    Directory for snakemake workflow outputs
+    """
     name: Optional[str] = None
+    """
+    Name of the workflow
+    """
     docker_metadata: Optional[DockerMetadata] = None
+    """
+    Credentials configuration for private docker repositories
+    """
+    env_config: Optional[EnvironmentConfig] = None
+    """
+    Environment configuration for spawning Snakemake tasks
+    """
     parameters: Dict[str, SnakemakeParameter] = field(default_factory=dict)
+    """
+    A dictionary mapping parameter names (strings) to `SnakemakeParameter` objects
+    """
 
     def __post_init__(self):
         if self.name is None:
diff --git a/latch_cli/snakemake/workflow.py b/latch_cli/snakemake/workflow.py
index 7b0ad50b..2b9daf02 100644
--- a/latch_cli/snakemake/workflow.py
+++ b/latch_cli/snakemake/workflow.py
@@ -668,6 +668,14 @@ def __init__(
         self._dag = dag
         self._cache_tasks = cache_tasks
         self._docker_metadata = metadata._snakemake_metadata.docker_metadata
+        self._use_conda = (
+            metadata._snakemake_metadata.env_config is not None
+            and metadata._snakemake_metadata.env_config.use_conda
+        )
+        self._use_containers = (
+            metadata._snakemake_metadata.env_config is not None
+            and metadata._snakemake_metadata.env_config.use_containers
+        )
         self.snakemake_tasks: List[SnakemakeJobTask] = []
 
         workflow_metadata = WorkflowMetadata(
@@ -1311,8 +1319,6 @@ def get_fn_code(
         if isinstance(self.job, GroupJob):
             jobs = self.job.jobs
 
-        need_conda = any(x.conda_env is not None for x in jobs)
-
         if non_blob_parameters is not None:
             for param, val in non_blob_parameters.items():
                 self.job.rule.workflow.globals["config"][param] = val
@@ -1322,8 +1328,8 @@ def get_fn_code(
             "latch_cli.snakemake.single_task_snakemake",
             "-s",
             snakefile_path_in_container,
-            *(["--use-conda"] if need_conda else []),
-            "--use-singularity",
+            *(["--use-conda"] if self.wf._use_conda else []),
+            *(["--use-singularity"] if self.wf._use_containers else []),
             "--target-jobs",
             *jobs_cli_args(jobs),
             "--allowed-rules",