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

kalavai_client/__init__.py

@@ -1,2 +1,2 @@
 
-__version__ = "0.5.30"
+__version__ = "0.6.0"

kalavai_client/assets/apps.yaml

@@ -28,7 +28,7 @@ releases:
   - name: lago
     namespace: kalavai
     chart: kalavai/lago
-    installed: true
+    installed: {{deploy_llm_sidecars}}
     set:
     - name: external.api.nodePort
       value: 32000
@@ -79,7 +79,7 @@ releases:
     namespace: kalavai
     chart: kalavai/kalavai-helios
     version: "0.1.11"
-    installed: {{not watcher_allow_unregistered_user}}
+    installed: false #{{deploy_llm_sidecars}}
     set:
     - name: deployment.watcher_endpoint
       value: "http://{{watcher_service}}"
@@ -94,11 +94,11 @@ releases:
   - name: opencost
     namespace: opencost
     chart: opencost-charts/opencost
-    installed: {{not watcher_allow_unregistered_user}}
+    installed: {{deploy_llm_sidecars}}
   - name: prometheus
     namespace: prometheus-system
     chart: prometheus/prometheus
-    installed: {{not watcher_allow_unregistered_user}}
+    installed: {{deploy_llm_sidecars}}
     set:
     - name: prometheus-pushgateway.enabled
       value: false
@@ -144,17 +144,19 @@ releases:
   - name: kalavai-watcher
     namespace: kalavai
     chart: kalavai/kalavai-watcher
-    version: "0.3.6"
+    version: "0.3.8"
     installed: true
     set:
     - name: namespace
       value: kalavai
     - name: replicas
-      value: 2
+      value: 1
     - name: image_tag
-      value: "v2025.03.12"
+      value: "v2025.03.18"
     - name: deployment.in_cluster
       value: "True"
+    - name: deployment.kalavai_username_key
+      value: "{{kalavai_username_key}}"
     - name: deployment.use_auth_key
       value: "True"
     - name: deployment.admin_key

kalavai_client/assets/apps_values.yaml

@@ -1,5 +1,3 @@
-# TODO: add helm versions here
-
 - name: kalavai_api_endpoint
   default: https://platform.kalavai.net/_/api
   description: ""

kalavai_client/assets/docker-compose-gui.yaml

@@ -10,6 +10,9 @@ services:
     environment:
     - KALAVAI_BRIDGE_URL=http://host.docker.internal
     - KALAVAI_BRIDGE_PORT={{bridge_port}}
+    {% if protected_access %}
+    - ACCESS_KEY={{protected_access}}
+    {% endif %}
     entrypoint: ["reflex"]
     command: >
       run

kalavai_client/assets/docker-compose-template.yaml

@@ -1,7 +1,7 @@
 services:
 {% if vpn %}
   {{vpn_name}}:
-    image: gravitl/netclient:v0.24.3
+    image: gravitl/netclient:v0.90.0
     container_name: {{vpn_name}}
     #privileged: true
     cap_add:

kalavai_client/auth.py

@@ -1,68 +1,110 @@
+from collections import defaultdict
 import os
-import pickle
+import json
+from typing import Optional, Any
 
-import anvil.server
-import anvil.users
+_USER_KEY = "id"
 
 
-AUTH_UPLINK_KEY = os.getenv("ANVIL_UPLINK_KEY", "client_AOPKTWK227ZV3R4ENTMOQRIY-ADMVMYW5OIRPH75P")
+class KalavaiAuth():
+    """Client for authenticating with the Kalavai watcher service."""
+    
+    def __init__(
+        self,
+        auth_service_url: str,
+        auth_service_key: str,
+        user_cookie_file: str
+    ):
+        """Initialize the Kalavai auth client.
+        
+        Args:
+            user_cookie_file: Path to store user session data
+        """
+        self.user_cookie = user_cookie_file
+        self.auth_service_url = auth_service_url
+        self.auth_service_key = auth_service_key
+        self.user_session = defaultdict(None)
+        self.load_user_session()
 
-class KalavaiAuthClient:
-    def __init__(self, user_cookie_file=None):
-        anvil.server.connect(AUTH_UPLINK_KEY, quiet=True)
-        self.user_cookie_file = user_cookie_file
-        user = self.load_user_session()
+    def save_auth(self, user_key: str) -> Optional[Any]:
+        """Login with user key.
+        
+        Args:
+            user_key: User's key (actual authentication key)
+            
+        Returns:
+            User data if successful, None otherwise
+        """
+        # Store the user key as the session
+        self.user_session[_USER_KEY] = user_key
+        with open(self.user_cookie, "w") as f:
+            json.dump(self.user_session, f)
+            print(f"User key securely saved to {self.user_cookie}")
 
-    def login(self, username, password):
+    def clear_auth(self) -> bool:
+        """Clear the current user and local session.
+        
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        self.user_session = defaultdict(None)
         try:
-            user = anvil.users.login_with_email(username, password, remember=True)
-        except:
-            return None
+            if os.path.exists(self.user_cookie):
+                os.remove(self.user_cookie)
+        except Exception as e:
+            print(f"Failed to clear user session: {str(e)}")
+
+    def is_authenticated(self) -> bool:
+        """Check if user is currently authenticated.
         
-        if self.user_cookie_file:
-            with open(self.user_cookie_file, "wb") as f:
-                    pickle.dump(user, f)  
-        return user
+        Returns:
+            bool: True if logged in, False otherwise
+        """
+        return self.user_session[_USER_KEY] is not None
 
-    def logout(self):
-        anvil.users.logout()
+    def load_user_session(self) -> Optional[Any]:
+        """Load user session from local storage.
+        
+        Returns:
+            User data if session exists, None otherwise
+        """
         try:
-            os.remove(self.user_cookie_file)
-        except:
-            pass
+            if not os.path.exists(self.user_cookie):
+                return None
+                
+            with open(self.user_cookie, "r") as f:
+                for key, value in json.load(f).items():
+                    self.user_session[key] = value
+            return self.user_session
+            
+        except Exception as e:
+            print(f"Failed to load user session: {str(e)}")
+            self.clear_auth()
+            return None
     
-    def is_logged_in(self):
-        user = self.load_user_session()
-        return user is not None
-
-    def load_user_session(self):
-        user = anvil.users.get_user()
-        if user:
-            return user
+    def get_user_id(self) -> Optional[Any]:
+        """Get the user key from the user session.
+        
+        Returns:
+            User key if session exists, None otherwise
+        """
         try:
-            with open(self.user_cookie_file, "rb") as f:
-                user = pickle.load(f)
-                return user
+            return self.user_session[_USER_KEY]
         except:
             return None
-        
-    def call_function(self, fn, *args):
-        return anvil.server.call(
-            fn,
-            *args
-        )
 
-    
 
+# Example usage
 if __name__ == "__main__":
-    auth = KalavaiAuthClient(
-        user_cookie_file="here.pickle"
+    # Initialize client
+    auth = KalavaiAuth(
+        auth_service_url="http://localhost:8000",
+        auth_service_key="example_auth_service_key",
+        user_cookie_file="user_cookie.json"
     )
-    user = auth.load_user_session()
-    if not user:
-        user = auth.login(username="carlos@kalavai.net", password="wrong_pass")
-
-    if user is None:
-        print("Failed to login")
-    else:
-        print(user)
+    
+    # Example of logging in with a user key
+    user_id = "example_user_id"
+    auth.save_auth(user_id)
+    
+    print(f"User id: {auth.get_user_id()}")

kalavai_client/bridge_api.py

@@ -2,7 +2,8 @@
 Core kalavai service.
 Used as a bridge between the kalavai-client app and the reflex frontend
 """
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, Depends
+from starlette.requests import Request
 import uvicorn
 
 from kalavai_client.bridge_models import (
@@ -47,11 +48,24 @@ from kalavai_client.core import (
     uncordon_nodes,
     TokenType
 )
+from kalavai_client.utils import load_user_id
 
 app = FastAPI()
 
+################################
+## API Key Validation methods ##
+################################
+async def verify_api_key(request: Request):
+    user_id = load_user_id()
+    if user_id is None:
+        return None
+    api_key = request.headers.get("X-API-KEY")
+    if api_key != user_id:
+        raise HTTPException(status_code=401, detail="Request requires API Key")
+    return api_key
+
 @app.post("/create_pool")
-def pool_create(request: CreatePoolRequest):
+def pool_create(request: CreatePoolRequest, api_key: str = Depends(verify_api_key)):
     result = create_pool(
         cluster_name=request.cluster_name,
         ip_address=request.ip_address,
@@ -66,7 +80,7 @@ def pool_create(request: CreatePoolRequest):
     return result
 
 @app.post("/join_pool")
-def pool_join(request: JoinPoolRequest):
+def pool_join(request: JoinPoolRequest, api_key: str = Depends(verify_api_key)):
     result = join_pool(
         token=request.token,
         num_gpus=request.num_gpus,
@@ -76,7 +90,7 @@ def pool_join(request: JoinPoolRequest):
     return result
 
 @app.post("/attach_to_pool")
-def pool_attach(request: JoinPoolRequest):
+def pool_attach(request: JoinPoolRequest, api_key: str = Depends(verify_api_key)):
     result = attach_to_pool(
         token=request.token,
         node_name=request.node_name
@@ -84,64 +98,64 @@ def pool_attach(request: JoinPoolRequest):
     return result
 
 @app.post("/stop_pool")
-def pool_stop(request: StopPoolRequest):
+def pool_stop(request: StopPoolRequest, api_key: str = Depends(verify_api_key)):
     result = stop_pool(
         skip_node_deletion=request.skip_node_deletion
     )
     return result
 
 @app.post("/delete_nodes")
-def device_delete(request: NodesActionRequest):
+def device_delete(request: NodesActionRequest, api_key: str = Depends(verify_api_key)):
     result = delete_nodes(
         nodes=request.nodes
     )
     return result
 
 @app.post("/cordon_nodes")
-def device_cordon(request: NodesActionRequest):
+def device_cordon(request: NodesActionRequest, api_key: str = Depends(verify_api_key)):
     result = cordon_nodes(
         nodes=request.nodes
     )
     return result
 
 @app.post("/uncordon_nodes")
-def device_uncordon(request: NodesActionRequest):
+def device_uncordon(request: NodesActionRequest, api_key: str = Depends(verify_api_key)):
     result = uncordon_nodes(
         nodes=request.nodes
     )
     return result
 
 @app.get("/get_pool_token")
-def get_token(mode: int):
+def get_token(mode: int, api_key: str = Depends(verify_api_key)):
 
     return get_pool_token(mode=TokenType(mode))
 
 @app.get("/fetch_devices")
-def get_devices():
+def get_devices(api_key: str = Depends(verify_api_key)):
     return fetch_devices()
 
 @app.post("/send_pool_invites")
-def send_pool_invites(request: InvitesRequest):
+def send_pool_invites(request: InvitesRequest, api_key: str = Depends(verify_api_key)):
     return send_invites(invitees=request.invitees)
 
 @app.get("/fetch_resources")
-def resources():
+def resources(api_key: str = Depends(verify_api_key)):
     return fetch_resources()
 
 @app.get("/fetch_job_names")
-def job_names():
+def job_names(api_key: str = Depends(verify_api_key)):
     return fetch_job_names()
 
 @app.get("/fetch_gpus")
-def gpus(available: bool = False):
+def gpus(available: bool = False, api_key: str = Depends(verify_api_key)):
     return fetch_gpus(available=available)
 
 @app.post("/fetch_job_details")
-def job_details(request: JobDetailsRequest):
+def job_details(request: JobDetailsRequest, api_key: str = Depends(verify_api_key)):
     return fetch_job_details(jobs=request.jobs)
 
 @app.get("/fetch_job_logs")
-def job_logs(job_name: str, force_namespace: str=None, pod_name: str=None, tail: int=100):
+def job_logs(job_name: str, force_namespace: str=None, pod_name: str=None, tail: int=100, api_key: str = Depends(verify_api_key)):
     return fetch_job_logs(
         job_name=job_name,
         force_namespace=force_namespace,
@@ -150,15 +164,15 @@ def job_logs(job_name: str, force_namespace: str=None, pod_name: str=None, tail:
     )
 
 @app.get("/fetch_job_templates")
-def job_templates():
+def job_templates(api_key: str = Depends(verify_api_key)):
     return fetch_job_templates()
 
 @app.get("/fetch_job_defaults")
-def job_templates(name: str):
+def job_templates(name: str, api_key: str = Depends(verify_api_key)):
     return fetch_job_defaults(name=name)
 
 @app.post("/deploy_job")
-def job_deploy(request: DeployJobRequest):
+def job_deploy(request: DeployJobRequest, api_key: str = Depends(verify_api_key)):
     result = deploy_job(
         template_name=request.template_name,
         values_dict=request.values,
@@ -168,7 +182,7 @@ def job_deploy(request: DeployJobRequest):
     return result
 
 @app.post("/delete_job")
-def job_delete(request: DeleteJobRequest):
+def job_delete(request: DeleteJobRequest, api_key: str = Depends(verify_api_key)):
     result = delete_job(
         name=request.name,
         force_namespace=request.force_namespace
@@ -176,15 +190,14 @@ def job_delete(request: DeleteJobRequest):
     return result
 
 @app.get("/authenticate_user")
-def user_authenticate(username: str, password: str):
+def user_authenticate(user_id: str, api_key: str = Depends(verify_api_key)):
     result = authenticate_user(
-        username=username,
-        password=password
+        user_id=user_id
     )
     return result
 
 @app.get("/load_user_session")
-def user_session():
+def user_session(api_key: str = Depends(verify_api_key)):
     result = load_user_session()
     return result
 
@@ -219,12 +232,12 @@ def agent_resume():
     return result
 
 @app.get("/get_ip_addresses")
-def ip_addresses(subnet: str=None):
+def ip_addresses(subnet: str=None, api_key: str = Depends(verify_api_key)):
     result = get_ip_addresses(subnet=subnet)
     return result
 
 @app.get("/list_available_pools")
-def pool_connected(user_only: bool=False):
+def pool_connected(user_only: bool=False, api_key: str = Depends(verify_api_key)):
     result = list_available_pools(user_only=user_only)
     return result