Repository: jupyterhub/kubespawner

Issue:
Support setting affinity items
We should support setting nodeAffinity, podAffinity, podAntiAffinity, and preferred/required for each.
```yaml
  affinity:
    nodeAffinity:
      requiredDuringSchedulingIgnoredDuringExecution: # nodeSelector
        nodeSelectorTerms: #
        # logical AND of the items below
        - matchExpressions:
          # local OR of the items below
          - key: kubernetes.io/e2e-az-name
            operator: In
            values:
            - e2e-az1
            - e2e-az2
      preferredDuringSchedulingIgnoredDuringExecution:
      # fullfilled weights summed over and compared in between nodes
      - weight: 1
        preference:
          matchExpressions:
          # local OR of the items below
          - key: another-node-label-key
            operator: In
            values:
            - another-node-label-value

    podAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
      # logical AND of the items below
      - labelSelector:
          matchExpressions:
          # logical OR of the items below
          - key: security
            operator: In
            values:
            - S1
        topologyKey: failure-domain.beta.kubernetes.io/zone
      preferredDuringSchedulingIgnoredDuringExecution:
      # sum the weights for each node and compare
      - weight: 100
        podAffinityTerm:
          labelSelector:
            matchExpressions:
            - key: security
              operator: In
              values:
              - S2
          topologyKey: kubernetes.io/hostname

    # Exactly the same fields etc as podAffinity
    podAntiAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
      - labelSelector:
          matchExpressions:
          - key: app
            operator: In
            values:
            - store
        topologyKey: "kubernetes.io/hostname"
      preferredDuringSchedulingIgnoredDuringExecution:
      - weight: 100
        podAffinityTerm:
          labelSelector:
            matchExpressions:
            - key: security
              operator: In
              values:
              - S2
          topologyKey: kubernetes.io/hostname
```


Comments:

Patch:
diff --git a/kubespawner/objects.py b/kubespawner/objects.py
index f29fe55..1c1944f 100644
--- a/kubespawner/objects.py
+++ b/kubespawner/objects.py
@@ -21,6 +21,9 @@ from kubernetes.client.models import (
     V1beta1HTTPIngressRuleValue, V1beta1HTTPIngressPath,
     V1beta1IngressBackend,
     V1Toleration,
+    V1Affinity,
+    V1NodeAffinity, V1NodeSelector, V1NodeSelectorTerm, V1PreferredSchedulingTerm, V1NodeSelectorRequirement,
+    V1PodAffinity, V1PodAntiAffinity, V1WeightedPodAffinityTerm, V1PodAffinityTerm,
 )
 
 def make_pod(
@@ -56,6 +59,12 @@ def make_pod(
     extra_containers=None,
     scheduler_name=None,
     tolerations=None,
+    node_affinity_preferred=None,
+    node_affinity_required=None,
+    pod_affinity_preferred=None,
+    pod_affinity_required=None,
+    pod_anti_affinity_preferred=None,
+    pod_anti_affinity_required=None,
     priority_class_name=None,
     logger=None,
 ):
@@ -156,6 +165,54 @@ def make_pod(
 
         Pass this field an array of "Toleration" objects.*
         * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.10/#nodeselectorterm-v1-core
+    node_affinity_preferred:
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "PreferredSchedulingTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.10/#preferredschedulingterm-v1-core
+    node_affinity_required:
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "NodeSelectorTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.10/#nodeselectorterm-v1-core
+    pod_affinity_preferred:
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "WeightedPodAffinityTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.10/#weightedpodaffinityterm-v1-core
+    pod_affinity_required:
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "PodAffinityTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.10/#podaffinityterm-v1-core
+    pod_anti_affinity_preferred:
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "WeightedPodAffinityTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.10/#weightedpodaffinityterm-v1-core
+    pod_anti_affinity_required:
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "PodAffinityTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.10/#podaffinityterm-v1-core
     priority_class_name:
         The name of the PriorityClass to be assigned the pod. This feature is Beta available in K8s 1.11.
     """
@@ -260,8 +317,63 @@ def make_pod(
     if scheduler_name:
         pod.spec.scheduler_name = scheduler_name
 
+    node_affinity = None
+    if node_affinity_preferred or node_affinity_required:
+        node_selector = None
+        if node_affinity_required:
+            node_selector = V1NodeSelector(
+                node_selector_terms=[get_k8s_model(V1NodeSelectorTerm, obj) for obj in node_affinity_required],
+            )
+
+        preferred_scheduling_terms = None
+        if node_affinity_preferred:
+            preferred_scheduling_terms = [get_k8s_model(V1PreferredSchedulingTerm, obj) for obj in node_affinity_preferred]
+
+        node_affinity = V1NodeAffinity(
+            preferred_during_scheduling_ignored_during_execution=preferred_scheduling_terms,
+            required_during_scheduling_ignored_during_execution=node_selector,
+        )
+
+    pod_affinity = None
+    if pod_affinity_preferred or pod_affinity_required:
+        weighted_pod_affinity_terms = None
+        if pod_affinity_preferred:
+            weighted_pod_affinity_terms = [get_k8s_model(V1WeightedPodAffinityTerm, obj) for obj in pod_affinity_preferred]
+
+        pod_affinity_terms = None
+        if pod_affinity_required:
+            pod_affinity_terms = [get_k8s_model(V1PodAffinityTerm, obj) for obj in pod_affinity_required]
 
+        pod_affinity = V1PodAffinity(
+            preferred_during_scheduling_ignored_during_execution=weighted_pod_affinity_terms,
+            required_during_scheduling_ignored_during_execution=pod_affinity_terms,
+        )
+
+    pod_anti_affinity = None
+    if pod_anti_affinity_preferred or pod_anti_affinity_required:
+        weighted_pod_affinity_terms = None
+        if pod_anti_affinity_preferred:
+            weighted_pod_affinity_terms = [get_k8s_model(V1WeightedPodAffinityTerm, obj) for obj in pod_anti_affinity_preferred]
+
+        pod_affinity_terms = None
+        if pod_anti_affinity_required:
+            pod_affinity_terms = [get_k8s_model(V1PodAffinityTerm, obj) for obj in pod_anti_affinity_required]
+
+        pod_anti_affinity = V1PodAffinity(
+            preferred_during_scheduling_ignored_during_execution=weighted_pod_affinity_terms,
+            required_during_scheduling_ignored_during_execution=pod_affinity_terms,
+        )
+
+    affinity = None
+    if (node_affinity or pod_affinity or pod_anti_affinity):
+        affinity = V1Affinity(
+            node_affinity=node_affinity,
+            pod_affinity=pod_affinity,
+            pod_anti_affinity=pod_anti_affinity,
+        )
 
+    if affinity:
+        pod.spec.affinity = affinity
 
     if priority_class_name:
         pod.spec.priority_class_name = priority_class_name
diff --git a/kubespawner/spawner.py b/kubespawner/spawner.py
index fb7deaa..39adce9 100644
--- a/kubespawner/spawner.py
+++ b/kubespawner/spawner.py
@@ -900,6 +900,79 @@ class KubeSpawner(Spawner):
         """
     )
 
+    node_affinity_preferred = List(
+        config=True,
+        help="""
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "PreferredSchedulingTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.11/#preferredschedulingterm-v1-core
+        """
+    )
+    node_affinity_required = List(
+        config=True,
+        help="""
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "NodeSelectorTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.11/#nodeselectorterm-v1-core
+        """
+    )
+    pod_affinity_preferred = List(
+        config=True,
+        help="""
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "WeightedPodAffinityTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.11/#weightedpodaffinityterm-v1-core
+        """
+    )
+    pod_affinity_required = List(
+        config=True,
+        help="""
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "PodAffinityTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.11/#podaffinityterm-v1-core
+        """
+    )
+    pod_anti_affinity_preferred = List(
+        config=True,
+        help="""
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "WeightedPodAffinityTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.11/#weightedpodaffinityterm-v1-core
+        """
+    )
+    pod_anti_affinity_required = List(
+        config=True,
+        help="""
+        Affinities describe where pods prefer or require to be scheduled, they
+        may prefer or require a node to have a certain label or be in proximity
+        / remoteness to another pod. To learn more visit
+        https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
+
+        Pass this field an array of "PodAffinityTerm" objects.*
+        * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.11/#podaffinityterm-v1-core
+        """
+    )
+
     extra_resource_guarantees = Dict(
         config=True,
         help="""
@@ -1286,6 +1359,12 @@ class KubeSpawner(Spawner):
             extra_containers=self.extra_containers,
             scheduler_name=self.scheduler_name,
             tolerations=self.tolerations,
+            node_affinity_preferred=self.node_affinity_preferred,
+            node_affinity_required=self.node_affinity_required,
+            pod_affinity_preferred=self.pod_affinity_preferred,
+            pod_affinity_required=self.pod_affinity_required,
+            pod_anti_affinity_preferred=self.pod_anti_affinity_preferred,
+            pod_anti_affinity_required=self.pod_anti_affinity_required,
             priority_class_name=self.priority_class_name,
             logger=self.log,
         )


Test patch:
diff --git a/tests/test_objects.py b/tests/test_objects.py
index 53bcef9..5cb5326 100644
--- a/tests/test_objects.py
+++ b/tests/test_objects.py
@@ -1075,6 +1075,371 @@ def test_make_pod_with_tolerations():
     }
 
 
+def test_make_pod_with_node_affinity_preferred():
+    """
+    Test specification of the simplest possible pod specification with non-empty node_affinity_preferred
+    """
+    node_affinity_preferred = [{
+        "weight": 1,
+        "preference": {
+            "matchExpressions": [{
+                "key": "hub.jupyter.org/node-purpose",
+                "operator": "In",
+                "values": ["user"],
+            }],
+        }
+    }]
+    assert api_client.sanitize_for_serialization(make_pod(
+        name='test',
+        image_spec='jupyter/singleuser:latest',
+        cmd=['jupyterhub-singleuser'],
+        port=8888,
+        image_pull_policy='IfNotPresent',
+        node_affinity_preferred=node_affinity_preferred
+    )) == {
+        "metadata": {
+            "name": "test",
+            "labels": {},
+            "annotations": {}
+        },
+        "spec": {
+            "automountServiceAccountToken": False,
+            "securityContext": {},
+            "containers": [
+                {
+                    "env": [],
+                    "name": "notebook",
+                    "image": "jupyter/singleuser:latest",
+                    "imagePullPolicy": "IfNotPresent",
+                    "args": ["jupyterhub-singleuser"],
+                    "ports": [{
+                        "name": "notebook-port",
+                        "containerPort": 8888
+                    }],
+                    'volumeMounts': [],
+                    "resources": {
+                        "limits": {},
+                        "requests": {}
+                    }
+                }
+            ],
+            "volumes": [],
+            "affinity": {
+                "nodeAffinity": {
+                    "preferredDuringSchedulingIgnoredDuringExecution": node_affinity_preferred
+                }
+            }
+        },
+        "kind": "Pod",
+        "apiVersion": "v1"
+    }
+
+
+def test_make_pod_with_node_affinity_required():
+    """
+    Test specification of the simplest possible pod specification with non-empty node_affinity_required
+    """
+    node_affinity_required = [{
+        "matchExpressions": [{
+            "key": "hub.jupyter.org/node-purpose",
+            "operator": "In",
+            "values": ["user"],
+        }]
+    }]
+    assert api_client.sanitize_for_serialization(make_pod(
+        name='test',
+        image_spec='jupyter/singleuser:latest',
+        cmd=['jupyterhub-singleuser'],
+        port=8888,
+        image_pull_policy='IfNotPresent',
+        node_affinity_required=node_affinity_required
+    )) == {
+        "metadata": {
+            "name": "test",
+            "labels": {},
+            "annotations": {}
+        },
+        "spec": {
+            "automountServiceAccountToken": False,
+            "securityContext": {},
+            "containers": [
+                {
+                    "env": [],
+                    "name": "notebook",
+                    "image": "jupyter/singleuser:latest",
+                    "imagePullPolicy": "IfNotPresent",
+                    "args": ["jupyterhub-singleuser"],
+                    "ports": [{
+                        "name": "notebook-port",
+                        "containerPort": 8888
+                    }],
+                    'volumeMounts': [],
+                    "resources": {
+                        "limits": {},
+                        "requests": {}
+                    }
+                }
+            ],
+            "volumes": [],
+            "affinity": {
+                "nodeAffinity": {
+                    "requiredDuringSchedulingIgnoredDuringExecution": {
+                        "nodeSelectorTerms": node_affinity_required
+                    }
+                }
+            }
+        },
+        "kind": "Pod",
+        "apiVersion": "v1"
+    }
+
+
+def test_make_pod_with_pod_affinity_preferred():
+    """
+    Test specification of the simplest possible pod specification with non-empty pod_affinity_preferred
+    """
+    pod_affinity_preferred = [{
+        "weight": 100,
+        "podAffinityTerm": {
+            "labelSelector": {
+                "matchExpressions": [{
+                    "key": "hub.jupyter.org/pod-kind",
+                    "operator": "In",
+                    "values": ["user"],
+                }]
+            },
+            "topologyKey": "kubernetes.io/hostname"
+        }
+    }]
+    assert api_client.sanitize_for_serialization(make_pod(
+        name='test',
+        image_spec='jupyter/singleuser:latest',
+        cmd=['jupyterhub-singleuser'],
+        port=8888,
+        image_pull_policy='IfNotPresent',
+        pod_affinity_preferred=pod_affinity_preferred
+    )) == {
+        "metadata": {
+            "name": "test",
+            "labels": {},
+            "annotations": {}
+        },
+        "spec": {
+            "automountServiceAccountToken": False,
+            "securityContext": {},
+            "containers": [
+                {
+                    "env": [],
+                    "name": "notebook",
+                    "image": "jupyter/singleuser:latest",
+                    "imagePullPolicy": "IfNotPresent",
+                    "args": ["jupyterhub-singleuser"],
+                    "ports": [{
+                        "name": "notebook-port",
+                        "containerPort": 8888
+                    }],
+                    'volumeMounts': [],
+                    "resources": {
+                        "limits": {},
+                        "requests": {}
+                    }
+                }
+            ],
+            "volumes": [],
+            "affinity": {
+                "podAffinity": {
+                    "preferredDuringSchedulingIgnoredDuringExecution": pod_affinity_preferred
+                }
+            }
+        },
+        "kind": "Pod",
+        "apiVersion": "v1"
+    }
+
+
+def test_make_pod_with_pod_affinity_required():
+    """
+    Test specification of the simplest possible pod specification with non-empty pod_affinity_required
+    """
+    pod_affinity_required = [{
+        "labelSelector": {
+            "matchExpressions": [{
+                "key": "security",
+                "operator": "In",
+                "values": ["S1"],
+            }]
+        },
+        "topologyKey": "failure-domain.beta.kubernetes.io/zone"
+    }]
+    assert api_client.sanitize_for_serialization(make_pod(
+        name='test',
+        image_spec='jupyter/singleuser:latest',
+        cmd=['jupyterhub-singleuser'],
+        port=8888,
+        image_pull_policy='IfNotPresent',
+        pod_affinity_required=pod_affinity_required
+    )) == {
+        "metadata": {
+            "name": "test",
+            "labels": {},
+            "annotations": {}
+        },
+        "spec": {
+            "automountServiceAccountToken": False,
+            "securityContext": {},
+            "containers": [
+                {
+                    "env": [],
+                    "name": "notebook",
+                    "image": "jupyter/singleuser:latest",
+                    "imagePullPolicy": "IfNotPresent",
+                    "args": ["jupyterhub-singleuser"],
+                    "ports": [{
+                        "name": "notebook-port",
+                        "containerPort": 8888
+                    }],
+                    'volumeMounts': [],
+                    "resources": {
+                        "limits": {},
+                        "requests": {}
+                    }
+                }
+            ],
+            "volumes": [],
+            "affinity": {
+                "podAffinity": {
+                    "requiredDuringSchedulingIgnoredDuringExecution": pod_affinity_required
+                }
+            }
+        },
+        "kind": "Pod",
+        "apiVersion": "v1"
+    }
+
+
+def test_make_pod_with_pod_anti_affinity_preferred():
+    """
+    Test specification of the simplest possible pod specification with non-empty pod_anti_affinity_preferred
+    """
+    pod_anti_affinity_preferred = [{
+        "weight": 100,
+        "podAffinityTerm": {
+            "labelSelector": {
+                "matchExpressions": [{
+                    "key": "hub.jupyter.org/pod-kind",
+                    "operator": "In",
+                    "values": ["user"],
+                }]
+            },
+            "topologyKey": "kubernetes.io/hostname"
+        }
+    }]
+    assert api_client.sanitize_for_serialization(make_pod(
+        name='test',
+        image_spec='jupyter/singleuser:latest',
+        cmd=['jupyterhub-singleuser'],
+        port=8888,
+        image_pull_policy='IfNotPresent',
+        pod_anti_affinity_preferred=pod_anti_affinity_preferred
+    )) == {
+        "metadata": {
+            "name": "test",
+            "labels": {},
+            "annotations": {}
+        },
+        "spec": {
+            "automountServiceAccountToken": False,
+            "securityContext": {},
+            "containers": [
+                {
+                    "env": [],
+                    "name": "notebook",
+                    "image": "jupyter/singleuser:latest",
+                    "imagePullPolicy": "IfNotPresent",
+                    "args": ["jupyterhub-singleuser"],
+                    "ports": [{
+                        "name": "notebook-port",
+                        "containerPort": 8888
+                    }],
+                    'volumeMounts': [],
+                    "resources": {
+                        "limits": {},
+                        "requests": {}
+                    }
+                }
+            ],
+            "volumes": [],
+            "affinity": {
+                "podAntiAffinity": {
+                    "preferredDuringSchedulingIgnoredDuringExecution": pod_anti_affinity_preferred
+                }
+            }
+        },
+        "kind": "Pod",
+        "apiVersion": "v1"
+    }
+
+
+def test_make_pod_with_pod_anti_affinity_required():
+    """
+    Test specification of the simplest possible pod specification with non-empty pod_anti_affinity_required
+    """
+    pod_anti_affinity_required = [{
+        "labelSelector": {
+            "matchExpressions": [{
+                "key": "security",
+                "operator": "In",
+                "values": ["S1"],
+            }]
+        },
+        "topologyKey": "failure-domain.beta.kubernetes.io/zone"
+    }]
+    assert api_client.sanitize_for_serialization(make_pod(
+        name='test',
+        image_spec='jupyter/singleuser:latest',
+        cmd=['jupyterhub-singleuser'],
+        port=8888,
+        image_pull_policy='IfNotPresent',
+        pod_anti_affinity_required=pod_anti_affinity_required
+    )) == {
+        "metadata": {
+            "name": "test",
+            "labels": {},
+            "annotations": {}
+        },
+        "spec": {
+            "automountServiceAccountToken": False,
+            "securityContext": {},
+            "containers": [
+                {
+                    "env": [],
+                    "name": "notebook",
+                    "image": "jupyter/singleuser:latest",
+                    "imagePullPolicy": "IfNotPresent",
+                    "args": ["jupyterhub-singleuser"],
+                    "ports": [{
+                        "name": "notebook-port",
+                        "containerPort": 8888
+                    }],
+                    'volumeMounts': [],
+                    "resources": {
+                        "limits": {},
+                        "requests": {}
+                    }
+                }
+            ],
+            "volumes": [],
+            "affinity": {
+                "podAntiAffinity": {
+                    "requiredDuringSchedulingIgnoredDuringExecution": pod_anti_affinity_required
+                }
+            }
+        },
+        "kind": "Pod",
+        "apiVersion": "v1"
+    }
+
+
 def test_make_pod_with_priority_class_name():
     """
     Test specification of the simplest possible pod specification with non-default priorityClassName set