kalavai-client 0.5.30__py3-none-any.whl → 0.6.1__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

@@ -3,10 +3,15 @@ from pathlib import Path
 import importlib.resources
 
 
+DEFAULT_KALAVAI_PATH = os.getenv("KALAVAI_PATH", None)
+
 def user_path(relative_path, create_path=False):
     """Transform a relative path into the user's cache folder path"""
-    base = os.path.expanduser("~")
-    kalavai_user_path = os.path.join(base, ".cache/kalavai")
+    if DEFAULT_KALAVAI_PATH is not None:
+        kalavai_user_path = DEFAULT_KALAVAI_PATH
+    else:
+        base = os.path.expanduser("~")
+        kalavai_user_path = os.path.join(base, ".cache/kalavai")
     full_path = os.path.join(kalavai_user_path, relative_path)
     if create_path:
         Path(full_path).mkdir(parents=True, exist_ok=True)
@@ -34,7 +39,7 @@ KALAVAI_PLATFORM_ENDPOINT = "https://platform.kalavai.net/_/api"
 DEFAULT_CONTAINER_NAME = "kalavai"
 DEFAULT_VPN_CONTAINER_NAME = "kalavai-vpn"
 CONTAINER_HOST_PATH = user_path("pool/", create_path=True)
-DEFAULT_FLANNEL_IFACE = os.getenv("KALAVAI_FLANNEL_IFACE", "netmaker-1")
+DEFAULT_FLANNEL_IFACE = os.getenv("KALAVAI_FLANNEL_IFACE", "netmaker")
 DEFAULT_WATCHER_PORT = 30001
 KUBE_VERSION = os.getenv("KALAVAI_KUBE_VERSION", "v1.31.1+k3s1")
 FORBIDEDEN_IPS = ["127.0.0.1"]
@@ -50,7 +55,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")