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/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 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):