kalavai-client 0.5.14__py3-none-any.whl → 0.5.16__py3-none-any.whl

kalavai_client/cli.py

@@ -7,19 +7,30 @@ import time
 import socket
 from pathlib import Path
 from getpass import getpass
-import ipaddress
 from sys import exit
 
 import yaml
-import netifaces as ni
+
 import arguably
 from rich.console import Console
 
+from kalavai_client.cluster import CLUSTER
 from kalavai_client.env import (
     USER_COOKIE,
     USER_LOCAL_SERVER_FILE,
     TEMPLATE_LABEL,
-    user_path
+    KALAVAI_PLATFORM_URL,
+    DEFAULT_VPN_CONTAINER_NAME,
+    CONTAINER_HOST_PATH,
+    USER_COMPOSE_FILE,
+    USER_HELM_APPS_FILE,
+    USER_KUBECONFIG_FILE,
+    USER_VPN_COMPOSE_FILE,
+    USER_TEMPLATES_FOLDER,
+    DOCKER_COMPOSE_GUI,
+    USER_GUI_COMPOSE_FILE,
+    user_path,
+    resource_path,
 )
 from kalavai_client.core import (
     fetch_resources,
@@ -28,22 +39,29 @@ from kalavai_client.core import (
     fetch_devices,
     fetch_job_logs,
     fetch_gpus,
-    load_gpu_models
+    load_gpu_models,
+    fetch_job_templates,
+    fetch_job_defaults,
+    deploy_job,
+    delete_job,
+    check_token,
+    attach_to_pool,
+    join_pool,
+    create_pool,
+    get_ip_addresses,
+    pause_agent,
+    resume_agent
 )
 from kalavai_client.utils import (
     check_gpu_drivers,
+    load_template,
     run_cmd,
-    decode_dict,
     generate_join_token,
     user_confirm,
-    load_template,
-    store_server_info,
     generate_table,
     request_to_server,
-    resource_path,
     safe_remove,
     leave_vpn,
-    get_vpn_details,
     load_server_info,
     user_login,
     user_logout,
@@ -51,9 +69,6 @@ from kalavai_client.utils import (
     register_cluster,
     unregister_cluster,
     get_public_seeds,
-    validate_join_public_seed,
-    is_storage_compatible,
-    is_watcher_alive,
     load_user_session,
     SERVER_IP_KEY,
     AUTH_KEY,
@@ -62,112 +77,30 @@ from kalavai_client.utils import (
     WRITE_AUTH_KEY,
     PUBLIC_LOCATION_KEY,
     NODE_NAME_KEY,
-    CLUSTER_NAME_KEY,
-    CLUSTER_IP_KEY,
-    CLUSTER_TOKEN_KEY,
-    WATCHER_PORT_KEY,
-    MANDATORY_TOKEN_FIELDS,
-    USER_NODE_LABEL_KEY,
-    ALLOW_UNREGISTERED_USER_KEY
-)
-from kalavai_client.cluster import (
-    dockerCluster
+    CLUSTER_NAME_KEY
 )
 
 
-KALAVAI_PLATFORM_URL = os.getenv("KALAVAI_PLATFORM_URL", "https://platform.kalavai.net")
 LOCAL_TEMPLATES_DIR = os.getenv("LOCAL_TEMPLATES_DIR", None)
 VERSION = 1
 RESOURCE_EXCLUDE = ["ephemeral-storage", "hugepages-1Gi", "hugepages-2Mi", "pods"]
 CORE_NAMESPACES = ["lws-system", "kube-system", "gpu-operator", "kalavai"]
 RAY_LABEL = "kalavai.ray.name"
 PVC_NAME_LABEL = "kalavai.storage.name"
-DOCKER_COMPOSE_TEMPLATE = resource_path("kalavai_client/assets/docker-compose-template.yaml")
 VPN_COMPOSE_TEMPLATE = resource_path("kalavai_client/assets/vpn-template.yaml")
-POOL_CONFIG_TEMPLATE = resource_path("kalavai_client/assets/pool_config_template.yaml")
-POOL_CONFIG_DEFAULT_VALUES = resource_path("kalavai_client/assets/pool_config_values.yaml")
-USER_WORKSPACE_TEMPLATE = resource_path("kalavai_client/assets/user_workspace.yaml")
-DEFAULT_USER_WORKSPACE_VALUES = resource_path("kalavai_client/assets/user_workspace_values.yaml")
 STORAGE_CLASS_NAME = "local-path"
 STORAGE_ACCESS_MODE = ["ReadWriteOnce"]
-STORAGE_CLASS_LABEL = "kalavai.storage.enabled"
 DEFAULT_STORAGE_NAME = "pool-cache"
 DEFAULT_STORAGE_SIZE = 20
-DEFAULT_WATCHER_PORT = 30001
-USER_NODE_LABEL = "kalavai.cluster.user"
-KUBE_VERSION = os.getenv("KALAVAI_KUBE_VERSION", "v1.31.1+k3s1")
-DEFAULT_FLANNEL_IFACE = os.getenv("KALAVAI_FLANNEL_IFACE", "netmaker-1")
-FORBIDEDEN_IPS = ["127.0.0.1"]
-# kalavai templates
-HELM_APPS_FILE = resource_path("kalavai_client/assets/apps.yaml")
-HELM_APPS_VALUES = resource_path("kalavai_client/assets/apps_values.yaml")
-# user specific config files
-DEFAULT_CONTAINER_NAME = "kalavai"
-DEFAULT_VPN_CONTAINER_NAME = "kalavai-vpn"
-CONTAINER_HOST_PATH = user_path("pool/", create_path=True)
-USER_COMPOSE_FILE = user_path("docker-compose-worker.yaml")
-USER_VPN_COMPOSE_FILE = user_path("docker-compose-vpn.yaml")
-USER_HELM_APPS_FILE = user_path("apps.yaml")
-USER_KUBECONFIG_FILE = user_path("kubeconfig")
-USER_TEMPLATES_FOLDER = user_path("templates", create_path=True)
-
 
+    
 console = Console()
-CLUSTER = dockerCluster(
-    container_name=DEFAULT_CONTAINER_NAME,
-    kube_version=KUBE_VERSION,
-    flannel_iface=DEFAULT_FLANNEL_IFACE,
-    compose_file=USER_COMPOSE_FILE,
-    kubeconfig_file=USER_KUBECONFIG_FILE,
-    poolconfig_file=USER_LOCAL_SERVER_FILE,
-    dependencies_file=USER_HELM_APPS_FILE
-)
 
 
 ######################
 ## HELPER FUNCTIONS ##
 ######################
 
-def check_seed_compatibility():
-    """Check required packages to start pools"""
-    logs = []
-    console.log("[white]Checking system requirements...")
-    # docker
-    try:
-        run_cmd("docker version >/dev/null 2>&1")
-    except:
-        logs.append("[red]Docker not installed. Install instructions:\n")
-        logs.append("   Linux: https://docs.docker.com/engine/install/\n")
-        logs.append("   Windows/MacOS: https://docs.docker.com/desktop/\n")
-    
-    if len(logs) == 0:
-        console.log("[green]System is ready to start a pool")
-        return True
-    else:
-        for log in logs:
-            console.log(log)
-        return False
-
-def check_worker_compatibility():
-    """Check required packages to join pools"""
-    logs = []
-    console.log("[white]Checking system requirements...")
-    # docker
-    try:
-        run_cmd("docker version >/dev/null 2>&1")
-    except:
-        logs.append("[red]Docker not installed. Install instructions:\n")
-        logs.append("   Linux: https://docs.docker.com/engine/install/\n")
-        logs.append("   Windows/MacOS: https://docs.docker.com/desktop/\n")
-    
-    if len(logs) == 0:
-        console.log("[green]System is ready to join a pool")
-        return True
-    else:
-        for log in logs:
-            console.log(log)
-        return False
-
 
 def cleanup_local():
     console.log("Removing local cache files...")
@@ -178,6 +111,7 @@ def cleanup_local():
     safe_remove(USER_KUBECONFIG_FILE)
     safe_remove(USER_LOCAL_SERVER_FILE)
     safe_remove(USER_TEMPLATES_FOLDER)
+    safe_remove(USER_GUI_COMPOSE_FILE)
 
 def pre_join_check(node_name, server_url, server_key):
     # check with the server that we can connect
@@ -217,75 +151,9 @@ def set_schedulable(schedulable, node_name=load_server_info(data_key=NODE_NAME_K
     except Exception as e:
         console.log(f"[red]Error when connecting to kalavai service: {str(e)}")
 
-
-def init_user_workspace(force_namespace=None):
-    
-    # load template config and populate with values
-    sidecar_template_yaml = load_template(
-        template_path=USER_WORKSPACE_TEMPLATE,
-        values={},
-        default_values_path=DEFAULT_USER_WORKSPACE_VALUES)
-
-    try:
-        data = {"config": sidecar_template_yaml}
-        if force_namespace is not None:
-            data["force_namespace"] = force_namespace
-        result = request_to_server(
-            method="post",
-            endpoint="/v1/create_user_space",
-            data=data,
-            server_creds=USER_LOCAL_SERVER_FILE,
-            user_cookie=USER_COOKIE
-        )
-        console.log(f"Workspace creation (ignore already created warnings): {result}" )
-    except Exception as e:
-        console.log(f"[red]Error when connecting to kalavai service: {str(e)}")
-
-def pool_init(pool_config_values_path=None):
-    """Deploy configured objects to initialise pool"""
-    if pool_config_values_path is None:
-        return
-    
-    # load template config and populate with values
-    sidecar_template_yaml = load_template(
-        template_path=POOL_CONFIG_TEMPLATE,
-        values={},
-        default_values_path=pool_config_values_path)
-
-    try:
-        result = request_to_server(
-            method="post",
-            endpoint="/v1/deploy_generic_model",
-            data={"config": sidecar_template_yaml},
-            server_creds=USER_LOCAL_SERVER_FILE,
-            user_cookie=USER_COOKIE
-        )
-        if 'failed' in result and len(result['failed']) > 0:
-            console.log(f"[red]Error when deploying pool config\n\n{result['failed']}")
-        if len(result['successful']) > 0:
-            console.log(f"[green]Deployed pool config!")
-    except Exception as e:
-        console.log(f"[red]Error when connecting to kalavai service: {str(e)}")
-
 def select_ip_address(subnet=None):
-    ips = []
-    retry = 3
-    while len(ips) == 0:
-        for iface in ni.interfaces():
-            try:
-                ip = ni.ifaddresses(iface)[ni.AF_INET][0]['addr']
-                if ip in FORBIDEDEN_IPS:
-                    continue
-                if subnet is None or ipaddress.ip_address(ip) in ipaddress.ip_network(subnet):
-                    ips.append(ip)
-            except:
-                pass
-        if len(ips) == 1:
-            return ips[0]
-        time.sleep(2)
-        retry -= 1
-        if retry < 0:
-            raise ValueError(f"No IPs available on subnet {subnet}")
+    ips = get_ip_addresses(subnet=subnet)
+    
     while True:
         option = user_confirm(
             question="Select IP to advertise the node (needs to be visible to other nodes)",
@@ -336,50 +204,50 @@ def select_token_type():
             break
     return {"admin": choice == 0, "user": choice == 1, "worker": choice == 2}
 
-def generate_compose_config(role, node_name, is_public, use_gpus=True, node_labels=None, pool_ip=None, vpn_token=None, pool_token=None):
+def input_gpus():
     num_gpus = 0
-    if use_gpus:
-        try:
-            has_gpus = check_gpu_drivers()
-            if has_gpus:
-                max_gpus = int(run_cmd("nvidia-smi -L | wc -l").decode())
-                num_gpus = user_confirm(
-                    question=f"{max_gpus} NVIDIA GPU(s) detected. How many GPUs would you like to include?",
-                    options=range(max_gpus+1)
-                ) 
-        except:
-            console.log(f"[red]WARNING: error when fetching NVIDIA GPU info. GPUs will not be used on this local machine")
-    if node_labels is not None:
-        node_labels = " ".join([f"--node-label {key}={value}" for key, value in node_labels.items()])
-    compose_values = {
-        "user_path": user_path(""),
-        "service_name": DEFAULT_CONTAINER_NAME,
-        "vpn": is_public,
-        "vpn_name": DEFAULT_VPN_CONTAINER_NAME,
-        "pool_ip": pool_ip,
-        "pool_token": pool_token,
-        "vpn_token": vpn_token,
-        "node_name": node_name,
-        "command": role,
-        "storage_enabled": "True",
-        "num_gpus": num_gpus,
-        "k3s_path": f"{CONTAINER_HOST_PATH}/k3s",
-        "etc_path": f"{CONTAINER_HOST_PATH}/etc",
-        "node_labels": node_labels,
-        "flannel_iface": DEFAULT_FLANNEL_IFACE if is_public else ""
-    }
-    # generate local config files
-    compose_yaml = load_template(
-        template_path=DOCKER_COMPOSE_TEMPLATE,
-        values=compose_values)
-    with open(USER_COMPOSE_FILE, "w") as f:
-        f.write(compose_yaml)
-    return compose_yaml
+    try:
+        has_gpus = check_gpu_drivers()
+        if has_gpus:
+            max_gpus = int(run_cmd("nvidia-smi -L | wc -l").decode())
+            num_gpus = user_confirm(
+                question=f"{max_gpus} NVIDIA GPU(s) detected. How many GPUs would you like to include?",
+                options=range(max_gpus+1)
+            ) 
+    except:
+        console.log(f"[red]WARNING: error when fetching NVIDIA GPU info. GPUs will not be used on this local machine")
+    return num_gpus
 
 ##################
 ## CLI COMMANDS ##
 ##################
 
+@arguably.command
+def gui__start(*others, gui_port=3000, backend_port=8000):
+    """Run GUI"""
+    values = {
+        "path": user_path(""),
+        "gui_port": gui_port,
+        "backend_port": backend_port
+    }
+    compose_yaml = load_template(
+        template_path=DOCKER_COMPOSE_GUI,
+        values=values)
+    with open(USER_GUI_COMPOSE_FILE, "w") as f:
+        f.write(compose_yaml)
+    
+    run_cmd(f"docker compose --file {USER_GUI_COMPOSE_FILE} up -d")
+
+    console.log(f"[green]Loading GUI, may take a few minutes. It will be available at http://localhost:{gui_port}")
+
+@arguably.command
+def gui__stop(*others):
+    """Stop GUI"""
+    run_cmd(f"docker compose --file {USER_GUI_COMPOSE_FILE} down")
+
+    console.log("[green]Kalavai GUI has been stopped")
+
+
 @arguably.command
 def login(*others,  username: str=None):
     """
@@ -461,8 +329,9 @@ def pool__publish(*others, description=None):
     description = description
     
     try:
-        if not pool__check_token(token=token, public=True):
-            raise ValueError("[red]Cluster must be started with a valid vpn_location to publish")
+        valid = check_token(token=token, public=True)
+        if "error" in valid:
+            raise ValueError(f"[red]Cluster must be started with a valid vpn_location to publish: {valid}")
         cluster_name = load_server_info(data_key=CLUSTER_NAME_KEY, file=USER_LOCAL_SERVER_FILE)
         
         register_cluster(
@@ -523,7 +392,7 @@ def pool__list(*others, user_only=False):
 
 
 @arguably.command
-def pool__start(cluster_name, *others,  only_registered_users: bool=False, ip_address: str=None, location: str=None, app_values: str=HELM_APPS_VALUES, pool_config_values: str=POOL_CONFIG_DEFAULT_VALUES):
+def pool__start(cluster_name, *others,  only_registered_users: bool=False, ip_address: str=None, location: str=None, app_values: str=None, pool_config_values: str=None):
     """
     Start Kalavai pool and start/resume sharing resources.
 
@@ -531,9 +400,6 @@ def pool__start(cluster_name, *others,  only_registered_users: bool=False, ip_ad
         *others: all the other positional arguments go here
     """
 
-    if not check_seed_compatibility():
-        return
-
     if CLUSTER.is_cluster_init():
         console.log(f"[white] You are already connected to {load_server_info(data_key=CLUSTER_NAME_KEY, file=USER_LOCAL_SERVER_FILE)}. Enter [yellow]kalavai pool stop[white] to exit and join another one.")
         return
@@ -547,127 +413,25 @@ def pool__start(cluster_name, *others,  only_registered_users: bool=False, ip_ad
         console.log("Installation was cancelled and did not complete.")
         return
     
-    # if only registered users are allowed, check user has logged in
-    user = defaultdict(lambda: None)
-    if only_registered_users or location is not None:
-        user = user_login(user_cookie=USER_COOKIE)
-        if user is None:
-            console.log("[white]--only-registered-users [red]or [white]--location[red] can only be used if the host is authenticated. Run [yellow]kalavai login[red] to authenticate")
-            exit()
-
-    # join private network if provided
-    vpn = defaultdict(lambda: None)
-    node_labels = {
-        STORAGE_CLASS_LABEL: is_storage_compatible()
-    }
-    if location is not None:
-        console.log("Fetching VPN credentials")
-        try:
-            vpn = get_vpn_details(
-                location=location,
-                user_cookie=USER_COOKIE)
-            node_labels[USER_NODE_LABEL] = user["username"]
-        except Exception as e:
-            console.log(f"[red]Error when joining network: {str(e)}")
-            return
-    
-    # Generate docker compose recipe
-    generate_compose_config(
-        role="server",
-        vpn_token=vpn["key"],
-        node_name=socket.gethostname(),
-        node_labels=node_labels,
-        is_public=location is not None
-    )
-    
-    # start server
-    console.log("Deploying seed...")
-    CLUSTER.start_seed_node()
-    
-    while not CLUSTER.is_agent_running():
-        console.log("Waiting for seed to start...")
-        time.sleep(10)
-    
     # select IP address (for external discovery)
     if ip_address is None and location is None:
         # local IP
         console.log(f"Scanning for valid IPs")
         ip_address = select_ip_address()
-    else:
-        # load VPN ip
-        ip_address = CLUSTER.get_vpn_ip()
+    
     console.log(f"Using {ip_address} address for server")
 
-    # populate local cred files
-    auth_key = str(uuid.uuid4())
-    write_auth_key = str(uuid.uuid4())
-    readonly_auth_key = str(uuid.uuid4())
-    
-    watcher_service = f"{ip_address}:{DEFAULT_WATCHER_PORT}"
-    values = {
-        CLUSTER_NAME_KEY: cluster_name,
-        CLUSTER_IP_KEY: ip_address,
-        AUTH_KEY: auth_key,
-        READONLY_AUTH_KEY: readonly_auth_key,
-        WRITE_AUTH_KEY: write_auth_key,
-        WATCHER_PORT_KEY: DEFAULT_WATCHER_PORT,
-        WATCHER_SERVICE_KEY: watcher_service,
-        USER_NODE_LABEL_KEY: USER_NODE_LABEL,
-        ALLOW_UNREGISTERED_USER_KEY: not only_registered_users
-    }
+    console.log(f"[green]Creating {cluster_name} pool, this may take a few minutes...")
 
-    store_server_info(
-        server_ip=ip_address,
-        auth_key=auth_key,
-        readonly_auth_key=readonly_auth_key,
-        write_auth_key=write_auth_key,
-        file=USER_LOCAL_SERVER_FILE,
-        watcher_service=watcher_service,
-        node_name=socket.gethostname(),
+    create_pool(
         cluster_name=cluster_name,
-        public_location=location,
-        user_api_key=user["api_key"])
-    
-    # Generate helmfile recipe
-    helm_yaml = load_template(
-        template_path=HELM_APPS_FILE,
-        values=values,
-        default_values_path=app_values,
-        force_defaults=True)
-    with open(USER_HELM_APPS_FILE, "w") as f:
-        f.write(helm_yaml)
-    
-    console.log("[green]Config files have been generated in your local machine\n")
-    
-    console.log("Setting pool dependencies...")
-    # set template values in helmfile
-    try:
-        CLUSTER.update_dependencies(
-            dependencies_file=USER_HELM_APPS_FILE
-        )
-    except Exception as e:
-        console.log(f"Error: {str(e)}")
-        exit()
-    console.log("[green]Your pool is ready! Grow it by sharing your joining token with others. Run [yellow]kalavai pool token[green] to generate one.")
-
-    if location is not None:
-        # register with kalavai if it's a public cluster
-        console.log("Registering public cluster with Kalavai...")
-        pool__publish()
-    
-    # wait until the server is ready to create objects
-    while True:
-        console.log("Waiting for core services to be ready, may take a few minutes...")
-        time.sleep(30)
-        if is_watcher_alive(server_creds=USER_LOCAL_SERVER_FILE, user_cookie=USER_COOKIE):
-            break
-    console.log("Initialise user workspace...")
-    pool_init(pool_config_values_path=pool_config_values)
-    # init default namespace
-    init_user_workspace(force_namespace="default")
-    if only_registered_users:
-        # init user namespace
-        init_user_workspace()
+        ip_address=ip_address,
+        app_values=app_values,
+        pool_config_values=pool_config_values,
+        num_gpus=input_gpus(),
+        only_registered_users=only_registered_users,
+        location=location
+    )
 
     return None
 
@@ -720,20 +484,13 @@ def pool__check_token(token, *others, public=False):
     """
     Utility to check the validity of a join token
     """
-    try:
-        data = decode_dict(token)
-        for field in MANDATORY_TOKEN_FIELDS:
-            assert field in data
-        if public:
-            if data[PUBLIC_LOCATION_KEY] is None:
-                raise ValueError("Token is not valid for public pools. Did you start the cluster with a public_location?")
-        console.log("[green]Token format is correct")
-        return True
-    except Exception as e:
-        console.log(f"[white]{str(e)}")
-        console.log("[red]Token is invalid.")
+    result = check_token(token=token, public=public)
+    if "error" in result:
+        console.log(f"[red]Error in token: {result}")
         return False
-
+    
+    console.log("[green]Token format is correct")
+    return True
 
 @arguably.command
 def pool__join(token, *others, node_name=None):
@@ -743,9 +500,6 @@ def pool__join(token, *others, node_name=None):
     Args:
         *others: all the other positional arguments go here
     """
-
-    if not check_worker_compatibility():
-        return
     
     # check that k3s is not running already in the host
     # k3s service running or preinstalled
@@ -763,132 +517,26 @@ def pool__join(token, *others, node_name=None):
             console.log("[green]Nothing happened.")
             return
     
-    if node_name is None:
-        node_name = socket.gethostname()
-    
-    # check token
-    if not pool__check_token(token):
-        return
-
-    try:
-        data = decode_dict(token)
-        kalavai_seed_ip = data[CLUSTER_IP_KEY]
-        kalavai_token = data[CLUSTER_TOKEN_KEY]
-        cluster_name = data[CLUSTER_NAME_KEY]
-        auth_key = data[AUTH_KEY]
-        watcher_service = data[WATCHER_SERVICE_KEY]
-        public_location = data[PUBLIC_LOCATION_KEY]
-        vpn = defaultdict(lambda: None)
-    except Exception as e:
-        console.log(str(e))
-        console.log("[red] Invalid token")
-        return
-    
-    # join private network if provided
-    node_labels = {
-        STORAGE_CLASS_LABEL: is_storage_compatible()
-    }
-    user = defaultdict(lambda: None)
-    if public_location is not None:
-        user = user_login(user_cookie=USER_COOKIE)
-        if user is None:
-            console.log("[red]Must be logged in to join public pools. Run [yellow]kalavai login[red] to authenticate")
-            exit()
-        console.log("Fetching VPN credentials")
-        try:
-            vpn = get_vpn_details(
-                location=public_location,
-                user_cookie=USER_COOKIE)
-            node_labels[USER_NODE_LABEL] = user["username"]
-        except Exception as e:
-            console.log(f"[red]Error when joining network: {str(e)}")
-            console.log("Are you authenticated? Try [yellow]kalavai login")
-            return
-        try:
-            validate_join_public_seed(
-                cluster_name=cluster_name,
-                join_key=token,
-                user_cookie=USER_COOKIE
-            )
-        except Exception as e:
-            console.log(f"[red]Error when joining network: {str(e)}")
-            return
-
-    # send note to server to let them know the node is coming online
-    # TODO: won't be able to check for VPN pools...
-    # if not pre_join_check(node_name=node_name, server_url=watcher_service, server_key=auth_key):
-    #     console.log(f"[red] Failed pre join checks. Server offline or node '{node_name}' may already exist. Please specify a different one with '--node-name'")
-    #     leave_vpn(container_name=DEFAULT_VPN_CONTAINER_NAME)
-    #     return
-        
-    # local agent join
-    # 1. Generate local cache files
-    console.log("Generating config files...")
-    
-    # Generate docker compose recipe
-    generate_compose_config(
-        role="agent",
-        pool_ip=f"https://{kalavai_seed_ip}:6443",
-        pool_token=kalavai_token,
-        vpn_token=vpn["key"],
-        node_name=node_name,
-        node_labels=node_labels,
-        is_public=public_location is not None)
-    
-    store_server_info(
-        server_ip=kalavai_seed_ip,
-        auth_key=auth_key,
-        file=USER_LOCAL_SERVER_FILE,
-        watcher_service=watcher_service,
-        node_name=node_name,
-        cluster_name=cluster_name,
-        public_location=public_location,
-        user_api_key=user["api_key"])
+    num_gpus = input_gpus()
 
     option = user_confirm(
         question="Docker compose ready. Would you like Kalavai to deploy it?",
         options=["no", "yes"]
     )
     if option == 0:
-        console.log("Manually deploy the worker with the following command:\n")
-        print(f"docker compose -f {USER_COMPOSE_FILE} up -d")
+        console.log("[red]Installation aborted")
         return
     
-    console.log(f"[white] Connecting to {cluster_name} @ {kalavai_seed_ip} (this may take a few minutes)...")
-    try:
-        CLUSTER.start_worker_node()
-    except Exception as e:
-        console.log(f"[red] Error connecting to {cluster_name} @ {kalavai_seed_ip}. Check with the admin if the token is still valid.")
-        pool__stop()
-        exit()
-
-    # ensure we are connected
-    while True:
-        console.log("Waiting for core services to be ready, may take a few minutes...")
-        time.sleep(30)
-        if is_watcher_alive(server_creds=USER_LOCAL_SERVER_FILE, user_cookie=USER_COOKIE):
-            break
-
-    # send note to server to let them know the node is coming online
-    if not pre_join_check(node_name=node_name, server_url=watcher_service, server_key=auth_key):
-        console.log(f"[red] Failed pre join checks. Server offline or node '{node_name}' may already exist. Please specify a different one with [yellow]--node-name'")
-        pool__stop()
-        return
-    
-    # check the node has connected successfully
-    try:
-        while not CLUSTER.is_agent_running():
-            console.log("waiting for runner, may take a few minutes... Press <ctrl+c> to stop")
-            time.sleep(30)
-    except KeyboardInterrupt:
-        console.log("[red]Installation aborted. Leaving pool.")
-        pool__stop()
-        return
-    
-    init_user_workspace()
-    
-    # set status to schedulable
-    console.log(f"[green] You are connected to {cluster_name}")
+    console.log("Connecting worker to the pool...")
+    result = join_pool(
+        token=token,
+        node_name=node_name,
+        num_gpus=num_gpus
+    )
+    if "error" in result:
+        console.log(f"[red]Error when connecting: {result}")
+    else:
+        console.log(f"[green] You are connected to {result}")
 
 @arguably.command
 def pool__stop(*others, skip_node_deletion=False):
@@ -942,11 +590,12 @@ def pool__pause(*others):
     """
     # k3s stop locally
     console.log("[white] Pausing kalavai app...")
-    success = CLUSTER.pause_agent()
-    if success:
-        console.log("[white] Kalava sharing paused. Resume with [yellow]kalavai pool resume")
+    success = pause_agent()
+    if "error" in success:
+        console.log(f"[red] Error when stopping. {success['error']}")
     else:
-        console.log("[red] Error when stopping. Please run [yellow]kalavai pool pause[red] again.")
+        console.log("[white] Kalava sharing paused. Resume with [yellow]kalavai pool resume")
+        
 
 @arguably.command
 def pool__resume(*others):
@@ -961,10 +610,12 @@ def pool__resume(*others):
         console.log("[red] Kalavai app was not started before, please run [yellow]kalavai pool start[red] to start a pool or [yellow]kalavai pool join[red] to join one first")
         return
     console.log("[white] Restarting sharing (may take a few minutes)...")
-    if CLUSTER.restart_agent():
-        console.log("[white] Kalava sharing resumed")
+    success = resume_agent()
+    if "error" in success:
+        console.log(f"[red] Error when restarting. {success['error']}")
     else:
-        console.log("[red] Error when restarting. Please run [yellow]kalavai pool resume[white] again.")
+        console.log("[white] Kalava sharing resumed")
+        
 
 
 @arguably.command
@@ -1103,7 +754,7 @@ def pool__attach(token, *others, node_name=None):
     """
 
     if node_name is None:
-        node_name = socket.gethostname()
+        node_name = f"{socket.gethostname()}-{uuid.uuid4().hex[:6]}"
     
     # check that is not attached to another instance
     if os.path.exists(USER_LOCAL_SERVER_FILE):
@@ -1115,70 +766,6 @@ def pool__attach(token, *others, node_name=None):
             console.log("[green]Nothing happened.")
             return
     
-    # check token
-    if not pool__check_token(token):
-        return
-
-    try:
-        data = decode_dict(token)
-        kalavai_seed_ip = data[CLUSTER_IP_KEY]
-        cluster_name = data[CLUSTER_NAME_KEY]
-        auth_key = data[AUTH_KEY]
-        watcher_service = data[WATCHER_SERVICE_KEY]
-        public_location = data[PUBLIC_LOCATION_KEY]
-        vpn = defaultdict(lambda: None)
-    except Exception as e:
-        console.log(str(e))
-        console.log("[red] Invalid token")
-        return
-    
-    user = defaultdict(lambda: None)
-    if public_location is not None:
-        user = user_login(user_cookie=USER_COOKIE)
-        if user is None:
-            console.log("[red]Must be logged in to join public pools. Run [yellow]kalavai login[red] to authenticate")
-            exit()
-        console.log("Fetching VPN credentials")
-        try:
-            vpn = get_vpn_details(
-                location=public_location,
-                user_cookie=USER_COOKIE)
-        except Exception as e:
-            console.log(f"[red]Error when joining network: {str(e)}")
-            console.log("Are you authenticated? Try [yellow]kalavai login")
-            return
-        try:
-            validate_join_public_seed(
-                cluster_name=cluster_name,
-                join_key=token,
-                user_cookie=USER_COOKIE
-            )
-        except Exception as e:
-            console.log(f"[red]Error when joining network: {str(e)}")
-            return
-        
-    # local agent join
-    # 1. Generate local cache files
-    console.log("Generating config files...")
-    
-    # Generate docker compose recipe
-    generate_compose_config(
-        use_gpus=False,
-        role="",
-        vpn_token=vpn["key"],
-        node_name=node_name,
-        is_public=public_location is not None)
-    
-    store_server_info(
-        server_ip=kalavai_seed_ip,
-        auth_key=auth_key,
-        file=USER_LOCAL_SERVER_FILE,
-        watcher_service=watcher_service,
-        node_name=node_name,
-        cluster_name=cluster_name,
-        public_location=public_location,
-        user_api_key=user["api_key"])
-
     option = user_confirm(
         question="Docker compose ready. Would you like Kalavai to deploy it?",
         options=["no", "yes"]
@@ -1188,17 +775,13 @@ def pool__attach(token, *others, node_name=None):
         print(f"docker compose -f {USER_COMPOSE_FILE} up -d")
         return
     
-    console.log(f"[white] Connecting to {cluster_name} @ {kalavai_seed_ip} (this may take a few minutes)...")
-    run_cmd(f"docker compose -f {USER_COMPOSE_FILE} up -d")
-    # ensure we are connected
-    while True:
-        console.log("Waiting for core services to be ready, may take a few minutes...")
-        time.sleep(30)
-        if is_watcher_alive(server_creds=USER_LOCAL_SERVER_FILE, user_cookie=USER_COOKIE):
-            break
+    result = attach_to_pool(token=token, node_name=node_name)
     
+    if "error" in result:
+        console.log(f"[red]Error when attaching to pool: {result}")
+        return
     # set status to schedulable
-    console.log(f"[green] You are connected to {cluster_name}")
+    console.log(f"[green] You are connected to {result}")
 
 
 @arguably.command
@@ -1422,18 +1005,13 @@ def job__templates(*others):
         console.log(f"[red]Problems with your pool: {str(e)}")
         return
     
-    try:
-        result = request_to_server(
-            method="get",
-            endpoint="/v1/get_job_templates",
-            server_creds=USER_LOCAL_SERVER_FILE,
-            data=None,
-            user_cookie=USER_COOKIE
-        )
-        console.log("Templates available in the pool")
-        console.log(result)
-    except Exception as e:
-        console.log(f"[red]Error when connecting to kalavai service: {str(e)}")
+    templates = fetch_job_templates()
+    if "error" in templates:
+        console.log(f"[red]Error when fetching templates: {str(e)}")
+        return
+    
+    console.log("Templates available in the pool")
+    console.log(templates)
 
 
 @arguably.command
@@ -1489,26 +1067,16 @@ def job__run(template_name, *others, values: str=None, force_namespace: str=None
         annotation_key="nvidia.com/nouse-gputype"
     )
 
-    # deploy template with kube-watcher
-    data = {
-        "template": template_name,
-        "template_values": values_dict
-    }
-    if force_namespace is not None:
-        data["force_namespace"] = force_namespace
+    result = deploy_job(
+        template_name=template_name,
+        values_dict=values_dict,
+        force_namespace=force_namespace
+    )
 
-    try:
-        result = request_to_server(
-            method="post",
-            endpoint="/v1/deploy_job",
-            data=data,
-            server_creds=USER_LOCAL_SERVER_FILE,
-            user_cookie=USER_COOKIE
-        )
+    if "error" in result:
+        console.log(f"[red]Error when deploying job: {str(e)}")
+    else:
         console.log(f"[green]{template_name} job deployed")
-    except Exception as e:
-        console.log(f"[red]Error when connecting to kalavai service: {str(e)}")
-        return
 
 @arguably.command
 def job__test(local_template_dir, *others, values, defaults, force_namespace: str=None):
@@ -1581,22 +1149,12 @@ def job__defaults(template_name, *others):
         return
     
     # deploy template with kube-watcher
-    data = {
-        "template": template_name
-    }
-    try:
-        result = request_to_server(
-            method="get",
-            endpoint="/v1/job_defaults",
-            data=data,
-            server_creds=USER_LOCAL_SERVER_FILE,
-            user_cookie=USER_COOKIE
-        )
-        print(
-            json.dumps(result,indent=3)
-        )
-    except Exception as e:
-        console.log(f"[red]Error when connecting to kalavai service: {str(e)}")
+    defaults = fetch_job_defaults(name=template_name)
+    if "error" in defaults:
+        console.log(f"[red]Error when fetching job defaults: {defaults}")
+    print(
+        json.dumps(defaults, indent=3)
+    )
 
 
 @arguably.command
@@ -1614,23 +1172,11 @@ def job__delete(name, *others, force_namespace: str=None):
         console.log("[WARNING][yellow]--force-namespace [white]requires an admin key. Request will fail if you are not an admin.")
     
     # deploy template with kube-watcher
-    data = {
-        "label": TEMPLATE_LABEL, # this ensures that both lws template and services are deleted
-        "value": name
-    }
-    if force_namespace is not None:
-        data["force_namespace"] = force_namespace
-    try:
-        result = request_to_server(
-            method="post",
-            endpoint="/v1/delete_labeled_resources",
-            data=data,
-            server_creds=USER_LOCAL_SERVER_FILE,
-            user_cookie=USER_COOKIE
-        )
+    result = delete_job(name=name, force_namespace=force_namespace)
+    if "error" in result:
+        console.log(f"[red]Error when deleting job: {str(e)}")
+    else:
         console.log(f"{result}")
-    except Exception as e:
-        console.log(f"[red]Error when connecting to kalavai service: {str(e)}")
 
 
 @arguably.command