diff --git a/pkg/scheduler/failover_test.go b/pkg/scheduler/failover_test.go index b176f0d3c..211389075 100644 --- a/pkg/scheduler/failover_test.go +++ b/pkg/scheduler/failover_test.go @@ -322,6 +322,50 @@ func TestFailoverRankOverridesPoolPriority(t *testing.T) { assert.Equal(t, first.Id, worker.Id) } +func TestExplicitlyRequestedGPUIsNotDemotedByAnotherGPUFailoverChain(t *testing.T) { + scheduler := failoverSchedulerForTest(t, nil) + configured := failoverWorker("configured-4090", "beta9-4090", "RTX4090", 1) + configured.Priority = 2 + legacyOutsideChain := failoverWorker("legacy-failover-4090", "legacy-failover", "RTX4090", 1) + legacyOutsideChain.Priority = -1 + + request := &types.ContainerRequest{ + GpuRequest: []string{"RTX4090", "A10G"}, + GpuCount: 1, + Cpu: 1000, + Memory: 1000, + } + worker, err := scheduler.selectWorkerFromWorkers([]*types.Worker{legacyOutsideChain, configured}, request) + assert.Nil(t, err) + assert.Equal(t, configured.Id, worker.Id) +} + +func TestStorageBackedRequestPrefersManagedAgentCapacity(t *testing.T) { + scheduler := failoverSchedulerForTest(t, nil) + agent := failoverWorker("managed-agent", "agent-4090", "RTX4090", 1) + agent.ControlPlaneManaged = true + agent.PoolSelector = "agent-4090" + agent.Priority = 2 + legacy := failoverWorker("legacy-4090", "legacy-4090", "RTX4090", 1) + legacy.Priority = 3 + + request := &types.ContainerRequest{ + GpuRequest: []string{"RTX4090", "A10G"}, + GpuCount: 1, + Cpu: 1000, + Memory: 1000, + Workspace: testWorkspaceWithStorage(), + } + worker, err := scheduler.selectWorkerFromWorkers([]*types.Worker{legacy, agent}, request) + assert.Nil(t, err) + assert.Equal(t, agent.Id, worker.Id) + + request.Workspace = types.Workspace{} + worker, err = scheduler.selectWorkerFromWorkers([]*types.Worker{legacy, agent}, request) + assert.Nil(t, err) + assert.Equal(t, legacy.Id, worker.Id) +} + // TestFailoverDemandRecordedOnEstateExhaustion covers the only trigger for // on-demand hardware: every pool the request could provision into, primary and // failover alike, has recently refused for capacity. diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index bbf62f6c3..de6860147 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -1147,6 +1147,7 @@ type scoredWorker struct { worker *types.Worker score int32 failoverRank int32 + storageRank int32 } // Constants used for scoring workers @@ -1199,15 +1200,22 @@ func (s *Scheduler) selectWorkerFromWorkersByStatus(workers []*types.Worker, req scoredWorkers = append(scoredWorkers, scoredWorker{ worker: worker, score: score, - failoverRank: chain.rank(worker.PoolName), + failoverRank: failoverRankForWorker(chain, worker, request), + storageRank: s.storagePreferenceRank(worker, request), }) } - // Chain rank takes precedence over pool priority. + // Native requested GPUs take precedence over failover capacity. Within the + // same tier, storage-backed work prefers agent capacity so legacy workers + // remain available to workspaces that cannot run on agents. Explicit chain + // order still takes precedence over ordinary pool priority. sort.Slice(scoredWorkers, func(i, j int) bool { if scoredWorkers[i].failoverRank != scoredWorkers[j].failoverRank { return scoredWorkers[i].failoverRank < scoredWorkers[j].failoverRank } + if scoredWorkers[i].storageRank != scoredWorkers[j].storageRank { + return scoredWorkers[i].storageRank < scoredWorkers[j].storageRank + } if scoredWorkers[i].score != scoredWorkers[j].score { return scoredWorkers[i].score > scoredWorkers[j].score } @@ -1217,6 +1225,36 @@ func (s *Scheduler) selectWorkerFromWorkersByStatus(workers []*types.Worker, req return scoredWorkers[0].worker, nil } +func failoverRankForWorker(chain *failoverChain, worker *types.Worker, request *types.ContainerRequest) int32 { + if chain == nil || worker == nil || request == nil { + return 0 + } + + // A pool can be both native capacity for one explicitly requested GPU and + // a configured failover target for another. Treat it as native in that + // case; otherwise an unrelated pool outside the chain gets rank zero and + // incorrectly jumps ahead of it regardless of priority. + if slices.Contains(gpuRequestsForScheduling(request), worker.Gpu) { + return 0 + } + return chain.rank(worker.PoolName) +} + +func (s *Scheduler) storagePreferenceRank(worker *types.Worker, request *types.ContainerRequest) int32 { + if worker == nil || request == nil || !request.StorageAvailable() { + return 0 + } + if worker.ControlPlaneManaged { + return 0 + } + if s != nil && s.workerPoolManager != nil { + if pool, ok := s.workerPoolManager.GetPool(workerPoolSelector(worker)); ok && pool.Config.AgentHosted() { + return 0 + } + } + return 1 +} + func (s *Scheduler) filterAgentWorkersByStorage(workers []*types.Worker, request *types.ContainerRequest) []*types.Worker { if len(workers) == 0 || request == nil || request.StorageAvailable() { return workers