From 1243130261843a09e78d6908397ecc62104d7a2c Mon Sep 17 00:00:00 2001
From: Ben Sherman <bentshermann@gmail.com>
Date: Sat, 26 Oct 2024 16:08:41 -0500
Subject: [PATCH] Fix support for micromamba (#4302) [ci fast]

Signed-off-by: Ben Sherman <bentshermann@gmail.com>
Signed-off-by: jorgee <jorge.ejarque@seqera.io>
Signed-off-by: Paolo Di Tommaso <paolo.ditommaso@gmail.com>
Co-authored-by: Jorge Ejarque <jorgee@users.noreply.github.com>
Co-authored-by: jorgee <jorge.ejarque@seqera.io>
Co-authored-by: Paolo Di Tommaso <paolo.ditommaso@gmail.com>
---
 .../groovy/nextflow/conda/CondaCache.groovy   |   3 +-
 .../executor/BashWrapperBuilder.groovy        |  11 +-
 .../groovy/nextflow/processor/TaskBean.groovy |   3 +
 .../groovy/nextflow/processor/TaskRun.groovy  |   6 +
 .../nextflow/conda/CondaCacheTest.groovy      | 107 ++++++++++++++++++
 .../executor/BashWrapperBuilderTest.groovy    |  17 +++
 .../nextflow/processor/TaskBeanTest.groovy    |   4 +
 7 files changed, 146 insertions(+), 5 deletions(-)

diff --git a/modules/nextflow/src/main/groovy/nextflow/conda/CondaCache.groovy b/modules/nextflow/src/main/groovy/nextflow/conda/CondaCache.groovy
index 6a008a6b6a..7d1ff901f5 100644
--- a/modules/nextflow/src/main/groovy/nextflow/conda/CondaCache.groovy
+++ b/modules/nextflow/src/main/groovy/nextflow/conda/CondaCache.groovy
@@ -285,7 +285,8 @@ class CondaCache {
         def cmd
         if( isYamlFilePath(condaEnv) ) {
             final target = isYamlUriPath(condaEnv) ? condaEnv : Escape.path(makeAbsolute(condaEnv))
-            cmd = "${binaryName} env create --prefix ${Escape.path(prefixPath)} --file ${target}"
+            final yesOpt = binaryName == 'micromamba' ? '--yes ' : ''
+            cmd = "${binaryName} env create ${yesOpt}--prefix ${Escape.path(prefixPath)} --file ${target}"
         }
         else if( isTextFilePath(condaEnv) ) {
             cmd = "${binaryName} create ${opts}--yes --quiet --prefix ${Escape.path(prefixPath)} --file ${Escape.path(makeAbsolute(condaEnv))}"
diff --git a/modules/nextflow/src/main/groovy/nextflow/executor/BashWrapperBuilder.groovy b/modules/nextflow/src/main/groovy/nextflow/executor/BashWrapperBuilder.groovy
index 3d190b63d3..6637d7a9b7 100644
--- a/modules/nextflow/src/main/groovy/nextflow/executor/BashWrapperBuilder.groovy
+++ b/modules/nextflow/src/main/groovy/nextflow/executor/BashWrapperBuilder.groovy
@@ -533,10 +533,13 @@ class BashWrapperBuilder {
     private String getCondaActivateSnippet() {
         if( !condaEnv )
             return null
-        def result = "# conda environment\n"
-        result += 'source $(conda info --json | awk \'/conda_prefix/ { gsub(/"|,/, "", $2); print $2 }\')'
-        result += "/bin/activate ${Escape.path(condaEnv)}\n"
-        return result
+        final command = useMicromamba
+            ? 'eval "$(micromamba shell hook --shell bash)" && micromamba activate'
+            : 'source $(conda info --json | awk \'/conda_prefix/ { gsub(/"|,/, "", $2); print $2 }\')/bin/activate'
+        return """\
+            # conda environment
+            ${command} ${Escape.path(condaEnv)}
+            """.stripIndent()
     }
 
     private String getSpackActivateSnippet() {
diff --git a/modules/nextflow/src/main/groovy/nextflow/processor/TaskBean.groovy b/modules/nextflow/src/main/groovy/nextflow/processor/TaskBean.groovy
index d1be0cf236..5d4175aeff 100644
--- a/modules/nextflow/src/main/groovy/nextflow/processor/TaskBean.groovy
+++ b/modules/nextflow/src/main/groovy/nextflow/processor/TaskBean.groovy
@@ -48,6 +48,8 @@ class TaskBean implements Serializable, Cloneable {
 
     Path condaEnv
 
+    Boolean useMicromamba
+
     Path spackEnv
 
     List<String> moduleNames
@@ -131,6 +133,7 @@ class TaskBean implements Serializable, Cloneable {
         this.environment = task.getEnvironment()
 
         this.condaEnv = task.getCondaEnv()
+        this.useMicromamba = task.getCondaConfig()?.useMicromamba()
         this.spackEnv = task.getSpackEnv()
         this.moduleNames = task.config.getModule()
         this.shell = task.config.getShell() ?: BashWrapperBuilder.BASH
diff --git a/modules/nextflow/src/main/groovy/nextflow/processor/TaskRun.groovy b/modules/nextflow/src/main/groovy/nextflow/processor/TaskRun.groovy
index cd72f69379..f3926c0b60 100644
--- a/modules/nextflow/src/main/groovy/nextflow/processor/TaskRun.groovy
+++ b/modules/nextflow/src/main/groovy/nextflow/processor/TaskRun.groovy
@@ -16,6 +16,8 @@
 
 package nextflow.processor
 
+import nextflow.conda.CondaConfig
+
 import java.nio.file.FileSystems
 import java.nio.file.NoSuchFileException
 import java.nio.file.Path
@@ -967,5 +969,9 @@ class TaskRun implements Cloneable {
     TaskBean toTaskBean() {
         return new TaskBean(this)
     }
+
+    CondaConfig getCondaConfig() {
+        return processor.session.getCondaConfig()
+    }
 }
 
diff --git a/modules/nextflow/src/test/groovy/nextflow/conda/CondaCacheTest.groovy b/modules/nextflow/src/test/groovy/nextflow/conda/CondaCacheTest.groovy
index 9e1f0c5ad4..637ae5623a 100644
--- a/modules/nextflow/src/test/groovy/nextflow/conda/CondaCacheTest.groovy
+++ b/modules/nextflow/src/test/groovy/nextflow/conda/CondaCacheTest.groovy
@@ -240,6 +240,33 @@ class CondaCacheTest extends Specification {
 
     }
 
+    def 'should create a conda environment - using micromamba' () {
+
+        given:
+        def ENV = 'bwa=1.1.1'
+        def PREFIX = Files.createTempDirectory('foo')
+        def cache = Spy(new CondaCache(useMicromamba: true))
+
+        when:
+        // the prefix directory exists ==> no mamba command is executed
+        def result = cache.createLocalCondaEnv(ENV)
+        then:
+        1 * cache.condaPrefixPath(ENV) >> PREFIX
+        0 * cache.isYamlFilePath(ENV)
+        0 * cache.runCommand(_)
+        result == PREFIX
+
+        when:
+        PREFIX.deleteDir()
+        result = cache.createLocalCondaEnv0(ENV, PREFIX)
+        then:
+        1 * cache.isYamlFilePath(ENV)
+        0 * cache.makeAbsolute(_)
+        1 * cache.runCommand("micromamba create --yes --quiet --prefix $PREFIX $ENV") >> null
+        result == PREFIX
+
+    }
+
     def 'should create a conda environment using mamba and remote lock file' () {
 
         given:
@@ -265,6 +292,32 @@ class CondaCacheTest extends Specification {
         1 * cache.runCommand("mamba env create --prefix $PREFIX --file $ENV") >> null
         result == PREFIX
 
+    }
+    def 'should create a conda environment using micromamba and remote lock file' () {
+
+        given:
+        def ENV = 'http://foo.com/some/file-lock.yml'
+        def PREFIX = Files.createTempDirectory('foo')
+        def cache = Spy(new CondaCache(useMicromamba: true))
+
+        when:
+        // the prefix directory exists ==> no mamba command is executed
+        def result = cache.createLocalCondaEnv(ENV)
+        then:
+        1 * cache.condaPrefixPath(ENV) >> PREFIX
+        0 * cache.isYamlFilePath(ENV)
+        0 * cache.runCommand(_)
+        result == PREFIX
+
+        when:
+        PREFIX.deleteDir()
+        result = cache.createLocalCondaEnv0(ENV, PREFIX)
+        then:
+        1 * cache.isYamlFilePath(ENV)
+        0 * cache.makeAbsolute(_)
+        1 * cache.runCommand("micromamba env create --yes --prefix $PREFIX --file $ENV") >> null
+        result == PREFIX
+
     }
 
     def 'should create conda env with options' () {
@@ -301,6 +354,23 @@ class CondaCacheTest extends Specification {
         result == PREFIX
     }
 
+    def 'should create conda env with options - using micromamba' () {
+        given:
+        def ENV = 'bwa=1.1.1'
+        def PREFIX = Paths.get('/foo/bar')
+        and:
+        def cache = Spy(new CondaCache(useMicromamba: true, createOptions: '--this --that'))
+
+        when:
+        def result = cache.createLocalCondaEnv0(ENV, PREFIX)
+        then:
+        1 * cache.isYamlFilePath(ENV)
+        1 * cache.isTextFilePath(ENV)
+        0 * cache.makeAbsolute(_)
+        1 * cache.runCommand("micromamba create --this --that --yes --quiet --prefix $PREFIX $ENV") >> null
+        result == PREFIX
+    }
+
     def 'should create conda env with channels' () {
         given:
         def ENV = 'bwa=1.1.1'
@@ -336,6 +406,24 @@ class CondaCacheTest extends Specification {
 
     }
 
+    def 'should create a conda env with a yaml file - using micromamba' () {
+
+        given:
+        def ENV = 'foo.yml'
+        def PREFIX = Paths.get('/conda/envs/my-env')
+        def cache = Spy(new CondaCache(useMicromamba: true))
+
+        when:
+        def result = cache.createLocalCondaEnv0(ENV, PREFIX)
+        then:
+        1 * cache.isYamlFilePath(ENV)
+        0 * cache.isTextFilePath(ENV)
+        1 * cache.makeAbsolute(ENV) >> Paths.get('/usr/base').resolve(ENV)
+        1 * cache.runCommand( "micromamba env create --yes --prefix $PREFIX --file /usr/base/foo.yml" ) >> null
+        result == PREFIX
+
+    }
+
     def 'should create a conda env with a text file' () {
 
         given:
@@ -355,6 +443,25 @@ class CondaCacheTest extends Specification {
 
     }
 
+    def 'should create a conda env with a text file - using micromamba' () {
+
+        given:
+        def ENV = 'foo.txt'
+        def PREFIX = Paths.get('/conda/envs/my-env')
+        and:
+        def cache = Spy(new CondaCache(useMicromamba: true, createOptions: '--this --that'))
+
+        when:
+        def result = cache.createLocalCondaEnv0(ENV, PREFIX)
+        then:
+        1 * cache.isYamlFilePath(ENV)
+        1 * cache.isTextFilePath(ENV)
+        1 * cache.makeAbsolute(ENV) >> Paths.get('/usr/base').resolve(ENV)
+        1 * cache.runCommand( "micromamba create --this --that --yes --quiet --prefix $PREFIX --file /usr/base/foo.txt" ) >> null
+        result == PREFIX
+
+    }
+
     def 'should get options from the config' () {
 
         when:
diff --git a/modules/nextflow/src/test/groovy/nextflow/executor/BashWrapperBuilderTest.groovy b/modules/nextflow/src/test/groovy/nextflow/executor/BashWrapperBuilderTest.groovy
index b9d9b8c3e1..5df66b6369 100644
--- a/modules/nextflow/src/test/groovy/nextflow/executor/BashWrapperBuilderTest.groovy
+++ b/modules/nextflow/src/test/groovy/nextflow/executor/BashWrapperBuilderTest.groovy
@@ -781,7 +781,24 @@ class BashWrapperBuilderTest extends Specification {
                 # conda environment
                 source $(conda info --json | awk '/conda_prefix/ { gsub(/"|,/, "", $2); print $2 }')/bin/activate /some/conda/env/foo
                 '''.stripIndent()
+    }
+
+    def 'should create micromamba activate snippet' () {
+
+        when:
+        def binding = newBashWrapperBuilder().makeBinding()
+        then:
+        binding.conda_activate == null
+        binding.containsKey('conda_activate')
 
+        when:
+        def CONDA = Paths.get('/some/conda/env/foo')
+        binding = newBashWrapperBuilder([condaEnv: CONDA, 'useMicromamba': true]).makeBinding()
+        then:
+        binding.conda_activate == '''\
+                # conda environment
+                eval "$(micromamba shell hook --shell bash)" && micromamba activate /some/conda/env/foo
+                '''.stripIndent()
     }
 
     def 'should create spack activate snippet' () {
diff --git a/modules/nextflow/src/test/groovy/nextflow/processor/TaskBeanTest.groovy b/modules/nextflow/src/test/groovy/nextflow/processor/TaskBeanTest.groovy
index c576bef6bb..4cfcb2a1fd 100644
--- a/modules/nextflow/src/test/groovy/nextflow/processor/TaskBeanTest.groovy
+++ b/modules/nextflow/src/test/groovy/nextflow/processor/TaskBeanTest.groovy
@@ -19,6 +19,7 @@ package nextflow.processor
 import java.nio.file.Paths
 
 import nextflow.Session
+import nextflow.conda.CondaConfig
 import nextflow.container.ContainerConfig
 import nextflow.executor.Executor
 import nextflow.script.ProcessConfig
@@ -69,6 +70,7 @@ class TaskBeanTest extends Specification {
         task.getEnvironment() >> [alpha: 'one', beta: 'xxx', gamma: 'yyy']
         task.getContainer() >> 'busybox:latest'
         task.getContainerConfig() >> [docker: true, registry: 'x']
+        task.getCondaConfig() >> new CondaConfig([useMicromamba:true], [:])
 
         when:
         def bean = new TaskBean(task)
@@ -99,6 +101,8 @@ class TaskBeanTest extends Specification {
         bean.stageInMode == 'link'
         bean.stageOutMode == 'rsync'
 
+        bean.useMicromamba == true
+
     }
 
     def 'should clone task bean' () {