update_dask_k8s.py

#!/usr/bin/env python3
"""
Update Dask configuration based on the configuration
of the running Pod.

To be run at startup.
"""

import os

import dask
import yaml

from kubernetes import client, config

user = os.environ.get('USER')
jhubuser = os.environ.get('JUPYTERHUB_USER')


def get_volumes():
    """Get volume mounts and volumnes defined in the
    running Pod.
    """
    config.load_incluster_config()
    v1 = client.CoreV1Api()
    namespace = open('/var/run/secrets/kubernetes.io/serviceaccount/namespace').read()
    pod_list = v1.list_namespaced_pod(namespace)
    this_pod = [pod for pod in pod_list.items if jhubuser in pod.metadata.name]
    this_pod = this_pod[0]
    volume_mounts = this_pod.spec.containers[0].volume_mounts
    volumes = this_pod.spec.volumes
    return volumes, volume_mounts


def save_config(config):
    """Save K8s Dask configuration to file.
    """
    res = yaml.safe_dump({'kubernetes': config})
    open(config['worker-template-path'], 'w').write(res)


volumes, volume_mounts = get_volumes()

# Only using hostPath here
dc = dask.config.config['kubernetes']
dc['worker-template']['spec']['volumes'] = [
    {'name': vol.name, 'hostPath': vol.host_path.to_dict()}
    for vol in volumes
    if vol.host_path
]

# Mounted in /data directory
dc['worker-template']['spec']['containers'][0]['volumeMounts'] = [
    {'name': vol.name, 'mountPath': vol.mount_path}
    for vol in volume_mounts
    if 'data' in vol.mount_path
]

# Propagate as well env variables
dc['worker-template']['spec']['containers'][0]['env'] = [
    {'name': 'EXTRA_PIP_PACKAGES', 'value': '${EXTRA_PIP_PACKAGES}'},
    {'name': 'OMP_NUM_THREADS', 'value': 1}
]

dc['worker-template-path'] = f'/home/{user}/.config/dask/kubernetes.yaml'

save_config(dc)