feat: actualize ydb_configure for k8s + add missing features

ydb/tools/cfg/base.py

@@ -3,6 +3,7 @@
 import collections
 import copy
 import os
+import logging
 from concurrent.futures import ThreadPoolExecutor
 
 from ydb.tools.cfg import types, validation, walle
@@ -22,6 +23,8 @@
 
 DEFAULT_BOX_ID = 1
 
+logger = logging.getLogger()
+
 # MOVE THESE DEFAULTS TO YDB CODE
 VDISKS_DEFAULT_CONFIG = {
     "vdisk_kinds": [
@@ -60,10 +63,11 @@ def merge_with_default(dft, override):
 
 
 class KiKiMRDrive(object):
-    def __init__(self, type, path, shared_with_os=False, expected_slot_count=None, kind=None):
+    def __init__(self, type, path, shared_with_os=False, expected_slot_count=None, kind=None, pdisk_config=None):
         self.type = type
         self.path = path
         self.shared_with_os = shared_with_os
+        self.pdisk_config = pdisk_config
         self.expected_slot_count = expected_slot_count
         self.kind = kind
 
@@ -74,10 +78,11 @@ def __eq__(self, other):
             and self.shared_with_os == other.shared_with_os
             and self.expected_slot_count == other.expected_slot_count
             and self.kind == other.kind
+            and self.pdisk_config == other.pdisk_config
         )
 
     def __hash__(self):
-        return hash("\0".join(map(str, (self.type, self.path, self.shared_with_os, self.expected_slot_count, self.kind))))
+        return hash("\0".join(map(str, (self.type, self.path, self.shared_with_os, self.expected_slot_count, self.kind, self.pdisk_config))))
 
 
 Domain = collections.namedtuple(
@@ -264,7 +269,7 @@ def normalize_domain(domain_name):
 
 
 class ClusterDetailsProvider(object):
-    def __init__(self, template, walle_provider, validator=None, database=None, use_new_style_cfg=False):
+    def __init__(self, template, host_info_provider, validator=None, database=None, use_new_style_cfg=False):
         if not validator:
             validator = validation.default_validator()
 
@@ -278,9 +283,13 @@ def __init__(self, template, walle_provider, validator=None, database=None, use_
             self.__cluster_description = self.get_subjective_description(self.__cluster_description, database, self.__validator)
 
         self._use_walle = self.__cluster_description.get("use_walle", True)
-        if not walle_provider:
-            walle_provider = walle.NopHostsInformationProvider()
-        self._walle = walle_provider
+        self._k8s_settings = self.__cluster_description.get("k8s_settings", {"use": False})
+
+        if host_info_provider is not None:
+            self._host_info_provider = host_info_provider
+        else:
+            self._host_info_provider = walle.NopHostsInformationProvider
+
         self.__translated_storage_pools_deprecated = None
         self.__translated_hosts = None
         self.__racks = {}
@@ -299,6 +308,7 @@ def __init__(self, template, walle_provider, validator=None, database=None, use_
         self.blob_storage_config = self.__cluster_description.get("blob_storage_config")
         self.memory_controller_config = self.__cluster_description.get("memory_controller_config", {})
         self.channel_profile_config = self.__cluster_description.get("channel_profile_config")
+        self.immediate_controls_config = self.__cluster_description.get("immediate_controls_config")
         self.pdisk_key_config = self.__cluster_description.get("pdisk_key_config", {})
         if not self.need_txt_files and not self.use_new_style_kikimr_cfg:
             assert "cannot remove txt files without new style kikimr cfg!"
@@ -344,6 +354,18 @@ def storage_config_generation(self):
     def use_walle(self):
         return self._use_walle
 
+    @property
+    def use_k8s_api(self):
+        return self._k8s_settings.get("use")
+
+    @property
+    def k8s_rack_label(self):
+        return self._k8s_settings.get("k8s_rack_label")
+
+    @property
+    def k8s_dc_label(self):
+        return self._k8s_settings.get("k8s_dc_label")
+
     @property
     def security_settings(self):
         return self.__cluster_description.get("security_settings", {})
@@ -358,23 +380,23 @@ def _get_datacenter(self, host_description):
         dc = host_description.get("location", {}).get("data_center", None)
         if dc:
             return str(dc)
-        return str(self._walle.get_datacenter(host_description.get("name", host_description.get("host"))))
+        return str(self._host_info_provider.get_datacenter(host_description.get("name", host_description.get("host"))))
 
     def _get_rack(self, host_description):
         if host_description.get("rack") is not None:
             return str(host_description.get("rack"))
         rack = host_description.get("location", {}).get("rack", None)
         if rack:
             return str(rack)
-        return str(self._walle.get_rack(host_description.get("name", host_description.get("host"))))
+        return str(self._host_info_provider.get_rack(host_description.get("name", host_description.get("host"))))
 
     def _get_body(self, host_description):
         if host_description.get("body") is not None:
             return str(host_description.get("body"))
         body = host_description.get("location", {}).get("body", None)
         if body:
             return str(body)
-        return str(self._walle.get_body(host_description.get("name", host_description.get("host"))))
+        return str(self._host_info_provider.get_body(host_description.get("name", host_description.get("host"))))
 
     def _collect_drives_info(self, host_description):
         host_config_id = host_description.get("host_config_id", None)
@@ -411,6 +433,7 @@ def hosts(self):
         if self._hosts is not None:
             return self._hosts
         futures = []
+
         for node_id, host_description in enumerate(self.__cluster_description.get("hosts"), 1):
             futures.append(self._thread_pool.submit(self.__collect_host_info, node_id, host_description))
 

ydb/tools/cfg/bin/__main__.py

@@ -12,6 +12,7 @@
 from ydb.tools.cfg.static import StaticConfigGenerator
 from ydb.tools.cfg.utils import write_to_file
 from ydb.tools.cfg.walle import NopHostsInformationProvider, WalleHostsInformationProvider
+from ydb.tools.cfg.k8s_api import K8sApiHostsInformationProvider
 
 logging_config.dictConfig(
     {
@@ -48,19 +49,21 @@ def cfg_generate(args):
     with open(args.cluster_description, "r") as yaml_template:
         cluster_template = yaml.safe_load(yaml_template)
 
-    hosts_provider = NopHostsInformationProvider()
+    host_info_provider = NopHostsInformationProvider()
     if args.hosts_provider_url:
-        hosts_provider = WalleHostsInformationProvider(args.hosts_provider_url)
+        host_info_provider = WalleHostsInformationProvider(args.hosts_provider_url)
+    elif args.hosts_provider_k8s:
+        host_info_provider = K8sApiHostsInformationProvider(args.kubeconfig)
 
-    generator = cfg_cls(cluster_template, args.binary_path, args.output_dir, walle_provider=hosts_provider, **kwargs)
+    generator = cfg_cls(cluster_template, args.binary_path, args.output_dir, host_info_provider=host_info_provider, **kwargs)
 
     all_configs = generator.get_all_configs()
     for cfg_name, cfg_value in all_configs.items():
         write_to_file(os.path.join(args.output_dir, cfg_name), cfg_value)
 
 
 def main():
-    parser = get_parser(cfg_generate, [{"name": "--hosts-provider-url", "help": "URL from which information about hosts can be obtained."}])
+    parser = get_parser(cfg_generate)
     args = parser.parse_args()
     args.func(args)
 

ydb/tools/cfg/configurator_setup.py

@@ -1,6 +1,8 @@
 import argparse
 import random
 
+from pathlib import Path
+
 
 def parse_optional_arguments(args):
     kwargs = {}
@@ -73,6 +75,31 @@ def get_parser(generate_func, extra_cfg_arguments=[]):
             help=v['help'],
         )
 
+    host_location_group = parser_cfg.add_mutually_exclusive_group()
+    host_location_group.add_argument(
+        "--hosts-provider-k8s", 
+        action='store_true', 
+        help="""If specified, K8s api will be used to query information 
+        about host location. Mutually exclusive with --host-location-url.
+        Usage example in template:
+
+        k8s_settings:
+            use: true
+            k8s_rack_label: <label>
+            k8s_dc_label: <label>""")
+    host_location_group.add_argument(
+        "--hosts-provider-url", 
+        type=str, 
+        help="""URL from which information about hosts can be obtained. 
+        Mutually exclusive with --hosts-provider-k8s""")
+
+    home_directory = str(Path.home())
+    defaultKubeconfigLocation = "{0}/.kube/config".format(home_directory)
+    parser_cfg.add_argument("--kubeconfig",
+                            type=str,
+                            help="path to the kubeconfig file. Default `$HOME/.kube/config`, also see --hosts-provider-k8s",
+                            default=defaultKubeconfigLocation)
+
     argument_group = parser_cfg.add_mutually_exclusive_group()
 
     argument_group.add_argument(

ydb/tools/cfg/dynamic.py

@@ -20,19 +20,19 @@ def __init__(
         output_dir,
         grpc_endpoint=None,
         local_binary_path=None,
-        walle_provider=None,
+        host_info_provider=None,
         **kwargs
     ):
         self._template = template
         self._binary_path = binary_path
         self._local_binary_path = local_binary_path or binary_path
         self._output_dir = output_dir
-        self._walle_provider = walle_provider
-        self._cluster_details = base.ClusterDetailsProvider(template, walle_provider=self._walle_provider)
+        self._host_info_provider = host_info_provider
+        self._cluster_details = base.ClusterDetailsProvider(template, host_info_provider=self._host_info_provider)
         self._grpc_endpoint = grpc_endpoint
         self.__configure_request = None
         self.__static_config = static.StaticConfigGenerator(
-            template, binary_path, output_dir, walle_provider=walle_provider, local_binary_path=local_binary_path
+            template, binary_path, output_dir, host_info_provider=host_info_provider, local_binary_path=local_binary_path
         )
 
     @property

ydb/tools/cfg/k8s_api/__init__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from .k8s_api import K8sApiHostsInformationProvider
+
+__all__ = ("k8s_api",)

ydb/tools/cfg/k8s_api/k8s_api.py

@@ -0,0 +1,68 @@
+import logging
+import threading
+import hashlib
+import logging
+
+from kubernetes import client, config
+
+from ydb.tools.cfg.walle import HostsInformationProvider
+
+logger = logging.getLogger()
+
+class K8sApiHostsInformationProvider(HostsInformationProvider):
+    def __init__(self, kubeconfig):
+        self._kubeconfig = kubeconfig
+        self._cache = {}
+        self._timeout_seconds = 5
+        self._retry_count = 10
+        self._lock = threading.Lock()
+        self._k8s_rack_label = None
+        self._k8s_dc_label = None
+
+    def _init_k8s_labels(self, rack_label, dc_label):
+        self._k8s_rack_label = rack_label
+        self._k8s_dc_label = dc_label
+        logger.info(f"initialized rack with {rack_label}, dc with {dc_label}")
+        self._populate_cache()
+
+    def _populate_cache(self):
+        try:
+            config.load_kube_config(config_file=self._kubeconfig)
+            with client.ApiClient() as api_client:
+                v1 = client.CoreV1Api(api_client)
+                nodes = v1.list_node().items
+
+            with self._lock:
+                for node in nodes:
+                    hostname = node.metadata.name
+                    self._cache[hostname] = node.metadata.labels
+        except client.exceptions.ApiException as e:
+            print(f"Failed to fetch node labels: {e}")
+
+    def get_rack(self, hostname):
+        if self._k8s_rack_label is None:
+            return "defaultRack"
+
+        labels = self._cache.get(hostname)
+        if labels and self._k8s_rack_label in labels:
+            logging.info(f"get_rack invoked on {hostname}, value {labels[self._k8s_rack_label]}")
+            return labels[self._k8s_rack_label]
+        logging.info(f"rack not found for hostname {hostname}")
+        return ""
+
+    def get_datacenter(self, hostname):
+        logging.info(f"get_datacenter invoked on {hostname}")
+
+        if self._k8s_dc_label is None:
+            return "defaultDC"
+
+        labels = self._cache.get(hostname)
+        if labels and self._k8s_dc_label in labels:
+            return labels[self._k8s_dc_label]
+        return ""
+
+    def get_body(self, hostname):
+        # Just something for now, please present better ideas
+        hex_digest = hashlib.md5(hostname.encode()).hexdigest()
+        decimal_value = int(hex_digest, 16) % (1 << 31)
+        return decimal_value

ydb/tools/cfg/k8s_api/ya.make

@@ -0,0 +1,12 @@
+PY3_LIBRARY()
+
+PY_SRCS(
+    __init__.py
+    k8s_api.py
+)
+
+PEERDIR(
+    contrib/python/kubernetes
+)
+
+END()