ci: dynamic partition + LPT from live sglang-ci-stats model (#25263)

python/sglang/test/ci/ci_register.py

@@ -287,17 +287,30 @@ def ut_parse_one_file(filename: str) -> Tuple[List[CIRegistry], bool]:
     return visitor.registries, visitor.has_main_entry
 
 
-def auto_partition(files: List[CIRegistry], rank: int, size: int) -> List[CIRegistry]:
+def auto_partition(
+    files: List[CIRegistry],
+    rank: int,
+    size: int,
+    live_est: Optional[dict] = None,
+) -> List[CIRegistry]:
     """Partition files into `size` sublists with approximately equal sums of
     estimated times using a greedy algorithm (LPT heuristic), and return the
     partition for the specified rank.
+
+    `live_est`: optional `filename -> est seconds` overrides; missing
+    files fall back to in-source `est_time`.
     """
     if not files or size <= 0:
         return []
 
+    def est_of(f: CIRegistry) -> float:
+        if live_est is not None and f.filename in live_est:
+            return live_est[f.filename]
+        return f.est_time
+
     # Sort by estimated_time descending; filename as tie-breaker for
     # deterministic partitioning regardless of glob ordering.
-    sorted_files = sorted(files, key=lambda f: (-f.est_time, f.filename))
+    sorted_files = sorted(files, key=lambda f: (-est_of(f), f.filename))
 
     partitions: List[List[CIRegistry]] = [[] for _ in range(size)]
     partition_sums = [0.0] * size
@@ -306,7 +319,7 @@ def auto_partition(files: List[CIRegistry], rank: int, size: int) -> List[CIRegi
     for file in sorted_files:
         min_sum_idx = min(range(size), key=partition_sums.__getitem__)
         partitions[min_sum_idx].append(file)
-        partition_sums[min_sum_idx] += file.est_time
+        partition_sums[min_sum_idx] += est_of(file)
 
     if rank < size:
         return partitions[rank]