perf(scheduler): cap the seats a shared claim may hold

scheduler/controller/cmd/manager/main.go

@@ -7,6 +7,7 @@ package main
 
 import (
 	"flag"
+	"fmt"
 	"os"
 	"strconv"
 	"time"
@@ -47,6 +48,7 @@ func main() {
 		retryInterval    time.Duration
 		placementTimeout time.Duration
 		strategy         string
+		maxSeats         int
 		maxConcurrent    int
 	)
 
@@ -62,6 +64,8 @@ func main() {
 		"How long a worker may go unplaced before the request is declared unsatisfiable. 0 waits forever.")
 	flag.StringVar(&strategy, "placement-strategy", env("OPEN_RL_PLACEMENT_STRATEGY", string(placement.StrategyBinPack)),
 		"binpack shares eligible FFT claims first; spread requests a GPU first, then falls back to FFT sharing. LoRA never shares.")
+	flag.IntVar(&maxSeats, "max-seats-per-claim", envInt("OPEN_RL_MAX_SEATS_PER_CLAIM", 0),
+		"Most workers that may share one claim; 0 is no limit, 1 is exclusive placement. A worker that finds every eligible claim full cuts its own and waits Pending until a GPU is added or a seat frees, or fails at the placement timeout.")
 	flag.IntVar(&maxConcurrent, "max-concurrent-reconciles", envInt("OPEN_RL_MAX_CONCURRENT_RECONCILES", 4),
 		"How many workers place at once. Seat booking is CAS-arbitrated, so concurrency risks only transient over-cut claims, which the sharing fallback drains.")
 
@@ -80,6 +84,10 @@ func main() {
 		setupLog.Error(err, "invalid placement strategy")
 		os.Exit(1)
 	}
+	if maxSeats < 0 {
+		setupLog.Error(fmt.Errorf("max-seats-per-claim is %d", maxSeats), "the seat limit must be 0 (no limit) or positive")
+		os.Exit(1)
+	}
 
 	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
 		Scheme:                 scheme,
@@ -116,6 +124,7 @@ func main() {
 		RetryInterval:           retryInterval,
 		PlacementTimeout:        placementTimeout,
 		PlacementStrategy:       parsedStrategy,
+		MaxSeatsPerClaim:        maxSeats,
 		MaxConcurrentReconciles: maxConcurrent,
 	}
 	if err := reconciler.SetupWithManager(mgr); err != nil {
@@ -133,7 +142,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	setupLog.Info("placing workers", "namespace", namespace, "deviceClass", deviceClass, "deviceDriver", deviceDriver, "strategy", strategy)
+	setupLog.Info("placing workers", "namespace", namespace, "deviceClass", deviceClass, "deviceDriver", deviceDriver, "strategy", strategy, "maxSeatsPerClaim", maxSeats)
 	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
 		setupLog.Error(err, "manager exited")
 		os.Exit(1)

scheduler/controller/internal/controller/ledger.go

@@ -45,8 +45,9 @@ func newSeat(worker *openrlv1alpha1.Workload, request placement.Request) openrlv
 
 // ensureSeat books the seat on the claim's ClaimLedger via one CAS loop. A seat
 // held by the same incarnation is adopted and a predecessor's is replaced.
-// There is no seat ceiling: host memory is the limit on parked workers,
-// checked at selection and verified against a fresh fleet read after booking.
+// Host memory is the node-wide limit on parked workers, checked at selection
+// and verified against a fresh fleet read after booking; MaxSeatsPerClaim,
+// when set, is the per-claim limit and is re-checked here.
 //
 // createMissing is true only for a worker's own claim: the founder books its
 // ClaimLedger before the claim exists, and a re-book heals the ledger under a
@@ -97,6 +98,13 @@ func (r *WorkloadReconciler) ensureSeat(ctx context.Context, claimName string, s
 				return nil, nil, errBookingContended
 			}
 		}
+		// The seat limit is re-checked here for the same reason: two joiners
+		// selecting from one snapshot must not both land the last seat. It
+		// gates every write that adds a seat, a re-book whose seat is gone
+		// included; adopting or replacing a seat adds none.
+		if existing == nil && r.MaxSeatsPerClaim > 0 && len(claimLedger.Spec.Seats) >= r.MaxSeatsPerClaim {
+			return nil, nil, errBookingContended
+		}
 		if existing != nil {
 			*existing = seat
 		} else {

scheduler/controller/internal/controller/ledger_test.go

@@ -13,6 +13,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/event"
 
 	openrlv1alpha1 "github.com/gke-labs/open-rl/scheduler/controller/api/v1alpha1"
 	"github.com/gke-labs/open-rl/scheduler/controller/internal/placement"
@@ -123,6 +124,95 @@ func TestRebookingAdoptsTheRecordedSeat(t *testing.T) {
 	}
 }
 
+// The seat limit is enforced inside the ledger's CAS loop, so a joiner
+// selecting from a snapshot that predates the last booking is refused.
+func TestBookingRefusesASeatPastTheLimit(t *testing.T) {
+	r := newReconciler(t, append(enabledNode(), trainerWorker("w-a", "model-a"), trainerWorker("w-b", "model-b"), trainerWorker("w-c", "model-c"))...)
+	r.PlacementStrategy = placement.StrategyBinPack
+	r.MaxSeatsPerClaim = 2
+	settle(t, r, "w-a")
+	claim := claimOf(t, r, "w-a")
+	allocateClaim(t, r, claim)
+	settle(t, r, "w-b")
+	if got := claimOf(t, r, "w-b"); got != claim {
+		t.Fatalf("w-b on %q, want %q", got, claim)
+	}
+	c := getWorker(t, r, "w-c")
+	if _, _, err := r.ensureSeat(context.Background(), claim, newSeat(c, requestFrom(c)), false); err != errBookingContended {
+		t.Fatalf("booking a third seat under a limit of two returned %v, want errBookingContended", err)
+	}
+	// A re-book that would add a seat is refused too: the limit holds on
+	// every write path, not only on joins.
+	if _, _, err := r.ensureSeat(context.Background(), claim, newSeat(c, requestFrom(c)), true); err != errBookingContended {
+		t.Fatalf("re-booking an unseated worker onto a full ledger returned %v, want errBookingContended", err)
+	}
+	// The founder's own re-book adopts its seat and is never refused.
+	a := getWorker(t, r, "w-a")
+	if _, _, err := r.ensureSeat(context.Background(), claim, newSeat(a, requestFrom(a)), true); err != nil {
+		t.Fatalf("re-booking the founder: %v", err)
+	}
+}
+
+// A booking refused on the selected ledger moves on to the next eligible one
+// instead of cutting a dedicated claim another ledger could have spared.
+func TestJoinRetriesTheNextLedgerWhenTheFirstIsFull(t *testing.T) {
+	ctx := context.Background()
+	r := newReconciler(t, append(enabledNode(),
+		trainerWorker("w-a", "model-a"), trainerWorker("w-b", "model-b"), trainerWorker("w-c", "model-c"),
+		trainerWorker("w-d", "model-d"), trainerWorker("w-e", "model-e"), trainerWorker("w-f", "model-f"))...)
+	r.PlacementStrategy = placement.StrategyBinPack
+	// Two allocated ledgers: a limit of one makes w-c cut a claim of its own.
+	r.MaxSeatsPerClaim = 1
+	settle(t, r, "w-a")
+	first := claimOf(t, r, "w-a")
+	allocateClaim(t, r, first)
+	settle(t, r, "w-c")
+	second := claimOf(t, r, "w-c")
+	if second == first {
+		t.Fatal("w-c joined under a limit of one, want a claim of its own")
+	}
+	allocateClaim(t, r, second)
+
+	r.MaxSeatsPerClaim = 3
+	book := func(name, claim string) {
+		t.Helper()
+		w := getWorker(t, r, name)
+		if _, _, err := r.ensureSeat(ctx, claim, newSeat(w, requestFrom(w)), false); err != nil {
+			t.Fatalf("booking %s on %s: %v", name, claim, err)
+		}
+	}
+	// Snapshot: the first ledger holds one seat and the second two, so
+	// selection prefers the first.
+	book("w-b", second)
+	fleet, err := r.readFleet(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	f := getWorker(t, r, "w-f")
+	request := requestFrom(f)
+	request.MaxSeats = 3
+	if got := placement.SelectClaim(request, fleet); got == nil || got.Name != first {
+		t.Fatalf("snapshot selects %v, want %q", got, first)
+	}
+	// Behind the snapshot, the first ledger fills to the limit.
+	book("w-d", first)
+	book("w-e", first)
+
+	joined, _, err := r.joinExistingClaim(ctx, f, request, fleet, "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if joined == nil || joined.Name != second {
+		t.Fatalf("joined %v, want the ledger with room %q", joined, second)
+	}
+	if got := len(getLedger(t, r, ledgerNameFor(second)).Spec.Seats); got != 3 {
+		t.Fatalf("second ledger holds %d seats, want 3", got)
+	}
+	if got := len(getLedger(t, r, ledgerNameFor(first)).Spec.Seats); got != 3 {
+		t.Fatalf("first ledger holds %d seats after the refused join, want 3", got)
+	}
+}
+
 func TestSharingRechecksTheLedgerAfterSelection(t *testing.T) {
 	resident, incoming := trainerWorker("resident", "model-a"), trainerWorker("incoming", "model-b")
 	r := newReconciler(t, append(enabledNode(), resident, incoming)...)
@@ -578,3 +668,28 @@ func TestWaitsForATerminatingPredecessorClaim(t *testing.T) {
 		t.Fatalf("placed on %q, want the fresh %q", got, claimNameFor(w))
 	}
 }
+
+// Only a ledger event that frees a seat wakes the fleet: a booking cannot
+// unblock anyone, a release or a retirement can.
+func TestSeatFreedWakesOnlyOnReleasedSeats(t *testing.T) {
+	seat := func(n int) *openrlv1alpha1.ClaimLedger {
+		l := &openrlv1alpha1.ClaimLedger{}
+		for i := 0; i < n; i++ {
+			l.Spec.Seats = append(l.Spec.Seats, openrlv1alpha1.Seat{Workload: fmt.Sprintf("w-%d", i)})
+		}
+		return l
+	}
+	p := seatFreed()
+	if p.Create(event.CreateEvent{Object: seat(1)}) {
+		t.Error("a new ledger woke the fleet")
+	}
+	if p.Update(event.UpdateEvent{ObjectOld: seat(1), ObjectNew: seat(2)}) {
+		t.Error("a booking woke the fleet")
+	}
+	if !p.Update(event.UpdateEvent{ObjectOld: seat(2), ObjectNew: seat(1)}) {
+		t.Error("a released seat did not wake the fleet")
+	}
+	if !p.Delete(event.DeleteEvent{Object: seat(1)}) {
+		t.Error("a retired ledger did not wake the fleet")
+	}
+}

scheduler/controller/internal/controller/workload_controller.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"strings"
 	"time"
 
@@ -54,6 +55,14 @@ type WorkloadReconciler struct {
 	// PlacementStrategy orders the two placement moves; the zero value reads
 	// as binpack.
 	PlacementStrategy placement.Strategy
+	// MaxSeatsPerClaim caps how many workers share one claim; zero means no
+	// limit, one is exclusive placement. A worker that finds every eligible
+	// ledger at the limit cuts its own claim and, if no device is free, waits
+	// on capacity like any other: Pending until one is added or a seat
+	// frees, Failed if PlacementTimeout runs out first. It bounds a
+	// time-sliced step at the limit times the single-worker time, and it
+	// turns a full fleet into standing claims an autoscaler can see.
+	MaxSeatsPerClaim int
 	// MaxConcurrentReconciles is how many workers place at once.
 	MaxConcurrentReconciles int
 
@@ -109,7 +118,9 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		return ctrl.Result{}, fmt.Errorf("cannot read the fleet, placing nothing this pass: %w", err)
 	}
 
-	return r.place(ctx, &worker, requestFrom(&worker), fleet)
+	request := requestFrom(&worker)
+	request.MaxSeats = r.MaxSeatsPerClaim
+	return r.place(ctx, &worker, request, fleet)
 }
 
 type placementScope struct {
@@ -518,23 +529,35 @@ func (r *WorkloadReconciler) teardown(ctx context.Context, worker *openrlv1alpha
 
 // joinExistingClaim books a seat on the claim SelectClaim proposes, re-reads
 // the fleet to verify concurrent bookings did not overcommit the node, and
-// keeps the caller's fleet snapshot current.
+// keeps the caller's fleet snapshot current. A ledger that refuses the
+// booking -- full, retiring, or holding an exclusive resident -- is dropped
+// from the candidates and selection runs again, so losing one race does not
+// cut a dedicated claim while another ledger has room.
 //
 // Because both racers write their seat before either re-reads, at least the
 // second reader sees the union across all claims on the node and yields.
 // Simultaneous racers may both observe the overcommit and both release their
 // seats; that double yield is benign and intentional.
 func (r *WorkloadReconciler) joinExistingClaim(ctx context.Context, worker *openrlv1alpha1.Workload, request placement.Request, fleet *placement.Fleet, current string) (*placement.Claim, *openrlv1alpha1.Seat, error) {
-	target := placement.SelectClaim(request, fleet)
-	if target == nil || target.Name == current {
-		return nil, nil, nil
-	}
-	_, seat, err := r.ensureSeat(ctx, target.Name, newSeat(worker, request), false)
-	if err == errBookingContended {
-		return nil, nil, nil
-	}
-	if err != nil {
-		return nil, nil, err
+	candidates := *fleet
+	candidates.Claims = maps.Clone(fleet.Claims)
+	var target *placement.Claim
+	var seat *openrlv1alpha1.Seat
+	for {
+		target = placement.SelectClaim(request, &candidates)
+		if target == nil || target.Name == current {
+			return nil, nil, nil
+		}
+		var err error
+		_, seat, err = r.ensureSeat(ctx, target.Name, newSeat(worker, request), false)
+		if err == errBookingContended {
+			delete(candidates.Claims, target.Name)
+			continue
+		}
+		if err != nil {
+			return nil, nil, err
+		}
+		break
 	}
 	fresh, err := r.readFleet(ctx)
 	if err != nil {
@@ -789,6 +812,20 @@ func managedClaims() predicate.Predicate {
 	})
 }
 
+// seatFreed passes ledger events that free a seat: a seat dropped from the
+// list, or the ledger retired. A seat added never unblocks a waiting worker,
+// so bookings and new ledgers do not wake the fleet.
+func seatFreed() predicate.Predicate {
+	return predicate.Funcs{
+		CreateFunc: func(event.CreateEvent) bool { return false },
+		UpdateFunc: func(e event.UpdateEvent) bool {
+			before, okBefore := e.ObjectOld.(*openrlv1alpha1.ClaimLedger)
+			after, okAfter := e.ObjectNew.(*openrlv1alpha1.ClaimLedger)
+			return okBefore && okAfter && len(after.Spec.Seats) < len(before.Spec.Seats)
+		},
+	}
+}
+
 // SetupWithManager registers the reconciler.
 func (r *WorkloadReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	r.reader = mgr.GetAPIReader()
@@ -826,6 +863,7 @@ func (r *WorkloadReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		For(&openrlv1alpha1.Workload{}).
 		Owns(&corev1.Pod{}).
 		Watches(&resourcev1.ResourceClaim{}, wakeAll, builder.WithPredicates(managedClaims())).
+		Watches(&openrlv1alpha1.ClaimLedger{}, wakeAll, builder.WithPredicates(managedClaims(), seatFreed())).
 		Watches(&corev1.Node{}, wakeAll, builder.WithPredicates(nodeCapacityChanged)).
 		WithOptions(controller.Options{MaxConcurrentReconciles: max(1, r.MaxConcurrentReconciles)}).
 		Complete(r)

scheduler/controller/internal/controller/workload_controller_test.go

@@ -687,6 +687,69 @@ func TestReconcileBinPacksOntoExistingClaims(t *testing.T) {
 	}
 }
 
+// With a seat limit, binpack fills a claim to the limit, moves to a free
+// device, and once every device is at the limit the next worker waits on a
+// standing claim of its own until a seat frees. Two devices, limit two: a and
+// b share the first, c and d the second, e stands pending; when b leaves, e
+// takes its seat.
+func TestReconcileStopsPackingAtTheSeatLimit(t *testing.T) {
+	r := newReconciler(t, append(enabledNode(),
+		trainerWorker("w-a", "model-a"), trainerWorker("w-b", "model-b"), trainerWorker("w-c", "model-c"),
+		trainerWorker("w-d", "model-d"), trainerWorker("w-e", "model-e"))...)
+	r.PlacementStrategy = placement.StrategyBinPack
+	r.MaxSeatsPerClaim = 2
+
+	settle(t, r, "w-a")
+	allocateClaim(t, r, claimOf(t, r, "w-a"))
+	settle(t, r, "w-b")
+	first := claimOf(t, r, "w-a")
+	if got := claimOf(t, r, "w-b"); got != first {
+		t.Fatalf("w-b landed on %q, want w-a's claim %q (under the limit)", got, first)
+	}
+
+	settle(t, r, "w-c")
+	second := claimOf(t, r, "w-c")
+	if second == first {
+		t.Fatal("w-c joined a claim at the limit, want a claim of its own")
+	}
+	allocateClaim(t, r, second)
+	settle(t, r, "w-d")
+	if got := claimOf(t, r, "w-d"); got != second {
+		t.Fatalf("w-d landed on %q, want w-c's claim %q (under the limit)", got, second)
+	}
+
+	// Both devices are claimed and both ledgers are full: w-e cuts a claim
+	// that cannot allocate, kube-scheduler declines its pod, and the sharing
+	// fallback finds no ledger under the limit. It stays Pending on its own
+	// claim with the limit in the reason.
+	settle(t, r, "w-e")
+	own := claimOf(t, r, "w-e")
+	if own == first || own == second {
+		t.Fatalf("w-e landed on %q, want a standing claim of its own", own)
+	}
+	fallBackToSharing(t, r, "w-e")
+	if got := claimOf(t, r, "w-e"); got != own {
+		t.Fatalf("w-e moved to %q past the limit, want it waiting on %q", got, own)
+	}
+	w := getWorker(t, r, "w-e")
+	if w.Status.Phase != openrlv1alpha1.PhasePending || !strings.Contains(w.Status.Reason, "at most 2 workers") {
+		t.Fatalf("w-e is %s %q, want Pending with the limit named", w.Status.Phase, w.Status.Reason)
+	}
+
+	// A seat frees: w-b leaves the first claim and the next pass seats w-e.
+	if err := r.releaseSeatAndReclaim(context.Background(), first, "w-b"); err != nil {
+		t.Fatal(err)
+	}
+	fallBackToSharing(t, r, "w-e")
+	if got := claimOf(t, r, "w-e"); got != first {
+		t.Fatalf("w-e is on %q after a seat freed, want %q", got, first)
+	}
+	ledger := getLedger(t, r, ledgerNameFor(first))
+	if len(ledger.Spec.Seats) != 2 {
+		t.Fatalf("first ledger holds %d seats, want 2: %+v", len(ledger.Spec.Seats), ledger.Spec.Seats)
+	}
+}
+
 // A trainer and a sampler on one accelerator, once the devices are contended.
 // Role selects which node pools may host a worker and nothing else -- in
 // particular it does not partition claims -- so on a node labelled for both,