Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
#!/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.
v1 = client.CoreV1Api()
namespace = open('/var/run/secrets/').read()
pod_list = v1.list_namespaced_pod(namespace)
this_pod = [pod for pod in pod_list.items if jhubuser in]
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':, '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':, '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'
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment