Add resilient streaming (#107)

* ft: add resilient watcher class * ft: fixes existing tests * ft: support custom reconnect exceptions * ft: add tests * ft: updates logs and docstrings * ft: add changelog entry * Update tests/test_events_replicator.py Co-authored-by: Alexander Streed <[email protected]> * Update tests/test_events_replicator.py Co-authored-by: Alexander Streed <[email protected]> * Update tests/test_worker.py Co-authored-by: Alexander Streed <[email protected]> * Update prefect_kubernetes/utilities.py Co-authored-by: Alexander Streed <[email protected]> * fix: add infrastructure kill tests * ft: update _stream form * ft: make _stream public, remove log_stream * ft: use a capped set for caching * Update prefect_kubernetes/utilities.py Co-authored-by: Alexander Streed <[email protected]> * Update prefect_kubernetes/utilities.py Co-authored-by: Alexander Streed <[email protected]> * Update prefect_kubernetes/utilities.py Co-authored-by: Alexander Streed <[email protected]> --------- Co-authored-by: Alexander Streed <[email protected]>
PrefectHQ · Jan 19, 2024 · f2e4d88 · f2e4d88
1 parent baeb9b4
commit f2e4d88
Show file tree

Hide file tree

Showing 10 changed files with 366 additions and 87 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- A `ResilientWatcher` utility class to reconnect Kubernetes client streams on `ProtocolErrors` - [#107](https://github.com/PrefectHQ/prefect-kubernetes/pull/107)
+
 ### Changed
 
 ### Deprecated

diff --git a/prefect_kubernetes/events.py b/prefect_kubernetes/events.py
@@ -1,10 +1,13 @@
 import atexit
+import logging
 import threading
 from typing import TYPE_CHECKING, Dict, List, Optional
 
 from prefect.events import Event, RelatedResource, emit_event
 from prefect.utilities.importtools import lazy_import
 
+from prefect_kubernetes.utilities import ResilientStreamWatcher
+
 if TYPE_CHECKING:
     import kubernetes
     import kubernetes.client
@@ -38,11 +41,13 @@ def __init__(
         worker_resource: Dict[str, str],
         related_resources: List[RelatedResource],
         timeout_seconds: int,
+        logger: Optional[logging.Logger] = None,
     ):
         self._client = client
         self._job_name = job_name
         self._namespace = namespace
         self._timeout_seconds = timeout_seconds
+        self._logger = logger
 
         # All events emitted by this replicator have the pod itself as the
         # resource. The `worker_resource` is what the worker uses when it's
@@ -52,7 +57,7 @@ def __init__(
         worker_related_resource = RelatedResource(__root__=worker_resource)
         self._related_resources = related_resources + [worker_related_resource]
 
-        self._watch = kubernetes.watch.Watch()
+        self._watch = ResilientStreamWatcher(logger=self._logger)
         self._thread = threading.Thread(target=self._replicate_pod_events)
 
         self._state = "READY"
@@ -90,7 +95,7 @@ def _replicate_pod_events(self):
 
         try:
             core_client = kubernetes.client.CoreV1Api(api_client=self._client)
-            for event in self._watch.stream(
+            for event in self._watch.api_object_stream(
                 func=core_client.list_namespaced_pod,
                 namespace=self._namespace,
                 label_selector=f"job-name={self._job_name}",

diff --git a/prefect_kubernetes/pods.py b/prefect_kubernetes/pods.py
@@ -2,11 +2,11 @@
 from typing import Any, Callable, Dict, Optional, Union
 
 from kubernetes.client.models import V1DeleteOptions, V1Pod, V1PodList
-from kubernetes.watch import Watch
 from prefect import task
 from prefect.utilities.asyncutils import run_sync_in_worker_thread
 
 from prefect_kubernetes.credentials import KubernetesCredentials
+from prefect_kubernetes.utilities import ResilientStreamWatcher
 
 
 @task
@@ -45,7 +45,6 @@ def kubernetes_orchestrator():
         ```
     """
     with kubernetes_credentials.get_client("core") as core_v1_client:
-
         return await run_sync_in_worker_thread(
             core_v1_client.create_namespaced_pod,
             namespace=namespace,
@@ -93,7 +92,6 @@ def kubernetes_orchestrator():
         ```
     """
     with kubernetes_credentials.get_client("core") as core_v1_client:
-
         return await run_sync_in_worker_thread(
             core_v1_client.delete_namespaced_pod,
             pod_name,
@@ -135,7 +133,6 @@ def kubernetes_orchestrator():
         ```
     """
     with kubernetes_credentials.get_client("core") as core_v1_client:
-
         return await run_sync_in_worker_thread(
             core_v1_client.list_namespaced_pod, namespace=namespace, **kube_kwargs
         )
@@ -180,7 +177,6 @@ def kubernetes_orchestrator():
         ```
     """
     with kubernetes_credentials.get_client("core") as core_v1_client:
-
         return await run_sync_in_worker_thread(
             core_v1_client.patch_namespaced_pod,
             name=pod_name,
@@ -224,7 +220,6 @@ def kubernetes_orchestrator():
         ```
     """
     with kubernetes_credentials.get_client("core") as core_v1_client:
-
         return await run_sync_in_worker_thread(
             core_v1_client.read_namespaced_pod,
             name=pod_name,
@@ -281,11 +276,11 @@ def kubernetes_orchestrator():
         ```
     """
     with kubernetes_credentials.get_client("core") as core_v1_client:
-
         if print_func is not None:
             # should no longer need to manually refresh on ApiException.status == 410
             # as of https://github.com/kubernetes-client/python-base/pull/133
-            for log_line in Watch().stream(
+            watcher = ResilientStreamWatcher()
+            for log_line in watcher.stream(
                 core_v1_client.read_namespaced_pod_log,
                 name=pod_name,
                 namespace=namespace,
@@ -341,7 +336,6 @@ def kubernetes_orchestrator():
         ```
     """
     with kubernetes_credentials.get_client("core") as core_v1_client:
-
         return await run_sync_in_worker_thread(
             core_v1_client.replace_namespaced_pod,
             body=new_pod,

diff --git a/prefect_kubernetes/utilities.py b/prefect_kubernetes/utilities.py
@@ -1,7 +1,11 @@
 """ Utilities for working with the Python Kubernetes API. """
+import logging
+import time
 from pathlib import Path
-from typing import Optional, TypeVar, Union
+from typing import Callable, List, Optional, Set, Type, TypeVar, Union
 
+import urllib3
+from kubernetes import watch
 from kubernetes.client import models as k8s_models
 from prefect.infrastructure.kubernetes import KubernetesJob, KubernetesManifest
 from slugify import slugify
@@ -13,6 +17,24 @@
 V1KubernetesModel = TypeVar("V1KubernetesModel")
 
 
+class _CappedSet(set):
+    """
+    A set with a bounded size.
+    """
+
+    def __init__(self, maxsize):
+        super().__init__()
+        self.maxsize = maxsize
+
+    def add(self, value):
+        """
+        Add to the set and maintain its max size.
+        """
+        if len(self) >= self.maxsize:
+            self.pop()
+        super().add(value)
+
+
 def convert_manifest_to_model(
     manifest: Union[Path, str, KubernetesManifest], v1_model_name: str
 ) -> V1KubernetesModel:
@@ -180,3 +202,123 @@ def _slugify_label_value(value: str, max_length: int = 63) -> str:
     # Kubernetes to throw the validation error
 
     return slug
+
+
+class ResilientStreamWatcher:
+    """
+    A wrapper class around kuberenetes.watch.Watch that will reconnect on
+    certain exceptions.
+    """
+
+    DEFAULT_RECONNECT_EXCEPTIONS = (urllib3.exceptions.ProtocolError,)
+
+    def __init__(
+        self,
+        logger: Optional[logging.Logger] = None,
+        max_cache_size: int = 50000,
+        reconnect_exceptions: Optional[List[Type[Exception]]] = None,
+    ) -> None:
+        """
+        A utility class for managing streams of Kuberenetes API objects and logs
+
+        Attributes:
+            logger: A logger which will be used interally to log errors
+            max_cache_size: The maximum number of API objects to track in an
+                internal cache to help deduplicate results on stream reconnects
+            reconnect_exceptions: A list of exceptions that will cause the stream
+                to reconnect.
+        """
+
+        self.max_cache_size = max_cache_size
+        self.logger = logger
+        self.watch = watch.Watch()
+
+        reconnect_exceptions = (
+            reconnect_exceptions
+            if reconnect_exceptions is not None
+            else self.DEFAULT_RECONNECT_EXCEPTIONS
+        )
+        self.reconnect_exceptions = tuple(reconnect_exceptions)
+
+    def stream(self, func: Callable, *args, cache: Optional[Set] = None, **kwargs):
+        """
+        A method for streaming API objects or logs from a Kubernetes
+        client function. This method will reconnect the stream on certain
+        configurable exceptions and deduplicate results on reconnects if
+        streaming API objects and a cache is provided.
+
+        Note that client functions that produce a stream will
+        restart a stream from the beginning of the log's history on reconnect.
+        If a cache is not provided, it is possible for duplicate entries to be yielded.
+
+        Args:
+            func: A Kubernetes client function to call which produces a stream
+                of logs
+            *args: Positional arguments to pass to `func`
+            cache: A keyward argument that provides a way to deduplicate
+                results on reconnects and bound
+            **kwargs: Keyword arguments to pass to `func`
+
+        Returns:
+            An iterator of log
+        """
+        keep_streaming = True
+        while keep_streaming:
+            try:
+                for event in self.watch.stream(func, *args, **kwargs):
+                    # check that we want to and can track this object
+                    if (
+                        cache is not None
+                        and isinstance(event, dict)
+                        and "object" in event
+                    ):
+                        uid = event["object"].metadata.uid
+                        if uid not in cache:
+                            cache.add(uid)
+                            yield event
+                    else:
+                        yield event
+                else:
+                    # Case: we've finished iterating
+                    keep_streaming = False
+            except self.reconnect_exceptions:
+                # Case: We've hit an exception we're willing to retry on
+                if self.logger:
+                    self.logger.error("Unable to connect, retrying...", exc_info=True)
+                time.sleep(1)
+            except Exception:
+                # Case: We hit an exception we're unwilling to retry on
+                if self.logger:
+                    self.logger.exception(
+                        f"Unexpected error while streaming {func.__name__}"
+                    )
+                keep_streaming = False
+                self.stop()
+                raise
+
+        self.stop()
+
+    def api_object_stream(self, func: Callable, *args, **kwargs):
+        """
+        Create a cache to maintain a record of API objects that have been
+        seen. This is useful because `stream` will reconnect a stream on
+        `self.reconnect_exceptions` and on reconnect it will restart streaming all
+        objects. This cache prevents the same object from being yielded twice.
+
+        Args:
+            func: A Kubernetes client function to call which produces a stream of API o
+                bjects
+            *args: Positional arguments to pass to `func`
+            **kwargs: Keyword arguments to pass to `func`
+
+        Returns:
+            An iterator of API objects
+        """
+        cache = _CappedSet(self.max_cache_size)
+        yield from self.stream(func, *args, cache=cache, **kwargs)
+
+    def stop(self):
+        """
+        Shut down the internal Watch object.
+        """
+        self.watch.stop()
diff --git a/prefect_kubernetes/worker.py b/prefect_kubernetes/worker.py
@@ -144,6 +144,7 @@
 
 from prefect_kubernetes.events import KubernetesEventsReplicator
 from prefect_kubernetes.utilities import (
+    ResilientStreamWatcher,
     _slugify_label_key,
     _slugify_label_value,
     _slugify_name,
@@ -575,7 +576,6 @@ async def run(
                 task_status.started(pid)
 
             # Monitor the job until completion
-
             events_replicator = KubernetesEventsReplicator(
                 client=client,
                 job_name=job.metadata.name,
@@ -585,6 +585,7 @@ async def run(
                     configuration=configuration
                 ),
                 timeout_seconds=configuration.pod_watch_timeout_seconds,
+                logger=logger,
             )
 
             with events_replicator:
@@ -910,15 +911,16 @@ def _watch_job(
 
         if configuration.stream_output:
             with self._get_core_client(client) as core_client:
-                logs = core_client.read_namespaced_pod_log(
-                    pod.metadata.name,
-                    configuration.namespace,
-                    follow=True,
-                    _preload_content=False,
-                    container="prefect-job",
-                )
+                watch = ResilientStreamWatcher(logger=logger)
                 try:
-                    for log in logs.stream():
+                    for log in watch.stream(
+                        core_client.read_namespaced_pod_log,
+                        pod.metadata.name,
+                        configuration.namespace,
+                        follow=True,
+                        _preload_content=False,
+                        container="prefect-job",
+                    ):
                         print(log.decode().rstrip())
 
                         # Check if we have passed the deadline and should stop streaming
@@ -928,7 +930,6 @@ def _watch_job(
                         )
                         if deadline and remaining_time <= 0:
                             break
-
                 except Exception:
                     logger.warning(
                         (
@@ -957,15 +958,15 @@ def _watch_job(
                     )
                     return -1
 
-                watch = kubernetes.watch.Watch()
+                watch = ResilientStreamWatcher(logger=logger)
                 # The kubernetes library will disable retries if the timeout kwarg is
                 # present regardless of the value so we do not pass it unless given
                 # https://github.com/kubernetes-client/python/blob/84f5fea2a3e4b161917aa597bf5e5a1d95e24f5a/kubernetes/base/watch/watch.py#LL160
                 timeout_seconds = (
                     {"timeout_seconds": remaining_time} if deadline else {}
                 )
 
-                for event in watch.stream(
+                for event in watch.api_object_stream(
                     func=batch_client.list_namespaced_job,
                     field_selector=f"metadata.name={job_name}",
                     namespace=configuration.namespace,
@@ -1065,12 +1066,12 @@ def _get_job_pod(
         """Get the first running pod for a job."""
         from kubernetes.client.models import V1Pod
 
-        watch = kubernetes.watch.Watch()
+        watch = ResilientStreamWatcher(logger=logger)
         logger.debug(f"Job {job_name!r}: Starting watch for pod start...")
         last_phase = None
         last_pod_name: Optional[str] = None
         with self._get_core_client(client) as core_client:
-            for event in watch.stream(
+            for event in watch.api_object_stream(
                 func=core_client.list_namespaced_pod,
                 namespace=configuration.namespace,
                 label_selector=f"job-name={job_name}",