Skip to content

Commit

Permalink
Added affinity option to pod directive (#2522) [ci fast]
Browse files Browse the repository at this point in the history 
Signed-off-by: Ben Sherman <bentshermann@gmail.com>
  • Loading branch information
@bentsherman
bentsherman committed Jan 14, 2022
1 parent 13f3e03 commit 640cbed
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 34 deletions.
3 changes: 2 additions & 1 deletion docs/process.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1979,7 +1979,8 @@ The ``pod`` directive allows the definition of the following options:
``imagePullPolicy: <V>`` Specifies the strategy to be used to pull the container image e.g. ``imagePullPolicy: 'Always'``.
``imagePullSecret: <V>`` Specifies the secret name to access a private container image registry. See `Kubernetes documentation <https://kubernetes.io/docs/concepts/containers/images/#specifying-imagepullsecrets-on-a-pod>`_ for details.
``runAsUser: <UID>`` Specifies the user ID to be used to run the container.
``nodeSelector: <V>`` Specifies which node the process will run on. See `Kubernetes nodeSelector <https://kubernetes.io/docs/concepts/configuration/assign-pod-node/#nodeselector>`_ for details.
``nodeSelector: <V>`` Specifies which node the process will run on. See `Kubernetes nodeSelector <https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#nodeselector>`_ for details.
``affinity: <V>`` Specifies affinity for which nodes the process should run on. See `Kubernetes affinity <https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#affinity-and-anti-affinity>`_ for details.
================================================= =================================================

When defined in the Nextflow configuration file, a pod setting can be defined using the canonical
Expand Down
17 changes: 14 additions & 3 deletions modules/nextflow/src/main/groovy/nextflow/k8s/model/PodOptions.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ class PodOptions {

private PodNodeSelector nodeSelector

private Map affinity

private PodSecurityContext securityContext

PodOptions( List<Map> options=null ) {
Expand Down Expand Up @@ -109,6 +111,9 @@ class PodOptions {
else if( entry.nodeSelector ) {
this.nodeSelector = new PodNodeSelector(entry.nodeSelector)
}
else if( entry.affinity instanceof Map ) {
this.affinity = entry.affinity as Map
}
else if( entry.annotation && entry.value ) {
this.annotations.put(entry.annotation as String, entry.value as String)
}
Expand All @@ -129,15 +134,17 @@ class PodOptions {

Map<String,String> getAnnotations() { annotations }

PodSecurityContext getSecurityContext() { securityContext }

PodNodeSelector getNodeSelector() { nodeSelector }

PodOptions setNodeSelector( PodNodeSelector sel ) {
nodeSelector = sel
return this
}

Map getAffinity() { affinity }

PodSecurityContext getSecurityContext() { securityContext }

PodOptions setSecurityContext( PodSecurityContext ctx ) {
this.securityContext = ctx
return this
Expand All @@ -159,6 +166,7 @@ class PodOptions {

PodOptions plus( PodOptions other ) {
def result = new PodOptions()

// env vars
result.envVars.addAll(envVars)
result.envVars.addAll( other.envVars )
Expand All @@ -181,9 +189,12 @@ class PodOptions {
else
result.securityContext = securityContext

// node select
// node selector
result.nodeSelector = other.nodeSelector ?: this.nodeSelector

// affinity
result.affinity = other.affinity ?: this.affinity

// pull policy
if (other.imagePullPolicy)
result.imagePullPolicy = other.imagePullPolicy
Expand Down
9 changes: 9 additions & 0 deletions modules/nextflow/src/main/groovy/nextflow/k8s/model/PodSpecBuilder.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,8 @@ class PodSpecBuilder {

PodNodeSelector nodeSelector

Map affinity

/**
* @return A sequential volume unique identifier
*/
Expand Down Expand Up @@ -252,8 +254,12 @@ class PodSpecBuilder {
// -- security context
if( opts.securityContext )
securityContext = opts.securityContext
// -- node selector
if( opts.nodeSelector )
nodeSelector = opts.nodeSelector
// -- affinity
if( opts.affinity )
affinity = opts.affinity

return this
}
Expand Down Expand Up @@ -309,6 +315,9 @@ class PodSpecBuilder {
if( nodeSelector )
spec.nodeSelector = nodeSelector.toSpec()

if( affinity )
spec.affinity = affinity

if( this.serviceAccount )
spec.serviceAccountName = this.serviceAccount

Expand Down
73 changes: 43 additions & 30 deletions modules/nextflow/src/test/groovy/nextflow/k8s/model/PodSpecBuilderTest.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -534,6 +534,16 @@ class PodSpecBuilderTest extends Specification {
def opts = Mock(PodOptions)
def builder = new PodSpecBuilder(podName: 'foo', imageName: 'image', command: ['echo'], labels: [runName: 'crazy_john'], annotations: [evict: 'false'])

def affinity = [
nodeAffinity: [
requiredDuringSchedulingIgnoredDuringExecution: [
nodeSelectorTerms: [
[key: "foo", operator: "In", values: ["bar", "baz"]]
]
]
]
]

when:
def spec = builder.withPodOptions(opts).build()
then:
Expand All @@ -547,41 +557,44 @@ class PodSpecBuilderTest extends Specification {
_ * opts.getAnnotations() >> [OMEGA:'zzz', SIGMA:'www']
_ * opts.getSecurityContext() >> new PodSecurityContext(1000)
_ * opts.getNodeSelector() >> new PodNodeSelector(gpu:true, queue: 'fast')
_ * opts.getAffinity() >> affinity

spec == [
apiVersion: 'v1',
kind: 'Pod',
metadata: [
apiVersion: 'v1',
kind: 'Pod',
metadata: [
name:'foo',
namespace:'default',
labels:[runName:'crazy_john', ALPHA:'xxx', GAMMA:'yyy'],
annotations: [evict: 'false', OMEGA:'zzz', SIGMA:'www']
],
spec: [
restartPolicy:'Never',
securityContext: [ runAsUser: 1000 ],
imagePullSecrets: [[ name: 'myPullSecret' ]],
nodeSelector: [gpu: 'true', queue: 'fast'],
affinity: affinity,

containers:[
[
name:'foo',
namespace:'default',
labels:[runName:'crazy_john', ALPHA:'xxx', GAMMA:'yyy'],
annotations: [evict: 'false', OMEGA:'zzz', SIGMA:'www']
],
spec: [
restartPolicy:'Never',
securityContext: [ runAsUser: 1000 ],
imagePullSecrets: [[ name: 'myPullSecret' ]],
nodeSelector: [gpu: 'true', queue: 'fast'],

containers:[
[name:'foo',
image:'image',
imagePullPolicy: 'always',
command:['echo'],
env:[[name:'HELLO', value:'WORLD']],
volumeMounts:[
[name:'vol-1', mountPath:'/work'],
[name:'vol-2', mountPath:'/home/user'],
[name:'vol-3', mountPath:'/etc/secret.txt']
],
]
image:'image',
imagePullPolicy: 'always',
command:['echo'],
env:[[name:'HELLO', value:'WORLD']],
volumeMounts:[
[name:'vol-1', mountPath:'/work'],
[name:'vol-2', mountPath:'/home/user'],
[name:'vol-3', mountPath:'/etc/secret.txt']
],
volumes:[
[name:'vol-1', persistentVolumeClaim:[claimName:'pvc1']],
[name:'vol-2', configMap:[name:'data']],
[name:'vol-3', secret:[secretName:'blah']] ]
]
],
volumes:[
[name:'vol-1', persistentVolumeClaim:[claimName:'pvc1']],
[name:'vol-2', configMap:[name:'data']],
[name:'vol-3', secret:[secretName:'blah']]
]

]
]

}
Expand Down

0 comments on commit 640cbed

Please sign in to comment.