Improve failover worker selection

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.

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