Skip to content

Instantly share code, notes, and snippets.

Embed
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.
"""
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.