Source code for dask_kubernetes.common.objects

"""
Convenience functions for creating pod templates.
"""
from collections import namedtuple
import copy
from kubernetes import client
import json

from kubernetes.client.configuration import Configuration

from dask_kubernetes.constants import KUBECLUSTER_WORKER_CONTAINER_NAME

_FakeResponse = namedtuple("_FakeResponse", ["data"])


class DummyApiClient(client.ApiClient):
    """A Dummy API client that is to be used solely for serialization/deserialization.

    This is to avoid starting a threadpool at initialization and for adapting the
    deserialize method to accept a python dictionary instead of a Response-like
    interface.
    """

    def __init__(self):
        self.configuration = Configuration.get_default_copy()

    def deserialize(self, dict_, klass):
        return super().deserialize(_FakeResponse(json.dumps(dict_)), klass)


SERIALIZATION_API_CLIENT = DummyApiClient()


def _set_k8s_attribute(obj, attribute, value):
    """
    Set a specific value on a kubernetes object's attribute

    obj
        an object from Kubernetes Python API client
    attribute
        Should be a Kubernetes API style attribute (with camelCase)
    value
        Can be anything (string, list, dict, k8s objects) that can be
        accepted by the k8s python client
    """
    current_value = None
    attribute_name = None
    # All k8s python client objects have an 'attribute_map' property
    # which has as keys python style attribute names (api_client)
    # and as values the kubernetes JSON API style attribute names
    # (apiClient). We want to allow users to use the JSON API style attribute
    # names only.
    for python_attribute, json_attribute in obj.attribute_map.items():
        if json_attribute == attribute:
            attribute_name = python_attribute
            break
    else:
        raise ValueError(
            "Attribute must be one of {}".format(obj.attribute_map.values())
        )

    if hasattr(obj, attribute_name):
        current_value = getattr(obj, attribute_name)

    if current_value is not None:
        # This will ensure that current_value is something JSONable,
        # so a dict, list, or scalar
        current_value = SERIALIZATION_API_CLIENT.sanitize_for_serialization(
            current_value
        )

    if isinstance(current_value, dict):
        # Deep merge our dictionaries!
        setattr(obj, attribute_name, merge_dictionaries(current_value, value))
    elif isinstance(current_value, list):
        # Just append lists
        setattr(obj, attribute_name, current_value + value)
    else:
        # Replace everything else
        setattr(obj, attribute_name, value)


def merge_dictionaries(a, b, path=None, update=True):
    """
    Merge two dictionaries recursively.

    From https://stackoverflow.com/a/25270947
    """
    if path is None:
        path = []
    for key in b:
        if key in a:
            if isinstance(a[key], dict) and isinstance(b[key], dict):
                merge_dictionaries(a[key], b[key], path + [str(key)])
            elif a[key] == b[key]:
                pass  # same leaf value
            elif isinstance(a[key], list) and isinstance(b[key], list):
                for idx, _ in enumerate(b[key]):
                    a[key][idx] = merge_dictionaries(
                        a[key][idx],
                        b[key][idx],
                        path + [str(key), str(idx)],
                        update=update,
                    )
            elif update:
                a[key] = b[key]
            else:
                raise Exception("Conflict at %s" % ".".join(path + [str(key)]))
        else:
            a[key] = b[key]
    return a


[docs]def make_pod_spec( image, labels={}, threads_per_worker=1, env={}, extra_container_config={}, extra_pod_config={}, resources=None, memory_limit=None, memory_request=None, cpu_limit=None, cpu_request=None, gpu_limit=None, annotations={}, ): """ Create generic pod template from input parameters Parameters ---------- image : str Docker image name labels : dict Dict of labels to pass to ``V1ObjectMeta`` threads_per_worker : int Number of threads per each worker env : dict Dict of environment variables to pass to ``V1Container`` extra_container_config : dict Extra config attributes to set on the container object extra_pod_config : dict Extra config attributes to set on the pod object resources : str Resources for task constraints like "GPU=2 MEM=10e9". Resources are applied separately to each worker process (only relevant when starting multiple worker processes. Passed to the `--resources` option in ``dask-worker``. memory_limit : int, float, or str Bytes of memory per process that the worker can use (applied to both ``dask-worker --memory-limit`` and ``spec.containers[].resources.limits.memory``). This can be: - an integer (bytes), note 0 is a special case for no memory management. - a float (bytes). Note: fraction of total system memory is not supported by k8s. - a string (like 5GiB or 5000M). Note: 'GB' is not supported by k8s. - 'auto' for automatically computing the memory limit. [default: auto] memory_request : int, float, or str Like ``memory_limit`` (applied only to ``spec.containers[].resources.requests.memory`` and ignored by ``dask-worker``). cpu_limit : float or str CPU resource limits (applied to ``spec.containers[].resources.limits.cpu``). cpu_request : float or str CPU resource requests (applied to ``spec.containers[].resources.requests.cpu``). gpu_limit : int GPU resource limits (applied to ``spec.containers[].resources.limits."nvidia.com/gpu"``). annotations : dict Dict of annotations passed to ``V1ObjectMeta`` Returns ------- pod : V1PodSpec Examples -------- >>> make_pod_spec(image='ghcr.io/dask/dask:latest', memory_limit='4G', memory_request='4G') """ args = [ "dask-worker", "$(DASK_SCHEDULER_ADDRESS)", "--nthreads", str(threads_per_worker), "--death-timeout", "60", ] if memory_limit: args.extend(["--memory-limit", str(memory_limit)]) if resources: args.extend(["--resources", str(resources)]) pod = client.V1Pod( metadata=client.V1ObjectMeta(labels=labels, annotations=annotations), spec=client.V1PodSpec( restart_policy="Never", containers=[ client.V1Container( name=KUBECLUSTER_WORKER_CONTAINER_NAME, image=image, args=args, env=[client.V1EnvVar(name=k, value=v) for k, v in env.items()], ) ], ), ) resources = client.V1ResourceRequirements(limits={}, requests={}) if cpu_request: resources.requests["cpu"] = cpu_request if memory_request: resources.requests["memory"] = memory_request if cpu_limit: resources.limits["cpu"] = cpu_limit if gpu_limit: resources.limits["nvidia.com/gpu"] = gpu_limit if memory_limit: resources.limits["memory"] = memory_limit pod.spec.containers[0].resources = resources for key, value in extra_container_config.items(): _set_k8s_attribute(pod.spec.containers[0], key, value) for key, value in extra_pod_config.items(): _set_k8s_attribute(pod.spec, key, value) return pod
def make_pod_from_dict(dict_): containers = dict_.get("spec", {}).get("containers", []) for i, container in enumerate(containers): container.setdefault("name", f"dask-{i}") return SERIALIZATION_API_CLIENT.deserialize(dict_, client.V1Pod) def make_service_from_dict(dict_): return SERIALIZATION_API_CLIENT.deserialize(dict_, client.V1Service) def make_pdb_from_dict(dict_): return SERIALIZATION_API_CLIENT.deserialize( dict_, client.V1beta1PodDisruptionBudget ) def clean_pod_template(pod_template, match_node_purpose="prefer", pod_type="worker"): """Normalize pod template""" pod_template = copy.deepcopy(pod_template) # Make sure metadata / labels / env objects exist, so they can be modified # later without a lot of `is None` checks if pod_template.metadata is None: pod_template.metadata = client.V1ObjectMeta() if pod_template.metadata.labels is None: pod_template.metadata.labels = {} if pod_template.spec.containers[0].env is None: pod_template.spec.containers[0].env = [] # add default tolerations tolerations = [ client.V1Toleration( key="k8s.dask.org/dedicated", operator="Equal", value=pod_type, effect="NoSchedule", ), # GKE currently does not permit creating taints on a node pool # with a `/` in the key field client.V1Toleration( key="k8s.dask.org_dedicated", operator="Equal", value=pod_type, effect="NoSchedule", ), ] if pod_template.spec.tolerations is None: pod_template.spec.tolerations = tolerations else: pod_template.spec.tolerations.extend(tolerations) # add default node affinity to k8s.dask.org/node-purpose=worker if match_node_purpose != "ignore": # for readability affinity = pod_template.spec.affinity if affinity is None: affinity = client.V1Affinity() if affinity.node_affinity is None: affinity.node_affinity = client.V1NodeAffinity() # a common object for both a preferred and a required node affinity node_selector_term = client.V1NodeSelectorTerm( match_expressions=[ client.V1NodeSelectorRequirement( key="k8s.dask.org/node-purpose", operator="In", values=[pod_type] ) ] ) if match_node_purpose == "require": if ( affinity.node_affinity.required_during_scheduling_ignored_during_execution is None ): affinity.node_affinity.required_during_scheduling_ignored_during_execution = client.V1NodeSelector( node_selector_terms=[] ) affinity.node_affinity.required_during_scheduling_ignored_during_execution.node_selector_terms.append( node_selector_term ) elif match_node_purpose == "prefer": if ( affinity.node_affinity.preferred_during_scheduling_ignored_during_execution is None ): affinity.node_affinity.preferred_during_scheduling_ignored_during_execution = ( [] ) preferred_scheduling_terms = [ client.V1PreferredSchedulingTerm( preference=node_selector_term, weight=100 ) ] affinity.node_affinity.preferred_during_scheduling_ignored_during_execution.extend( preferred_scheduling_terms ) else: raise ValueError( 'Attribute must be one of "ignore", "prefer", or "require".' ) pod_template.spec.affinity = affinity return pod_template def clean_service_template(service_template): """Normalize service template and check for type errors""" service_template = copy.deepcopy(service_template) # Make sure metadata / labels objects exist, so they can be modified # later without a lot of `is None` checks if service_template.metadata is None: service_template.metadata = client.V1ObjectMeta() if service_template.metadata.labels is None: service_template.metadata.labels = {} return service_template def clean_pdb_template(pdb_template): """Normalize pdb template and check for type errors""" pdb_template = copy.deepcopy(pdb_template) # Make sure metadata / labels objects exist, so they can be modified # later without a lot of `is None` checks if pdb_template.metadata is None: pdb_template.metadata = client.V1ObjectMeta() if pdb_template.metadata.labels is None: pdb_template.metadata.labels = {} if pdb_template.spec.selector is None: pdb_template.spec.selector = client.V1LabelSelector() return pdb_template