<commit_before>"""
JupyterHub Spawner to spawn user notebooks on a Kubernetes cluster.

This module exports `KubeSpawner` class, which is the actual spawner
implementation that should be used by JupyterHub.
"""
import os
import json
import time
import string
import threading
from urllib.parse import urlparse, urlunparse
from concurrent.futures import ThreadPoolExecutor
import multiprocessing


from tornado import gen
from tornado.concurrent import run_on_executor
from traitlets.config import SingletonConfigurable
from traitlets import Type, Unicode, List, Integer, Union, Dict, Bool, Any
from jupyterhub.spawner import Spawner
from jupyterhub.traitlets import Command
from kubernetes.client.models.v1_volume import V1Volume
from kubernetes.client.models.v1_volume_mount import V1VolumeMount
from kubernetes.client.rest import ApiException
from kubernetes import client, config, watch
import escapism

from kubespawner.traitlets import Callable
from kubespawner.utils import request_maker, k8s_url
from kubespawner.objects import make_pod, make_pvc

class PodReflector(SingletonConfigurable):
    """
    Local up-to-date copy of a set of kubernetes pods
    """
    labels = Dict(
        {
            'heritage': 'jupyterhub',
            'component': 'singleuser-server',
        },
        config=True,
        help="""
        Labels to reflect onto local cache
        """
    )

    namespace = Unicode(
        None,
        allow_none=True,
        help="""
        Namespace to watch for resources in
        """
    )

    pods = Dict(
        {},
        help="""
        Dictionary of pod names to Pod objects.

        This can be directly accessed from multiple threads.
        """
    )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Load kubernetes config here, since this is a Singleton and
        # so this __init__ will be run way before anything else gets run.
        try:
            config.load_incluster_config()
        except config.ConfigException:
            config.load_kube_config()
        self.api = client.CoreV1Api()

        # FIXME: Protect against malicious labels?
        self.label_selector = ','.join(['{}={}'.format(k, v) for k, v in self.labels.items()])

        self.start()

    def _list_and_update(self):
        """
        Update current list of pods by doing a full fetch of pod information.

        Overwrites all current pod info.
        """
        initial_pods = self.api.list_namespaced_pod(
            self.namespace,
            label_selector=self.label_selector
        )
        # This is an atomic operation on the dictionary!
        self.pods = {p.metadata.name: p for p in initial_pods.items}

    def _watch_and_update(self):
        """
        Keeps the current list of pods up-to-date

        This method is to be run not on the main thread!

        We first fetch the list of current pods, and store that. Then we
        register to be notified of changes to those pods, and keep our
        local store up-to-date based on these notifications.

        We also perform exponential backoff, giving up after we hit 32s
        wait time. This should protect against network connections dropping
        and intermittent unavailability of the api-server. Every time we
        recover from an exception we also do a full fetch, to pick up
        changes that might've been missed in the time we were not doing
        a watch.

        Note that we're playing a bit with fire here, by updating a dictionary
        in this thread while it is probably being read in another thread
        without using locks! However, dictionary access itself is atomic,
        and as long as we don't try to mutate them (do a 'fetch / modify /
        update' cycle on them), we should be ok!
        """
        cur_delay = 0.1
        while True:
            self.log.info("watching for pods with label selector %s in namespace %s", self.label_selector, self.namespace)
            try:
                self._list_and_update()
                w = watch.Watch()
                for ev in w.stream(
                        self.api.list_namespaced_pod,
                        self.namespace,
                        label_selector=self.label_selector
                ):
                    cur_delay = 0.1
                    pod = ev['object']
                    if ev['type'] == 'DELETED':
                        try:
                            # This is an atomic operation on the dictionary!
                            del self.pods[pod.metadata.name]
                        except KeyError:
                            # Race, somehow another thread deleted this?
                            pass
                    else:
                        # This is an atomic operation on the dictionary!
                        self.pods[pod.metadata.name] = pod
            except:
                cur_delay = cur_delay * 2
                self.log.exception("Error when watching pods, retrying in %ss", cur_delay)
                time.sleep(cur_delay)
                if cur_delay > 30:
                    self.log.fatal('Unable to connect to APIServer')
                    sys.exit(1)
                continue
            finally:
                w.stop()

    def start(self):
        """
        Start the reflection process!

        We'll do a blocking read of all pods first, so that we don't
        race with any operations that are checking the state of the pod
        store - such as polls. This should be called only once at the
        start of program initialization (when the singleton is being created),
        and not afterwards!
        """
        if hasattr(self, 'watch_thread'):
            raise ValueError('Thread watching for pods is already running')

        self._list_and_update()
        self.watch_thread = threading.Thread(target=self._watch_and_update)
        # If the watch_thread is only thread left alive, exit app
        self.watch_thread.daemon = True
        self.watch_thread.start()

class KubeSpawner(Spawner):
    """
    Implement a JupyterHub spawner to spawn pods in a Kubernetes Cluster.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # By now, all the traitlets have been set, so we can use them to compute
        # other attributes
        self.executor = ThreadPoolExecutor(max_workers=self.k8s_api_threadpool_workers)

        # FIXME: Is this a memory leak? Will there be other side effects?
        self.pod_reflector = PodReflector.instance(parent=self, namespace=self.namespace)<commit_msg>Use one threadpool for all spawners<commit_after>"""
JupyterHub Spawner to spawn user notebooks on a Kubernetes cluster.

This module exports `KubeSpawner` class, which is the actual spawner
implementation that should be used by JupyterHub.
"""
import os
import json
import time
import string
import threading
from urllib.parse import urlparse, urlunparse
from concurrent.futures import ThreadPoolExecutor
import multiprocessing


from tornado import gen
from tornado.concurrent import run_on_executor
from traitlets.config import SingletonConfigurable
from traitlets import Type, Unicode, List, Integer, Union, Dict, Bool, Any
from jupyterhub.spawner import Spawner
from jupyterhub.traitlets import Command
from kubernetes.client.models.v1_volume import V1Volume
from kubernetes.client.models.v1_volume_mount import V1VolumeMount
from kubernetes.client.rest import ApiException
from kubernetes import client, config, watch
import escapism

from kubespawner.traitlets import Callable
from kubespawner.utils import request_maker, k8s_url
from kubespawner.objects import make_pod, make_pvc

class PodReflector(SingletonConfigurable):
    """
    Local up-to-date copy of a set of kubernetes pods
    """
    labels = Dict(
        {
            'heritage': 'jupyterhub',
            'component': 'singleuser-server',
        },
        config=True,
        help="""
        Labels to reflect onto local cache
        """
    )

    namespace = Unicode(
        None,
        allow_none=True,
        help="""
        Namespace to watch for resources in
        """
    )

    pods = Dict(
        {},
        help="""
        Dictionary of pod names to Pod objects.

        This can be directly accessed from multiple threads.
        """
    )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Load kubernetes config here, since this is a Singleton and
        # so this __init__ will be run way before anything else gets run.
        try:
            config.load_incluster_config()
        except config.ConfigException:
            config.load_kube_config()
        self.api = client.CoreV1Api()

        # FIXME: Protect against malicious labels?
        self.label_selector = ','.join(['{}={}'.format(k, v) for k, v in self.labels.items()])

        self.start()

    def _list_and_update(self):
        """
        Update current list of pods by doing a full fetch of pod information.

        Overwrites all current pod info.
        """
        initial_pods = self.api.list_namespaced_pod(
            self.namespace,
            label_selector=self.label_selector
        )
        # This is an atomic operation on the dictionary!
        self.pods = {p.metadata.name: p for p in initial_pods.items}

    def _watch_and_update(self):
        """
        Keeps the current list of pods up-to-date

        This method is to be run not on the main thread!

        We first fetch the list of current pods, and store that. Then we
        register to be notified of changes to those pods, and keep our
        local store up-to-date based on these notifications.

        We also perform exponential backoff, giving up after we hit 32s
        wait time. This should protect against network connections dropping
        and intermittent unavailability of the api-server. Every time we
        recover from an exception we also do a full fetch, to pick up
        changes that might've been missed in the time we were not doing
        a watch.

        Note that we're playing a bit with fire here, by updating a dictionary
        in this thread while it is probably being read in another thread
        without using locks! However, dictionary access itself is atomic,
        and as long as we don't try to mutate them (do a 'fetch / modify /
        update' cycle on them), we should be ok!
        """
        cur_delay = 0.1
        while True:
            self.log.info("watching for pods with label selector %s in namespace %s", self.label_selector, self.namespace)
            try:
                self._list_and_update()
                w = watch.Watch()
                for ev in w.stream(
                        self.api.list_namespaced_pod,
                        self.namespace,
                        label_selector=self.label_selector
                ):
                    cur_delay = 0.1
                    pod = ev['object']
                    if ev['type'] == 'DELETED':
                        try:
                            # This is an atomic operation on the dictionary!
                            del self.pods[pod.metadata.name]
                        except KeyError:
                            # Race, somehow another thread deleted this?
                            pass
                    else:
                        # This is an atomic operation on the dictionary!
                        self.pods[pod.metadata.name] = pod
            except:
                cur_delay = cur_delay * 2
                self.log.exception("Error when watching pods, retrying in %ss", cur_delay)
                time.sleep(cur_delay)
                if cur_delay > 30:
                    self.log.fatal('Unable to connect to APIServer')
                    sys.exit(1)
                continue
            finally:
                w.stop()

    def start(self):
        """
        Start the reflection process!

        We'll do a blocking read of all pods first, so that we don't
        race with any operations that are checking the state of the pod
        store - such as polls. This should be called only once at the
        start of program initialization (when the singleton is being created),
        and not afterwards!
        """
        if hasattr(self, 'watch_thread'):
            raise ValueError('Thread watching for pods is already running')

        self._list_and_update()
        self.watch_thread = threading.Thread(target=self._watch_and_update)
        # If the watch_thread is only thread left alive, exit app
        self.watch_thread.daemon = True
        self.watch_thread.start()

class SingletonExecutor(SingletonConfigurable, ThreadPoolExecutor):
    """
    Simple wrapper to ThreadPoolExecutor that is also a singleton.

    We want one ThreadPool that is used by all the spawners, rather
    than one ThreadPool per spawner!
    """
    pass

class KubeSpawner(Spawner):
    """
    Implement a JupyterHub spawner to spawn pods in a Kubernetes Cluster.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # By now, all the traitlets have been set, so we can use them to compute
        # other attributes
        self.executor = SingletonExecutor.instance(max_workers=self.k8s_api_threadpool_workers)

        # FIXME: Is this a memory leak? Will there be other side effects?
        self.pod_reflector = PodReflector.instance(parent=self, namespace=self.namespace)
{'fraction_non_alphanumeric': 0.06673064147043625, 'fraction_numerical': 0.0025097807632686203, 'mean_word_length': 2.0748978665456197, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4000687', 'n_tokens_mistral': 3600, 'n_tokens_neox': 3444, 'n_words': 1326}