From 77dd3bf9ed9d0bc96adb4da5478726a7b27fd916 Mon Sep 17 00:00:00 2001 From: ShubyM Date: Thu, 29 Oct 2026 13:50:38 -0700 Subject: [PATCH 1/2] feat(scheduler): stamp GPU UUIDs and placement times on Workload status The scheduler now records which GPUs a Workload holds and when. It resolves the claim's allocated devices to the UUIDs their ResourceSlices publish, sets placedAt when it sets nodeName, and sets releasedAt when the node clears, the claim changes or teardown frees the seat. The API server can then tell which run held which GPU from Workload status alone. --- .../controller/api/v1alpha1/workload_types.go | 14 ++++++ .../api/v1alpha1/zz_generated.deepcopy.go | 13 ++++++ .../controller/internal/controller/fleet.go | 40 +++++++++++++++++ .../controller/workload_controller.go | 43 +++++++++++++++++++ .../controller/workload_controller_test.go | 34 +++++++++++++++ scheduler/deploy/base/00-workload-crd.yaml | 18 ++++++++ 6 files changed, 162 insertions(+) diff --git a/scheduler/controller/api/v1alpha1/workload_types.go b/scheduler/controller/api/v1alpha1/workload_types.go index 57c3d8cbf..9496116a9 100644 --- a/scheduler/controller/api/v1alpha1/workload_types.go +++ b/scheduler/controller/api/v1alpha1/workload_types.go @@ -171,6 +171,20 @@ type WorkloadStatus struct { // the node is genuinely unknown. NodeName string `json:"nodeName,omitempty"` + // GPUUUIDs are the UUIDs of the claim's allocated devices, which GPU + // metrics are keyed by. + // +optional + GPUUUIDs []string `json:"gpuUUIDs,omitempty"` + + // PlacedAt is when the controller observed the allocation and set NodeName. + // +optional + PlacedAt *metav1.Time `json:"placedAt,omitempty"` + + // ReleasedAt is when the worker stopped holding its devices. Cleared by + // the next placement. + // +optional + ReleasedAt *metav1.Time `json:"releasedAt,omitempty"` + // Reason is the most recent human-readable explanation of Phase. Reason string `json:"reason,omitempty"` diff --git a/scheduler/controller/api/v1alpha1/zz_generated.deepcopy.go b/scheduler/controller/api/v1alpha1/zz_generated.deepcopy.go index ef6ccd2c7..e5c84407a 100644 --- a/scheduler/controller/api/v1alpha1/zz_generated.deepcopy.go +++ b/scheduler/controller/api/v1alpha1/zz_generated.deepcopy.go @@ -200,6 +200,19 @@ func (in *WorkloadSpec) DeepCopy() *WorkloadSpec { // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *WorkloadStatus) DeepCopyInto(out *WorkloadStatus) { *out = *in + if in.GPUUUIDs != nil { + in, out := &in.GPUUUIDs, &out.GPUUUIDs + *out = make([]string, len(*in)) + copy(*out, *in) + } + if in.PlacedAt != nil { + in, out := &in.PlacedAt, &out.PlacedAt + *out = (*in).DeepCopy() + } + if in.ReleasedAt != nil { + in, out := &in.ReleasedAt, &out.ReleasedAt + *out = (*in).DeepCopy() + } if in.Conditions != nil { in, out := &in.Conditions, &out.Conditions *out = make([]v1.Condition, len(*in)) diff --git a/scheduler/controller/internal/controller/fleet.go b/scheduler/controller/internal/controller/fleet.go index f805aa46b..0fd79919a 100644 --- a/scheduler/controller/internal/controller/fleet.go +++ b/scheduler/controller/internal/controller/fleet.go @@ -3,11 +3,14 @@ package controller import ( "context" "fmt" + "path" "sort" + "strings" corev1 "k8s.io/api/core/v1" resourcev1 "k8s.io/api/resource/v1" "k8s.io/apimachinery/pkg/api/resource" + "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/log" @@ -250,6 +253,43 @@ func latestCompletePools(ctx context.Context, slices []resourcev1.ResourceSlice, return keep } +// gpuUUIDs resolves the devices allocated to a claim to the UUIDs their +// ResourceSlices publish, in allocation order. A device without a UUID +// attribute is left out; an empty result is retried on a later reconcile. +func (r *WorkloadReconciler) gpuUUIDs(ctx context.Context, claimName string) ([]string, error) { + var claim resourcev1.ResourceClaim + if err := r.Get(ctx, types.NamespacedName{Namespace: r.Namespace, Name: claimName}, &claim); err != nil { + return nil, client.IgnoreNotFound(err) + } + if claim.Status.Allocation == nil { + return nil, nil + } + var slices resourcev1.ResourceSliceList + if err := r.List(ctx, &slices); err != nil { + return nil, fmt.Errorf("list resourceslices: %w", err) + } + type deviceKey struct{ driver, pool, device string } + published := map[deviceKey]string{} + for _, i := range latestCompletePools(ctx, slices.Items, r.DeviceDriver) { + spec := slices.Items[i].Spec + for _, device := range spec.Devices { + for name, attr := range device.Attributes { + // NVIDIA's driver publishes "uuid"; match any domain and case. + if attr.StringValue != nil && strings.EqualFold(path.Base(string(name)), "uuid") { + published[deviceKey{spec.Driver, spec.Pool.Name, device.Name}] = *attr.StringValue + } + } + } + } + var uuids []string + for _, result := range claim.Status.Allocation.Devices.Results { + if uuid, ok := published[deviceKey{result.Driver, result.Pool, result.Device}]; ok { + uuids = append(uuids, uuid) + } + } + return uuids, nil +} + // claimFrom reads a managed claim back into the fleet: identity from // metadata, shape from DRA's allocation. An unallocated claim has no shape // yet -- the tiers were alternatives, and which one holds is DRA's answer. diff --git a/scheduler/controller/internal/controller/workload_controller.go b/scheduler/controller/internal/controller/workload_controller.go index d8dfe8856..04282fd08 100644 --- a/scheduler/controller/internal/controller/workload_controller.go +++ b/scheduler/controller/internal/controller/workload_controller.go @@ -201,6 +201,16 @@ func (r *WorkloadReconciler) repairPlacementAssignment(ctx context.Context, scop return placementPhaseResult{}, err } scope.claimName = "" + // Record the loss now. A re-cut reuses the claim's name, and the + // later Pending or Failed patches touch neither claim nor node, + // so this is the only transition that releases the placement. + if err := r.patchStatus(ctx, worker, func(status *openrlv1alpha1.WorkloadStatus) { + status.Phase = openrlv1alpha1.PhasePending + status.ClaimName, status.AssignmentID, status.NodeName = "", "", "" + status.Reason = "ReplacingLostClaim: the assigned claim no longer exists" + }); err != nil { + return placementPhaseResult{}, err + } } } // One seat per worker. Any other seat under this name is stale, left by @@ -512,6 +522,14 @@ func (r *WorkloadReconciler) teardown(ctx context.Context, worker *openrlv1alpha return ctrl.Result{}, err } } + if err := r.patchStatus(ctx, worker, func(status *openrlv1alpha1.WorkloadStatus) { + if status.PlacedAt != nil && status.ReleasedAt == nil { + now := metav1.Now() + status.ReleasedAt = &now + } + }); err != nil { + return ctrl.Result{}, err + } controllerutil.RemoveFinalizer(worker, workerFinalizer) return ctrl.Result{}, r.Update(ctx, worker) } @@ -694,9 +712,19 @@ func (r *WorkloadReconciler) reportPod(ctx context.Context, worker *openrlv1alph node = claim.Node } } + var uuids []string + if node != "" && len(worker.Status.GPUUUIDs) == 0 { + var err error + if uuids, err = r.gpuUUIDs(ctx, claimName); err != nil { + return err + } + } return r.patchStatus(ctx, worker, func(status *openrlv1alpha1.WorkloadStatus) { status.Phase, status.ClaimName, status.PodName, status.NodeName, status.Reason = phase, claimName, podName, node, reason + if len(uuids) > 0 { + status.GPUUUIDs = uuids + } // The footprint is DRA's answer, recordable only once the // allocation is observed: which tier held, and what each device // carries under it. @@ -715,6 +743,7 @@ func (r *WorkloadReconciler) reportPod(ctx context.Context, worker *openrlv1alph func (r *WorkloadReconciler) patchStatus(ctx context.Context, worker *openrlv1alpha1.Workload, mutate func(*openrlv1alpha1.WorkloadStatus)) error { before := worker.Status.DeepCopy() mutate(&worker.Status) + trackPlacement(before, &worker.Status, metav1.Now()) worker.Status.ObservedGeneration = worker.Generation // Skip no-op writes: they would re-trigger the watch forever. if apiequality.Semantic.DeepEqual(before, &worker.Status) { @@ -730,6 +759,20 @@ func (r *WorkloadReconciler) patchStatus(ctx context.Context, worker *openrlv1al return nil } +// trackPlacement stamps the placement interval on every status write. A worker +// is placed when NodeName is set and released when NodeName clears or its claim +// changes. Teardown releases it explicitly, since nothing else changes there. +func trackPlacement(before, after *openrlv1alpha1.WorkloadStatus, now metav1.Time) { + held := before.NodeName != "" && before.PlacedAt != nil && before.ReleasedAt == nil + if held && (after.NodeName == "" || after.ClaimName != before.ClaimName) { + after.NodeName, after.GPUUUIDs, after.ReleasedAt = "", nil, &now + return + } + if after.NodeName != "" && (after.NodeName != before.NodeName || after.PlacedAt == nil) { + after.PlacedAt, after.ReleasedAt = &now, nil + } +} + func (r *WorkloadReconciler) markPending(ctx context.Context, worker *openrlv1alpha1.Workload, reason string) error { return r.patchStatus(ctx, worker, func(status *openrlv1alpha1.WorkloadStatus) { status.Phase, status.Reason = openrlv1alpha1.PhasePending, reason diff --git a/scheduler/controller/internal/controller/workload_controller_test.go b/scheduler/controller/internal/controller/workload_controller_test.go index a9658d6db..2c78c9b5a 100644 --- a/scheduler/controller/internal/controller/workload_controller_test.go +++ b/scheduler/controller/internal/controller/workload_controller_test.go @@ -79,10 +79,12 @@ func enabledNode() []client.Object { nodeName := testNode product := "NVIDIA RTX PRO 6000 Blackwell" device := func(name string) resourcev1.Device { + uuid := "GPU-" + name return resourcev1.Device{ Name: name, Attributes: map[resourcev1.QualifiedName]resourcev1.DeviceAttribute{ "productName": {StringValue: &product}, + "uuid": {StringValue: &uuid}, }, Capacity: map[resourcev1.QualifiedName]resourcev1.DeviceCapacity{ "memory": {Value: resource.MustParse("96Gi")}, @@ -1248,3 +1250,35 @@ func TestMultiGPUWorkloadCutsAnExactCountExclusiveClaim(t *testing.T) { t.Fatalf("status = %+v, want no claim and a reason naming four devices", waiting.Status) } } + +func TestWorkerStatusRecordsItsPlacementInterval(t *testing.T) { + r := newReconciler(t, append(enabledNode(), trainerWorker("w-a", "model-a"))...) + settle(t, r, "w-a") + allocateClaim(t, r, claimOf(t, r, "w-a")) + runReconcile(t, r, "w-a") + + placed := getWorker(t, r, "w-a").Status + if placed.NodeName != testNode || placed.PlacedAt == nil || placed.ReleasedAt != nil { + t.Fatalf("status = %+v, want placed on %s and not released", placed, testNode) + } + if got := placed.GPUUUIDs; len(got) != 1 || got[0] != "GPU-gpu-0" { + t.Fatalf("gpuUUIDs = %v, want the allocated device's UUID", got) + } + + // A second finalizer keeps the worker readable after teardown lets go. + held := getWorker(t, r, "w-a") + held.Finalizers = append(held.Finalizers, "test.openrl.io/hold") + if err := r.Update(context.Background(), held); err != nil { + t.Fatal(err) + } + if err := r.Delete(context.Background(), held); err != nil { + t.Fatal(err) + } + runReconcile(t, r, "w-a") + runReconcile(t, r, "w-a") + + released := getWorker(t, r, "w-a").Status + if released.ReleasedAt == nil || !released.PlacedAt.Equal(placed.PlacedAt) { + t.Fatalf("status = %+v, want releasedAt stamped and placedAt kept", released) + } +} diff --git a/scheduler/deploy/base/00-workload-crd.yaml b/scheduler/deploy/base/00-workload-crd.yaml index 2042ce6bb..9eb73f66d 100644 --- a/scheduler/deploy/base/00-workload-crd.yaml +++ b/scheduler/deploy/base/00-workload-crd.yaml @@ -8773,6 +8773,13 @@ spec: the pools registered. format: int32 type: integer + gpuUUIDs: + description: |- + GPUUUIDs are the UUIDs of the claim's allocated devices, which GPU + metrics are keyed by. + items: + type: string + type: array memoryPerDevice: description: |- MemoryPerDevice is what each of those devices must provide, i.e. Memory @@ -8797,6 +8804,11 @@ spec: - Succeeded - Failed type: string + placedAt: + description: PlacedAt is when the controller observed the allocation + and set NodeName. + format: date-time + type: string podName: description: PodName is the worker pod, once created. type: string @@ -8804,6 +8816,12 @@ spec: description: Reason is the most recent human-readable explanation of Phase. type: string + releasedAt: + description: |- + ReleasedAt is when the worker stopped holding its devices. Cleared by + the next placement. + format: date-time + type: string type: object required: - spec From 2aa061a3475cc5cff730409a19bf583e49d788e9 Mon Sep 17 00:00:00 2001 From: ShubyM Date: Thu, 29 Oct 2026 13:55:02 -0700 Subject: [PATCH 2/2] feat(observability): list the GPUs each Workload holds The inventory now has placements: each Workload that holds GPUs, with its node, pod, GPU UUIDs, placement time and the runs it serves, read live from Workload status. When the API server releases a run's workers, a background thread saves each Workload's last placement to open_rl:placements: so it outlives the Workload. Nothing reads that history yet; the activity view will. --- docs/configuration.md | 7 ++++ src/server/observability/fleet.py | 23 ++++++++++- src/server/observability/kubernetes.py | 6 ++- src/server/scheduler_worker_manager.py | 55 ++++++++++++++++++++++++-- src/server/store.py | 15 ++++++- tests/test_observability.py | 50 +++++++++++++++++++++++ tests/test_scheduler_worker_manager.py | 55 +++++++++++++++++++++++++- 7 files changed, 202 insertions(+), 9 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 4abed401f..17584fc9a 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -129,6 +129,13 @@ The API server serves a read-only inventory of runs at `/api/v1/inspect/inventory` and a run's logs at `/api/v1/inspect/logs?run_id=`. See `/docs` for their fields. +`placements` lists the `Workload`s that hold GPUs now, from their status in +Kubernetes. Each entry has the node, pod, GPU UUIDs, `placed_at` in Unix +seconds, and the `run_ids` the `Workload` serves. When the server releases a +run's workers, it also saves each `Workload`'s last placement to the Redis list +`open_rl:placements:`. The list keeps the newest 200 records for seven +days. + | Env var | Default | What it does | | --- | --- | --- | | `OPEN_RL_LOG_STORE` | unset | Where it reads logs. Unset turns log reads off. `cloud-logging` reads Cloud Logging in the project of the server's Google credentials. `cloud-logging:` reads that Cloud Logging resource instead, such as `projects//locations/global/buckets//views/`. | diff --git a/src/server/observability/fleet.py b/src/server/observability/fleet.py index b5ec39756..1a998e48c 100644 --- a/src/server/observability/fleet.py +++ b/src/server/observability/fleet.py @@ -98,4 +98,25 @@ async def read() -> dict: runs.sort(key=lambda r: -r["created_at"]) pods = [{"name": p["name"], "problem": p["problem"], "restarts": p["restarts"]} for p in cluster["pods"] if p["problem"]] view = {"available": cluster["available"], "nodes": cluster["nodes"], "pods": pods} - return {"observed_at": observed_at, "errors": errors, "runs": runs, "cluster": view} + return {"observed_at": observed_at, "errors": errors, "runs": runs, "cluster": view, "placements": placements(cluster["workloads"], metas)} + + +def placements(workloads: list[dict] | None, metas: dict[str, str]) -> list[dict]: + """The Workloads holding GPUs now, each with the runs it serves: those whose runtime is its model ID and whose kind + matches. A shared runtime serves only its active runs.""" + served = [] + for name, raw in metas.items(): + run_id = name.removeprefix("open_rl:model_meta:") + try: + meta = decode_model_metadata(raw) + except Exception: + continue # read reports the unreadable record + if meta.status == "active" or not meta.shares_runtime(): + served.append((meta.runtime(run_id), "lora" if meta.fine_tuning_type == "lora" else "fft", run_id)) + fields = ("name", "role", "node", "pod", "gpu_uuids", "placed_at") + return [ + {field: w[field] for field in fields} + | {"run_ids": sorted(r for runtime, kind, r in served if (runtime, kind) == (w["model_id"], w["training_kind"]))} + for w in workloads or [] + if w["holds_gpus"] + ] diff --git a/src/server/observability/kubernetes.py b/src/server/observability/kubernetes.py index f3cfebce7..8e4b1b2e9 100644 --- a/src/server/observability/kubernetes.py +++ b/src/server/observability/kubernetes.py @@ -59,7 +59,7 @@ def kubernetes_api() -> tuple: async def kubernetes() -> Cluster: try: - from server.scheduler_worker_manager import worker_state + from server.scheduler_worker_manager import holds_gpus, placement_of, worker_state core, custom = kubernetes_api() except Exception as exc: @@ -97,5 +97,7 @@ async def items(name: str, call, *args) -> list[dict] | None: 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}) + placement = placement_of(item) + row = {"model_id": spec["modelID"], "training_kind": spec.get("trainingKind"), "state": state, "holds_gpus": holds_gpus(placement)} + result["workloads"].append(row | placement) return result diff --git a/src/server/scheduler_worker_manager.py b/src/server/scheduler_worker_manager.py index cbead33e1..8949c0e99 100644 --- a/src/server/scheduler_worker_manager.py +++ b/src/server/scheduler_worker_manager.py @@ -21,8 +21,10 @@ import dataclasses import datetime +import json import logging import os +import threading import time from dataclasses import dataclass, replace from typing import Any @@ -32,6 +34,7 @@ from server.estimator import Footprint, footprint from server.lora_snapshots import snapshot_root from server.sampling_dispatcher import revision_file, tmp_root +from server.store import get_state_store from server.vllm_options import ENGINE_ENV from server.worker_manager import ( FAILED, @@ -58,6 +61,7 @@ ROUTER_CONFIG = "/etc/sampler-router" # On a pool's Deployment and NetworkPolicy, so release_owner finds them after the pool's Workloads are gone. OWNER_LABEL = "openrl.io/owner" +PLACEMENT_HISTORY, PLACEMENT_HISTORY_SECONDS = 200, 7 * 24 * 3600 def workload_name(role: str, owner: str, is_lora: bool, index: int = 0) -> str: @@ -460,6 +464,7 @@ def release_owner(self, owner: str) -> set[str]: selector = "app.kubernetes.io/managed-by=open-rl-api-server" found = self.custom_api.list_namespaced_custom_object(GROUP, VERSION, self.namespace, PLURAL, label_selector=selector) ours = [item for item in found["items"] if item["spec"]["ownerID"] == owner] + record_placements(ours) pools = {item["metadata"]["labels"].get("app.kubernetes.io/instance") for item in ours} - {None} owned = f"{OWNER_LABEL}={owner}" pools |= {d.metadata.name for d in self.apps_api.list_namespaced_deployment(self.namespace, label_selector=owned).items} @@ -611,8 +616,50 @@ def worker_state(phase: str, reason: str, node: str | None, deleting: bool) -> t return STARTING, facts +def epoch_seconds(timestamp: str | None) -> float | None: + return datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00")).timestamp() if timestamp else None + + def age_seconds(timestamp: str | None, now: float) -> float | None: - if not timestamp: - return None - created = datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00")).timestamp() - return max(0.0, now - created) + created = epoch_seconds(timestamp) + return None if created is None else max(0.0, now - created) + + +def placement_of(workload: dict[str, Any]) -> dict[str, Any]: + """Where a Workload runs and the GPUs it holds, as the scheduler stamps them on its status.""" + meta, spec, status = workload["metadata"], workload["spec"], workload.get("status") or {} + return { + "name": meta["name"], + "role": spec.get("role"), + "node": status.get("nodeName"), + "pod": status.get("podName"), + "gpu_uuids": status.get("gpuUUIDs") or [], + "placed_at": epoch_seconds(status.get("placedAt")), + "released_at": epoch_seconds(status.get("releasedAt")), + "deleting": bool(meta.get("deletionTimestamp")), + } + + +def holds_gpus(placement: dict[str, Any]) -> bool: + return placement["placed_at"] is not None and placement["released_at"] is None and not placement["deleting"] + + +def record_placements(workloads: list[dict[str, Any]]) -> None: + """Keep the final placement of each Workload still holding GPUs in open_rl:placements:, newest first, so it + outlives the Workload. A daemon thread writes it, so a slow or failing Redis neither delays nor fails the release.""" + records: dict[str, list[str]] = {} + for item in workloads: + if holds_gpus(placement := placement_of(item)): + del placement["deleting"] + placement.update(training_kind=item["spec"].get("trainingKind"), released_at=time.time()) + records.setdefault(f"open_rl:placements:{item['spec']['modelID']}", []).append(json.dumps(placement)) + if records: + threading.Thread(target=push_placements, args=(records,), daemon=True).start() + + +def push_placements(records: dict[str, list[str]]) -> None: + try: + for key, values in records.items(): + get_state_store().push_recent_sync(key, values, PLACEMENT_HISTORY, PLACEMENT_HISTORY_SECONDS) + except Exception as exc: + logger.warning("Could not record the placements of released workloads: %s", exc) diff --git a/src/server/store.py b/src/server/store.py index a37b8ed86..f2d5a6a60 100644 --- a/src/server/store.py +++ b/src/server/store.py @@ -274,12 +274,19 @@ async def scan_values(self, prefix: str) -> dict[str, str]: """Every live value whose key starts with prefix.""" pass + @abstractmethod + def push_recent_sync(self, key: str, values: list[str], keep: int, ttl_seconds: float) -> None: + """Synchronously put values at the head of the list at key and keep its newest `keep`. Redis drops the list + ttl_seconds after the last push.""" + pass + class InMemoryStateStore(StateStore): def __init__(self): self.kv_store: dict[str, str] = {} self.expiries: dict[str, float] = {} self.sets: dict[str, set[str]] = {} + self.lists: dict[str, list[str]] = {} async def set_value(self, key: str, value: str, ttl_seconds: float | None = None) -> None: self.kv_store[key] = value @@ -314,6 +321,9 @@ async def set_members(self, key: str) -> set[str]: async def scan_values(self, prefix: str) -> dict[str, str]: return {key: value for key in list(self.kv_store) if key.startswith(prefix) and (value := self.get_value_sync(key)) is not None} + def push_recent_sync(self, key: str, values: list[str], keep: int, ttl_seconds: float) -> None: + self.lists[key] = (values[::-1] + self.lists.get(key, []))[:keep] + class RedisStateStore(StateStore): def __init__(self, client: redis.Redis, sync_client: sync_redis.Redis): @@ -346,6 +356,9 @@ async def scan_values(self, prefix: str) -> dict[str, str]: keys = [key async for key in self.redis.scan_iter(match=prefix + "*", count=1000)] return {key: value for key, value in zip(keys, await self.redis.mget(keys)) if value is not None} if keys else {} + def push_recent_sync(self, key: str, values: list[str], keep: int, ttl_seconds: float) -> None: + self.sync_redis.pipeline().lpush(key, *values).ltrim(key, 0, keep - 1).expire(key, int(ttl_seconds)).execute() + @cache def redis_client(redis_url: str) -> redis.Redis: @@ -364,5 +377,5 @@ def get_store() -> RequestStore: def get_state_store() -> StateStore: redis_url = os.environ.get("REDIS_URL") if redis_url: - return RedisStateStore(redis_client(redis_url), sync_redis.Redis.from_url(redis_url, decode_responses=True)) + return RedisStateStore(redis_client(redis_url), sync_redis.Redis.from_url(redis_url, decode_responses=True, socket_timeout=5)) return InMemoryStateStore() diff --git a/tests/test_observability.py b/tests/test_observability.py index 93130e385..9b0db3401 100644 --- a/tests/test_observability.py +++ b/tests/test_observability.py @@ -1,7 +1,10 @@ import asyncio import copy +import json import os import unittest +from datetime import UTC, datetime +from types import SimpleNamespace from unittest import mock import httpx @@ -20,6 +23,12 @@ def meta(**fields) -> str: return TrainingModelMetadata(base_model="Qwen/Qwen3-8B", **fields).model_dump_json() +def workload(name: str, model_id: str, kind: str, deleting: bool = False, **status: str) -> dict: + metadata = {"name": name} | ({"deletionTimestamp": "2026-10-29T12:05:00Z"} if deleting else {}) + status = {"phase": "Running", "nodeName": "n1", "podName": f"orw-{name}", "gpuUUIDs": [f"GPU-{name}"]} | status + return {"metadata": metadata, "spec": {"role": "trainer", "modelID": model_id, "trainingKind": kind}, "status": status} + + def cluster(*states: str) -> dict: return {"pods": [], "workloads": [{"model_id": "r1", "training_kind": "fft", "state": s} for s in states]} @@ -71,6 +80,47 @@ def test_a_malformed_model_record_is_skipped_with_an_error(self): self.assertEqual([r["run_id"] for r in body["runs"]], ["good"]) self.assertTrue(any("bad" in error for error in body["errors"]), body["errors"]) + def test_placements_list_the_workloads_holding_gpus_with_the_runs_they_serve(self): + def listing(items) -> object: + return lambda *args, **kwargs: SimpleNamespace(data=json.dumps({"items": list(items)})) + + placed, released = "2026-10-29T12:00:00Z", "2026-10-29T12:10:00Z" + workloads = ( + workload("fft-a", "fft-1", "fft", placedAt=placed), + workload("lora-a", "Qwen/Qwen3-8B", "lora", placedAt=placed), + workload("fft-gone", "fft-2", "fft", placedAt=placed, releasedAt=released, nodeName=""), + workload("fft-deleting", "fft-3", "fft", deleting=True, placedAt=placed), + workload("fft-waiting", "fft-4", "fft", nodeName=""), + ) + api = SimpleNamespace(list_namespaced_pod=listing([]), list_node=listing([]), list_cluster_custom_object=listing([])) + api.list_namespaced_custom_object = listing(workloads) + self.enterContext(mock.patch("server.observability.kubernetes.kubernetes_api", return_value=(api, api))) + records = { + "model_meta:fft-1": meta(fine_tuning_type="full", status="completed"), + "model_meta:lora-live": meta(), + "model_meta:lora-done": meta(status="completed"), + } + for key, value in records.items(): + asyncio.run(self.store.set_value(f"open_rl:{key}", value)) + body = TestClient(api_server.app).get("/api/v1/inspect/inventory").json() + + at = datetime(2026, 10, 29, 12, tzinfo=UTC).timestamp() + self.assertEqual( + body["placements"], + [ + {"name": "fft-a", "role": "trainer", "node": "n1", "pod": "orw-fft-a", "gpu_uuids": ["GPU-fft-a"], "placed_at": at, "run_ids": ["fft-1"]}, + { + "name": "lora-a", + "role": "trainer", + "node": "n1", + "pod": "orw-lora-a", + "gpu_uuids": ["GPU-lora-a"], + "placed_at": at, + "run_ids": ["lora-live"], + }, + ], + ) + class JoinTest(unittest.TestCase): def test_lifecycle_follows_the_runs_workloads(self): diff --git a/tests/test_scheduler_worker_manager.py b/tests/test_scheduler_worker_manager.py index cb5618cc9..39729146b 100644 --- a/tests/test_scheduler_worker_manager.py +++ b/tests/test_scheduler_worker_manager.py @@ -1,10 +1,12 @@ import json import os import tempfile +import threading +import time import unittest from types import SimpleNamespace from typing import Any -from unittest.mock import patch +from unittest.mock import Mock, patch from server import api_server from server.estimator import footprint @@ -83,6 +85,16 @@ def call(*args, **kwargs): return call +class HistoryStore(InMemoryStateStore): + def __init__(self): + super().__init__() + self.pushed = threading.Event() + + def push_recent_sync(self, *args) -> None: + super().push_recent_sync(*args) + self.pushed.set() + + class SchedulerWorkerManagerTest(unittest.TestCase): def setUp(self) -> None: # The API server Deployment sets OPEN_RL_TMP_DIR to the shared volume every worker pod mounts. A temp dir stands in for it. @@ -404,6 +416,47 @@ def test_release_owner_finds_workloads_by_their_spec_not_their_labels(self) -> N self.assertEqual(self.manager.release_owner("old"), {"old"}) self.assertEqual(self.api.deleted, ["fft-old-trainer"]) + def place_lora_pair_and_other(self) -> str: + s = self.store_with("adapter", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora"}) + s.kv_store["open_rl:model_meta:other"] = json.dumps({"base_model": "Qwen/Qwen3-0.6B", "fine_tuning_type": "lora"}) + with patch("server.worker_manager.get_state_store", return_value=s): + self.manager.ensure("adapter", "trainer") + self.manager.ensure("adapter", "sampler") + self.manager.ensure("other", "trainer") + for name, body in self.api.existing.items(): + body["status"] = {"nodeName": "n1", "podName": f"orw-{name}", "gpuUUIDs": [f"GPU-{name}"], "placedAt": "2026-10-29T12:00:00Z"} + return owner_id("Qwen/Qwen2.5-0.5B") + + def test_release_records_the_final_placement_of_each_released_workload(self) -> None: + owner, store = self.place_lora_pair_and_other(), HistoryStore() + with patch("server.scheduler_worker_manager.get_state_store", return_value=store): + self.manager.release_owner(owner) + self.assertTrue(store.pushed.wait(5)) + + self.assertEqual(list(store.lists), ["open_rl:placements:Qwen/Qwen2.5-0.5B"]) + records = [json.loads(r) for r in store.lists["open_rl:placements:Qwen/Qwen2.5-0.5B"]] + self.assertEqual( + sorted((r["name"], r["role"], r["node"], r["gpu_uuids"][0]) for r in records), + [ + (f"lora-{owner}-0-sampler", "sampler", "n1", f"GPU-lora-{owner}-0-sampler"), + (f"lora-{owner}-0-trainer", "trainer", "n1", f"GPU-lora-{owner}-0-trainer"), + ], + ) + self.assertTrue(all(r["training_kind"] == "lora" and r["released_at"] > r["placed_at"] for r in records), records) + + def test_a_failing_or_hanging_history_store_does_not_hold_up_the_release(self) -> None: + unblock = threading.Event() + self.addCleanup(unblock.set) + for push in (RuntimeError("redis down"), lambda *args: unblock.wait(5)): + self.api.deleted.clear() + owner = self.place_lora_pair_and_other() + store = SimpleNamespace(push_recent_sync=Mock(side_effect=push)) + with patch("server.scheduler_worker_manager.get_state_store", return_value=store): + started = time.monotonic() + self.assertEqual(self.manager.release_owner(owner), {"Qwen/Qwen2.5-0.5B"}) + self.assertLess(time.monotonic() - started, 1) + self.assertEqual(len(self.api.deleted), 2) + def test_ensure_waits_for_a_terminating_workload_before_recreating_it(self) -> None: s = self.store_with("adapter", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora"}) with patch("server.worker_manager.get_state_store", return_value=s):