`ding.utils.k8s_helper`¶

`ding.utils.k8s_helper` ¶

`K8sLauncher` ¶

Bases: object

Overview

object to manage the K8s cluster

Interfaces: __init__, _load, create_cluster, _check_k3d_tools, delete_cluster, preload_images

`init(config_path)` ¶

Overview

Initialize the K8sLauncher object.

Arguments: - config_path (:obj:str): The path of the config file.

`create_cluster()` ¶

Overview

Create the k8s cluster.

`delete_cluster()` ¶

Overview

Delete the k8s cluster.

`preload_images(images)` ¶

Overview

Preload images.

`get_operator_server_kwargs(cfg)` ¶

Overview

Get kwarg dict from config file

Arguments: - cfg (:obj:EasyDict) System config Returns: - result (:obj:dict) Containing api_version, namespace, name, port, host.

`exist_operator_server()` ¶

Overview

Check if the 'KUBERNETES_SERVER_URL' environment variable exists.

`pod_exec_command(kubeconfig, name, namespace, cmd)` ¶

Overview

Execute command in pod

Arguments: - kubeconfig (:obj:str) The path of kubeconfig file - name (:obj:str) The name of pod - namespace (:obj:str) The namespace of pod

Full Source Code

../ding/utils/k8s_helper.py

import osimport jsonfrom typing import Tuplefrom easydict import EasyDictimport yamlimport subprocessfrom enum import Enum, uniquefrom ding.interaction.base import split_http_addressfrom .default_helper import one_time_warningDEFAULT_NAMESPACE = 'default'DEFAULT_POD_NAME = 'dijob-example-coordinator'DEFAULT_API_VERSION = '/v1alpha1'DEFAULT_K8S_COLLECTOR_PORT = 22270DEFAULT_K8S_LEARNER_PORT = 22271DEFAULT_K8S_AGGREGATOR_SLAVE_PORT = 22272DEFAULT_K8S_COORDINATOR_PORT = 22273DEFAULT_K8S_AGGREGATOR_MASTER_PORT = 22273def get_operator_server_kwargs(cfg: EasyDict) -> dict:    """    Overview:        Get kwarg dict from config file    Arguments:        - cfg (:obj:`EasyDict`) System config    Returns:        - result (:obj:`dict`) Containing ``api_version``,  ``namespace``, ``name``, ``port``, ``host``.    """    namespace = os.environ.get('KUBERNETES_POD_NAMESPACE', DEFAULT_NAMESPACE)    name = os.environ.get('KUBERNETES_POD_NAME', DEFAULT_POD_NAME)    url = cfg.get('system_addr', None) or os.environ.get('KUBERNETES_SERVER_URL', None)    assert url, 'please set environment variable KUBERNETES_SERVER_URL in Kubenetes platform.'    api_version = cfg.get('api_version', None) \        or os.environ.get('KUBERNETES_SERVER_API_VERSION', DEFAULT_API_VERSION)    try:        host, port = url.split(":")[0], int(url.split(":")[1])    except Exception as e:        host, port, _, _ = split_http_address(url)    return {        'api_version': api_version,        'namespace': namespace,        'name': name,        'host': host,        'port': port,    }def exist_operator_server() -> bool:    """    Overview:        Check if the 'KUBERNETES_SERVER_URL' environment variable exists.    """    return 'KUBERNETES_SERVER_URL' in os.environdef pod_exec_command(kubeconfig: str, name: str, namespace: str, cmd: str) -> Tuple[int, str]:    """    Overview:        Execute command in pod    Arguments:        - kubeconfig (:obj:`str`) The path of kubeconfig file        - name (:obj:`str`) The name of pod        - namespace (:obj:`str`) The namespace of pod    """    try:        from kubernetes import config        from kubernetes.client import CoreV1Api        from kubernetes.client.rest import ApiException        from kubernetes.stream import stream    except ModuleNotFoundError as e:        one_time_warning("You have not installed kubernetes package! Please try 'pip install DI-engine[k8s]'.")        exit(-1)    config.load_kube_config(config_file=kubeconfig)    core_v1 = CoreV1Api()    resp = None    try:        resp = core_v1.read_namespaced_pod(name=name, namespace=namespace)    except ApiException as e:        if e.status != 404:            return -1, "Unknown error: %s" % e    if not resp:        return -1, f"Pod {name} does not exist."    if resp.status.phase != 'Running':        return -1, f"Pod {name} is not in Running."    exec_command = ['/bin/sh', '-c', cmd]    resp = stream(        core_v1.connect_get_namespaced_pod_exec,        name,        namespace,        command=exec_command,        stderr=False,        stdin=False,        stdout=True,        tty=False    )    resp = resp.replace("\'", "\"") \        .replace('None', 'null') \        .replace(': False', ': 0') \        .replace(': True', ': 1') \        .replace('"^([+-]?[0-9.]+)([eEinumkKMGTP]*[-+]?[0-9]*)$"', '\\"^([+-]?[0-9.]+)([eEinumkKMGTP]*[-+]?[0-9]*)$\\"')    resp = json.loads(resp)    return resp['code'], resp['message']@uniqueclass K8sType(Enum):    Local = 1    K3s = 2class K8sLauncher(object):    """    Overview:        object to manage the K8s cluster    Interfaces:        ``__init__``, ``_load``, ``create_cluster``, ``_check_k3d_tools``, ``delete_cluster``, ``preload_images``    """    def __init__(self, config_path: str) -> None:        """        Overview:            Initialize the K8sLauncher object.        Arguments:            - config_path (:obj:`str`): The path of the config file.        """        self.name = None        self.servers = 1        self.agents = 0        self.type = K8sType.Local        self._images = []        self._load(config_path)        self._check_k3d_tools()    def _load(self, config_path: str) -> None:        """        Overview:            Load the config file.        Arguments:            - config_path (:obj:`str`): The path of the config file.        """        with open(config_path, 'r') as f:            data = yaml.safe_load(f)            self.name = data.get('name') if data.get('name') else self.name            if data.get('servers'):                if type(data.get('servers')) is not int:                    raise TypeError(f"servers' type is expected int, actual {type(data.get('servers'))}")                self.servers = data.get('servers')            if data.get('agents'):                if type(data.get('agents')) is not int:                    raise TypeError(f"agents' type is expected int, actual {type(data.get('agents'))}")                self.agents = data.get('agents')            if data.get('type'):                if data.get('type') == 'k3s':                    self.type = K8sType.K3s                elif data.get('type') == 'local':                    self.type = K8sType.Local                else:                    raise ValueError(f"no type found for {data.get('type')}")            if data.get('preload_images'):                if type(data.get('preload_images')) is not list:                    raise TypeError(f"preload_images' type is expected list, actual {type(data.get('preload_images'))}")                self._images = data.get('preload_images')    def _check_k3d_tools(self) -> None:        """        Overview:            Check if the k3d tools exist.        """        if self.type != K8sType.K3s:            return        args = ['which', 'k3d']        proc = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)        out, _ = proc.communicate()        if out.decode('utf-8') == '':            raise FileNotFoundError(                "No k3d tools found, please install by executing ./ding/scripts/install-k8s-tools.sh"            )    def create_cluster(self) -> None:        """        Overview:            Create the k8s cluster.        """        print('Creating k8s cluster...')        if self.type != K8sType.K3s:            return        args = ['k3d', 'cluster', 'create', f'{self.name}', f'--servers={self.servers}', f'--agents={self.agents}']        proc = subprocess.Popen(args, stderr=subprocess.PIPE)        _, err = proc.communicate()        err_str = err.decode('utf-8').strip()        if err_str != '' and 'WARN' not in err_str:            if 'already exists' in err_str:                print('K8s cluster already exists')            else:                raise RuntimeError(f'Failed to create cluster {self.name}: {err_str}')        # preload images        self.preload_images(self._images)    def delete_cluster(self) -> None:        """        Overview:            Delete the k8s cluster.        """        print('Deleting k8s cluster...')        if self.type != K8sType.K3s:            return        args = ['k3d', 'cluster', 'delete', f'{self.name}']        proc = subprocess.Popen(args, stderr=subprocess.PIPE)        _, err = proc.communicate()        err_str = err.decode('utf-8').strip()        if err_str != '' and 'WARN' not in err_str and \                'NotFound' not in err_str:            raise RuntimeError(f'Failed to delete cluster {self.name}: {err_str}')    def preload_images(self, images: list) -> None:        """        Overview:            Preload images.        """        if self.type != K8sType.K3s or len(images) == 0:            return        args = ['k3d', 'image', 'import', f'--cluster={self.name}']        args += images        proc = subprocess.Popen(args, stderr=subprocess.PIPE)        _, err = proc.communicate()        err_str = err.decode('utf-8').strip()        if err_str != '' and 'WARN' not in err_str:            raise RuntimeError(f'Failed to preload images: {err_str}')

ding.utils.k8s_helper¶