account for in-flight capacity in launch scheduling

agent/nix_builder_autoscaler/scheduler.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import logging
 import time
+from collections import Counter
 from datetime import datetime
 from typing import TYPE_CHECKING
 
@@ -146,6 +147,25 @@ def _launch_for_unmet_demand(
     if not pending:
         return
 
+    demand_by_system = Counter(str(resv["system"]) for resv in pending)
+    in_flight_slots = (
+        db.list_slots(SlotState.LAUNCHING)
+        + db.list_slots(SlotState.BOOTING)
+        + db.list_slots(SlotState.BINDING)
+    )
+    in_flight_by_system = Counter(str(slot["system"]) for slot in in_flight_slots)
+
+    leases_per_slot = max(1, config.capacity.max_leases_per_slot)
+    for system, in_flight_count in in_flight_by_system.items():
+        in_flight_capacity = in_flight_count * leases_per_slot
+        if in_flight_capacity <= 0:
+            continue
+        current_demand = demand_by_system.get(system, 0)
+        demand_by_system[system] = max(0, current_demand - in_flight_capacity)
+
+    if sum(demand_by_system.values()) <= 0:
+        return
+
     active = _count_active_slots(db)
     if active >= config.capacity.max_slots:
         return
@@ -154,12 +174,20 @@ def _launch_for_unmet_demand(
     if not empty_slots:
         return
 
-    for launched, slot in enumerate(empty_slots):
-        if launched >= len(pending):
-            break
+    launched = 0
+    for slot in empty_slots:
         if active + launched >= config.capacity.max_slots:
             break
+
+        system = str(slot["system"])
+        if demand_by_system.get(system, 0) <= 0:
+            continue
+
         _launch_slot(db, runtime, config, metrics, slot)
+        launched += 1
+        demand_by_system[system] = max(0, demand_by_system[system] - leases_per_slot)
+        if sum(demand_by_system.values()) <= 0:
+            break
 
 
 def _ensure_min_and_warm(