fix: cache conda envs to fix performance regression introduced in #1300 (#3093)

johanneskoester · web-flow · commit 66600c4a2934 · 2024-09-19T16:16:53.000+02:00
&lt;!--Add a description of your PR here--&gt;

### QC
&lt;!-- Make sure that you can tick the boxes below. --&gt;

* [x] The PR contains a test case for the changes or the changes are
already covered by an existing test case.
* [x] The documentation (`docs/`) is updated to reflect the changes or
this is not necessary (e.g. if the change does neither modify the
language nor the behavior or functionalities of Snakemake).


&lt;!-- This is an auto-generated comment: release notes by coderabbit.ai
--&gt;

## Summary by CodeRabbit

- **New Features**
- Introduced a caching mechanism for the conda environment expansion,
improving performance by reducing redundant computations.
- **Bug Fixes**
- Enhanced logic for determining cacheability of the conda environment,
ensuring accurate evaluations.

&lt;!-- end of auto-generated comment: release notes by coderabbit.ai --&gt;
diff --git a/snakemake/rules.py b/snakemake/rules.py
@@ -25,6 +25,7 @@
     _IOFile,
     Namedlist,
     AnnotatedString,
+    contains_wildcard,
     contains_wildcard_constraints,
     update_wildcard_constraints,
     flag,
@@ -66,6 +67,9 @@
 from snakemake_interface_common.rules import RuleInterface
 
 
+_NOT_CACHED = object()
+
+
 class Rule(RuleInterface):
     def __init__(self, name, workflow, lineno=None, snakefile=None):
         """
@@ -92,6 +96,7 @@ def __init__(self, name, workflow, lineno=None, snakefile=None):
         self._log = Log()
         self._benchmark = None
         self._conda_env = None
+        self._expanded_conda_env = _NOT_CACHED
         self._container_img = None
         self.is_containerized = False
         self.env_modules = None
@@ -1061,6 +1066,9 @@ def expand_group(self, wildcards):
             return self.group
 
     def expand_conda_env(self, wildcards, params=None, input=None):
+        if self._expanded_conda_env is not _NOT_CACHED:
+            return self._expanded_conda_env
+
         from snakemake.common import is_local_file
         from snakemake.sourcecache import SourceFile, infer_source_file
         from snakemake.deployment.conda import (
@@ -1070,12 +1078,18 @@ def expand_conda_env(self, wildcards, params=None, input=None):
         )
 
         conda_env = self._conda_env
-        if callable(conda_env):
-            conda_env, _ = self.apply_input_function(
-                conda_env, wildcards=wildcards, params=params, input=input
-            )
-
-        if conda_env is None:
+        if conda_env is not None:
+            if not callable(conda_env):
+                cacheable = not contains_wildcard(conda_env)
+            else:
+                conda_env, _ = self.apply_input_function(
+                    conda_env, wildcards=wildcards, params=params, input=input
+                )
+                cacheable = False
+                if conda_env is None:
+                    return None
+        else:
+            self._expanded_conda_env = None
             return None
 
         if is_conda_env_file(conda_env):
@@ -1096,6 +1110,9 @@ def expand_conda_env(self, wildcards, params=None, input=None):
         conda_env = conda_env.apply_wildcards(wildcards, self)
         conda_env.check()
 
+        if cacheable:
+            self._expanded_conda_env = conda_env
+
         return conda_env
 
     def is_producer(self, requested_output):
