konduktor-nightly 0.1.0.dev20250209104336__py3-none-any.whl → 0.1.0.dev20250313070642__py3-none-any.whl

konduktor/controller/node.py

@@ -6,7 +6,7 @@ from konduktor import kube_client
 from konduktor import logging as konduktor_logging
 
 # node taint/label
-NODE_HEALTH_LABEL = "trainy.konduktor.ai/faulty"
+NODE_HEALTH_LABEL = 'trainy.konduktor.ai/faulty'
 
 logger = konduktor_logging.get_logger(__name__)
 
@@ -68,7 +68,7 @@ def untaint(node_name: str):
         _request_timeout=kube_client.API_TIMEOUT,
     )
 
-    logger.info(f"Node {node_name} taint removed.")
+    logger.info(f'Node {node_name} taint removed.')
 
 
 def taint(node_name: str):
@@ -80,8 +80,8 @@ def taint(node_name: str):
     core_api = kube_client.core_api()
     taint = kubernetes.client.V1Taint(
         key=NODE_HEALTH_LABEL,
-        value="true",
-        effect="NoSchedule",
+        value='true',
+        effect='NoSchedule',
     )
     node = core_api.read_node(
         name=node_name,
@@ -103,7 +103,7 @@ def taint(node_name: str):
         _request_timeout=kube_client.API_TIMEOUT,
     )
 
-    logger.info(f"Node {node_name} tainted.")
+    logger.info(f'Node {node_name} tainted.')
 
 
 def list_nodes() -> List[str]:

konduktor/controller/parse.py

@@ -8,19 +8,19 @@ from konduktor import logging as konduktor_logging
 from konduktor.controller import constants
 
 # comma separated list of namespaces to watch for pod errors
-WATCHED_NAMESPACES: List[str] = os.environ.get("WATCHED_NAMESPACES", "default").split(
-    ","
+WATCHED_NAMESPACES: List[str] = os.environ.get('WATCHED_NAMESPACES', 'default').split(
+    ','
 )
 LOGS_SINCE: int = 10  # retrieves logs generated in the past 10 seconds
 LOG_ENDPOINT: str = os.environ.get(
-    "LOG_ENDPOINT",
+    'LOG_ENDPOINT',
     # this assumes you have access to this endpoint by
     # running as a deployment within the cluster
     # for local development use 'http://localhost:3100' and
     # kubectl port-forward svc/loki -n loki 3100:3100
-    "http://loki.loki.svc.cluster.local:3100",
+    'http://loki.loki.svc.cluster.local:3100',
 )
-QUERY_URL: str = "/loki/api/v1/query_range"
+QUERY_URL: str = '/loki/api/v1/query_range'
 
 logger = konduktor_logging.get_logger(__name__)
 
@@ -35,32 +35,32 @@ def _query_range(pattern: str, **label_filters) -> List[Dict[str, Any]]:
     Returns:
         List[Dict[str, Any]]: List of loglines
     """
-    url = f"{LOG_ENDPOINT}{QUERY_URL}"
-    formatted_filters = ", ".join(
+    url = f'{LOG_ENDPOINT}{QUERY_URL}'
+    formatted_filters = ', '.join(
         f'{key}="{value}"' for key, value in label_filters.items()
     )
-    query = r"{" f"{formatted_filters}" r"}" f"|~ {pattern}"
-    params = {"query": query, "since": f"{LOGS_SINCE}s"}
+    query = r'{' f'{formatted_filters}' r'}' f'|~ {pattern}'
+    params = {'query': query, 'since': f'{LOGS_SINCE}s'}
     response = requests.get(url, params=params)
     if response.status_code == 200:
         data = response.json()
-        return data["data"]["result"]
+        return data['data']['result']
     elif response.status_code == 400:
-        logger.error(f"Bad Request: {response.status_code}")
+        logger.error(f'Bad Request: {response.status_code}')
         logger.error(response.json())  # Optionally print the error details
     else:
-        logger.error(f"loki query failed {params}")
+        logger.error(f'loki query failed {params}')
     return []
 
 
 def pod_errors() -> Set[str]:
-    logger.info("querying pod logs")
+    logger.info('querying pod logs')
     bad_nodes = set()
     for regex in constants.POD_LOG_ERROR_REGEXES:
         for namespace in WATCHED_NAMESPACES:
             log_lines = _query_range(regex, k8s_namespace_name=namespace)
             for line in log_lines:
-                log_node = line["stream"]["k8s_node_name"]
+                log_node = line['stream']['k8s_node_name']
                 bad_nodes.add(log_node)
     return bad_nodes
 
@@ -82,28 +82,28 @@ def sxid_error(pattern: str, log_content: str) -> int:
 
 def is_sxid_error(log_content: str) -> bool:
     """Returns (S)Xid error code, zero otherwise"""
-    error_code = sxid_error(r"SXid.*?: (\d+),", log_content) or sxid_error(
-        r"NVRM: Xid.*?: (\d+),", log_content
+    error_code = sxid_error(r'SXid.*?: (\d+),', log_content) or sxid_error(
+        r'NVRM: Xid.*?: (\d+),', log_content
     )
     return error_code not in constants.ALLOWLISTED_NVSWITCH_SXID_ERRORS
 
 
 def dmesg_errors() -> Set[str]:
-    logger.info("checking dmesg logs")
-    pattern = " or ".join(constants.DMESG_ERROR_REGEXES)
-    log_lines = _query_range(pattern, k8s_daemonset_name="dmesg")
+    logger.info('checking dmesg logs')
+    pattern = ' or '.join(constants.DMESG_ERROR_REGEXES)
+    log_lines = _query_range(pattern, k8s_daemonset_name='dmesg')
     bad_nodes = set()
     for line in log_lines:
-        log_node, log_content = line["stream"]["k8s_node_name"], line["values"][0][1]
+        log_node, log_content = line['stream']['k8s_node_name'], line['values'][0][1]
         if is_sxid_error(log_content):
-            logger.info(f"node `{log_node}` has (S)Xid error: {log_content}")
+            logger.info(f'node `{log_node}` has (S)Xid error: {log_content}')
         else:
-            logger.info(f"dmesg error on node `{log_node}`: {log_content}")
+            logger.info(f'dmesg error on node `{log_node}`: {log_content}')
         bad_nodes.add(log_node)
     return bad_nodes
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     import time
 
     while True:

konduktor/dashboard/backend/main.py

@@ -21,10 +21,10 @@ app = FastAPI()
 # CORS Configuration
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allow all origins
+    allow_origins=['*'],  # Allow all origins
     allow_credentials=True,
-    allow_methods=["*"],  # Allow all methods
-    allow_headers=["*"],  # Allow all headers
+    allow_methods=['*'],  # Allow all methods
+    allow_headers=['*'],  # Allow all headers
 )
 
 # Use Kubernetes API clients
@@ -35,44 +35,44 @@ core_client = core_api()
 crd_client = crd_api()
 
 
-@app.get("/")
+@app.get('/')
 async def home():
-    return JSONResponse({"home": "/"})
+    return JSONResponse({'home': '/'})
 
 
-@app.delete("/deleteJob")
+@app.delete('/deleteJob')
 async def delete_job(request: Request):
     data = await request.json()
-    name = data.get("name", "")
-    namespace = data.get("namespace", "default")
+    name = data.get('name', '')
+    namespace = data.get('namespace', 'default')
 
     try:
-        delete_options = client.V1DeleteOptions(propagation_policy="Background")
+        delete_options = client.V1DeleteOptions(propagation_policy='Background')
 
         crd_client.delete_namespaced_custom_object(
-            group="kueue.x-k8s.io",
-            version="v1beta1",
+            group='kueue.x-k8s.io',
+            version='v1beta1',
             namespace=namespace,
-            plural="workloads",
+            plural='workloads',
             name=name,
             body=delete_options,
         )
         logger.debug(f"Kueue Workload '{name}' deleted successfully.")
 
-        return JSONResponse({"success": True, "status": 200})
+        return JSONResponse({'success': True, 'status': 200})
 
     except ApiException as e:
-        logger.debug(f"Exception: {e}")
-        return JSONResponse({"error": str(e)}, status_code=e.status)
+        logger.debug(f'Exception: {e}')
+        return JSONResponse({'error': str(e)}, status_code=e.status)
 
 
-@app.get("/getJobs")
+@app.get('/getJobs')
 async def get_jobs():
     rows = fetch_jobs()
     return JSONResponse(rows)
 
 
-@app.get("/getNamespaces")
+@app.get('/getNamespaces')
 async def get_namespaces():
     try:
         # Get the list of namespaces
@@ -81,50 +81,50 @@ async def get_namespaces():
         namespace_list = [ns.metadata.name for ns in namespaces.items]
         return JSONResponse(namespace_list)
     except ApiException as e:
-        logger.debug(f"Exception: {e}")
-        return JSONResponse({"error": str(e)}, status_code=e.status)
+        logger.debug(f'Exception: {e}')
+        return JSONResponse({'error': str(e)}, status_code=e.status)
 
 
-@app.put("/updatePriority")
+@app.put('/updatePriority')
 async def update_priority(request: Request):
     data = await request.json()
-    name = data.get("name", "")
-    namespace = data.get("namespace", "default")
-    priority = data.get("priority", 0)
+    name = data.get('name', '')
+    namespace = data.get('namespace', 'default')
+    priority = data.get('priority', 0)
 
     try:
         job = crd_client.get_namespaced_custom_object(
-            group="kueue.x-k8s.io",
-            version="v1beta1",
+            group='kueue.x-k8s.io',
+            version='v1beta1',
             namespace=namespace,
-            plural="workloads",
+            plural='workloads',
             name=name,
         )
 
-        job["spec"]["priority"] = priority
+        job['spec']['priority'] = priority
 
         crd_client.patch_namespaced_custom_object(
-            group="kueue.x-k8s.io",
-            version="v1beta1",
+            group='kueue.x-k8s.io',
+            version='v1beta1',
             namespace=namespace,
-            plural="workloads",
+            plural='workloads',
             name=name,
             body=job,
         )
-        return JSONResponse({"success": True, "status": 200})
+        return JSONResponse({'success': True, 'status': 200})
 
     except ApiException as e:
-        logger.debug(f"Exception: {e}")
-        return JSONResponse({"error": str(e)}, status_code=e.status)
+        logger.debug(f'Exception: {e}')
+        return JSONResponse({'error': str(e)}, status_code=e.status)
 
 
 # Get a listing of workloads in kueue
 def fetch_jobs():
     listing = crd_client.list_namespaced_custom_object(
-        group="kueue.x-k8s.io",
-        version="v1beta1",
-        namespace="default",
-        plural="workloads",
+        group='kueue.x-k8s.io',
+        version='v1beta1',
+        namespace='default',
+        plural='workloads',
     )
 
     return format_workloads(listing)
@@ -136,30 +136,30 @@ def format_workloads(listing: Dict[str, Any]) -> List[Dict[str, Any]]:
 
     res = []
 
-    for job in listing["items"]:
-        id = job["metadata"]["uid"]
-        name = job["metadata"]["name"]
-        created_at = job["metadata"]["creationTimestamp"]
-        namespace = job["metadata"]["namespace"]
-        localQueueName = job["spec"].get("queueName", "Unknown")
-        priority = job["spec"]["priority"]
-        active = job["spec"].get("active", 0)
-        status = "ADMITTED" if "admission" in job.get("status", {}) else "PENDING"
-
-        statusVal = 1 if "admission" in job.get("status", {}) else 0
+    for job in listing['items']:
+        id = job['metadata']['uid']
+        name = job['metadata']['name']
+        created_at = job['metadata']['creationTimestamp']
+        namespace = job['metadata']['namespace']
+        localQueueName = job['spec'].get('queueName', 'Unknown')
+        priority = job['spec']['priority']
+        active = job['spec'].get('active', 0)
+        status = 'ADMITTED' if 'admission' in job.get('status', {}) else 'PENDING'
+
+        statusVal = 1 if 'admission' in job.get('status', {}) else 0
         order = (statusVal * 10) + priority
 
         res.append(
             {
-                "id": id,
-                "name": name,
-                "namespace": namespace,
-                "localQueueName": localQueueName,
-                "priority": priority,
-                "status": status,
-                "active": active,
-                "created_at": created_at,
-                "order": order,
+                'id': id,
+                'name': name,
+                'namespace': namespace,
+                'localQueueName': localQueueName,
+                'priority': priority,
+                'status': status,
+                'active': active,
+                'created_at': created_at,
+                'order': order,
             }
         )
 

konduktor/dashboard/backend/sockets.py

@@ -11,7 +11,7 @@ from konduktor import logging as konduktor_logging
 
 # SocketIO configuration
 socketio = AsyncServer(
-    cors_allowed_origins="*", ping_interval=25, ping_timeout=60, async_mode="asgi"
+    cors_allowed_origins='*', ping_interval=25, ping_timeout=60, async_mode='asgi'
 )
 
 logger = konduktor_logging.get_logger(__name__)
@@ -25,7 +25,7 @@ SELECTED_NAMESPACES: list[str] = []
 
 # "http://loki.loki.svc.cluster.local:3100/loki/api/v1/query_range" for prod
 # "http://localhost:3100/loki/api/v1/query_range" for local
-LOGS_URL = os.environ.get("LOGS_URL", "http://localhost:3100/loki/api/v1/query_range")
+LOGS_URL = os.environ.get('LOGS_URL', 'http://localhost:3100/loki/api/v1/query_range')
 
 
 def format_log_entry(entry: List[str], namespace: str) -> Dict[str, str]:
@@ -44,11 +44,11 @@ def format_log_entry(entry: List[str], namespace: str) -> Dict[str, str]:
     log_message = entry[1]
     timestamp_s = int(timestamp_ns) / 1e9
     dt = datetime.datetime.utcfromtimestamp(timestamp_s)
-    human_readable_time = dt.strftime("%Y-%m-%d %H:%M:%S")
+    human_readable_time = dt.strftime('%Y-%m-%d %H:%M:%S')
     formatted_log = {
-        "timestamp": human_readable_time,
-        "log": log_message,
-        "namespace": namespace,
+        'timestamp': human_readable_time,
+        'log': log_message,
+        'namespace': namespace,
     }
     return formatted_log
 
@@ -56,15 +56,15 @@ def format_log_entry(entry: List[str], namespace: str) -> Dict[str, str]:
 def get_logs(FIRST_RUN: bool) -> List[Dict[str, str]]:
     global LOG_CHECKPOINT_TIME
 
-    logger.debug(f"Selected namespaces: {SELECTED_NAMESPACES}")
+    logger.debug(f'Selected namespaces: {SELECTED_NAMESPACES}')
 
     # Use the selected namespaces in the query
     namespace_filter = (
-        "|".join(SELECTED_NAMESPACES) if SELECTED_NAMESPACES else "default"
+        '|'.join(SELECTED_NAMESPACES) if SELECTED_NAMESPACES else 'default'
     )
     query = f'{{k8s_namespace_name=~"{namespace_filter}"}}'
 
-    logger.debug(f"Loki logs query: {query}")
+    logger.debug(f'Loki logs query: {query}')
 
     if FIRST_RUN:
         # Calculate how many nanoseconds to look back when first time looking at logs
@@ -78,7 +78,7 @@ def get_logs(FIRST_RUN: bool) -> List[Dict[str, str]]:
             LOG_CHECKPOINT_TIME = 0
         start_time = str(int(LOG_CHECKPOINT_TIME) + 1)
 
-    params = {"query": query, "start": start_time, "limit": "300"}
+    params = {'query': query, 'start': start_time, 'limit': '300'}
 
     url = LOGS_URL
     response = requests.get(url, params=params)
@@ -88,11 +88,11 @@ def get_logs(FIRST_RUN: bool) -> List[Dict[str, str]]:
 
     if response.status_code == 200:
         data = response.json()
-        rows = data["data"]["result"]
+        rows = data['data']['result']
 
         for row in rows:
-            namespace = row["stream"]["k8s_namespace_name"]
-            for value in row["values"]:
+            namespace = row['stream']['k8s_namespace_name']
+            for value in row['values']:
                 last = max(int(value[0]), last)
                 formatted_logs.append(format_log_entry(value, namespace))
 
@@ -100,12 +100,12 @@ def get_logs(FIRST_RUN: bool) -> List[Dict[str, str]]:
         # sort because sometimes loki API is wrong and logs are out of order
         formatted_logs.sort(
             key=lambda log: datetime.datetime.strptime(
-                log["timestamp"], "%Y-%m-%d %H:%M:%S"
+                log['timestamp'], '%Y-%m-%d %H:%M:%S'
             )
         )
         LOG_CHECKPOINT_TIME = last
 
-    logger.debug(f"Formatted logs length: {len(formatted_logs)}")
+    logger.debug(f'Formatted logs length: {len(formatted_logs)}')
 
     return formatted_logs
 
@@ -117,7 +117,7 @@ async def send_logs():
 
         FIRST_RUN = False  # After the first successful fetch, set to False
         if logs:
-            await socketio.emit("log_data", logs)
+            await socketio.emit('log_data', logs)
 
         await asyncio.sleep(5)
 
@@ -130,7 +130,7 @@ async def connect(sid, environ):
     global CLIENT_CONNECTED, FIRST_RUN, BACKGROUND_TASK_RUNNING
     CLIENT_CONNECTED = True
     FIRST_RUN = True
-    logger.debug("Client connected")
+    logger.debug('Client connected')
 
     # Start the background task only if it's not already running
     if not BACKGROUND_TASK_RUNNING:
@@ -142,7 +142,7 @@ async def connect(sid, environ):
 async def update_namespaces(sid, namespaces):
     global SELECTED_NAMESPACES
     SELECTED_NAMESPACES = namespaces
-    logger.debug("Updated namespaces")
+    logger.debug('Updated namespaces')
 
 
 @socketio.event
@@ -151,4 +151,4 @@ async def disconnect(sid):
     CLIENT_CONNECTED = False
     FIRST_RUN = True
     BACKGROUND_TASK_RUNNING = False
-    logger.debug("Client disconnected")
+    logger.debug('Client disconnected')

konduktor/data/__init__.py

@@ -0,0 +1,9 @@
+"""Data sync between workstation <--> blob (s3, gcs, etc.) <--> worker pods"""
+
+from konduktor.data.storage import Storage, StorageMode, StoreType
+
+__all__ = [
+    'Storage',
+    'StorageMode',
+    'StoreType',
+]

konduktor/data/constants.py

@@ -0,0 +1,12 @@
+from typing import Any, Union
+
+Path = str
+SourceType = Union[Path]
+StorageHandle = Any
+
+# TODO(asaiacai) This should match the cloud store
+# classes in cloud_stores.py,
+# should honestly just use one or the other instead of both
+STORE_ENABLED_CLOUDS = ['gs']
+
+_STORAGE_LOG_FILE_NAME = 'storage.log'