add enterprise folder with progress page backend features

lynxkite-enterprise/README.md

@@ -0,0 +1,3 @@
+# lynxkite-enterprise
+
+Enterprise backend extensions for LynxKite.

lynxkite-enterprise/pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "lynxkite-enterprise"
+version = "0.3.0"
+description = "Enterprise backend extensions for LynxKite"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "fastapi[standard]>=0.115.6",
+    "pycrdt-websocket>=0.16",
+    "pycrdt>=0.12.26",
+]
+license = "AGPL-3.0-or-later"
+
+[project.urls]
+Homepage = "https://github.com/lynxkite/lynxkite-2000/"
+
+[tool.uv.sources]
+lynxkite-core = { workspace = true }
+
+[build-system]
+requires = ["setuptools", "wheel", "setuptools-scm"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+namespaces = true
+where = ["src"]

lynxkite-enterprise/src/lynxkite_enterprise/__init__.py

@@ -0,0 +1 @@
+"""Enterprise backend extensions for LynxKite."""

lynxkite-enterprise/src/lynxkite_enterprise/backend.py

@@ -0,0 +1,267 @@
+"""Enterprise backend routes and background refresh jobs."""
+
+import asyncio
+import contextlib
+import importlib
+from typing import Any
+import fastapi
+import pycrdt.websocket
+from lynxkite_core import workspace
+from . import progress_crdt
+
+_k8s_unavailable_logged = False
+_progress_server = None  # Set by lifespan_context
+
+
+def _log_k8s_unavailable_once() -> None:
+    global _k8s_unavailable_logged
+    if _k8s_unavailable_logged:
+        return
+    print("Kubernetes package is not installed; GPU service/K8s progress features are disabled.")
+    _k8s_unavailable_logged = True
+
+
+def _k8s_client_module():
+    try:
+        return importlib.import_module("kubernetes.client")
+    except ModuleNotFoundError:
+        _log_k8s_unavailable_once()
+        return None
+
+
+def _k8s_config_module():
+    try:
+        return importlib.import_module("kubernetes.config")
+    except ModuleNotFoundError:
+        _log_k8s_unavailable_once()
+        return None
+
+
+def _load_k8s_config() -> bool:
+    """Load Kubernetes configuration (in-cluster or from kubeconfig)."""
+    config = _k8s_config_module()
+    if config is None:
+        return False
+
+    try:
+        config.load_kube_config()
+    except Exception:
+        try:
+            config.load_incluster_config()
+        except Exception:
+            return False
+    return True
+
+
+def _get_active_workspaces(ws_server) -> list[tuple[str, workspace.Workspace]]:
+    """Return list of (name, workspace) tuples for all active workspaces."""
+    active_workspaces = []
+    for name, room in getattr(ws_server, "rooms", {}).items():
+        try:
+            ws = workspace.Workspace.model_validate(room.ws.to_py())
+            active_workspaces.append((name, ws))
+        except Exception as e:
+            print(f"Error loading workspace {name}: {e}")
+    return active_workspaces
+
+
+def _compute_gpu_services(ws_server) -> list[dict[str, Any]]:
+    """Build a list of GPU service status dicts from live Kubernetes deployments."""
+    client = _k8s_client_module()
+    if client is None:
+        return []
+
+    try:
+        if not _load_k8s_config():
+            return []
+        system_namespaces = {"kube-system", "kube-public", "kube-node-lease"}
+        active_workspaces = _get_active_workspaces(ws_server)
+        gpu_services = []
+        for d in client.AppsV1Api().list_deployment_for_all_namespaces().items:
+            if (d.metadata.namespace or "") in system_namespaces:
+                continue
+            labels = dict(d.metadata.labels) if d.metadata.labels else {}
+            used_by = [w for w in [labels.get("workspace")] if w]
+            for ws_name, ws in active_workspaces:
+                for node in getattr(ws, "nodes", []):
+                    node_name = getattr(node.data, "op_id", None) or getattr(
+                        node.data, "title", None
+                    )
+                    if node_name and node_name == d.metadata.name and ws_name not in used_by:
+                        used_by.append(ws_name)
+                        break
+            gpu_services.append(
+                {
+                    "publisher": d.metadata.namespace or "",
+                    "name": d.metadata.name,
+                    "status": "running" if (d.status.available_replicas or 0) > 0 else "stopped",
+                    "replicasHealthy": d.status.available_replicas or 0,
+                    "replicasRequested": d.spec.replicas or 0,
+                    "usedByWorkspaces": used_by,
+                }
+            )
+        return gpu_services
+    except Exception as e:
+        print(f"GPU service refresh error: {e}")
+        return []
+
+
+def _get_k8s_workspace_gpus() -> dict[str, int]:
+    """Get workspace GPU counts from Kubernetes deployment labels."""
+    k8s_workspace_gpus: dict[str, int] = {}
+    client = _k8s_client_module()
+    if client is None:
+        return k8s_workspace_gpus
+    try:
+        if not _load_k8s_config():
+            return k8s_workspace_gpus
+        for d in client.AppsV1Api().list_deployment_for_all_namespaces().items:
+            labels = dict(d.metadata.labels) if d.metadata.labels else {}
+            ws_label = labels.get("workspace")
+            if ws_label:
+                k8s_workspace_gpus[ws_label] = k8s_workspace_gpus.get(ws_label, 0) + (
+                    d.spec.replicas or 0
+                )
+    except Exception as e:
+        print(f"K8s GPU info refresh error: {e}")
+    return k8s_workspace_gpus
+
+
+async def _progress_refresh_loop(ws_server):
+    """Background task: refresh K8s GPU counts and GPU service data into the progress CRDT doc."""
+    while True:
+        progress_crdt.update_progress_workspaces(ws_server, _get_k8s_workspace_gpus())
+        progress_crdt.update_progress_gpu_services(_compute_gpu_services(ws_server))
+        await asyncio.sleep(15)
+
+
+@contextlib.asynccontextmanager
+async def progress_refresh_context(ws_server):
+    task = asyncio.create_task(_progress_refresh_loop(ws_server))
+    try:
+        yield
+    finally:
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+
+
+@contextlib.asynccontextmanager
+async def lifespan_context(ws_server):
+    """Async context manager for enterprise lifecycle: starts the progress WebSocket server
+    and the background K8s refresh loop."""
+    global _progress_server
+    server = progress_crdt.ProgressWebsocketServer(auto_clean_rooms=False)
+    _progress_server = server
+    async with server:
+        await server.get_room("progress")
+        progress_crdt.update_progress_workspaces(ws_server)
+        progress_crdt.update_progress_gpu_services([])
+        async with progress_refresh_context(ws_server):
+            yield
+    _progress_server = None
+
+
+def on_workspace_connection_open(room_name: str, ws_server) -> None:
+    """Track workspace connection and update progress."""
+    progress_crdt.on_workspace_connection_open(room_name, ws_server)
+
+
+def on_workspace_connection_close(room_name: str, ws_server) -> None:
+    """Track workspace disconnection and update progress."""
+    progress_crdt.on_workspace_connection_close(room_name, ws_server)
+
+
+def on_workspace_deleted(name: str) -> None:
+    progress_crdt.delete_workspace_entry(name)
+
+
+def on_workspace_changed(ws_server) -> None:
+    """Push an immediate workspace snapshot to the progress CRDT doc."""
+    progress_crdt.update_progress_workspaces(ws_server)
+
+
+def register_routes(app: fastapi.FastAPI, crdt_module) -> None:
+    @app.websocket("/ws/progress/crdt/{room_name:path}")
+    async def progress_crdt_websocket(websocket: fastapi.WebSocket, room_name: str):
+        room_name = crdt_module.sanitize_path(room_name)
+        if _progress_server is None:
+            await websocket.close(code=1008)
+            return
+        server = pycrdt.websocket.ASGIServer(_progress_server)
+        await server({"path": room_name, "type": "websocket"}, websocket._receive, websocket._send)
+
+    @app.post("/api/pause_workspace")
+    async def pause_workspace(req: dict):
+        """Pause or resume a workspace."""
+        room_name = req.get("room_name")
+        if not isinstance(room_name, str) or not room_name:
+            raise fastapi.HTTPException(status_code=400, detail="Missing or invalid room_name")
+        room = await crdt_module.get_room(room_name)
+        paused = req.get("paused", True)
+        with room.ws.doc.transaction():
+            room.ws["paused"] = paused
+        return {"status": "ok", "room_name": room_name, "paused": paused}
+
+    @app.post("/api/stop_workspace")
+    async def stop_workspace(req: dict):
+        """Stop and reset all nodes in a workspace."""
+        room_name = req.get("room_name")
+        if not isinstance(room_name, str) or not room_name:
+            raise fastapi.HTTPException(status_code=400, detail="Missing or invalid room_name")
+        room = await crdt_module.get_room(room_name)
+        with room.ws.doc.transaction():
+            room.ws["paused"] = True
+            for node in room.ws["nodes"]:
+                node["data"]["status"] = "planned"
+                node["data"]["message"] = None
+        return {"status": "ok", "room_name": room_name}
+
+    @app.post("/api/scale_gpu_service")
+    async def scale_gpu_service(req: dict):
+        """Scale a GPU service deployment to the specified replica count."""
+        client = _k8s_client_module()
+        if client is None:
+            raise fastapi.HTTPException(
+                status_code=503,
+                detail="Kubernetes integration is unavailable: python package 'kubernetes' is not installed.",
+            )
+
+        name = req.get("name")
+        namespace = req.get("namespace") or req.get("publisher")
+        replicas = req.get("replicas")
+
+        print(f"Scaling GPU service {namespace}/{name} to {replicas} replicas")
+        if not isinstance(name, str) or not name:
+            raise fastapi.HTTPException(status_code=400, detail="Missing or invalid name")
+        if not isinstance(namespace, str) or not namespace:
+            raise fastapi.HTTPException(status_code=400, detail="Missing or invalid namespace")
+        if not isinstance(replicas, int) or replicas < 0:
+            raise fastapi.HTTPException(
+                status_code=400, detail="Replicas must be a non-negative integer"
+            )
+
+        if not _load_k8s_config():
+            raise fastapi.HTTPException(
+                status_code=503,
+                detail="Kubernetes integration is unavailable: configuration could not be loaded.",
+            )
+        apps = client.AppsV1Api()
+        apps.patch_namespaced_deployment_scale(
+            name=name,
+            namespace=namespace,
+            body={"spec": {"replicas": replicas}},
+        )
+
+        # Push a fresh snapshot right away so the UI updates without waiting for poll loop.
+        progress_crdt.update_progress_workspaces(
+            crdt_module.ws_websocket_server,
+            _get_k8s_workspace_gpus(),
+        )
+        progress_crdt.update_progress_gpu_services(
+            _compute_gpu_services(crdt_module.ws_websocket_server)
+        )
+
+        return {"status": "ok", "name": name, "namespace": namespace, "replicas": replicas}

lynxkite-enterprise/src/lynxkite_enterprise/progress_crdt.py

@@ -0,0 +1,173 @@
+"""Enterprise CRDT helpers for progress-page data and sockets."""
+
+import json
+import pathlib
+import re
+from typing import Any
+import pycrdt
+import pycrdt.websocket
+from lynxkite_core import workspace
+
+# Tracks currently open workspace websocket connections by room name.
+_active_workspace_ws_connections: dict[str, int] = {}
+
+# Progress CRDT - singleton ephemeral room for the progress page.
+_progress_ydoc: pycrdt.Doc | None = None
+
+
+def ws_exception_handler(exception, log):
+    for ex in exception.exceptions if isinstance(exception, ExceptionGroup) else [exception]:
+        log.exception(ex)
+    return True
+
+
+class ProgressWebsocketServer(pycrdt.websocket.WebsocketServer):
+    """WebSocket server for the singleton progress CRDT room."""
+
+    async def init_room(self, name: str) -> pycrdt.websocket.YRoom:
+        global _progress_ydoc
+        ydoc = pycrdt.Doc()
+        ydoc["workspaces"] = pycrdt.Map()
+        ydoc["gpu_services"] = pycrdt.Text()
+        _progress_ydoc = ydoc
+        return pycrdt.websocket.YRoom(ydoc=ydoc, exception_handler=ws_exception_handler)
+
+    async def get_room(self, name: str) -> pycrdt.websocket.YRoom:
+        """Get or initialize a progress room with the expected schema."""
+        if name not in self.rooms:
+            self.rooms[name] = await self.init_room(name)
+        room = self.rooms[name]
+        await self.start_room(room)
+        return room
+
+
+def _workspace_display_name(room_name: str) -> str:
+    return re.sub(r"[^a-zA-Z0-9]", "-", pathlib.Path(room_name).name.removesuffix(".lynxkite.json"))
+
+
+def _extract_active_node(nodes: list) -> tuple[dict[str, Any] | None, float | None]:
+    for node in nodes:
+        if node.data.status != "active":
+            continue
+        telemetry = node.data.telemetry or {}
+        tqdm_n = telemetry.get("n")
+        tqdm_total = telemetry.get("total")
+        tqdm_rate = telemetry.get("rate")
+        active_node = {
+            "id": node.id,
+            "title": node.data.title,
+            "tqdm": {"n": tqdm_n, "total": tqdm_total, "rate": tqdm_rate} if tqdm_total else None,
+        }
+        eta_seconds = None
+        if tqdm_total and tqdm_rate and tqdm_rate > 0 and tqdm_n is not None:
+            eta_seconds = (tqdm_total - tqdm_n) / tqdm_rate
+        return active_node, eta_seconds
+    return None, None
+
+
+def _workspace_status(*, total: int, done: int, active: int, paused: bool) -> str:
+    if total == 0:
+        return "idle"
+    if done == total:
+        return "done"
+    return "paused" if paused else ("active" if active > 0 else "running")
+
+
+def _build_workspace_entry(room_name: str, room, k8s_workspace_gpus: dict[str, int]) -> str:
+    display_name = _workspace_display_name(room_name)
+    ws = workspace.Workspace.model_validate(room.ws.to_py())
+    nodes = ws.nodes or []
+    total = len(nodes)
+    done = sum(1 for n in nodes if n.data.status == "done")
+    active = sum(1 for n in nodes if n.data.status == "active")
+    paused = bool(ws.paused)
+    active_node, eta_seconds = _extract_active_node(nodes)
+    status = _workspace_status(total=total, done=done, active=active, paused=paused)
+
+    return json.dumps(
+        {
+            "name": display_name,
+            "room_name": room_name,
+            "status": status,
+            "boxes_done": done,
+            "boxes_total": total,
+            "active_node": active_node,
+            "eta_seconds": eta_seconds,
+            "gpus": (ws.execution_options or {}).get("gpus")
+            or k8s_workspace_gpus.get(display_name, 0),
+            "paused": paused,
+        }
+    )
+
+
+def _connected_workspace_rooms(server) -> list[tuple[str, Any]]:
+    return [
+        (room_name, room)
+        for room_name in sorted(_active_workspace_ws_connections)
+        if (room := server.rooms.get(room_name)) is not None
+    ]
+
+
+def _update_workspace_ws_connection(room_name: str, delta: int) -> None:
+    count = _active_workspace_ws_connections.get(room_name, 0) + delta
+    if count <= 0:
+        _active_workspace_ws_connections.pop(room_name, None)
+    else:
+        _active_workspace_ws_connections[room_name] = count
+
+
+def on_workspace_connection_open(room_name: str, ws_websocket_server) -> None:
+    _update_workspace_ws_connection(room_name, +1)
+    update_progress_workspaces(ws_websocket_server)
+
+
+def on_workspace_connection_close(room_name: str, ws_websocket_server) -> None:
+    _update_workspace_ws_connection(room_name, -1)
+    update_progress_workspaces(ws_websocket_server)
+
+
+def update_progress_workspaces(ws_websocket_server, k8s_workspace_gpus: dict | None = None):
+    """Recompute workspace status entries and push them into the progress CRDT doc."""
+    if _progress_ydoc is None or not hasattr(ws_websocket_server, "rooms"):
+        return
+    k8s_workspace_gpus = k8s_workspace_gpus or {}
+    ws_map: pycrdt.Map = _progress_ydoc["workspaces"]
+    connected_rooms = _connected_workspace_rooms(ws_websocket_server)
+    entries_by_room: dict[str, str] = {}
+    for room_name, room in connected_rooms:
+        try:
+            entries_by_room[room_name] = _build_workspace_entry(room_name, room, k8s_workspace_gpus)
+        except Exception as e:
+            print(f"Error updating progress for workspace {room_name}: {e}")
+
+    # Sync the map in one transaction to avoid transient partial snapshots.
+    with _progress_ydoc.transaction():
+        for room_name, entry in entries_by_room.items():
+            ws_map[room_name] = entry
+        for name in list(ws_map.keys()):
+            if name not in entries_by_room:
+                del ws_map[name]
+
+
+def update_progress_gpu_services(gpu_services_data: list):
+    """Update the GPU services entry in the progress CRDT doc with fresh Kubernetes data."""
+    if _progress_ydoc is None:
+        return
+    gpu_services_text: pycrdt.Text = _progress_ydoc["gpu_services"]
+    new_content = json.dumps(gpu_services_data)
+    if str(gpu_services_text) == new_content:
+        return
+    with _progress_ydoc.transaction():
+        if len(gpu_services_text) > 0:
+            del gpu_services_text[0 : len(gpu_services_text)]
+        gpu_services_text += new_content
+
+
+def delete_workspace_entry(name: str) -> None:
+    _active_workspace_ws_connections.pop(name, None)
+    if _progress_ydoc is None:
+        return
+    ws_map: pycrdt.Map = _progress_ydoc["workspaces"]
+    if name in ws_map:
+        with _progress_ydoc.transaction():
+            del ws_map[name]

ruff.toml

@@ -1 +1,2 @@
 line-length = 100
+target-version = "py311"