pilot/api/dask_k8s_utils.py

0001 #!/usr/bin/env python
0002 # Licensed under the Apache License, Version 2.0 (the "License");
0003 # you may not use this file except in compliance with the License.
0004 # You may obtain a copy of the License at
0005 # http://www.apache.org/licenses/LICENSE-2.0
0006 #
0007 # Authors:
0008 # - Paul Nilsson, paul.nilsson@cern.ch, 2021
0009
0010 # based on Harvester k8s_utils
0011
0012 import os
0013
0014 from kubernetes import client, config
0015 # from kubernetes.client.rest import ApiException
0016
0017
0018 class K8sClient(object):
0019
0020     def __init__(self, namespace, config_file=None):
0021         """
0022         Init.
0023         """
0024
0025         if not os.path.isfile(config_file):
0026             raise RuntimeError('Cannot find k8s config file: {0}'.format(config_file))
0027
0028         config.load_kube_config(config_file=config_file)
0029         self.namespace = namespace if namespace else 'default'
0030         self.core_v1 = client.CoreV1Api()
0031         self.batch_v1 = client.BatchV1Api()
0032         self.apps_v1 = client.AppsV1Api()
0033         self.delete_v1 = client.V1DeleteOptions(propagation_policy='Background')
0034
0035     def create_ssh_keys_secret(self, work_spec):
0036         return True
0037
0038     def create_dask_head(self, work_spec, panda_queue, evaluation_image, pilot_image, cert, dfs_claim_name,
0039                          cpu_adjust_ratio=100, memory_adjust_ratio=100, max_time=None):
0040         """
0041
0042         """
0043
0044         #worker_id = str(work_spec.workerID)
0045         #tmp_log = core_utils.make_logger(base_logger, 'workerID={0}'.format(worker_id),
0046         #                                 method_name='create_dask_head')
0047
0048         if not self.create_configmap_dask(work_spec, panda_queue):
0049             return False
0050
0051         if not self.create_host_discovery_configmap(work_spec):
0052             return False
0053
0054     def create_dask_formation(self, work_spec, panda_queue, evaluation_image, pilot_image,
0055                               worker_command, cert, dfs_claim_name, cpu_adjust_ratio=100, memory_adjust_ratio=100,
0056                               max_time=None):
0057
0058         rsp = self.create_ssh_keys_secret(work_spec)
0059         if not rsp:
0060             return rsp
0061
0062         rsp = self.create_dask_head(work_spec, panda_queue, evaluation_image, pilot_image, cert, dfs_claim_name,
0063                                     cpu_adjust_ratio, memory_adjust_ratio, max_time)
0064         if not rsp:
0065             return rsp
0066
0067         rsp = self.create_dask_workers(work_spec, evaluation_image, worker_command, dfs_claim_name,
0068                                        cpu_adjust_ratio, memory_adjust_ratio, max_time)
0069
0070         if not rsp:
0071             return rsp
0072
0073         return True