Repository: mars-project/mars

Issue:
Make labels of Mars pods configurable
Sometimes we need to add customized labels to pods in kubernetes, for instance, we may define groups by specifying different labels. This functionality is not implemented in ``mars.deploy.kubernetes``.


Discussion:

Resolving patch:
diff --git a/mars/deploy/kubernetes/client.py b/mars/deploy/kubernetes/client.py
--- a/mars/deploy/kubernetes/client.py
+++ b/mars/deploy/kubernetes/client.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import logging
 import random
 import time
@@ -52,7 +53,7 @@ def session(self):
 
     def start(self):
         self._endpoint = self._cluster.start()
-        self._session = new_session(self._endpoint)
+        self._session = new_session(self._endpoint).as_default()
 
     def stop(self, wait=False, timeout=0):
         self._cluster.stop(wait=wait, timeout=timeout)
@@ -88,6 +89,7 @@ def __init__(self, kube_api_client=None, image=None, namespace=None,
         extra_modules = extra_modules.split(',') if isinstance(extra_modules, str) \
             else extra_modules
         extra_envs = kwargs.pop('extra_env', None) or dict()
+        extra_labels = kwargs.pop('extra_labels', None) or dict()
         service_port = kwargs.pop('service_port', None) or self._default_service_port
 
         def _override_modules(updates):
@@ -97,16 +99,21 @@ def _override_modules(updates):
             modules.update(updates)
             return sorted(modules)
 
-        def _override_envs(updates):
-            ret = extra_envs.copy()
+        def _override_dict(d, updates):
+            updates = updates or dict()
+            ret = d.copy()
             ret.update(updates)
             return ret
 
+        _override_envs = functools.partial(_override_dict, extra_envs)
+        _override_labels = functools.partial(_override_dict, extra_labels)
+
         self._scheduler_num = scheduler_num
         self._scheduler_cpu = scheduler_cpu
         self._scheduler_mem = scheduler_mem
         self._scheduler_extra_modules = _override_modules(kwargs.pop('scheduler_extra_modules', []))
-        self._scheduler_extra_env = _override_envs(kwargs.pop('scheduler_extra_env', None) or dict())
+        self._scheduler_extra_env = _override_envs(kwargs.pop('scheduler_extra_env', None))
+        self._scheduler_extra_labels = _override_labels(kwargs.pop('scheduler_extra_labels', None))
         self._scheduler_service_port = kwargs.pop('scheduler_service_port', None) or service_port
 
         self._worker_num = worker_num
@@ -116,14 +123,16 @@ def _override_envs(updates):
         self._worker_cache_mem = worker_cache_mem
         self._min_worker_num = min_worker_num
         self._worker_extra_modules = _override_modules(kwargs.pop('worker_extra_modules', []))
-        self._worker_extra_env = _override_envs(kwargs.pop('worker_extra_env', None) or dict())
+        self._worker_extra_env = _override_envs(kwargs.pop('worker_extra_env', None))
+        self._worker_extra_labels = _override_labels(kwargs.pop('worker_extra_labels', None))
         self._worker_service_port = kwargs.pop('worker_service_port', None) or service_port
 
         self._web_num = web_num
         self._web_cpu = web_cpu
         self._web_mem = web_mem
         self._web_extra_modules = _override_modules(kwargs.pop('web_extra_modules', []))
-        self._web_extra_env = _override_envs(kwargs.pop('web_extra_env', None) or dict())
+        self._web_extra_env = _override_envs(kwargs.pop('web_extra_env', None))
+        self._web_extra_labels = _override_labels(kwargs.pop('web_extra_labels', None))
         self._web_service_port = kwargs.pop('web_service_port', None) or service_port
 
     @property
@@ -149,7 +158,7 @@ def _create_kube_service(self):
 
         service_config = ServiceConfig(
             'marsservice', service_type='NodePort', port=self._web_service_port,
-            selector=dict(name=MarsWebsConfig.rc_name),
+            selector={'mars/service-type': MarsWebsConfig.rc_name},
         )
         self._core_api.create_namespaced_service(self._namespace, service_config.build())
 
@@ -159,7 +168,7 @@ def _create_kube_service(self):
         port = svc_data['spec']['ports'][0]['node_port']
 
         web_pods = self._core_api.list_namespaced_pod(
-            self._namespace, label_selector='name=' + MarsWebsConfig.rc_name).to_dict()
+            self._namespace, label_selector='mars/service-type=' + MarsWebsConfig.rc_name).to_dict()
         host_ip = random.choice(web_pods['items'])['status']['host_ip']
 
         # docker desktop use a VM to hold docker processes, hence
@@ -213,6 +222,7 @@ def _create_schedulers(self):
             pre_stop_command=self._pre_stop_command,
         )
         schedulers_config.add_simple_envs(self._scheduler_extra_env)
+        schedulers_config.add_labels(self._scheduler_extra_labels)
         self._core_api.create_namespaced_replication_controller(
             self._namespace, schedulers_config.build())
 
@@ -226,6 +236,7 @@ def _create_workers(self):
             pre_stop_command=self._pre_stop_command,
         )
         workers_config.add_simple_envs(self._worker_extra_env)
+        workers_config.add_labels(self._worker_extra_labels)
         self._core_api.create_namespaced_replication_controller(
             self._namespace, workers_config.build())
 
@@ -236,6 +247,7 @@ def _create_webs(self):
             service_port=self._web_service_port, pre_stop_command=self._pre_stop_command,
         )
         webs_config.add_simple_envs(self._web_extra_env)
+        webs_config.add_labels(self._web_extra_labels)
         self._core_api.create_namespaced_replication_controller(
             self._namespace, webs_config.build())
 
@@ -248,9 +260,9 @@ def _wait_services_ready(self):
         min_worker_num = int(self._min_worker_num or self._worker_num)
         limits = [self._scheduler_num, min_worker_num, self._web_num]
         selectors = [
-            'name=' + MarsSchedulersConfig.rc_name,
-            'name=' + MarsWorkersConfig.rc_name,
-            'name=' + MarsWebsConfig.rc_name,
+            'mars/service-type=' + MarsSchedulersConfig.rc_name,
+            'mars/service-type=' + MarsWorkersConfig.rc_name,
+            'mars/service-type=' + MarsWebsConfig.rc_name,
             ]
         wait_services_ready(selectors, limits,
                             lambda sel: self._get_ready_pod_count(sel),
@@ -277,7 +289,7 @@ def start(self):
         except:  # noqa: E722
             if self._log_when_fail:  # pargma: no cover
                 logger.error('Error when creating cluster')
-                for name, log in self._load_cluster_logs():
+                for name, log in self._load_cluster_logs().items():
                     logger.error('Error logs for %s:\n%s', name, log)
             self.stop()
             raise
diff --git a/mars/deploy/kubernetes/config.py b/mars/deploy/kubernetes/config.py
--- a/mars/deploy/kubernetes/config.py
+++ b/mars/deploy/kubernetes/config.py
@@ -290,6 +290,7 @@ def __init__(self, name, image, replicas, resource_request=None, resource_limit=
         self._ports = []
         self._volumes = []
         self._envs = dict()
+        self._labels = dict()
 
         self.add_default_envs()
 
@@ -308,6 +309,9 @@ def add_simple_envs(self, envs):
         for k, v in envs.items() or ():
             self.add_env(k, v)
 
+    def add_labels(self, labels):
+        self._labels.update(labels)
+
     def add_port(self, container_port):
         self._ports.append(PortConfig(container_port))
 
@@ -358,10 +362,9 @@ def build(self):
             },
             'spec': {
                 'replicas': int(self._replicas),
-                'selector': {'name': self._name},
                 'template': {
                     'metadata': {
-                        'labels': {'name': self._name},
+                        'labels': _remove_nones(self._labels) or None,
                     },
                     'spec': self.build_template_spec()
                 }
@@ -401,6 +404,8 @@ def __init__(self, replicas, cpu=None, memory=None, limit_resources=False,
         for vol in volumes or ():
             self.add_volume(vol)
 
+        self.add_labels({'mars/service-type': self.rc_name})
+
     def add_default_envs(self):
         self.add_env('MARS_K8S_POD_NAME', field_path='metadata.name')
         self.add_env('MARS_K8S_POD_NAMESPACE', field_path='metadata.namespace')
@@ -430,6 +435,11 @@ def config_readiness_probe(self):
     def get_local_app_module(mod_name):
         return __name__.rsplit('.', 1)[0] + '.' + mod_name
 
+    def build(self):
+        result = super().build()
+        result['spec']['selector'] = {'mars/service-type': self.rc_name}
+        return result
+
 
 class MarsSchedulersConfig(MarsReplicationControllerConfig):
     """
diff --git a/mars/deploy/kubernetes/core.py b/mars/deploy/kubernetes/core.py
--- a/mars/deploy/kubernetes/core.py
+++ b/mars/deploy/kubernetes/core.py
@@ -43,6 +43,7 @@ def __init__(self, k8s_config=None, k8s_namespace=None, label_selector=None):
 
         self._k8s_namespace = k8s_namespace or os.environ.get('MARS_K8S_POD_NAMESPACE') or 'default'
         self._label_selector = label_selector
+        self._full_label_selector = None
         self._client = client.CoreV1Api(client.ApiClient(self._k8s_config))
         self._pool = ThreadPool(1)
 
@@ -51,12 +52,32 @@ def __init__(self, k8s_config=None, k8s_namespace=None, label_selector=None):
     def __reduce__(self):
         return type(self), (self._k8s_config, self._k8s_namespace, self._label_selector)
 
+    def _get_label_selector(self):
+        if self._full_label_selector is not None:
+            return self._full_label_selector
+
+        selectors = [self._label_selector]
+        if 'MARS_K8S_GROUP_LABELS' in os.environ:
+            group_labels = os.environ['MARS_K8S_GROUP_LABELS'].split(',')
+            cur_pod_info = self._pool.spawn(self._client.read_namespaced_pod,
+                                            os.environ['MARS_K8S_POD_NAME'],
+                                            namespace=self._k8s_namespace).result().to_dict()
+            for label in group_labels:
+                label_val = cur_pod_info['metadata']['labels'][label]
+                selectors.append(f'{label}={label_val}')
+        self._full_label_selector = ','.join(selectors)
+        logger.debug('Using pod selector %s', self._full_label_selector)
+        return self._full_label_selector
+
     def _extract_pod_name_ep(self, pod_data):
+        pod_ip = pod_data["status"]["pod_ip"]
         svc_port = pod_data['spec']['containers'][0]['ports'][0]['container_port']
-        return pod_data['metadata']['name'], f'{pod_data["status"]["pod_ip"]}:{svc_port}'
+        return pod_data['metadata']['name'], f'{pod_ip}:{svc_port}'
 
     @staticmethod
     def _extract_pod_ready(obj_data):
+        if obj_data['status']['phase'] != 'Running':
+            return False
         # if conditions not supported, always return True
         if 'status' not in obj_data or 'conditions' not in obj_data['status']:
             return True
@@ -66,7 +87,7 @@ def _extract_pod_ready(obj_data):
     def _get_pod_to_ep(self):
         query = self._pool.spawn(self._client.list_namespaced_pod,
                                  namespace=self._k8s_namespace,
-                                 label_selector=self._label_selector).result().to_dict()
+                                 label_selector=self._get_label_selector()).result().to_dict()
         result = dict()
         for el in query['items']:
             name, pod_ep = self._extract_pod_name_ep(el)
@@ -82,6 +103,8 @@ def get(self, update=False):
 
     def is_all_ready(self):
         self.get(True)
+        if not self._pod_to_ep:
+            return False
         return all(a is not None for a in self._pod_to_ep.values())
 
     def watch(self):
@@ -96,7 +119,7 @@ def watch(self):
             linger = 10 if self.is_all_ready() else 1
             streamer = w.stream(self._client.list_namespaced_pod,
                                 namespace=self._k8s_namespace,
-                                label_selector=self._label_selector,
+                                label_selector=self._get_label_selector(),
                                 timeout_seconds=linger)
             while True:
                 try:
@@ -151,9 +174,14 @@ def wait_all_schedulers_ready(self):
         # check if all schedulers are registered in ClusterInfoActor
         actor_client = new_client()
         while True:
-            cluster_info = actor_client.actor_ref(
-                SchedulerClusterInfoActor.default_uid(), address=random.choice(kube_schedulers))
-            cluster_info_schedulers = cluster_info.get_schedulers()
+            try:
+                cluster_info = actor_client.actor_ref(
+                    SchedulerClusterInfoActor.default_uid(), address=random.choice(kube_schedulers))
+                cluster_info_schedulers = cluster_info.get_schedulers()
+            except ConnectionError:  # pragma: no cover
+                time.sleep(0.1)
+                continue
+
             if set(cluster_info_schedulers) == set(kube_schedulers):
                 from ...cluster_info import INITIAL_SCHEDULER_FILE
                 with open(INITIAL_SCHEDULER_FILE, 'w') as scheduler_file:
@@ -164,4 +192,5 @@ def wait_all_schedulers_ready(self):
             sleep_fun(1)  # pragma: no cover
 
     def create_scheduler_discoverer(self):
-        self.scheduler_discoverer = K8SPodsIPWatcher(label_selector='name=marsscheduler')
+        self.scheduler_discoverer = K8SPodsIPWatcher(
+            label_selector='mars/service-type=marsscheduler')


Test patch:
diff --git a/mars/deploy/kubernetes/tests/test_config.py b/mars/deploy/kubernetes/tests/test_config.py
--- a/mars/deploy/kubernetes/tests/test_config.py
+++ b/mars/deploy/kubernetes/tests/test_config.py
@@ -36,7 +36,7 @@ def testSimpleObjects(self):
                          'mars-pod-reader-binding')
 
         service_config_dict = ServiceConfig(
-            'mars-test-service', 'NodePort', 'name=marsscheduler', 7103, 7103).build()
+            'mars-test-service', 'NodePort', 'mars/service-type=marsscheduler', 7103, 7103).build()
         self.assertEqual(service_config_dict['metadata']['name'], 'mars-test-service')
 
     def testSchedulerObject(self):
diff --git a/mars/deploy/kubernetes/tests/test_kubernetes.py b/mars/deploy/kubernetes/tests/test_kubernetes.py
--- a/mars/deploy/kubernetes/tests/test_kubernetes.py
+++ b/mars/deploy/kubernetes/tests/test_kubernetes.py
@@ -100,6 +100,8 @@ def testRunInKubernetes(self):
             extra_vol_config = HostPathVolumeConfig('mars-src-path', '/mnt/mars', MARS_ROOT)
             cluster_client = new_cluster(api_client, image=self._docker_image,
                                          worker_spill_paths=[temp_spill_dir],
+                                         extra_labels={'mars-test/group': 'test-label-name'},
+                                         extra_env={'MARS_K8S_GROUP_LABELS': 'mars-test/group'},
                                          extra_volumes=[extra_vol_config],
                                          pre_stop_command=['rm', '/tmp/stopping.tmp'],
                                          timeout=600, log_when_fail=True)