kalavai-client 0.5.25__py3-none-any.whl → 0.5.27__py3-none-any.whl

kalavai_client/__init__.py

@@ -1,2 +1,2 @@
 
-__version__ = "0.5.25"
+__version__ = "0.5.27"

kalavai_client/bridge_api.py

@@ -7,17 +7,19 @@ import uvicorn
 
 from kalavai_client.bridge_models import (
     CreatePoolRequest,
+    InvitesRequest,
     JoinPoolRequest,
     StopPoolRequest,
     DeployJobRequest,
     DeleteJobRequest,
     JobDetailsRequest,
-    DeleteNodesRequest
+    NodesActionRequest
 )
 from kalavai_client.core import (
     create_pool,
     join_pool,
     attach_to_pool,
+    send_invites,
     stop_pool,
     fetch_devices,
     fetch_resources,
@@ -41,6 +43,8 @@ from kalavai_client.core import (
     get_ip_addresses,
     get_pool_token,
     delete_nodes,
+    cordon_nodes,
+    uncordon_nodes,
     TokenType
 )
 
@@ -55,7 +59,9 @@ def pool_create(request: CreatePoolRequest):
         num_gpus=request.num_gpus,
         node_name=request.node_name,
         only_registered_users=request.only_registered_users,
-        location=request.location
+        location=request.location,
+        description=request.description,
+        token_mode=request.token_mode
     )
     return result
 
@@ -85,21 +91,39 @@ def pool_stop(request: StopPoolRequest):
     return result
 
 @app.post("/delete_nodes")
-def device_delete(request: DeleteNodesRequest):
+def device_delete(request: NodesActionRequest):
     result = delete_nodes(
         nodes=request.nodes
     )
     return result
 
+@app.post("/cordon_nodes")
+def device_cordon(request: NodesActionRequest):
+    result = cordon_nodes(
+        nodes=request.nodes
+    )
+    return result
+
+@app.post("/uncordon_nodes")
+def device_uncordon(request: NodesActionRequest):
+    result = uncordon_nodes(
+        nodes=request.nodes
+    )
+    return result
+
 @app.get("/get_pool_token")
-def devices(mode: int):
+def get_token(mode: int):
 
     return get_pool_token(mode=TokenType(mode))
 
 @app.get("/fetch_devices")
-def devices():
+def get_devices():
     return fetch_devices()
 
+@app.post("/send_pool_invites")
+def send_pool_invites(request: InvitesRequest):
+    return send_invites(invitees=request.invitees)
+
 @app.get("/fetch_resources")
 def resources():
     return fetch_resources()

kalavai_client/bridge_models.py

@@ -1,8 +1,11 @@
 from pydantic import BaseModel
 
-from kalavai_client.core import Job
+from kalavai_client.core import Job, TokenType
 
 
+class InvitesRequest(BaseModel):
+    invitees: list[str]
+
 class CreatePoolRequest(BaseModel):
     cluster_name: str
     ip_address: str
@@ -11,8 +14,10 @@ class CreatePoolRequest(BaseModel):
     node_name: str = None
     only_registered_users: bool = False
     location: str = None
+    token_mode: TokenType = TokenType.USER
+    description: str = ""
 
-class DeleteNodesRequest(BaseModel):
+class NodesActionRequest(BaseModel):
     nodes: list[str]
 
 class JoinPoolRequest(BaseModel):

kalavai_client/cli.py

@@ -51,10 +51,15 @@ from kalavai_client.core import (
     create_pool,
     get_ip_addresses,
     pause_agent,
+    register_pool,
     resume_agent,
     get_pool_token,
     delete_nodes,
-    TokenType
+    cordon_nodes,
+    stop_pool,
+    uncordon_nodes,
+    TokenType,
+    unregister_pool
 )
 from kalavai_client.utils import (
     check_gpu_drivers,
@@ -127,27 +132,6 @@ def pre_join_check(node_name, server_url, server_key):
     except Exception as e:
         console.log(f"[red]Error when connecting to kalavai service: {str(e)}")
         return False
-    
-def set_schedulable(schedulable, node_name=load_server_info(data_key=NODE_NAME_KEY, file=USER_LOCAL_SERVER_FILE)):
-    """
-    Delete job in the cluster
-    """
-    # deploy template with kube-watcher
-    data = {
-        "schedulable": str(schedulable),
-        "node_names": [node_name]
-    }
-    try:
-        res = request_to_server(
-            method="post",
-            endpoint="/v1/set_node_schedulable",
-            data=data,
-            server_creds=USER_LOCAL_SERVER_FILE,
-            user_cookie=USER_COOKIE
-        )
-        console.log(f"{res}")
-    except Exception as e:
-        console.log(f"[red]Error when connecting to kalavai service: {str(e)}")
 
 def select_ip_address(subnet=None):
     ips = get_ip_addresses(subnet=subnet)
@@ -223,7 +207,7 @@ def input_gpus():
 ##################
 
 @arguably.command
-def gui__start(*others, backend_only=False, gui_frontend_port=3000, gui_backend_port=8000, bridge_port=8001):
+def gui__start(*others, backend_only=False, gui_frontend_port=3000, gui_backend_port=8000, bridge_port=8001, log_level="critical"):
     """Run GUI (docker) and kalavai core backend (api)"""
 
     if not backend_only:
@@ -241,7 +225,7 @@ def gui__start(*others, backend_only=False, gui_frontend_port=3000, gui_backend_
         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_frontend_port}")
-    run_api(port=bridge_port)
+    run_api(port=bridge_port, log_level=log_level)
     
     if not backend_only:
         run_cmd(f"docker compose --file {USER_GUI_COMPOSE_FILE} down")
@@ -305,7 +289,7 @@ def location__list(*others):
     console.log(table)
 
 @arguably.command
-def pool__publish(*others, description=None):
+def pool__publish(*others, description=None, is_private=True):
     """
     [AUTH] Publish pool to Kalavai platform, where other users may be able to join
     """
@@ -319,26 +303,29 @@ def pool__publish(*others, description=None):
         console.log(f"[red]Problems with your pool: {str(e)}")
         return
     choices = select_token_type()
-    token = pool__token(**choices)["token"]
+    if choices["admin"]:
+        mode = TokenType.ADMIN
+    elif choices["user"]:
+        mode = TokenType.USER
+    else:
+        mode = TokenType.WORKER
     
     if description is None:
         console.log("[yellow] [Markdown] In a few words (max 500 chars), describe your goals with this cluster. Remember, this is what other users will see to decide whether to share their resources with you, [blue]so inspire them!")
-        description = input(f"(You can edit this later in {KALAVAI_PLATFORM_URL}\n")
-    
-    try:
-        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(
-            name=cluster_name,
-            token=token,
-            description=description,
-            user_cookie=USER_COOKIE)
-        console.log(f"[green]Your cluster is now public on {KALAVAI_PLATFORM_URL}")
-    except Exception as e:
-        console.log(f"[red]Error when publishing cluster. {str(e)}")
+        description = input(f"(You can edit this later at {KALAVAI_PLATFORM_URL}\n")
+    cluster_name = load_server_info(data_key=CLUSTER_NAME_KEY, file=USER_LOCAL_SERVER_FILE)
+
+    result = register_pool(
+        cluster_name=cluster_name,
+        token_mode=mode,
+        description=description,
+        is_private=is_private        
+    )
+
+    if "error" in result:
+        console.log(f"[red]Error when publishing cluster: {result['error']}")
+    else:
+        console.log(f"[green]Your cluster is now registered with {KALAVAI_PLATFORM_URL}")
 
 @arguably.command
 def pool__unpublish(cluster_name=None, *others):
@@ -354,15 +341,13 @@ def pool__unpublish(cluster_name=None, *others):
         console.log(f"[red]Problems with your pool: {str(e)}")
         return
     
-    try:
-        if cluster_name is None:
-            cluster_name = load_server_info(data_key=CLUSTER_NAME_KEY, file=USER_LOCAL_SERVER_FILE)
-        unregister_cluster(
-            name=cluster_name,
-            user_cookie=USER_COOKIE)
+    result = unregister_pool()
+    if "error" in result:
+        console.log(f"[red]{result['error']}")
+    elif "warning" in result:
+        console.log(f"[yellow]{result['warning']}")
+    else:
         console.log(f"[green]Your cluster has been removed from {KALAVAI_PLATFORM_URL}")
-    except Exception as e:
-        console.log(f"[red]Error when unpublishing cluster. {str(e)}")
 
 @arguably.command
 def pool__list(*others, user_only=False):
@@ -430,6 +415,9 @@ def pool__start(cluster_name, *others,  only_registered_users: bool=False, ip_ad
         location=location
     )
 
+    if "warning" in result:
+        console.log(f"[yellow]Warning: {result['warning']}")
+
     if "error" in result:
         console.log(f"[red]{result}")
 
@@ -539,38 +527,13 @@ def pool__stop(*others, skip_node_deletion=False):
         *others: all the other positional arguments go here
     """
     console.log("[white] Stopping kalavai app...")
-    # delete local node from server
-    if not skip_node_deletion:
-        node__delete(load_server_info(data_key=NODE_NAME_KEY, file=USER_LOCAL_SERVER_FILE))
-    # unpublish event (only if seed node)
-    # TODO: no, this should be done via the platform!!!
-    # try:
-    #     if CLUSTER.is_seed_node():
-    #         console.log("Unregistering pool...")
-    #         unregister_cluster(
-    #             name=load_server_info(data_key=CLUSTER_NAME_KEY, file=USER_LOCAL_SERVER_FILE),
-    #             user_cookie=USER_COOKIE)
-    # except Exception as e:
-    #     console.log(f"[red][WARNING]: (ignore if not a public pool) Error when unpublishing cluster. {str(e)}")
-    # remove local node agent
-    console.log("Removing agent and local cache")
-    
-    # disconnect from VPN first, then remove agent, then remove local files
-    console.log("Disconnecting from VPN...")
-    try:
-        vpns = leave_vpn(container_name=DEFAULT_VPN_CONTAINER_NAME)
-        if vpns is not None:
-            for vpn in vpns:
-                console.log(f"You have left {vpn} VPN")
-    except:
-        # no vpn
-        pass
 
-    CLUSTER.remove_agent()
-
-    # clean local files
-    cleanup_local()
-    console.log("[white] Kalavai has stopped sharing your resources. Use [yellow]kalavai pool start[white] or [yellow]kalavai pool join[white] to start again!")
+    result = stop_pool(skip_node_deletion=skip_node_deletion)
+    if "error" in result:
+        console.log(f"[red]{result['error']}")
+    else:
+        console.log(result)
+        console.log("[white] Kalavai has stopped sharing your resources. Use [yellow]kalavai pool start[white] or [yellow]kalavai pool join[white] to start again!")
 
 @arguably.command
 def pool__pause(*others):
@@ -958,8 +921,11 @@ def node__cordon(node_name, *others):
     except Exception as e:
         console.log(f"[red]Problems with your pool: {str(e)}")
         return
-    set_schedulable(schedulable=False, node_name=node_name)
-
+    result = cordon_nodes(nodes=[node_name])
+    if "error" in result:
+        console.log(f"[red]{result['error']}")
+    else:
+        console.log(result)
 
 @arguably.command
 def node__uncordon(node_name, *others):
@@ -971,8 +937,11 @@ def node__uncordon(node_name, *others):
     except Exception as e:
         console.log(f"[red]Problems with your pool: {str(e)}")
         return
-    set_schedulable(schedulable=True, node_name=node_name)
-
+    result = uncordon_nodes(nodes=[node_name])
+    if "error" in result:
+        console.log(f"[red]{result['error']}")
+    else:
+        console.log(result)
 
 @arguably.command
 def job__templates(*others):

kalavai_client/core.py

@@ -9,15 +9,19 @@ import netifaces as ni
 from typing import Optional
 from pydantic import BaseModel
 from enum import Enum
+import re
 
 from kalavai_client.cluster import CLUSTER
 from kalavai_client.utils import (
     check_gpu_drivers,
     generate_join_token,
+    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,
@@ -98,6 +102,30 @@ class TokenType(Enum):
     WORKER = 2
 
 
+def set_schedulable(schedulable, node_names):
+    """
+    Delete job in the cluster
+    """
+    # deploy template with kube-watcher
+    data = {
+        "schedulable": str(schedulable),
+        "node_names": node_names
+    }
+    try:
+        res = request_to_server(
+            method="post",
+            endpoint="/v1/set_node_schedulable",
+            data=data,
+            server_creds=USER_LOCAL_SERVER_FILE,
+            user_cookie=USER_COOKIE
+        )
+        if res is not None and "detail" in res:
+            return {"error": res["detail"]}
+        else:
+            return {"result": res}
+    except Exception as e:
+        return {"error": f"Error when connecting to kalavai service: {str(e)}"}
+
 def init_user_workspace(force_namespace=None):
     
     # load template config and populate with values
@@ -495,6 +523,12 @@ def delete_nodes(nodes):
     except Exception as e:
         return {"error": f"Error when removing nodes {nodes}: {str(e)}"}
 
+def cordon_nodes(nodes):
+    return set_schedulable(schedulable=False, node_names=nodes)
+
+def uncordon_nodes(nodes):
+    return set_schedulable(schedulable=True, node_names=nodes)
+
 def attach_to_pool(token, node_name=None):
     if node_name is None:
         node_name = f"{socket.gethostname()}-{uuid.uuid4().hex[:6]}"
@@ -673,7 +707,7 @@ def join_pool(token, num_gpus=None, node_name=None, ip_address=None):
     
     return cluster_name
 
-def create_pool(cluster_name: str, ip_address: str, 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, only_registered_users: bool=False, location: str=None):
 
     if not check_seed_compatibility():
         return {"error": "Requirements failed"}
@@ -799,6 +833,18 @@ def create_pool(cluster_name: str, ip_address: str, app_values: str=None, pool_c
         # init user namespace
         init_user_workspace()
     
+    # register cluster (if user is logged in)
+    result = register_pool(
+        cluster_name=cluster_name,
+        description=description,
+        is_private=True,
+        token_mode=token_mode)
+
+    if "error" in result:
+        return {"warning": result["error"]}
+    if "warning" in result:
+        return {"warning": result["warning"]}
+    
     return {"success"}
 
 def get_pool_token(mode: TokenType):
@@ -857,6 +903,33 @@ def pool_init(pool_config_values_path=None):
         return result
     except Exception as e:
         return {"error": f"[red]Error when connecting to kalavai service: {str(e)}"}
+    
+def register_pool(cluster_name, description, is_private=True, token_mode=TokenType.USER):
+    token = get_pool_token(mode=token_mode)["token"]
+    valid = check_token(token=token, public=not is_private)
+    if "error" in valid:
+        return {"error": valid}
+    try:
+        result = register_cluster(
+            name=cluster_name,
+            token=token,
+            description=description,
+            user_cookie=USER_COOKIE,
+            is_private=is_private)
+
+        return {"success": result}
+    except Exception as e:
+        return {"warning": str(e)}
+    
+def unregister_pool():
+    cluster_name = load_server_info(data_key=CLUSTER_NAME_KEY, file=USER_LOCAL_SERVER_FILE)
+    try:
+        unregister_cluster(
+            name=cluster_name,
+            user_cookie=USER_COOKIE)
+    except Exception as e:
+        return {"warning": str(e)}
+    return {"success"}
 
 def is_connected():
     if not os.path.isfile(USER_LOCAL_SERVER_FILE):
@@ -927,16 +1000,8 @@ def stop_pool(skip_node_deletion=False):
             delete_node(load_server_info(data_key=NODE_NAME_KEY, file=USER_LOCAL_SERVER_FILE))
         )
     # unpublish event (only if seed node)
-    # TODO: no, this should be done via the platform!!!
-    # try:
-    #     if CLUSTER.is_seed_node():
-    #         console.log("Unregistering pool...")
-    #         unregister_cluster(
-    #             name=load_server_info(data_key=CLUSTER_NAME_KEY, file=USER_LOCAL_SERVER_FILE),
-    #             user_cookie=USER_COOKIE)
-    # except Exception as e:
-    #     console.log(f"[red][WARNING]: (ignore if not a public pool) Error when unpublishing cluster. {str(e)}")
-    # remove local node agent
+    if CLUSTER.is_seed_node():
+        unregister_pool()
     
     # disconnect from VPN first, then remove agent, then remove local files
     try:
@@ -948,13 +1013,22 @@ def stop_pool(skip_node_deletion=False):
         # no vpn
         pass
 
+    # remove local node agent
     CLUSTER.remove_agent()
 
     # clean local files
     cleanup_local()
 
-    return logs
+    return {"logs": logs}
 
 def list_available_pools(user_only=False):
     pools = get_public_seeds(user_only=user_only, user_cookie=USER_COOKIE)
-    return pools
+    return pools
+
+def send_invites(invitees):
+    result = send_pool_invite(
+        cluster_name=load_server_info(data_key=CLUSTER_NAME_KEY, file=USER_LOCAL_SERVER_FILE),
+        invitee_addresses=invitees,
+        user_cookie=USER_COOKIE
+    )
+    return result

kalavai_client/utils.py

@@ -206,7 +206,7 @@ def get_vpn_details(location, user_cookie):
     )
     return vpn
 
-def register_cluster(name, token, description, user_cookie):
+def register_cluster(name, token, description, user_cookie, is_private=True):
     auth = KalavaiAuthClient(
         user_cookie_file=user_cookie
     )
@@ -218,7 +218,8 @@ def register_cluster(name, token, description, user_cookie):
         name,
         user,
         description,
-        token
+        token,
+        is_private
     )
     return seed
 
@@ -251,6 +252,21 @@ def validate_join_public_seed(cluster_name, join_key, user_cookie):
     )
     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
+
 def validate_poolconfig(poolconfig_file):
     if not Path(poolconfig_file).is_file():
         return False

{kalavai_client-0.5.25.dist-info → kalavai_client-0.5.27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kalavai-client
-Version: 0.5.25
+Version: 0.5.27
 Summary: Client app for kalavai platform
 License: Apache-2.0
 Keywords: LLM,platform

{kalavai_client-0.5.25.dist-info → kalavai_client-0.5.27.dist-info}/RECORD

@@ -1,4 +1,4 @@
-kalavai_client/__init__.py,sha256=5jLxgxx_gD3kMCIx_aZdf95uxR2xc87aRxAL6dfWZqA,23
+kalavai_client/__init__.py,sha256=Wek6gRlCzMJXNHhlUntIAj24wWPV3k64gGvcy5WL7rU,23
 kalavai_client/__main__.py,sha256=WQUfxvRsBJH5gsCJg8pLz95QnZIj7Ol8psTO77m0QE0,73
 kalavai_client/assets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kalavai_client/assets/apps.yaml,sha256=d13TzkWtqdwpWOxuoG7eG0Jp0UhVUsboS28496H8iH4,5982
@@ -11,15 +11,15 @@ kalavai_client/assets/pool_config_values.yaml,sha256=VrM3XHQfQo6QLZ68qvagooUptaY
 kalavai_client/assets/user_workspace.yaml,sha256=wDvlMYknOPABAEo0dsQwU7bac8iubjAG9tdkFbJZ5Go,476
 kalavai_client/assets/user_workspace_values.yaml,sha256=G0HOzQUxrDMCwuW9kbWUZaKMzDDPVwDwzBHCL2Xi2ZM,542
 kalavai_client/auth.py,sha256=QsBh28L2LwjBBK6pTUE4Xu36lLDTyetyU1YfS1Hbb6g,1717
-kalavai_client/bridge_api.py,sha256=lFS1fvH7UQxZ24vrHGeea5z8bnGWhOZhgU_zKIrnlHI,4891
-kalavai_client/bridge_models.py,sha256=hWxEODqHwc8Sj21zkIZoZmU9W3PvNdK-CA1tH6yMQ4s,798
-kalavai_client/cli.py,sha256=Ev9SsnbgCU21WnxOOpSJbvqfQ1pMhIqSlk3tiQIuz7w,48099
+kalavai_client/bridge_api.py,sha256=rXTz6WpzQtsDmBDUUOqPEjWx8vfiIWEfvP1iM4MYDGM,5501
+kalavai_client/bridge_models.py,sha256=WwGIaWBIk4s32YemgDB2CcrrCWC-KeZjTT3iBi-kaa0,936
+kalavai_client/cli.py,sha256=zQ205vqPW69oEt9EAO5wC8_yEvWbFEtLUnI90Oqt4t4,46642
 kalavai_client/cluster.py,sha256=gwjmdsd--YrffT0BmZDOEpbrdm3lPskUuN5jdgcrOR0,12947
-kalavai_client/core.py,sha256=c9UI_RJ7QmvWPfB9VQpc5QtNCbExOcPtGc-zed1ZgzI,30187
+kalavai_client/core.py,sha256=Trv2DDdlBAsBYUaTKnosWrfboYRrZnRcv-jfvmG7-LU,32288
 kalavai_client/env.py,sha256=Zg2pP-xGJpQumo56KMBxBLgIsBmcNN0S9R-ZP2-s630,2604
-kalavai_client/utils.py,sha256=rz5W9PRZrTpgdmOs6yeqUi4f_q_L-3BJ5g1o7Asgnyo,13386
-kalavai_client-0.5.25.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-kalavai_client-0.5.25.dist-info/METADATA,sha256=orePuUjefv9e6fOzA0jk0ioz4NluV7JjuIx6iaj4YOY,14443
-kalavai_client-0.5.25.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-kalavai_client-0.5.25.dist-info/entry_points.txt,sha256=9T6D45gxwzfVbglMm1r6XPdXuuZdHfy_7fCeu2jUphc,50
-kalavai_client-0.5.25.dist-info/RECORD,,
+kalavai_client/utils.py,sha256=OPmrsycNyrs2ZpTsjAzBuPN8hQNJtsYDLPKU13tnf-U,13862
+kalavai_client-0.5.27.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+kalavai_client-0.5.27.dist-info/METADATA,sha256=5-vseYyG7Ya57UyFAqdEabXKGoFrzTMRnJgPb6FRZ6Q,14443
+kalavai_client-0.5.27.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+kalavai_client-0.5.27.dist-info/entry_points.txt,sha256=9T6D45gxwzfVbglMm1r6XPdXuuZdHfy_7fCeu2jUphc,50
+kalavai_client-0.5.27.dist-info/RECORD,,