Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #!/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)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement