openbayes-cli 0.4.5__py3-none-any.whl

bayes/client/resource_client.py

@@ -0,0 +1,108 @@
+from gql.transport.exceptions import TransportQueryError
+
+from bayes.client.base import BayesGQLClient
+from bayes.error import Error
+from bayes.model.resource import ResourceData, Quota, Limitations
+
+
+def get_resources(client: BayesGQLClient, partyId, type, labels):
+    query = """
+    query NormalClusterResources($partyId: String, $type: [DeployType], $labels: [String]) {
+      normalClusterResources(partyId: $partyId, type: $type, labels: $labels) {
+        name
+        memory
+        type
+        cpu {
+          type
+          millicores
+          count
+        }
+        disk {
+          type
+          size
+        }
+        gpu {
+          verboseName
+          vendor
+          type
+          name
+          mode
+          memory
+          group
+          description
+          count
+        }
+        verboseName
+        gpuResource
+        labels
+      }
+    }
+    """
+    variables = {"partyId": partyId, "type": type, "labels": labels}
+    try:
+        response = client.exec(query, variables)
+    except TransportQueryError as e:
+        raise Error(e.errors[0]['message'] if e.errors else str(e))
+
+    resource_data_list = response.get("normalClusterResources", [])
+    if resource_data_list is None or not resource_data_list:
+        return None
+
+    return [ResourceData(**resource_data) for resource_data in resource_data_list]
+
+
+def get_resource_quota(client: BayesGQLClient, partyId):
+    query = """
+        query Quota($partyId: ID!) {
+          party(id: $partyId) {
+            quota {
+              computationQuota {
+                value {
+                  availableMinutes
+                }
+                key
+              }
+            }
+          }
+        }
+    """
+    variables = {"partyId": partyId}
+    try:
+        response = client.exec(query, variables)
+    except TransportQueryError as e:
+        raise Error(e.errors[0]['message'] if e.errors else str(e))
+
+    result = response.get("party", {}).get("quota", {})
+    if not result:
+        return None
+
+    return Quota(**result)
+
+
+def get_resource_limitation(client: BayesGQLClient, partyId):
+    query = """
+       query Resources($partyId: ID!) {
+          party(id: $partyId) {
+            limitations {
+              resources {
+                key
+                value {
+                  current
+                  limit
+                }
+              }
+            }
+          }
+        }
+    """
+    variables = {"partyId": partyId}
+    try:
+        response = client.exec(query, variables)
+    except TransportQueryError as e:
+        raise Error(e.errors[0]['message'] if e.errors else str(e))
+
+    result = response.get("party", {}).get("limitations", {})
+    if not result:
+        return None
+
+    return Limitations(**result)

bayes/client/runtime_client.py

@@ -0,0 +1,32 @@
+from gql.transport.exceptions import TransportQueryError
+
+from bayes.client.base import BayesGQLClient
+from bayes.error import Error
+from bayes.model.runtime import ClusterRuntime
+
+
+def get_runtimes(client: BayesGQLClient, partyId, type, labels, runtimeType):
+    query = """
+     query NormalClusterRuntimes($partyId: String, $type: [DeployType], $labels: [String!], $runtimeType: [RuntimeType]) {
+          normalClusterRuntimes(partyId: $partyId, type: $type, labels: $labels, runtimeType: $runtimeType) {
+            framework
+            name
+            version
+            type
+            device
+            deprecated
+            labels
+          }
+        }
+    """
+    variables = {"partyId": partyId, "type": type, "labels": labels, "runtimeType": runtimeType}
+    try:
+        response = client.exec(query, variables)
+    except TransportQueryError as e:
+        raise Error(e.errors[0]['message'] if e.errors else str(e))
+
+    runtime_data_list = response.get("normalClusterRuntimes", [])
+    if runtime_data_list is None or not runtime_data_list:
+        return None
+
+    return [ClusterRuntime(**runtime_data) for runtime_data in runtime_data_list]

bayes/client/ssh_client.py

@@ -0,0 +1,66 @@
+import sys
+from datetime import datetime
+
+from gql.transport.exceptions import TransportQueryError
+from pydantic import BaseModel
+
+from bayes.client.base import BayesGQLClient
+
+
+class SSHKey(BaseModel):
+    id: int
+    name: str
+    fingerprint: str
+    createdAt: datetime
+
+
+def get_keys(client: BayesGQLClient, username: str):
+    query = """
+    query Keys($userId: String!) {
+      keys(userId: $userId) {
+        id
+        name
+        fingerprint
+        createdAt
+      }
+    }  
+    """
+    variables = {"userId": username}
+    try:
+        response = client.exec(query, variables)
+    except TransportQueryError as e:
+        error_message = e.errors[0]['message'] if e.errors else str(e)
+        print(error_message, file=sys.stderr)
+        sys.exit(1)
+
+    keys_list = response.get("keys", [])
+    if keys_list is None or not keys_list:
+        return None
+
+    return [SSHKey(**key) for key in keys_list]
+
+
+def create_key(client: BayesGQLClient, userId, name, content):
+    query = """
+    mutation CreateSSHKey($userId: String!, $name: String!, $content: String!) {
+      createSSHKey(userId: $userId, name: $name, content: $content) {
+        id
+        name
+        fingerprint
+        createdAt
+      }
+    }
+    """
+    variables = {"userId": userId, "name": name, "content": content}
+    try:
+        response = client.exec(query, variables)
+    except TransportQueryError as e:
+        error_message = e.errors[0]['message'] if e.errors else str(e)
+        print(error_message, file=sys.stderr)
+        sys.exit(1)
+
+    result = response.get("createSSHKey")
+    if result is None:
+        raise Exception("create SSHKey result is none")
+
+    return SSHKey(**result)

bayes/client/status_client.py

@@ -0,0 +1,14 @@
+from .base import BayesGQLClient
+
+
+def health_check(endpoint):
+    client = BayesGQLClient(endpoint, None)
+    query = """
+    query { __typename }
+    """
+
+    try:
+        response = client.exec(query)
+        return response is not None
+    except Exception as e:
+        return False

bayes/client/user_client.py

@@ -0,0 +1,58 @@
+from gql.transport.exceptions import TransportQueryError
+
+from .base import BayesGQLClient
+
+from pydantic import BaseModel
+
+from ..error import Error
+
+
+class LoginModel(BaseModel):
+    email: str
+    token: str
+    username: str
+
+
+def login(client: BayesGQLClient, username: str, password: str) -> LoginModel:
+    query = """
+    mutation Login($username: String!, $password: String!) {
+      login(username: $username, password: $password) {
+        email
+        token
+        username
+      }
+    }
+    """
+    variables = {"username": username, "password": password}
+    try:
+        response = client.exec(query, variables)
+    except TransportQueryError as e:
+        raise Error(e.errors[0]['message'] if e.errors else str(e))
+
+    login_data = response.get("login")
+    if login_data is None:
+        raise Error("Login failed: Unexpected response format")
+
+    return LoginModel(**login_data)
+
+
+def login_with_token(client: BayesGQLClient, token):
+    query = """
+        query LoginWithToken($token: String!) {
+            loginWithToken(token: $token) {
+                email
+                username
+            }
+        }
+        """
+    variables = {"token": token}
+    try:
+        response = client.exec(query, variables)
+    except TransportQueryError as e:
+        raise Error(e.errors[0]['message'] if e.errors else str(e))
+
+    login_data = response.get("loginWithToken")
+    if login_data is None:
+        raise Error("Login failed: Unexpected response format")
+
+    return LoginModel(token=token, **login_data)