dask.py

import os
import boto3
import yaml

jupyter_image = os.environ['JUPYTER_IMAGE']

session = boto3.Session()
credentials = session.get_credentials()
aws_creds = {'key': credentials.access_key, 'secret': credentials.secret_key}
dask_worker_spec = f'''kind: Pod
spec:
  restartPolicy: Never
  affinity:
    nodeAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
        nodeSelectorTerms:
          - matchExpressions:
              - key: hub.jupyter.org/node-purpose
                operator: In
                values:
                  - cpu-worker
  tolerations:
    - effect: NoSchedule
      key: hub.jupyter.org_dedicated
      operator: Equal
      value: cpu-worker
  containers:
    - image: {jupyter_image}
      imagePullPolicy: IfNotPresent
      args: [conda, run, -n, geostack, dask-worker, --nprocs, '8', --nthreads, '1', --dashboard, --memory-limit, '63.75GB', --death-timeout, '60']
      name: dask-high-cpu-memory
      resources:
        limits:
          cpu: "8"
          memory: "63G"
        requests:
          cpu: "7.75"
          memory: "63G"'''
from dask_kubernetes import KubeCluster
from dask.distributed import Client
cluster = KubeCluster.from_dict(yaml.safe_load(dask_worker_spec), silence_logs='error', env=creds)
client = Client(cluster)
cluster