diff --git a/Makefile b/Makefile index b7fd838f..10c080e5 100644 --- a/Makefile +++ b/Makefile @@ -23,6 +23,7 @@ UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_for # Only forward BASE_URL to e2e when the user supplied it. The Makefile default # is for local CLI usage; e2e should start its own backend by default. TRAINING_TEST_BASE_URL ?= $(if $(filter environment command line,$(origin BASE_URL)),$(BASE_URL),) +UNIT_TESTS += tests.test_dashboard_kubernetes TRAINING_TEST_EXTRA ?= gpu TRAINING_TEST_ARGS ?= PIGLATIN_TEST_PYTHONPATH ?= examples/sft/pig-latin diff --git a/src/server/dashboard/__init__.py b/src/server/dashboard/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/server/dashboard/kubernetes.py b/src/server/dashboard/kubernetes.py new file mode 100644 index 00000000..f3cfebce --- /dev/null +++ b/src/server/dashboard/kubernetes.py @@ -0,0 +1,101 @@ +import asyncio +import json +import os +from functools import cache +from typing import TypedDict + +READ_TIMEOUT_SECONDS = 6 + + +class Cluster(TypedDict): + available: bool # whether pods could be listed + pods: list[dict] + nodes: list[dict] + workloads: list[dict] | None # None when the scheduler's Workloads cannot be listed + errors: list[str] + + +def namespace() -> str: + return os.getenv("OPEN_RL_WORKER_NAMESPACE") or os.getenv("POD_NAMESPACE") or "openrl-system" + + +def problem(status: dict) -> str | None: + if status.get("phase") == "Failed": + return f"Failed: {': '.join(p for p in (status.get('reason'), status.get('message')) if p) or 'see logs'}" + for cs in status.get("containerStatuses") or []: + waiting, ended = cs["state"].get("waiting"), cs["state"].get("terminated") + if waiting and waiting.get("reason") not in (None, "ContainerCreating", "PodInitializing"): + return f"{waiting['reason']}: {waiting.get('message') or ''}".strip(": ") + if ended and (ended.get("exitCode") or ended.get("reason") not in (None, "Completed")): + return f"{ended.get('reason') or 'Terminated'}: exit code {ended.get('exitCode')}" + (f" — {ended['message']}" if ended.get("message") else "") + if status.get("phase") == "Pending": + unscheduled = [c for c in status.get("conditions") or [] if c["type"] == "PodScheduled" and c["status"] != "True"] + return f"Unschedulable: {unscheduled[0].get('message') or 'no node available'}" if unscheduled else "Pending" + return None + + +def gpu_capacity(node: dict, slices: list[dict]) -> int: + """A node using DRA publishes no nvidia.com/gpu capacity; its GPUs are the devices in its NVIDIA ResourceSlices. Only + each pool's newest generation counts: while the driver restarts, old and new generations coexist.""" + name, capacity = node["metadata"]["name"], int((node["status"].get("capacity") or {}).get("nvidia.com/gpu", 0)) + mine = [s["spec"] for s in slices if s["spec"].get("nodeName") == name and s["spec"].get("driver") == "gpu.nvidia.com"] + newest: dict[str, int] = {} + for pool in (spec["pool"] for spec in mine): + newest[pool["name"]] = max(newest.get(pool["name"], 0), pool["generation"]) + devices = sum(len(spec.get("devices") or []) for spec in mine if spec["pool"]["generation"] == newest[spec["pool"]["name"]]) + return max(capacity, devices) + + +@cache +def kubernetes_api() -> tuple: + from kubernetes import client, config + + try: + config.load_incluster_config() + except Exception: + config.load_kube_config() + return client.CoreV1Api(), client.CustomObjectsApi() + + +async def kubernetes() -> Cluster: + try: + from server.scheduler_worker_manager import worker_state + + core, custom = kubernetes_api() + except Exception as exc: + return {"available": False, "pods": [], "nodes": [], "workloads": None, "errors": [f"Kubernetes unavailable: {exc}"]} + ns, errors = namespace(), [] + + async def items(name: str, call, *args) -> list[dict] | None: + try: + # Read as raw JSON, because the client's own decoding of a large pod list costs more than the request. + found = await asyncio.to_thread(call, *args, _preload_content=False, _request_timeout=READ_TIMEOUT_SECONDS) + return json.loads(found.data)["items"] + except Exception as exc: + if getattr(exc, "status", None) != 404: # no scheduler CRD or no DRA on a local stack + errors.append(f"{name} read failed: {exc}") + return None + + pods, nodes, workloads, slices = await asyncio.gather( + items("pods", core.list_namespaced_pod, ns), + items("nodes", core.list_node), + items("workloads", custom.list_namespaced_custom_object, "openrl.io", "v1alpha1", ns, "workloads"), + items("resourceslices", custom.list_cluster_custom_object, "resource.k8s.io", "v1", "resourceslices"), + ) + result: Cluster = {"available": pods is not None, "pods": [], "nodes": [], "workloads": None if workloads is None else [], "errors": errors} + for pod in pods or []: + labels, statuses = pod["metadata"].get("labels") or {}, pod["status"].get("containerStatuses") or [] + row = {"name": pod["metadata"]["name"], "node": pod["spec"].get("nodeName"), "owner": labels.get("timeslice.io/owner")} + result["pods"].append(row | {"problem": problem(pod["status"]), "restarts": sum(cs.get("restartCount") or 0 for cs in statuses)}) + for node in nodes or []: + labels, status = node["metadata"].get("labels") or {}, node["status"] + row = {"name": node["metadata"]["name"], "ready": any(c["type"] == "Ready" and c["status"] == "True" for c in status.get("conditions") or [])} + row.update(accelerator=labels.get("cloud.google.com/gke-accelerator") or labels.get("nvidia.com/gpu.product")) + row["gpu_capacity"] = gpu_capacity(node, slices or []) + result["nodes"].append(row | {"openrl_enabled": labels.get("openrl.io/enabled") == "true"}) + for item in workloads or []: + spec, status = item["spec"], item.get("status") or {} + deleting = bool(item["metadata"].get("deletionTimestamp")) + state, _ = worker_state(phase=status.get("phase") or "Pending", reason=status.get("reason") or "", node=status.get("nodeName"), deleting=deleting) + result["workloads"].append({"model_id": spec["modelID"], "training_kind": spec.get("trainingKind"), "state": state}) + return result diff --git a/tests/test_dashboard_kubernetes.py b/tests/test_dashboard_kubernetes.py new file mode 100644 index 00000000..92c12dcd --- /dev/null +++ b/tests/test_dashboard_kubernetes.py @@ -0,0 +1,40 @@ +import asyncio +import unittest +from unittest import mock + +from server.dashboard.kubernetes import gpu_capacity, kubernetes, problem + + +class NodeTest(unittest.TestCase): + def test_a_dra_node_counts_the_gpus_in_its_resource_slices(self): + node = {"metadata": {"name": "n1"}, "status": {"capacity": {"cpu": "200"}}} + + def slice_on(on: str, devices: int, generation: int = 2, driver: str = "gpu.nvidia.com") -> dict: + return {"spec": {"driver": driver, "nodeName": on, "pool": {"name": on, "generation": generation}, "devices": [{}] * devices}} + + slices = [slice_on("n1", 8), slice_on("n1", 8, generation=1), slice_on("n2", 4), slice_on("n1", 4, driver="tpu.google.com")] + self.assertEqual((gpu_capacity(node, []), gpu_capacity(node, slices)), (0, 8)) + + +class ProblemTest(unittest.TestCase): + def test_a_pod_status_names_what_keeps_it_from_running(self): + crashing = {"phase": "Running", "containerStatuses": [{"state": {"waiting": {"reason": "CrashLoopBackOff", "message": "back-off 5m0s"}}}]} + pending = {"phase": "Pending", "conditions": [{"type": "PodScheduled", "status": "False", "message": "0/3 nodes are available"}]} + completed = {"phase": "Succeeded", "containerStatuses": [{"state": {"terminated": {"reason": "Completed", "exitCode": 0}}}]} + + self.assertEqual(problem(crashing), "CrashLoopBackOff: back-off 5m0s") + self.assertEqual(problem(pending), "Unschedulable: 0/3 nodes are available") + self.assertIsNone(problem(completed)) + + +class KubernetesTest(unittest.TestCase): + def test_a_cluster_without_a_client_is_reported_as_unavailable(self): + with mock.patch("server.dashboard.kubernetes.kubernetes_api", side_effect=RuntimeError("no kubeconfig")): + cluster = asyncio.run(kubernetes()) + + self.assertFalse(cluster["available"]) + self.assertEqual(cluster["errors"], ["Kubernetes unavailable: no kubeconfig"]) + + +if __name__ == "__main__": + unittest.main()