kalavai-client 0.5.30__py3-none-any.whl → 0.6.0__py3-none-any.whl

kalavai_client/core.py

@@ -13,16 +13,17 @@ import re
 
 from kalavai_client.cluster import CLUSTER
 from kalavai_client.utils import (
+    DEPLOY_LLM_SIDECARS_KEY,
+    NODE_ROLE_LABEL,
     check_gpu_drivers,
     generate_join_token,
+    load_user_id,
     register_cluster,
     request_to_server,
     load_server_info,
     decode_dict,
-    get_vpn_details,
     send_pool_invite,
     unregister_cluster,
-    validate_join_public_seed,
     generate_compose_config,
     store_server_info,
     is_watcher_alive,
@@ -45,10 +46,8 @@ from kalavai_client.utils import (
     WATCHER_PORT_KEY,
     WATCHER_SERVICE_KEY,
     USER_NODE_LABEL_KEY,
-    ALLOW_UNREGISTERED_USER_KEY
-)
-from kalavai_client.auth import (
-    KalavaiAuthClient
+    ALLOW_UNREGISTERED_USER_KEY,
+    KALAVAI_AUTH
 )
 from kalavai_client.env import (
     USER_COOKIE,
@@ -126,7 +125,7 @@ def set_schedulable(schedulable, node_names):
     except Exception as e:
         return {"error": f"Error when connecting to kalavai service: {str(e)}"}
 
-def init_user_workspace(force_namespace=None):
+def init_user_workspace(user_id=None, node_name=None, force_namespace=None):
     
     # load template config and populate with values
     sidecar_template_yaml = load_template(
@@ -138,6 +137,10 @@ def init_user_workspace(force_namespace=None):
         data = {"config": sidecar_template_yaml}
         if force_namespace is not None:
             data["force_namespace"] = force_namespace
+        if user_id is not None:
+            data["user_id"] = user_id
+        if node_name is not None:
+            data["node_name"] = node_name
         result = request_to_server(
             method="post",
             endpoint="/v1/create_user_space",
@@ -495,31 +498,18 @@ def fetch_gpus(available=False):
 
     except Exception as e:
         return {"error": str(e)}
-
-def load_user_session():
-    auth = KalavaiAuthClient(
-        user_cookie_file=USER_COOKIE
-    )
-    return auth.load_user_session()
     
-def authenticate_user(username=None, password=None):
-    auth = KalavaiAuthClient(
-        user_cookie_file=USER_COOKIE
-    )
-    user = auth.load_user_session()
-    if user is None:
-        user = auth.login(username=username, password=password)
+def authenticate_user(user_key=None):
+    if user_key is None:
+        KALAVAI_AUTH.save_auth(user_key)
     
-    if user is None:
-        return {"error": "Username or password incorrect"}
-    return user
+    return KALAVAI_AUTH.load_user_session()
+
+def load_user_session():
+    return KALAVAI_AUTH.load_user_session()
 
 def user_logout():
-    auth = KalavaiAuthClient(
-        user_cookie_file=USER_COOKIE
-    )
-    auth.logout()
-    return True
+    return KALAVAI_AUTH.clear_auth()
 
 def check_token(token, public=False):
     try:
@@ -574,38 +564,16 @@ def attach_to_pool(token, node_name=None):
         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:
         return {"error": f"Invalid token. {str(e)}"} 
     
-    user = defaultdict(lambda: None)
-    if public_location is not None:
-        user = load_user_session()
-        if user is None:
-            return {"error ": "Must be logged in to join public pools"}
-        try:
-            vpn = get_vpn_details(
-                location=public_location,
-                user_cookie=USER_COOKIE)
-        except Exception as e:
-            return {"error": f"Are you authenticated? {str(e)}"}
-        try:
-            validate_join_public_seed(
-                cluster_name=cluster_name,
-                join_key=token,
-                user_cookie=USER_COOKIE
-            )
-        except Exception as e:
-            return {"error": f"Error when joining network: {str(e)}"}
-        
     # local agent join
     # 1. Generate local cache files
     # Generate docker compose recipe
     generate_compose_config(
         role="",
-        vpn_token=vpn["key"],
-        node_name=node_name,
-        is_public=public_location is not None)
+        vpn_token=public_location,
+        node_name=node_name)
     
     store_server_info(
         server_ip=kalavai_seed_ip,
@@ -615,7 +583,7 @@ def attach_to_pool(token, node_name=None):
         node_name=node_name,
         cluster_name=cluster_name,
         public_location=public_location,
-        user_api_key=user["api_key"])
+        user_api_key=None)
     
     run_cmd(f"docker compose -f {USER_COMPOSE_FILE} up -d")
     # ensure we are connected
@@ -636,6 +604,44 @@ def get_max_gpus():
     except:
         return 0
 
+def generate_worker_package(num_gpus=0, node_name=None, ip_address="0.0.0.0", storage_compatible=True):
+    # get pool data from token  
+    token = get_pool_token(mode=TokenType.WORKER)
+    if "error" in token:
+        return {"error": f"[red]Error when getting pool token: {token['error']}"}
+    
+    if node_name is None:
+        node_name = f"worker-{uuid.uuid4().hex[:6]}"
+    
+    # parse pool data
+    try:
+        data = decode_dict(token["token"])
+        kalavai_seed_ip = data[CLUSTER_IP_KEY]
+        kalavai_token = data[CLUSTER_TOKEN_KEY]
+        public_location = data[PUBLIC_LOCATION_KEY]
+    except Exception as e:
+        return {"error": f"Invalid token. {str(e)}"} 
+    
+    # join private network if provided
+    node_labels = {
+        STORAGE_CLASS_LABEL: storage_compatible,
+        NODE_ROLE_LABEL: "worker"
+    }
+    # Generate docker compose recipe
+    compose = generate_compose_config(
+        write_to_file=False,
+        role="agent",
+        node_ip_address=ip_address,
+        pool_ip=f"https://{kalavai_seed_ip}:6443",
+        pool_token=kalavai_token,
+        num_gpus=num_gpus,
+        vpn_token=public_location,
+        node_name=node_name,
+        node_labels=node_labels)
+    
+    return compose
+
+
 def join_pool(token, num_gpus=None, node_name=None, ip_address=None):
     compatibility = check_worker_compatibility()
     if len(compatibility["issues"]) > 0:
@@ -660,35 +666,14 @@ def join_pool(token, num_gpus=None, node_name=None, ip_address=None):
         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:
         return {"error": f"Invalid token. {str(e)}"} 
     
     # join private network if provided
     node_labels = {
-        STORAGE_CLASS_LABEL: is_storage_compatible()
-    }
-    user = defaultdict(lambda: None)
-    if public_location is not None:
-        user = authenticate_user()
-        if user is None:
-            return {"error": "Must be logged in to join public pools"}
-        try:
-            vpn = get_vpn_details(
-                location=public_location,
-                user_cookie=USER_COOKIE)
-            node_labels[USER_NODE_LABEL] = user["username"]
-        except Exception as e:
-            return {"error": f"Are you authenticated? Error: {str(e)}"}
-        try:
-            validate_join_public_seed(
-                cluster_name=cluster_name,
-                join_key=token,
-                user_cookie=USER_COOKIE
-            )
-        except Exception as e:
-            return {"error": f"Error when joining network: {str(e)}"}
-        
+        STORAGE_CLASS_LABEL: is_storage_compatible(),
+        NODE_ROLE_LABEL: "worker"
+    }  
     # local agent join
     # Generate docker compose recipe
     generate_compose_config(
@@ -697,10 +682,9 @@ def join_pool(token, num_gpus=None, node_name=None, ip_address=None):
         pool_ip=f"https://{kalavai_seed_ip}:6443",
         pool_token=kalavai_token,
         num_gpus=num_gpus,
-        vpn_token=vpn["key"],
+        vpn_token=public_location,
         node_name=node_name,
-        node_labels=node_labels,
-        is_public=public_location is not None)
+        node_labels=node_labels)
     
     store_server_info(
         server_ip=kalavai_seed_ip,
@@ -710,7 +694,7 @@ def join_pool(token, num_gpus=None, node_name=None, ip_address=None):
         node_name=node_name,
         cluster_name=cluster_name,
         public_location=public_location,
-        user_api_key=user["api_key"])
+        user_api_key=None)
     
     try:
         CLUSTER.start_worker_node()
@@ -730,13 +714,25 @@ def join_pool(token, num_gpus=None, node_name=None, ip_address=None):
     except KeyboardInterrupt:
         return {"error": "Installation aborted. Leaving pool."}
     
-    result = init_user_workspace()
+    result = init_user_workspace(
+        user_id=load_user_id(),
+        node_name=node_name)
     if "error" in result:
         return {"error": f"Error when creating user workspace: {result}"}
     
     return cluster_name
 
-def create_pool(cluster_name: str, ip_address: str, description: str="", token_mode: TokenType=TokenType.USER, app_values: str=None, pool_config_values: str=None, num_gpus: int=0, node_name: str=None, only_registered_users: bool=False, location: str=None):
+def create_pool(
+        cluster_name: str,
+        ip_address: str,
+        description: str="",
+        token_mode: TokenType=TokenType.USER,
+        app_values: str=None,
+        pool_config_values: str=None,
+        num_gpus: int=0,
+        node_name: str=None,
+        location: str=None
+    ):
 
     if not check_seed_compatibility():
         return {"error": "Requirements failed"}
@@ -747,28 +743,22 @@ def create_pool(cluster_name: str, ip_address: str, description: str="", token_m
     if pool_config_values is None:
         pool_config_values = POOL_CONFIG_DEFAULT_VALUES
 
-    node_name = f"{socket.gethostname()}-{uuid.uuid4().hex[:6]}" 
+    node_name = f"{socket.gethostname()}-{uuid.uuid4().hex[:6]}"
+    user_id = load_user_id()
     
-    # 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 = authenticate_user()
-        if user is None:
-            return {"error": "[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"}
-
-    # join private network if provided
-    vpn = defaultdict(lambda: None)
     node_labels = {
-        STORAGE_CLASS_LABEL: is_storage_compatible()
+        STORAGE_CLASS_LABEL: is_storage_compatible(),
+        NODE_ROLE_LABEL: "server"
     }
-    if location is not None:
-        try:
-            vpn = get_vpn_details(
-                location=location,
-                user_cookie=USER_COOKIE)
-            node_labels[USER_NODE_LABEL] = user["username"]
-        except Exception as e:
-            return {"error": f"[red]Error when joining network: {str(e)}"}
+
+    # if location is not None:
+    #     try:
+    #         vpn = get_vpn_details(
+    #             location=location,
+    #             user_cookie=USER_COOKIE)
+    #         node_labels[USER_NODE_LABEL] = user["username"]
+    #     except Exception as e:
+    #         return {"error": f"[red]Error when joining network: {str(e)}"}
         
     if num_gpus is None:
         num_gpus = get_max_gpus()
@@ -776,12 +766,11 @@ def create_pool(cluster_name: str, ip_address: str, description: str="", token_m
     # Generate docker compose recipe
     generate_compose_config(
         role="server",
-        vpn_token=vpn["key"],
+        vpn_token=location,
         node_ip_address=ip_address,
         num_gpus=num_gpus,
         node_name=node_name,
-        node_labels=node_labels,
-        is_public=location is not None
+        node_labels=node_labels
     )
     
     # start server
@@ -796,7 +785,8 @@ def create_pool(cluster_name: str, ip_address: str, description: str="", token_m
         ip_address = CLUSTER.get_vpn_ip()
 
     # populate local cred files
-    auth_key = str(uuid.uuid4())
+    
+    auth_key = user_id if user_id is not None else str(uuid.uuid4())
     write_auth_key = str(uuid.uuid4())
     readonly_auth_key = str(uuid.uuid4())
     
@@ -810,7 +800,8 @@ def create_pool(cluster_name: str, ip_address: str, description: str="", token_m
         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
+        ALLOW_UNREGISTERED_USER_KEY: True, # Change this if only registered users are allowed,
+        DEPLOY_LLM_SIDECARS_KEY: location is not None
     }
 
     store_server_info(
@@ -823,7 +814,7 @@ def create_pool(cluster_name: str, ip_address: str, description: str="", token_m
         node_name=node_name,
         cluster_name=cluster_name,
         public_location=location,
-        user_api_key=user["api_key"])
+        user_api_key=None)
     
     # Generate helmfile recipe
     helm_yaml = load_template(
@@ -857,11 +848,14 @@ def create_pool(cluster_name: str, ip_address: str, description: str="", token_m
     if "error" in result or ("failed" in result and len(result['failed']) > 0):
         return {"error": f"Error when initialising pool: {result}"}
     # init default namespace
-    init_user_workspace(force_namespace="default")
-    if only_registered_users:
-        # init user namespace
-        init_user_workspace()
-    
+    init_user_workspace(
+        user_id=user_id,
+        node_name=node_name,
+        force_namespace="default")
+    # if only_registered_users:
+    #     # init user namespace
+    #     init_user_workspace(
+    #         user_id=user_id)
     # register cluster (if user is logged in)
     result = register_pool(
         cluster_name=cluster_name,
@@ -876,6 +870,7 @@ def create_pool(cluster_name: str, ip_address: str, description: str="", token_m
     
     return {"success"}
 
+
 def get_pool_token(mode: TokenType):
 
     try:

kalavai_client/env.py

@@ -50,7 +50,7 @@ POOL_CONFIG_DEFAULT_VALUES = resource_path("kalavai_client/assets/pool_config_va
 # user specific config files
 USER_TEMPLATES_FOLDER = user_path("templates", create_path=True)
 USER_LOCAL_SERVER_FILE = user_path(".server")
-USER_COOKIE = user_path(".user_cookie.pkl")
+USER_COOKIE = user_path(".user_cookie.json")
 USER_COMPOSE_FILE = user_path("docker-compose-worker.yaml")
 USER_GUI_COMPOSE_FILE = user_path("docker-compose-gui.yaml")
 USER_HELM_APPS_FILE = user_path("apps.yaml")

kalavai_client/utils.py

@@ -12,8 +12,7 @@ from jinja2 import Template
 from rich.table import Table
 import yaml
 
-
-from kalavai_client.auth import KalavaiAuthClient
+from kalavai_client.auth import KalavaiAuth
 from kalavai_client.env import (
     SERVER_IP_KEY,
     DOCKER_COMPOSE_TEMPLATE,
@@ -22,6 +21,7 @@ from kalavai_client.env import (
     DEFAULT_VPN_CONTAINER_NAME,
     CONTAINER_HOST_PATH,
     USER_COMPOSE_FILE,
+    USER_COOKIE,
     user_path
 )
 
@@ -38,6 +38,8 @@ CLUSTER_NAME_KEY = "cluster_name"
 AUTH_KEY = "watcher_admin_key"
 WRITE_AUTH_KEY = "watcher_write_key"
 ALLOW_UNREGISTERED_USER_KEY = "watcher_allow_unregistered_user"
+DEPLOY_LLM_SIDECARS_KEY = "deploy_llm_sidecars"
+NODE_ROLE_LABEL = "kalavai.node_role"
 USER_API_KEY = "user_api_key"
 READONLY_AUTH_KEY = "watcher_readonly_key"
 WATCHER_SERVICE_KEY = "watcher_service"
@@ -62,6 +64,12 @@ MANDATORY_POOLCONFIG_FIELDS = [
     PUBLIC_LOCATION_KEY
 ]
 
+KALAVAI_AUTH = KalavaiAuth(
+    auth_service_url="dummy_url",
+    auth_service_key="dummy_key",
+    user_cookie_file=USER_COOKIE
+)
+
 
 ####### Methods to check OS compatibility ########
 def check_gpu_drivers():
@@ -92,7 +100,7 @@ def is_storage_compatible():
         return False
 ################
 
-def generate_compose_config(role, node_name, is_public, node_ip_address="0.0.0.0", num_gpus=0, node_labels=None, pool_ip=None, vpn_token=None, pool_token=None):
+def generate_compose_config(role, node_name, write_to_file=True, node_ip_address="0.0.0.0", num_gpus=0, node_labels=None, pool_ip=None, vpn_token=None, pool_token=None):
     
     if node_labels is not None:
         node_labels = " ".join([f"--node-label {key}={value}" for key, value in node_labels.items()])
@@ -100,7 +108,7 @@ def generate_compose_config(role, node_name, is_public, node_ip_address="0.0.0.0
     compose_values = {
         "user_path": user_path(""),
         "service_name": DEFAULT_CONTAINER_NAME,
-        "vpn": is_public,
+        "vpn": vpn_token is not None,
         "vpn_name": DEFAULT_VPN_CONTAINER_NAME,
         "node_ip_address": node_ip_address,
         "pool_ip": pool_ip,
@@ -113,14 +121,15 @@ def generate_compose_config(role, node_name, is_public, node_ip_address="0.0.0.0
         "k3s_path": f"{CONTAINER_HOST_PATH}/{rand_suffix}/k3s",
         "etc_path": f"{CONTAINER_HOST_PATH}/{rand_suffix}/etc",
         "node_labels": node_labels,
-        "flannel_iface": DEFAULT_FLANNEL_IFACE if is_public else ""
+        "flannel_iface": DEFAULT_FLANNEL_IFACE if vpn_token is not None 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)
+    if write_to_file:
+        with open(USER_COMPOSE_FILE, "w") as f:
+            f.write(compose_yaml)
     return compose_yaml
 
 def is_watcher_alive(server_creds, user_cookie, timeout=30):
@@ -146,126 +155,63 @@ def load_server_info(data_key, file):
             return json.load(f)[data_key]
     except:
         return None
-    
-def user_login(user_cookie, username=None, password=None):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    user = auth.load_user_session()
-    if user is None:
-        user = auth.login(username=username, password=password)
-    return user
-
-def user_logout(user_cookie):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    auth.logout()
 
-def load_user_session(user_cookie):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    if not auth.is_logged_in():
-        return None
-    return auth.load_user_session()
+def load_user_session():
+    return KALAVAI_AUTH.load_user_session()
 
-def get_public_vpns(user_cookie):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    if not auth.is_logged_in():
-        raise ValueError("Cannot access vpns, user is not authenticated")
-    seeds = auth.call_function(
-        "get_public_vpns"
-    )
-    return seeds
+def load_user_id():
+    return KALAVAI_AUTH.get_user_id()
 
 def get_public_seeds(user_only, user_cookie):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    if not auth.is_logged_in():
-        raise ValueError("Cannot access vpns, user is not authenticated")
-    user = auth.load_user_session() if user_only else None
-    seeds = auth.call_function(
-        "get_available_seeds",
-        user
-    )
-    return seeds
-
-def get_vpn_details(location, user_cookie):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    if not auth.is_logged_in():
-        raise ValueError("Cannot access vpns, user is not authenticated")
-    vpn = auth.call_function(
-        "get_vpn_details",
-        location
-    )
-    return vpn
+    return []
+    # if not auth_obj.is_logged_in():
+    #     raise ValueError("Cannot access vpns, user is not authenticated")
+    # user = auth_obj.load_user_session() if user_only else None
+    # seeds = auth_obj.call_function(
+    #     "get_available_seeds",
+    #     user
+    # )
+    # return seeds
 
 def register_cluster(name, token, description, user_cookie, is_private=True):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    if not auth.is_logged_in():
-        raise ValueError("Cannot register cluster, user is not authenticated")
-    user = auth.load_user_session()
-    seed = auth.call_function(
-        "register_seed",
-        name,
-        user,
-        description,
-        token,
-        is_private
-    )
-    return seed
+    return None
+    # if not auth_obj.is_logged_in():
+    #     raise ValueError("Cannot register cluster, user is not authenticated")
+    # user = auth_obj.load_user_session()
+    # seed = auth_obj.call_function(
+    #     "register_seed",
+    #     name,
+    #     user,
+    #     description,
+    #     token,
+    #     is_private
+    # )
+    # return seed
 
 def unregister_cluster(name, user_cookie):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    if not auth.is_logged_in():
-        raise ValueError("Cannot unregister cluster, user is not authenticated")
-    user = auth.load_user_session()
-    seed = auth.call_function(
-        "unregister_seed",
-        name,
-        user,
-    )
-    return seed
-
-def validate_join_public_seed(cluster_name, join_key, user_cookie):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    if not auth.is_logged_in():
-        raise ValueError("Cannot notify join cluster, user is not authenticated")
-    user = auth.load_user_session()
-    seed = auth.call_function(
-        "validate_join_public_seed",
-        cluster_name,
-        join_key,
-        user,
-    )
-    return seed
+    return None
+    # if not auth_obj.is_logged_in():
+    #     raise ValueError("Cannot unregister cluster, user is not authenticated")
+    # user = auth_obj.load_user_session()
+    # seed = auth_obj.call_function(
+    #     "unregister_seed",
+    #     name,
+    #     user,
+    # )
+    # return seed
 
 def send_pool_invite(cluster_name, invitee_addresses, user_cookie):
-    auth = KalavaiAuthClient(
-        user_cookie_file=user_cookie
-    )
-    if not auth.is_logged_in():
-        raise ValueError("Cannot notify join cluster, user is not authenticated")
-    user = auth.load_user_session()
-    result = auth.call_function(
-        "send_pool_invite",
-        user,
-        cluster_name,
-        invitee_addresses
-    )
-    return result
+    return None
+    # if not auth_obj.is_logged_in():
+    #     raise ValueError("Cannot notify join cluster, user is not authenticated")
+    # user = auth_obj.load_user_session()
+    # result = auth_obj.call_function(
+    #     "send_pool_invite",
+    #     user,
+    #     cluster_name,
+    #     invitee_addresses
+    # )
+    # return result
 
 def validate_poolconfig(poolconfig_file):
     if not Path(poolconfig_file).is_file():
@@ -284,20 +230,6 @@ def run_cmd(command):
     except OSError as error:
         return error # for exit code
 
-def join_vpn(location, user_cookie):
-    vpn = get_vpn_details(location=location, user_cookie=user_cookie)
-    token = vpn["key"]
-    if token is None:
-        raise ValueError(f"VPN location {location} not found or is private")
-    try:
-        data = decode_dict(token)
-        for field in ["server", "value"]:
-            assert field in data
-    except:
-        raise ValueError("Invalid net token")
-    run_cmd(f"sudo netclient join -t {token} >/dev/null 2>&1")
-    return vpn
-
 def leave_vpn(container_name):
     try:
         vpns = json.loads(run_cmd(f"docker exec {container_name} netclient list").decode())
@@ -332,9 +264,10 @@ def request_to_server(
         "X-API-KEY": auth_key
     }
 
-    user = load_user_session(user_cookie=user_cookie)
     headers["USER-KEY"] = load_server_info(data_key=USER_API_KEY, file=server_creds)
-    headers["USER"] = user["username"] if user is not None else None
+    user_id = load_user_id()
+    if user_id is not None:
+        headers["USER"] = user_id
 
     response = requests.request(
         method=method,