[MRG] add "sticky builds" functionality (#949)

[MRG] add "sticky builds" functionality

Author: betatim

binderhub/app.py

@@ -196,6 +196,19 @@ def _valid_badge_base_url(self, proposal):
         config=True,
     )
 
+    sticky_builds = Bool(
+        False,
+        help="""
+        Attempt to assign builds for the same repository to the same node.
+
+        In order to speed up re-builds of a repository all its builds will
+        be assigned to the same node in the cluster.
+
+        Note: This feature only works if you also enable docker-in-docker support.
+        """,
+        config=True,
+    )
+
     use_registry = Bool(
         True,
         help="""
@@ -556,6 +569,7 @@ def initialize(self, *args, **kwargs):
             "build_image": self.build_image,
             'build_node_selector': self.build_node_selector,
             'build_pool': self.build_pool,
+            "sticky_builds": self.sticky_builds,
             'log_tail_lines': self.log_tail_lines,
             'pod_quota': self.pod_quota,
             'per_repo_quota': self.per_repo_quota,

binderhub/build.py

@@ -12,6 +12,8 @@
 from tornado.ioloop import IOLoop
 from tornado.log import app_log
 
+from .utils import rendezvous_rank
+
 
 class Build:
     """Represents a build of a git repository into a docker image.
@@ -36,7 +38,7 @@ class Build:
     """
     def __init__(self, q, api, name, namespace, repo_url, ref, git_credentials, build_image,
                  image_name, push_secret, memory_limit, docker_host, node_selector,
-                 appendix='', log_tail_lines=100):
+                 appendix='', log_tail_lines=100, sticky_builds=False):
         self.q = q
         self.api = api
         self.repo_url = repo_url
@@ -56,6 +58,10 @@ def __init__(self, q, api, name, namespace, repo_url, ref, git_credentials, buil
         self.stop_event = threading.Event()
         self.git_credentials = git_credentials
 
+        self.sticky_builds = sticky_builds
+
+        self._component_label = "binderhub-build"
+
     def get_cmd(self):
         """Get the cmd to run to build the image"""
         cmd = [
@@ -144,8 +150,71 @@ def progress(self, kind, obj):
         """Put the current action item into the queue for execution."""
         self.main_loop.add_callback(self.q.put, {'kind': kind, 'payload': obj})
 
+    def get_affinity(self):
+        """Determine the affinity term for the build pod.
+
+        There are a two affinity strategies, which one is used depends on how
+        the BinderHub is configured.
+
+        In the default setup the affinity of each build pod is an "anti-affinity"
+        which causes the pods to prefer to schedule on separate nodes.
+
+        In a setup with docker-in-docker enabled pods for a particular
+        repository prefer to schedule on the same node in order to reuse the
+        docker layer cache of previous builds.
+        """
+        dind_pods = self.api.list_namespaced_pod(
+            self.namespace,
+            label_selector="component=dind,app=binder",
+        )
+
+        if self.sticky_builds and dind_pods:
+            node_names = [pod.spec.node_name for pod in dind_pods.items]
+            ranked_nodes = rendezvous_rank(node_names, self.repo_url)
+            best_node_name = ranked_nodes[0]
+
+            affinity = client.V1Affinity(
+                node_affinity=client.V1NodeAffinity(
+                    preferred_during_scheduling_ignored_during_execution=[
+                        client.V1PreferredSchedulingTerm(
+                            weight=100,
+                            preference=client.V1NodeSelectorTerm(
+                                match_fields=[
+                                    client.V1NodeSelectorRequirement(
+                                        key="metadata.name",
+                                        operator="In",
+                                        values=[best_node_name],
+                                    )
+                                ]
+                            ),
+                        )
+                    ]
+                )
+            )
+
+        else:
+            affinity = client.V1Affinity(
+                pod_anti_affinity=client.V1PodAntiAffinity(
+                    preferred_during_scheduling_ignored_during_execution=[
+                        client.V1WeightedPodAffinityTerm(
+                            weight=100,
+                            pod_affinity_term=client.V1PodAffinityTerm(
+                                topology_key="kubernetes.io/hostname",
+                                label_selector=client.V1LabelSelector(
+                                    match_labels=dict(
+                                        component=self._component_label
+                                    )
+                                )
+                            )
+                        )
+                    ]
+                )
+            )
+
+        return affinity
+
     def submit(self):
-        """Submit a image spec to openshift's s2i and wait for completion """
+        """Submit a build pod to create the image for the repository."""
         volume_mounts = [
             client.V1VolumeMount(mount_path="/var/run/docker.sock", name="docker-socket")
         ]
@@ -166,13 +235,12 @@ def submit(self):
         if self.git_credentials:
             env.append(client.V1EnvVar(name='GIT_CREDENTIAL_ENV', value=self.git_credentials))
 
-        component_label = "binderhub-build"
         self.pod = client.V1Pod(
             metadata=client.V1ObjectMeta(
                 name=self.name,
                 labels={
                     "name": self.name,
-                    "component": component_label,
+                    "component": self._component_label,
                 },
                 annotations={
                     "binder-repo": self.repo_url,
@@ -211,23 +279,7 @@ def submit(self):
                 node_selector=self.node_selector,
                 volumes=volumes,
                 restart_policy="Never",
-                affinity=client.V1Affinity(
-                    pod_anti_affinity=client.V1PodAntiAffinity(
-                        preferred_during_scheduling_ignored_during_execution=[
-                            client.V1WeightedPodAffinityTerm(
-                                weight=100,
-                                pod_affinity_term=client.V1PodAffinityTerm(
-                                    topology_key="kubernetes.io/hostname",
-                                    label_selector=client.V1LabelSelector(
-                                        match_labels=dict(
-                                            component=component_label
-                                        )
-                                    )
-                                )
-                            )
-                        ]
-                    )
-                )
+                affinity=self.get_affinity()
             )
         )
 

binderhub/builder.py

@@ -340,7 +340,8 @@ async def get(self, provider_prefix, _unescaped_spec):
             node_selector=self.settings['build_node_selector'],
             appendix=appendix,
             log_tail_lines=self.settings['log_tail_lines'],
-            git_credentials=provider.git_credentials
+            git_credentials=provider.git_credentials,
+            sticky_builds=self.settings['sticky_builds'],
         )
 
         with BUILDS_INPROGRESS.track_inprogress():

binderhub/tests/test_build.py

@@ -2,12 +2,15 @@
 
 import json
 import sys
+from collections import namedtuple
 from unittest import mock
 from urllib.parse import quote
 
 import pytest
 from tornado.httputil import url_concat
 
+from kubernetes import client
+
 from binderhub.build import Build
 from .utils import async_requests
 
@@ -51,6 +54,57 @@ async def test_build(app, needs_build, needs_launch, always_build, slug, pytestc
     assert r.url.startswith(final['url'])
 
 
+def test_default_affinity():
+    # check that the default affinity is a pod anti-affinity
+    build = Build(
+        mock.MagicMock(), api=mock.MagicMock(), name='test_build',
+        namespace='build_namespace', repo_url=mock.MagicMock(),
+        ref=mock.MagicMock(), build_image=mock.MagicMock(),
+        image_name=mock.MagicMock(), push_secret=mock.MagicMock(),
+        memory_limit=mock.MagicMock(), git_credentials=None,
+        docker_host='http://mydockerregistry.local',
+        node_selector=mock.MagicMock())
+
+    affinity = build.get_affinity()
+
+    assert isinstance(affinity, client.V1Affinity)
+    assert affinity.node_affinity is None
+    assert affinity.pod_affinity is None
+    assert affinity.pod_anti_affinity is not None
+
+
+def test_sticky_builds_affinity():
+    # Setup some mock objects for the response from the k8s API
+    Pod = namedtuple("Pod", "spec")
+    PodSpec = namedtuple("PodSpec", "node_name")
+    PodList = namedtuple("PodList", "items")
+
+    mock_k8s_api = mock.MagicMock()
+    mock_k8s_api.list_namespaced_pod.return_value = PodList(
+        [Pod(PodSpec("node-a")), Pod(PodSpec("node-b"))],
+    )
+
+    build = Build(
+        mock.MagicMock(), api=mock_k8s_api, name='test_build',
+        namespace='build_namespace', repo_url=mock.MagicMock(),
+        ref=mock.MagicMock(), build_image=mock.MagicMock(),
+        image_name=mock.MagicMock(), push_secret=mock.MagicMock(),
+        memory_limit=mock.MagicMock(), git_credentials=None,
+        docker_host='http://mydockerregistry.local',
+        node_selector=mock.MagicMock(),
+        sticky_builds=True)
+
+    affinity = build.get_affinity()
+
+    assert isinstance(affinity, client.V1Affinity)
+    assert affinity.node_affinity is not None
+    assert affinity.pod_affinity is None
+    assert affinity.pod_anti_affinity is None
+
+    # One of the two nodes we have in our mock should be the preferred node
+    assert affinity.node_affinity.preferred_during_scheduling_ignored_during_execution[0].preference.match_fields[0].values[0] in ("node-a", "node-b")
+
+
 def test_git_credentials_passed_to_podspec_upon_submit():
     git_credentials = {
         'client_id': 'my_username',

binderhub/tests/test_utils.py

@@ -0,0 +1,59 @@
+from binderhub import utils
+
+
+def test_rendezvous_rank():
+    # check that a key doesn't move if its assigned bucket remains but the
+    # other buckets are removed
+    key = "crazy frog is a crazy key"
+    first_round = utils.rendezvous_rank(["b1", "b2", "b3"], key)
+    second_round = utils.rendezvous_rank([first_round[0], first_round[1]], key)
+
+    assert first_round[0] == second_round[0], key
+
+
+def test_rendezvous_independence():
+    # check that the relative ranking of 80 buckets doesn't depend on the
+    # presence of 20 extra buckets
+    key = "k1"
+    eighty_buckets = utils.rendezvous_rank(["b%i" % i for i in range(80)], key)
+    hundred_buckets = utils.rendezvous_rank(["b%i" % i for i in range(100)], key)
+
+    for i in range(80, 100):
+        hundred_buckets.remove("b%i" % i)
+
+    assert eighty_buckets == hundred_buckets
+
+
+def test_rendezvous_redistribution():
+    # check that approximately a third of keys move to the new bucket
+    # when one is added
+    n_keys = 3000
+
+    # count how many keys were moved, which bucket a key started from and
+    # which bucket a key was moved from (to the new bucket)
+    n_moved = 0
+    from_bucket = {"b1": 0, "b2": 0}
+    start_in = {"b1": 0, "b2": 0}
+
+    for i in range(n_keys):
+        key = f"key-{i}"
+        two_buckets = utils.rendezvous_rank(["b1", "b2"], key)
+        start_in[two_buckets[0]] += 1
+        three_buckets = utils.rendezvous_rank(["b1", "b2", "b3"], key)
+
+        if two_buckets[0] != three_buckets[0]:
+            n_moved += 1
+            from_bucket[two_buckets[0]] += 1
+
+            # should always move to the newly added bucket
+            assert three_buckets[0] == "b3"
+
+    # because of statistical fluctuations we have to leave some room when
+    # making this comparison
+    assert 0.31 < n_moved / n_keys < 0.35
+    # keys should move from the two original buckets with approximately
+    # equal probability. We pick 30 because it is "about right"
+    assert abs(from_bucket["b1"] - from_bucket["b2"]) < 30
+    # the initial distribution of keys should be roughly the same
+    # We pick 30 because it is "about right"
+    assert abs(start_in["b1"] - start_in["b2"]) < 30