Merge pull request #644 from dolfinus/add_pod_post_create_hook

Add after_pod_created_hook

kubespawner/spawner.py

@@ -1022,6 +1022,26 @@ def _validate_image_pull_secrets(self, proposal):
         """,
     )
 
+    after_pod_created_hook = Callable(
+        None,
+        allow_none=True,
+        config=True,
+        help="""
+        Callable to augment the Pod object after launching.
+
+        Expects a callable that takes two parameters:
+
+           1. The spawner object that is doing the spawning
+           2. The Pod object that was launched
+
+        This can be a coroutine if necessary. When set to none, no augmenting is done.
+
+        This is very useful if you want to add some services or ingress to the pod after it is launched.
+        Note that the spawner object can change between versions of KubeSpawner and JupyterHub,
+        so be careful relying on this!
+        """,
+    )
+
     volumes = List(
         config=True,
         help="""
@@ -2656,6 +2676,7 @@ async def _start(self):
         # try again. We try 4 times, and if it still fails we give up.
         pod = await self.get_pod_manifest()
         if self.modify_pod_hook:
+            self.log.info('Pod is being modified via modify_pod_hook')
             pod = await maybe_future(self.modify_pod_hook(self, pod))
 
         ref_key = f"{self.namespace}/{self.pod_name}"
@@ -2666,7 +2687,7 @@ async def _start(self):
             timeout=self.k8s_api_request_retry_timeout,
         )
 
-        if self.internal_ssl or self.services_enabled:
+        if self.internal_ssl or self.services_enabled or self.after_pod_created_hook:
             try:
                 # wait for pod to have uid,
                 # required for creating owner reference
@@ -2702,21 +2723,28 @@ async def _start(self):
                         f"Failed to create secret {secret_manifest.metadata.name}",
                     )
 
-                service_manifest = self.get_service_manifest(owner_reference)
-                await exponential_backoff(
-                    partial(
-                        self._ensure_not_exists,
-                        "service",
-                        service_manifest.metadata.name,
-                    ),
-                    f"Failed to delete service {service_manifest.metadata.name}",
-                )
-                await exponential_backoff(
-                    partial(
-                        self._make_create_resource_request, "service", service_manifest
-                    ),
-                    f"Failed to create service {service_manifest.metadata.name}",
-                )
+                if self.internal_ssl or self.services_enabled:
+                    service_manifest = self.get_service_manifest(owner_reference)
+                    await exponential_backoff(
+                        partial(
+                            self._ensure_not_exists,
+                            "service",
+                            service_manifest.metadata.name,
+                        ),
+                        f"Failed to delete service {service_manifest.metadata.name}",
+                    )
+                    await exponential_backoff(
+                        partial(
+                            self._make_create_resource_request,
+                            "service",
+                            service_manifest,
+                        ),
+                        f"Failed to create service {service_manifest.metadata.name}",
+                    )
+
+                if self.after_pod_created_hook:
+                    self.log.info('Executing after_pod_created_hook')
+                    await maybe_future(self.after_pod_created_hook(self, pod))
             except Exception:
                 # cleanup on failure and re-raise
                 await self.stop(True)

tests/test_spawner.py

@@ -1,11 +1,13 @@
 import asyncio
 import json
 import os
+from functools import partial
 from unittest.mock import Mock
 
 import pytest
 from jupyterhub.objects import Hub, Server
 from jupyterhub.orm import Spawner
+from jupyterhub.utils import exponential_backoff
 from kubernetes_asyncio.client.models import (
     V1Capabilities,
     V1Container,
@@ -16,6 +18,7 @@
 from traitlets.config import Config
 
 from kubespawner import KubeSpawner
+from kubespawner.objects import make_owner_reference, make_service
 
 
 class MockUser(Mock):
@@ -298,6 +301,88 @@ async def test_spawn_internal_ssl(
     assert service_name not in service_names
 
 
+async def test_spawn_after_pod_created_hook(
+    kube_ns,
+    kube_client,
+    hub,
+    config,
+):
+    async def after_pod_created_hook(spawner: KubeSpawner, pod: dict):
+        owner_reference = make_owner_reference(spawner.pod_name, pod["metadata"]["uid"])
+
+        labels = spawner._build_common_labels(spawner._expand_all(spawner.extra_labels))
+        annotations = spawner._build_common_annotations(
+            spawner._expand_all(spawner.extra_annotations)
+        )
+
+        service_manifest = make_service(
+            name=spawner.pod_name + "-hook",
+            port=spawner.port,
+            servername=spawner.name,
+            owner_references=[owner_reference],
+            labels=labels,
+            annotations=annotations,
+        )
+
+        await exponential_backoff(
+            partial(
+                spawner._ensure_not_exists,
+                "service",
+                service_manifest.metadata.name,
+            ),
+            f"Failed to delete service {service_manifest.metadata.name}",
+        )
+        await exponential_backoff(
+            partial(spawner._make_create_resource_request, "service", service_manifest),
+            f"Failed to create service {service_manifest.metadata.name}",
+        )
+
+    spawner = KubeSpawner(
+        config=config,
+        hub=hub,
+        user=MockUser(name="ssl"),
+        api_token="abc123",
+        oauth_client_id="unused",
+        after_pod_created_hook=after_pod_created_hook,
+    )
+    # start the spawner
+    await spawner.start()
+    pod_name = "jupyter-%s" % spawner.user.name
+    # verify the pod exists
+    pods = (await kube_client.list_namespaced_pod(kube_ns)).items
+    pod_names = [p.metadata.name for p in pods]
+    assert pod_name in pod_names
+    # verify poll while running
+    status = await spawner.poll()
+    assert status is None
+
+    # verify service exist
+    service_name = pod_name + "-hook"
+    services = (await kube_client.list_namespaced_service(kube_ns)).items
+    service_names = [s.metadata.name for s in services]
+    assert service_name in service_names
+
+    # stop the pod
+    await spawner.stop()
+
+    # verify pod is gone
+    pods = (await kube_client.list_namespaced_pod(kube_ns)).items
+    pod_names = [p.metadata.name for p in pods]
+    assert "jupyter-%s" % spawner.user.name not in pod_names
+
+    # verify service is gone
+    # it may take a little while for them to get cleaned up
+    for _ in range(5):
+        services = (await kube_client.list_namespaced_service(kube_ns)).items
+        service_names = {s.metadata.name for s in services}
+        if service_name in service_names:
+            await asyncio.sleep(1)
+        else:
+            break
+
+    assert service_name not in service_names
+
+
 async def test_spawn_progress(kube_ns, kube_client, config, hub_pod, hub):
     spawner = KubeSpawner(
         hub=hub,