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

kalavai_client/__init__.py

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

kalavai_client/bridge_api.py

@@ -7,6 +7,7 @@ import uvicorn
 
 from kalavai_client.bridge_models import (
     CreatePoolRequest,
+    InvitesRequest,
     JoinPoolRequest,
     StopPoolRequest,
     DeployJobRequest,
@@ -18,6 +19,7 @@ from kalavai_client.core import (
     create_pool,
     join_pool,
     attach_to_pool,
+    send_invites,
     stop_pool,
     fetch_devices,
     fetch_resources,
@@ -57,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
 
@@ -108,14 +112,18 @@ def device_uncordon(request: NodesActionRequest):
     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,6 +14,8 @@ class CreatePoolRequest(BaseModel):
     node_name: str = None
     only_registered_users: bool = False
     location: str = None
+    token_mode: TokenType = TokenType.USER
+    description: str = ""
 
 class NodesActionRequest(BaseModel):
     nodes: list[str]

kalavai_client/cli.py

@@ -51,12 +51,15 @@ from kalavai_client.core import (
     create_pool,
     get_ip_addresses,
     pause_agent,
+    register_pool,
     resume_agent,
     get_pool_token,
     delete_nodes,
     cordon_nodes,
+    stop_pool,
     uncordon_nodes,
-    TokenType
+    TokenType,
+    unregister_pool
 )
 from kalavai_client.utils import (
     check_gpu_drivers,
@@ -286,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
     """
@@ -300,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):
@@ -335,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):
@@ -411,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}")
 
@@ -520,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):

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,
@@ -703,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"}
@@ -829,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):
@@ -887,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):
@@ -957,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:
@@ -978,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.26.dist-info → kalavai_client-0.5.27.dist-info}/METADATA

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

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

@@ -1,4 +1,4 @@
-kalavai_client/__init__.py,sha256=yEQn5h7HJ9p7TIniys_VwtONGK1a_jUSpvxigYZe4PI,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=vhkwgVOm0ZoKVz1dr-SGPQ7uMp-XPdhyzLJfKI2nsuc,5245
-kalavai_client/bridge_models.py,sha256=mEDdqUcK9BtijLpNo20HXuYdVccn4895hPzejSv666A,798
-kalavai_client/cli.py,sha256=U6Z4ESwV0L2jnKmpzEC4fKTcmrlYo4K0k9kSwUeiWL4,47700
+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=fhKRW5EAFwKUFdswgcuEzbSpyHAfIOrGkyghq04ckj8,31078
+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.26.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-kalavai_client-0.5.26.dist-info/METADATA,sha256=XlK_eljvXmKS6L2-XtvPzAuwCG6bVHZGPFEBifAOjlE,14443
-kalavai_client-0.5.26.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-kalavai_client-0.5.26.dist-info/entry_points.txt,sha256=9T6D45gxwzfVbglMm1r6XPdXuuZdHfy_7fCeu2jUphc,50
-kalavai_client-0.5.26.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,,