add constraint spec (#40)

Signed-off-by: Jorge Aguilera <[email protected]>
nextflow-io · Jun 13, 2024 · 419b14a · 419b14a
1 parent 737410d
commit 419b14a
Show file tree

Hide file tree

Showing 5 changed files with 237 additions and 28 deletions.
diff --git a/plugins/nf-nomad/src/main/nextflow/nomad/NomadConfig.groovy b/plugins/nf-nomad/src/main/nextflow/nomad/NomadConfig.groovy
@@ -19,6 +19,7 @@ package nextflow.nomad
 import groovy.transform.CompileStatic
 import groovy.util.logging.Slf4j
 import nextflow.nomad.config.AffinitySpec
+import nextflow.nomad.config.ConstraintSpec
 import nextflow.nomad.config.VolumeSpec
 
 /**
@@ -61,6 +62,7 @@ class NomadConfig {
         final String dockerVolume
         final VolumeSpec volumeSpec
         final AffinitySpec affinitySpec
+        final ConstraintSpec constraintSpec
 
         NomadJobOpts(Map nomadJobOpts){
             deleteOnCompletion = nomadJobOpts.containsKey("deleteOnCompletion") ?
@@ -78,25 +80,51 @@ class NomadConfig {
             if( dockerVolume ){
                 log.info "dockerVolume config will be deprecated, use volume type:'docker' name:'name' instead"
             }
+
+            this.volumeSpec = parseVolume(nomadJobOpts)
+            this.affinitySpec = parseAffinity(nomadJobOpts)
+            this.constraintSpec = parseConstraint(nomadJobOpts)
+        }
+
+        VolumeSpec parseVolume(Map nomadJobOpts){
             if( nomadJobOpts.volume && nomadJobOpts.volume instanceof Closure){
-                this.volumeSpec = new VolumeSpec()
+                def volumeSpec = new VolumeSpec()
                 def closure = (nomadJobOpts.volume as Closure)
-                def clone = closure.rehydrate(this.volumeSpec, closure.owner, closure.thisObject)
+                def clone = closure.rehydrate(volumeSpec, closure.owner, closure.thisObject)
                 clone.resolveStrategy = Closure.DELEGATE_FIRST
                 clone()
-                this.volumeSpec.validate()
+                volumeSpec.validate()
+                volumeSpec
             }else{
-                volumeSpec = null
+                null
             }
-            if( nomadJobOpts.affinity && nomadJobOpts.affinity instanceof Closure){
-                this.affinitySpec = new AffinitySpec()
+        }
+
+        AffinitySpec parseAffinity(Map nomadJobOpts) {
+            if (nomadJobOpts.affinity && nomadJobOpts.affinity instanceof Closure) {
+                def affinitySpec = new AffinitySpec()
                 def closure = (nomadJobOpts.affinity as Closure)
-                def clone = closure.rehydrate(this.affinitySpec, closure.owner, closure.thisObject)
+                def clone = closure.rehydrate(affinitySpec, closure.owner, closure.thisObject)
                 clone.resolveStrategy = Closure.DELEGATE_FIRST
                 clone()
-                this.affinitySpec.validate()
-            }else{
-                affinitySpec = null
+                affinitySpec.validate()
+                affinitySpec
+            } else {
+                null
+            }
+        }
+
+        ConstraintSpec parseConstraint(Map nomadJobOpts){
+            if (nomadJobOpts.constraint && nomadJobOpts.constraint instanceof Closure) {
+                def constraintSpec = new ConstraintSpec()
+                def closure = (nomadJobOpts.constraint as Closure)
+                def clone = closure.rehydrate(constraintSpec, closure.owner, closure.thisObject)
+                clone.resolveStrategy = Closure.DELEGATE_FIRST
+                clone()
+                constraintSpec.validate()
+                constraintSpec
+            } else {
+                null
             }
         }
     }

diff --git a/plugins/nf-nomad/src/main/nextflow/nomad/config/ConstraintSpec.groovy b/plugins/nf-nomad/src/main/nextflow/nomad/config/ConstraintSpec.groovy
@@ -0,0 +1,38 @@
+package nextflow.nomad.config
+
+class ConstraintSpec {
+
+    private String attribute
+    private String operator
+    private String value
+
+    String getOperator(){
+        return operator
+    }
+
+    String getAttribute() {
+        return attribute
+    }
+
+    String getValue() {
+        return value
+    }
+
+    ConstraintSpec attribute(String attribute){
+        this.attribute=attribute
+        this
+    }
+
+    ConstraintSpec operator(String operator){
+        this.operator = operator
+        this
+    }
+
+    ConstraintSpec value(String value){
+        this.value = value
+        this
+    }
+
+    void validate(){
+    }
+}
diff --git a/plugins/nf-nomad/src/main/nextflow/nomad/executor/NomadService.groovy b/plugins/nf-nomad/src/main/nextflow/nomad/executor/NomadService.groovy
@@ -23,6 +23,7 @@ import io.nomadproject.client.ApiClient
 import io.nomadproject.client.api.JobsApi
 import io.nomadproject.client.model.Affinity
 import io.nomadproject.client.model.AllocationListStub
+import io.nomadproject.client.model.Constraint
 import io.nomadproject.client.model.Job
 import io.nomadproject.client.model.JobRegisterRequest
 import io.nomadproject.client.model.JobRegisterResponse
@@ -198,6 +199,21 @@ class NomadService implements Closeable{
             }
             taskDef.affinities([affinity])
         }
+
+        if( config.jobOpts.constraintSpec ){
+            def constraint = new Constraint()
+            if(config.jobOpts.constraintSpec.attribute){
+                constraint.ltarget(config.jobOpts.constraintSpec.attribute)
+            }
+
+            constraint.operand(config.jobOpts.constraintSpec.operator ?: "=")
+
+            if(config.jobOpts.constraintSpec.value){
+                constraint.rtarget(config.jobOpts.constraintSpec.value)
+            }
+            taskDef.constraints([constraint])
+        }
+
         taskDef
     }
 

diff --git a/plugins/nf-nomad/src/test/nextflow/nomad/NomadConfigSpec.groovy b/plugins/nf-nomad/src/test/nextflow/nomad/NomadConfigSpec.groovy
@@ -181,4 +181,21 @@ class NomadConfigSpec extends Specification {
         config.jobOpts.affinitySpec.getValue() == '3'
         config.jobOpts.affinitySpec.getWeight() == 50
     }
+
+    void "should instantiate a constraint spec if specified"() {
+        when:
+        def config = new NomadConfig([
+                jobs: [constraint : {
+                    attribute '${meta.my_custom_value}'
+                    operator  ">"
+                    value     "3"
+                }]
+        ])
+
+        then:
+        config.jobOpts.constraintSpec
+        config.jobOpts.constraintSpec.getAttribute() == '${meta.my_custom_value}'
+        config.jobOpts.constraintSpec.getOperator() == '>'
+        config.jobOpts.constraintSpec.getValue() == '3'
+    }
 }