From 735ab2301d0905ea054ad6efa3150acb296d0e78 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Johannes=20K=C3=B6ster?= <johannes.koester@uni-due.de>
Date: Tue, 25 Jan 2022 11:19:07 +0100
Subject: [PATCH] feat: specify conda environments via their name (#1340)

* feat: specify conda environments via their name (preparation)

* Implement named env support

* docs

* fix

* fmt

* fix code smells

* Mention wildcard expansion.

* add assertion

* fix list envs

* fix attribute

* fix wildcard test

* dbg

* cleanup

* fix test case
---
 docs/snakefiles/deployment.rst                |  33 +++
 snakemake/dag.py                              |  19 +-
 snakemake/deployment/conda.py                 | 203 +++++++++++++++---
 snakemake/deployment/containerize.py          |   4 +-
 snakemake/executors/__init__.py               |   4 +-
 snakemake/jobs.py                             |  50 +++--
 snakemake/rules.py                            |  11 +-
 snakemake/workflow.py                         |  11 +-
 tests/test_conda_named/Snakefile              |  18 ++
 .../expected-results/test.out                 |   1 +
 tests/tests.py                                |   5 +
 11 files changed, 289 insertions(+), 70 deletions(-)
 create mode 100644 tests/test_conda_named/Snakefile
 create mode 100644 tests/test_conda_named/expected-results/test.out

diff --git a/docs/snakefiles/deployment.rst b/docs/snakefiles/deployment.rst
index 1ff1f2394..f6856ce45 100644
--- a/docs/snakefiles/deployment.rst
+++ b/docs/snakefiles/deployment.rst
@@ -246,6 +246,8 @@ with the following `environment definition <https://conda.io/projects/conda/en/l
 
 The path to the environment definition is interpreted as **relative to the Snakefile that contains the rule** (unless it is an absolute path, which is discouraged).
 
+Instead of using a concrete path, it is also possible to provide a path containing wildcards (which must also occur in the output files of the rule), analogous to the specification of input files.
+
 .. sidebar:: Note
 
    Note that conda environments are only used with ``shell``, ``script`` and the ``wrapper`` directive, not the ``run`` directive.
@@ -260,6 +262,37 @@ Note that you need to clean up environments manually for now. However, in many c
 
 Conda deployment also works well for offline or air-gapped environments. Running ``snakemake --use-conda --conda-create-envs-only`` will only install the required conda environments without running the full workflow. Subsequent runs with ``--use-conda`` will make use of the local environments without requiring internet access.
 
+
+.. _conda_named_env:
+
+Using already existing named conda environments
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Sometimes it can be handy to refer to an already existing named conda environment from a rule, instead of defining a new one from scratch.
+Importantly, one should be aware that this can **hamper reproducibility**, because the workflow then relies on this environment to be present
+**in exactly the same way** on any new system where the workflow is executed. Essentially, you will have to take care of this manually in such a case.
+Therefore, the approach using environment definition files described above is highly recommended and preferred.
+
+Nevertheless, in case you are still sure that you want to use an existing named environment, it can simply be put into the conda directive, e.g.
+
+.. code-block:: python
+    rule NAME:
+        input:
+            "table.txt"
+        output:
+            "plots/myplot.pdf"
+        conda:
+            "some-env-name"
+        script:
+            "scripts/plot-stuff.R"
+
+For such a rule, Snakemake will just activate the given environment, instead of automatically deploying anything.
+Instead of using a concrete name, it is also possible to provide a name containing wildcards (which must also occur in the output files of the rule), analogous to the specification of input files.
+
+Note that Snakemake distinguishes file based environments from named ones as follows: 
+if the given specification ends on ``.yaml`` or ``.yml``, Snakemake assumes it to be a path to an environment definition file; otherwise, it assumes the given specification
+to be the name of an existing environment.
+
 .. _singularity:
 
 
diff --git a/snakemake/dag.py b/snakemake/dag.py
index 3770e54f9..a50efc41c 100755
--- a/snakemake/dag.py
+++ b/snakemake/dag.py
@@ -275,37 +275,36 @@ def cleanup(self):
     def create_conda_envs(
         self, dryrun=False, forceall=False, init_only=False, quiet=False
     ):
-        # First deduplicate based on job.conda_env_file
+        # First deduplicate based on job.conda_env_spec
         jobs = self.jobs if forceall else self.needrun_jobs
         env_set = {
-            (job.conda_env_file, job.container_img_url)
+            (job.conda_env_spec, job.container_img_url)
             for job in jobs
-            if job.conda_env_file
+            if job.conda_env_spec
         }
         # Then based on md5sum values
         self.conda_envs = dict()
-        for (env_file, simg_url) in env_set:
+        for (env_spec, simg_url) in env_set:
             simg = None
             if simg_url and self.workflow.use_singularity:
                 assert (
                     simg_url in self.container_imgs
                 ), "bug: must first pull singularity images"
                 simg = self.container_imgs[simg_url]
-            env = conda.Env(
-                env_file,
+            env = env_spec.get_conda_env(
                 self.workflow,
                 container_img=simg,
                 cleanup=self.workflow.conda_cleanup_pkgs,
             )
-            self.conda_envs[(env_file, simg_url)] = env
+            self.conda_envs[(env_spec, simg_url)] = env
 
         if not init_only:
             for env in self.conda_envs.values():
-                if not dryrun or not quiet:
+                if (not dryrun or not quiet) and not env.is_named:
                     env.create(dryrun)
 
     def pull_container_imgs(self, dryrun=False, forceall=False, quiet=False):
-        # First deduplicate based on job.conda_env_file
+        # First deduplicate based on job.conda_env_spec
         jobs = self.jobs if forceall else self.needrun_jobs
         img_set = {
             (job.container_img_url, job.is_containerized)
@@ -2075,7 +2074,7 @@ def add(path):
                 logger.info("Archiving conda environments...")
                 envs = set()
                 for job in self.jobs:
-                    if job.conda_env_file:
+                    if job.conda_env_spec:
                         env_archive = job.archive_conda_env()
                         envs.add(env_archive)
                 for env in envs:
diff --git a/snakemake/deployment/conda.py b/snakemake/deployment/conda.py
index d3dbb94c9..59a91a556 100644
--- a/snakemake/deployment/conda.py
+++ b/snakemake/deployment/conda.py
@@ -22,6 +22,7 @@
 from enum import Enum
 import threading
 import shutil
+from abc import ABC, abstractmethod
 
 
 from snakemake.exceptions import CreateCondaEnvironmentException, WorkflowError
@@ -29,7 +30,13 @@
 from snakemake.common import is_local_file, parse_uri, strip_prefix, ON_WINDOWS
 from snakemake import utils
 from snakemake.deployment import singularity, containerize
-from snakemake.io import git_content
+from snakemake.io import (
+    IOFile,
+    apply_wildcards,
+    contains_wildcard,
+    git_content,
+    _IOFile,
+)
 
 
 class CondaCleanupMode(Enum):
@@ -45,9 +52,21 @@ class Env:
     """Conda environment from a given specification file."""
 
     def __init__(
-        self, env_file, workflow, env_dir=None, container_img=None, cleanup=None
+        self,
+        workflow,
+        env_file=None,
+        env_name=None,
+        env_dir=None,
+        container_img=None,
+        cleanup=None,
     ):
-        self.file = infer_source_file(env_file)
+        self.file = None
+        self.name = None
+        if env_file is not None:
+            self.file = infer_source_file(env_file)
+        if env_name is not None:
+            assert env_file is None, "bug: both env_file and env_name specified"
+            self.name = env_name
 
         self.frontend = workflow.conda_frontend
         self.workflow = workflow
@@ -68,9 +87,20 @@ def __init__(
         self._singularity_args = workflow.singularity_args
 
     def _get_content(self):
-        return self.workflow.sourcecache.open(self.file, "rb").read()
+        if self.is_named:
+            from snakemake.shell import shell
+
+            content = shell.check_output(
+                "conda env export {}".format(self.address_argument),
+                stderr=subprocess.STDOUT,
+                universal_newlines=True,
+            )
+            return content.encode()
+        else:
+            return self.workflow.sourcecache.open(self.file, "rb").read()
 
     def _get_content_deploy(self):
+        self.check_is_file_based()
         deploy_file = Path(self.file).with_suffix(".post-deploy.sh")
         return self.workflow.sourcecache.open(deploy_file, "rb").read()
 
@@ -123,29 +153,50 @@ def content_hash(self):
 
     @property
     def is_containerized(self):
-
         if not self._container_img:
             return False
         return self._container_img.is_containerized
 
     @property
-    def path(self):
+    def is_named(self):
+        return self.file is None
+
+    def check_is_file_based(self):
+        assert (
+            self.file is not None
+        ), "bug: trying to access conda env file based functionality for named environment"
+
+    @property
+    def address(self):
         """Path to directory of the conda environment.
 
         First tries full hash, if it does not exist, (8-prefix) is used
         as default.
 
         """
-        hash = self.hash
-        env_dir = self._env_dir
-        get_path = lambda h: os.path.join(env_dir, h)
-        hash_candidates = [hash[:8], hash]  # [0] is the old fallback hash (shortened)
-        exists = [os.path.exists(get_path(h)) for h in hash_candidates]
-        if self.is_containerized or exists[1] or (not exists[0]):
-            # containerizes, full hash exists or fallback hash does not exist: use full hash
-            return get_path(hash_candidates[1])
-        # use fallback hash
-        return get_path(hash_candidates[0])
+        if self.is_named:
+            return self.name
+        else:
+            hash = self.hash
+            env_dir = self._env_dir
+            get_path = lambda h: os.path.join(env_dir, h)
+            hash_candidates = [
+                hash[:8],
+                hash,
+            ]  # [0] is the old fallback hash (shortened)
+            exists = [os.path.exists(get_path(h)) for h in hash_candidates]
+            if self.is_containerized or exists[1] or (not exists[0]):
+                # containerizes, full hash exists or fallback hash does not exist: use full hash
+                return get_path(hash_candidates[1])
+            # use fallback hash
+            return get_path(hash_candidates[0])
+
+    @property
+    def address_argument(self):
+        if self.is_named:
+            return "--name '{}'".format(self.address)
+        else:
+            return "--prefix '{}'".format(self.address)
 
     @property
     def archive_file(self):
@@ -167,6 +218,8 @@ def create_archive(self):
         # importing requests locally because it interferes with instantiating conda environments
         import requests
 
+        self.check_is_file_based()
+
         env_archive = self.archive_file
         if os.path.exists(env_archive):
             return env_archive
@@ -181,7 +234,7 @@ def create_archive(self):
             os.makedirs(env_archive, exist_ok=True)
             try:
                 out = shell.check_output(
-                    "conda list --explicit --prefix '{}'".format(self.path),
+                    "conda list --explicit {}".format(self.address_argument),
                     stderr=subprocess.STDOUT,
                     universal_newlines=True,
                 )
@@ -240,7 +293,7 @@ def execute_deployment_script(self, env_file, deploy_file):
         )
         conda = Conda(self._container_img)
         shell.check_output(
-            conda.shellcmd(self.path, "sh {}".format(deploy_file)),
+            conda.shellcmd(self.address, "sh {}".format(deploy_file)),
             stderr=subprocess.STDOUT,
         )
 
@@ -248,6 +301,8 @@ def create(self, dryrun=False):
         """Create the conda enviroment."""
         from snakemake.shell import shell
 
+        self.check_is_file_based()
+
         # Read env file and create hash.
         env_file = self.file
         deploy_file = None
@@ -279,8 +334,7 @@ def create(self, dryrun=False):
             if Path(env_file).with_suffix(".post-deploy.sh").exists():
                 deploy_file = Path(env_file).with_suffix(".post-deploy.sh")
 
-        env_hash = self.hash
-        env_path = self.path
+        env_path = self.address
 
         if self.is_containerized:
             if not dryrun:
@@ -462,11 +516,17 @@ def get_singularity_envvars(self):
 
     def __hash__(self):
         # this hash is only for object comparison, not for env paths
-        return hash(self.file)
+        if self.is_named:
+            return hash(self.name)
+        else:
+            return hash(self.file)
 
     def __eq__(self, other):
         if isinstance(other, Env):
-            return self.file == other.file
+            if self.is_named:
+                return self.name == other.name
+            else:
+                return self.file == other.file
         return False
 
 
@@ -583,24 +643,111 @@ def bin_path(self):
         else:
             return os.path.join(self.prefix_path, "bin")
 
-    def shellcmd(self, env_path, cmd):
+    def shellcmd(self, env_address, cmd):
         # get path to activate script
         activate = os.path.join(self.bin_path(), "activate")
 
         if ON_WINDOWS:
             activate = activate.replace("\\", "/")
-            env_path = env_path.replace("\\", "/")
+            env_address = env_address.replace("\\", "/")
 
-        return "source {} '{}'; {}".format(activate, env_path, cmd)
+        return "source {} '{}'; {}".format(activate, env_address, cmd)
 
-    def shellcmd_win(self, env_path, cmd):
+    def shellcmd_win(self, env_address, cmd):
         """Prepend the windows activate bat script."""
         # get path to activate script
         activate = os.path.join(self.bin_path(), "activate.bat").replace("\\", "/")
-        env_path = env_path.replace("\\", "/")
+        env_address = env_address.replace("\\", "/")
 
-        return '"{}" "{}"&&{}'.format(activate, env_path, cmd)
+        return '"{}" "{}"&&{}'.format(activate, env_address, cmd)
 
 
 def is_mamba_available():
     return shutil.which("mamba") is not None
+
+
+class CondaEnvSpec(ABC):
+    @abstractmethod
+    def apply_wildcards(self, wildcards):
+        ...
+
+    @abstractmethod
+    def get_conda_env(self, workflow, env_dir=None, container_img=None, cleanup=None):
+        ...
+
+    @abstractmethod
+    def check(self):
+        ...
+
+    @property
+    def is_file(self):
+        return False
+
+    @property
+    @abstractmethod
+    def contains_wildcard(self):
+        ...
+
+
+class CondaEnvFileSpec(CondaEnvSpec):
+    def __init__(self, filepath: str, rule=None):
+        if isinstance(filepath, _IOFile):
+            self.file = filepath
+        else:
+            self.file = IOFile(filepath, rule=rule)
+
+    def apply_wildcards(self, wildcards):
+        filepath = self.file.apply_wildcards(wildcards)
+        if is_local_file(filepath):
+            # Normalize 'file:///my/path.yml' to '/my/path.yml'
+            filepath = parse_uri(filepath).uri_path
+        return CondaEnvFileSpec(filepath)
+
+    def check(self):
+        self.file.check()
+
+    def get_conda_env(self, workflow, env_dir=None, container_img=None, cleanup=None):
+        return Env(
+            workflow,
+            env_file=self.file,
+            env_dir=env_dir,
+            container_img=container_img,
+            cleanup=cleanup,
+        )
+
+    @property
+    def is_file(self):
+        return True
+
+    @property
+    def contains_wildcard(self):
+        return contains_wildcard(self.file)
+
+
+class CondaEnvNameSpec(CondaEnvSpec):
+    def __init__(self, name: str):
+        self.name = name
+
+    def apply_wildcards(self, wildcards):
+        return CondaEnvNameSpec(apply_wildcards(self.name, wildcards))
+
+    def get_conda_env(self, workflow, env_dir=None, container_img=None, cleanup=None):
+        return Env(
+            workflow,
+            env_name=self.name,
+            env_dir=env_dir,
+            container_img=container_img,
+            cleanup=cleanup,
+        )
+
+    def check(self):
+        # not a file, nothing to check here
+        pass
+
+    @property
+    def contains_wildcard(self):
+        return contains_wildcard(self.name)
+
+
+def is_conda_env_file(spec: str):
+    return spec.endswith(".yaml") or spec.endswith(".yml")
diff --git a/snakemake/deployment/containerize.py b/snakemake/deployment/containerize.py
index 755133826..850faaea4 100644
--- a/snakemake/deployment/containerize.py
+++ b/snakemake/deployment/containerize.py
@@ -15,7 +15,7 @@
 
 def containerize(workflow):
     if any(
-        contains_wildcard(rule.conda_env)
+        rule.conda_env.contains_wildcard
         for rule in workflow.rules
         if rule.conda_env is not None
     ):
@@ -31,7 +31,7 @@ def relfile(env):
 
     envs = sorted(
         set(
-            conda.Env(rule.conda_env, workflow, env_dir=CONDA_ENV_PATH)
+            rule.conda_env.get_conda_env(workflow, env_dir=CONDA_ENV_PATH)
             for rule in workflow.rules
             if rule.conda_env is not None
         ),
diff --git a/snakemake/executors/__init__.py b/snakemake/executors/__init__.py
index 5d13cc3a6..57c025f0c 100644
--- a/snakemake/executors/__init__.py
+++ b/snakemake/executors/__init__.py
@@ -485,7 +485,9 @@ def run(self, job, callback=None, submit_callback=None, error_callback=None):
     def job_args_and_prepare(self, job):
         job.prepare()
 
-        conda_env = job.conda_env_path if self.workflow.use_conda else None
+        conda_env = (
+            job.conda_env.address if self.workflow.use_conda and job.conda_env else None
+        )
         container_img = (
             job.container_img_path if self.workflow.use_singularity else None
         )
diff --git a/snakemake/jobs.py b/snakemake/jobs.py
index ae4558bd4..e7b0014d4 100644
--- a/snakemake/jobs.py
+++ b/snakemake/jobs.py
@@ -175,7 +175,7 @@ def __init__(
         self._log = None
         self._benchmark = None
         self._resources = None
-        self._conda_env_file = None
+        self._conda_env_spec = None
         self._conda_env = None
         self._group = None
 
@@ -316,34 +316,30 @@ def reset_params_and_resources(self):
         self._params = None
 
     @property
-    def conda_env_file(self):
-        if self._conda_env_file is None:
-            expanded_env = self.rule.expand_conda_env(self.wildcards_dict)
-            if expanded_env is not None:
-                # Normalize 'file:///my/path.yml' to '/my/path.yml'
-                if is_local_file(expanded_env):
-                    self._conda_env_file = parse_uri(expanded_env).uri_path
-                else:
-                    self._conda_env_file = expanded_env
-        return self._conda_env_file
+    def conda_env_spec(self):
+        if self._conda_env_spec is None:
+            self._conda_env_spec = self.rule.expand_conda_env(self.wildcards_dict)
+        return self._conda_env_spec
 
     @property
     def conda_env(self):
-        if self.conda_env_file:
+        if self.conda_env_spec:
             if self._conda_env is None:
                 self._conda_env = self.dag.conda_envs.get(
-                    (self.conda_env_file, self.container_img_url)
+                    (self.conda_env_spec, self.container_img_url)
                 )
             return self._conda_env
         return None
 
-    @property
-    def conda_env_path(self):
-        return self.conda_env.path if self.conda_env else None
-
     def archive_conda_env(self):
         """Archive a conda environment into a custom local channel."""
-        if self.conda_env_file:
+        if self.conda_env_spec:
+            if self.conda_env.is_named:
+                raise WorkflowError(
+                    "Workflow archives cannot be created for workflows using named conda environments."
+                    "Please use paths to YAML files for all your conda directives.",
+                    rule=self.rule,
+                )
             return self.conda_env.create_archive()
         return None
 
@@ -964,7 +960,7 @@ def log_error(self, msg=None, indent=False, **kwargs):
             jobid=self.dag.jobid(self),
             output=list(format_files(self, self.output, self.dynamic_output)),
             log=list(self.log),
-            conda_env=self.conda_env.path if self.conda_env else None,
+            conda_env=self.conda_env.address if self.conda_env else None,
             aux=kwargs,
             indent=indent,
             shellcmd=self.shellcmd,
@@ -984,8 +980,12 @@ def get_wait_for_files(self):
 
         if self.shadow_dir:
             wait_for_files.append(self.shadow_dir)
-        if self.dag.workflow.use_conda and self.conda_env:
-            wait_for_files.append(self.conda_env_path)
+        if (
+            self.dag.workflow.use_conda
+            and self.conda_env
+            and not self.conda_env.is_named
+        ):
+            wait_for_files.append(self.conda_env.address)
         return wait_for_files
 
     @property
@@ -1230,8 +1230,12 @@ def get_wait_for_files(self):
         for job in self.jobs:
             if job.shadow_dir:
                 wait_for_files.append(job.shadow_dir)
-            if self.dag.workflow.use_conda and job.conda_env:
-                wait_for_files.append(job.conda_env_path)
+            if (
+                self.dag.workflow.use_conda
+                and job.conda_env
+                and not job.conda_env.is_named
+            ):
+                wait_for_files.append(job.conda_env.address)
         return wait_for_files
 
     @property
diff --git a/snakemake/rules.py b/snakemake/rules.py
index 064ec5e38..1e6a4e19e 100644
--- a/snakemake/rules.py
+++ b/snakemake/rules.py
@@ -342,7 +342,16 @@ def conda_env(self):
 
     @conda_env.setter
     def conda_env(self, conda_env):
-        self._conda_env = IOFile(conda_env, rule=self)
+        from snakemake.deployment.conda import (
+            is_conda_env_file,
+            CondaEnvFileSpec,
+            CondaEnvNameSpec,
+        )
+
+        if is_conda_env_file(conda_env):
+            self._conda_env = CondaEnvFileSpec(conda_env, rule=self)
+        else:
+            self._conda_env = CondaEnvNameSpec(conda_env)
 
     @property
     def container_img(self):
diff --git a/snakemake/workflow.py b/snakemake/workflow.py
index 721fb3d24..3d5c6cbc5 100644
--- a/snakemake/workflow.py
+++ b/snakemake/workflow.py
@@ -89,7 +89,7 @@
     SourceFile,
     infer_source_file,
 )
-from snakemake.deployment.conda import Conda
+from snakemake.deployment.conda import Conda, is_conda_env_file
 from snakemake import sourcecache
 
 
@@ -353,8 +353,8 @@ def norm_rule_relpath(f, rule):
                     for dirpath, _, files in os.walk(script_dir)
                     for f in files
                 )
-            if rule.conda_env:
-                f = local_path(rule.conda_env)
+            if rule.conda_env and rule.conda_env.is_file:
+                f = local_path(rule.conda_env.file)
                 if f:
                     # url points to a local env file
                     env_path = norm_rule_relpath(f, rule)
@@ -957,11 +957,11 @@ def files(items):
         if list_conda_envs:
             print("environment", "container", "location", sep="\t")
             for env in set(job.conda_env for job in dag.jobs):
-                if env:
+                if env and not env.is_named:
                     print(
                         env.file.simplify_path(),
                         env.container_img_url or "",
-                        simplify_path(env.path),
+                        simplify_path(env.address),
                         sep="\t",
                     )
             return True
@@ -1481,6 +1481,7 @@ def decorate(ruleinfo):
 
                 if (
                     ruleinfo.conda_env is not None
+                    and is_conda_env_file(ruleinfo.conda_env)
                     and is_local_file(ruleinfo.conda_env)
                     and not os.path.isabs(ruleinfo.conda_env)
                 ):
diff --git a/tests/test_conda_named/Snakefile b/tests/test_conda_named/Snakefile
new file mode 100644
index 000000000..ba5ff78df
--- /dev/null
+++ b/tests/test_conda_named/Snakefile
@@ -0,0 +1,18 @@
+shell("source activate base; mamba create -y -n xxx-test-env -c conda-forge --override-channels ripgrep==13.0.0")
+
+
+onsuccess:
+    shell("conda env remove -y -n xxx-test-env")
+
+
+onerror:
+    shell("conda env remove -y -n xxx-test-env")
+
+
+rule a:
+    output:
+        "test.out"
+    conda:
+        "xxx-test-env"
+    shell:
+        r"rg --version | head -n1 | grep -o 'ripgrep [0-9]*\.[0-9]*\.[0-9]*' > {output}"
diff --git a/tests/test_conda_named/expected-results/test.out b/tests/test_conda_named/expected-results/test.out
new file mode 100644
index 000000000..fa1e10f0c
--- /dev/null
+++ b/tests/test_conda_named/expected-results/test.out
@@ -0,0 +1 @@
+ripgrep 13.0.0
diff --git a/tests/tests.py b/tests/tests.py
index 9727866bb..c7da69b8c 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -1392,3 +1392,8 @@ def test_checkpoint_allowed_rules():
 @skip_on_windows
 def test_modules_ruledeps_inheritance():
     run(dpath("test_modules_ruledeps_inheritance"))
+
+
+@skip_on_windows
+def test_conda_named():
+    run(dpath("test_conda_named"), use_conda=True)