kalavai-client 0.5.0__py2.py3-none-any.whl

kalavai_client/cluster.py

@@ -0,0 +1,308 @@
+import os
+import time
+from pathlib import Path
+from abc import ABC, abstractmethod
+
+from kalavai_client.utils import (
+    run_cmd,
+    check_gpu_drivers,
+    validate_poolconfig,
+    user_path
+)
+
+
+class Cluster(ABC):
+    @abstractmethod
+    def start_seed_node(self, ip_address, labels, flannel_iface):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def start_worker_node(self, url, token, node_name, auth_key, watcher_service, ip_address, labels, flannel_iface):
+        raise NotImplementedError()
+
+
+    @abstractmethod
+    def update_dependencies(self, dependencies_files):
+        raise NotImplementedError()
+
+
+    @abstractmethod
+    def remove_agent(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def is_agent_running(self) -> bool:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def is_seed_node(self) -> bool:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def is_cluster_init(self) -> bool:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def pause_agent(self) -> bool:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def restart_agent(self) -> bool:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_cluster_token(self) -> str:
+        raise NotImplementedError()
+    
+    @abstractmethod
+    def diagnostics(self) -> str:
+        raise NotImplementedError()
+    
+    @abstractmethod
+    def validate_cluster(self) -> bool:
+        raise NotImplementedError
+    
+class dockerCluster(Cluster):
+    def __init__(self, container_name, compose_file, kubeconfig_file, poolconfig_file, dependencies_file, kube_version="v1.31.1+k3s1", flannel_iface=None):
+        self.kube_version = kube_version
+        self.container_name = container_name
+        self.compose_file = compose_file
+        self.kubeconfig_file = kubeconfig_file
+        self.poolconfig_file = poolconfig_file
+        self.dependencies_file = dependencies_file
+
+        if flannel_iface is not None:
+            self.default_flannel_iface = flannel_iface
+        else:
+            self.default_flannel_iface = ""
+        
+    def start_seed_node(self):
+        
+        run_cmd(f"docker compose -f {self.compose_file} up -d")
+        time.sleep(5)
+        run_cmd(f"docker cp {self.container_name}:/etc/rancher/k3s/k3s.yaml {self.kubeconfig_file}")
+
+    def start_worker_node(self):
+        run_cmd(f"docker compose -f {self.compose_file} up -d")
+        
+
+    def update_dependencies(self, dependencies_file=None, debug=False, retries=3):
+        if dependencies_file is not None:
+            self.dependencies_file = dependencies_file
+        if debug:
+            output = ""
+        else:
+            output = " >/dev/null 2>&1"
+        while True:
+            try:
+                home = user_path("")
+                run_cmd(f"docker run --rm --net=host -v {home}:{home} ghcr.io/helmfile/helmfile:v0.169.2 helmfile sync --file {self.dependencies_file} --kubeconfig {self.kubeconfig_file} {output}")
+                #run_cmd(f"helmfile sync --file {self.dependencies_file} --kubeconfig {self.kubeconfig_file} {output}")
+                break
+            except Exception as e:
+                if retries > 0:
+                    retries -= 1
+                    print(f"[{retries}] Dependencies failed. Retrying...")
+                else:
+                    raise Exception(f"Dependencies failed. Are you connected to the internet?\n\nTrace: {str(e)}")
+
+    def remove_agent(self):
+        try:
+            run_cmd(f'docker compose -f {self.compose_file} down')
+            return True
+        except:
+            return False
+
+    def is_agent_running(self):
+        if not os.path.isfile(self.compose_file):
+            return False
+        status = self.container_name in run_cmd(f"docker compose -f {self.compose_file} ps --services --status=running").decode()
+        return status
+
+    def is_seed_node(self):
+        if not os.path.isfile(self.compose_file):
+            return False
+        status = "server" in run_cmd(f"docker compose -f {self.compose_file} ps --services --status=running").decode()
+        return status
+    
+    def is_cluster_init(self):
+        if not os.path.isfile(self.compose_file):
+            return False
+        status = self.container_name in run_cmd(f"docker compose -f {self.compose_file} ps --services --all").decode()
+        return status
+
+    def pause_agent(self):
+        status = False
+        try:
+            run_cmd(f'docker compose -f {self.compose_file} stop')
+            status = True
+        except:
+            pass
+        return status
+
+    def restart_agent(self):
+        try:
+            run_cmd(f'docker compose -f {self.compose_file} start')
+
+        except:
+            pass
+        return self.is_agent_running()
+
+    def get_cluster_token(self):
+        if self.is_seed_node():
+            return run_cmd(f"docker container exec {self.container_name} cat /var/lib/rancher/k3s/server/node-token").decode()
+            #return run_cmd("sudo k3s token create --kubeconfig /etc/rancher/k3s/k3s.yaml --ttl 0").decode()
+        else:
+            return None
+    
+    def diagnostics(self) -> str:
+        # TODO: check cache files are in order
+        # get cluster status
+        if self.is_seed_node():
+            return run_cmd(f"docker exec {self.container_name} kubectl get pods -A -o wide").decode() + "\n\n" + run_cmd(f"docker exec {self.container_name} kubectl get nodes").decode()
+        else:
+            return None
+        
+    def validate_cluster(self) -> bool:
+        # check if credentials are present
+        return os.path.isfile(self.poolconfig_file)
+
+
+class k3sCluster(Cluster):
+
+    def __init__(self, kubeconfig_file, poolconfig_file, dependencies_file, kube_version="v1.31.1+k3s1", flannel_iface=None):
+        self.kube_version = kube_version
+        self.kubeconfig_file = kubeconfig_file
+        self.poolconfig_file = poolconfig_file
+        self.dependencies_file = dependencies_file
+
+        if flannel_iface is not None:
+            self.default_flannel_iface = flannel_iface
+        else:
+            self.default_flannel_iface = ""
+        try:
+            if check_gpu_drivers():
+                self.node_labels = "--node-label gpu=on"
+        except:
+            print("[Warning] issues detected with nvidia, GPU has been disabled for this node")
+            self.node_labels = ""
+        
+    def start_seed_node(self, ip_address, labels=None, is_public=False):
+        node_labels = self.node_labels 
+        if labels is not None:
+            for key, value in labels.items():
+                node_labels += f" --node-label {key}={value}"
+        if is_public:
+            flannel_iface = f"--flannel-iface {self.default_flannel_iface}"
+        else:
+            flannel_iface = ""
+        run_cmd(f'curl -sfL https://get.k3s.io | INSTALL_K3S_VERSION="{self.kube_version}" INSTALL_K3S_EXEC="server --node-ip {ip_address} --node-external-ip {ip_address} {flannel_iface} --flannel-backend wireguard-native {node_labels}" sh - >/dev/null 2>&1')
+        run_cmd(f"sudo cp /etc/rancher/k3s/k3s.yaml {self.kubeconfig_file}")
+        run_cmd(f"sudo chown $USER {self.kubeconfig_file}")
+
+
+    def start_worker_node(self, url, token, node_name, ip_address, labels=None, is_public=False):
+        node_labels = self.node_labels 
+        if labels is not None:
+            for key, value in labels.items():
+                node_labels += f" --node-label {key}={value}"
+        if is_public:
+            flannel_iface = f"--flannel-iface {self.default_flannel_iface}"
+        else:
+            flannel_iface = ""
+        command = f'curl -sfL https://get.k3s.io | INSTALL_K3S_VERSION="{self.kube_version}" INSTALL_K3S_EXEC="agent --token {token} --server https://{url}:6443 --node-name {node_name} --node-ip {ip_address} --node-external-ip {ip_address} {flannel_iface} {node_labels}" sh - >/dev/null 2>&1'
+        run_cmd(command)
+        
+
+    def update_dependencies(self, dependencies_file=None, debug=False, retries=3):
+        if dependencies_file is not None:
+            self.dependencies_file = dependencies_file
+        if debug:
+            output = ""
+        else:
+            output = " >/dev/null 2>&1"
+        while True:
+            try:
+                run_cmd(f"helmfile sync --file {self.dependencies_file} --kubeconfig {self.kubeconfig_file} {output}")
+                break
+            except Exception as e:
+                if retries > 0:
+                    retries -= 1
+                    print(f"[{retries}] Dependencies failed. Retrying...")
+                else:
+                    raise Exception(f"Dependencies failed. Are you connected to the internet?\n\nTrace: {str(e)}")
+                
+
+    def remove_agent(self):
+        try:
+            run_cmd('/usr/local/bin/k3s-uninstall.sh >/dev/null 2>&1')
+            run_cmd('sudo rm -r /etc/rancher/node/ >/dev/null 2>&1')
+            return True
+        except:
+            pass
+        try:
+            run_cmd('/usr/local/bin/k3s-agent-uninstall.sh >/dev/null 2>&1')
+            return True
+        except:
+            pass
+        return False
+
+    def is_agent_running(self):
+        status = (0 == os.system('sudo systemctl is-active --quiet k3s-agent.service')) or (0 == os.system('sudo systemctl is-active --quiet k3s.service'))
+        return status
+
+    def is_seed_node(self):
+        return 0 == os.system('sudo systemctl is-active --quiet k3s.service')
+
+    def is_cluster_init(self):
+        status = Path("/usr/local/bin/k3s-agent-uninstall.sh").is_file() or Path("/usr/local/bin/k3s-uninstall.sh").is_file()
+        return status
+
+    def pause_agent(self):
+        status = False
+        try:
+            run_cmd('sudo systemctl stop k3s >/dev/null 2>&1')
+            status = True
+        except:
+            pass
+        try:
+            run_cmd('sudo systemctl stop k3s-agent >/dev/null 2>&1')
+            status = True
+        except:
+            pass
+        return status
+
+    def restart_agent(self):
+        try:
+            run_cmd('sudo systemctl start k3s >/dev/null 2>&1')
+        except:
+            pass
+        try:
+            run_cmd('sudo systemctl start k3s-agent >/dev/null 2>&1')
+        except:
+            pass
+        return self.is_agent_running()
+
+    def get_cluster_token(self):
+        if self.is_seed_node():
+            return run_cmd("sudo cat /var/lib/rancher/k3s/server/node-token").decode()
+            #return run_cmd("sudo k3s token create --kubeconfig /etc/rancher/k3s/k3s.yaml --ttl 0").decode()
+        else:
+            return None
+    
+    def diagnostics(self) -> str:
+        if self.is_seed_node():
+            return run_cmd(f"k3s kubectl get pods -A -o wide --kubeconfig {self.kubeconfig_file}").decode() + "\n\n" + run_cmd(f"k3s kubectl get nodes --kubeconfig {self.kubeconfig_file}").decode()
+        else:
+            return None
+        
+    def validate_cluster(self) -> bool:
+        if not self.is_cluster_init():
+            raise ValueError("Pool not initialised")
+        if not self.is_agent_running():
+            raise ValueError("Pool initialised but agent is not running")
+        # check cache files
+        if self.is_seed_node():
+            if not validate_poolconfig(self.poolconfig_file):
+                raise ValueError("Cache missconfigured. Run 'kalavai pool stop' to clear.")
+        return True