|
| 1 | +"""Build setup-aware consecutive batches and expose their timing plan.""" |
| 2 | +from __future__ import annotations |
| 3 | + |
| 4 | +from dataclasses import dataclass |
| 5 | + |
| 6 | + |
| 7 | +# Add a suite only after a full split run proves that parallel batches break |
| 8 | +# its order/state assumptions. Each entry needs a short reason. |
| 9 | +SERIAL_SUITES: dict[str, str] = { |
| 10 | + "daemon_application": "The 10-second default-capacity admission fixture timed out in the split tail batch under full-run load.", |
| 11 | +} |
| 12 | + |
| 13 | + |
| 14 | +@dataclass(frozen=True) |
| 15 | +class Batch: |
| 16 | + suite: str |
| 17 | + tests: tuple[str, ...] |
| 18 | + setup_ms: float |
| 19 | + work_ms: float |
| 20 | + |
| 21 | + @property |
| 22 | + def estimate_ms(self) -> float: |
| 23 | + return self.setup_ms + self.work_ms |
| 24 | + |
| 25 | + |
| 26 | +def plan_batches( |
| 27 | + names: list[str], timing, target_seconds: float, multiplier: float |
| 28 | +) -> list[Batch]: |
| 29 | + """Keep registration order within each suite and account for setup cost.""" |
| 30 | + by_suite: dict[str, list[str]] = {} |
| 31 | + for name in names: |
| 32 | + suite = name.split(":", 1)[0] |
| 33 | + by_suite.setdefault(suite, []).append(name) |
| 34 | + |
| 35 | + batches = [] |
| 36 | + target_ms = target_seconds * multiplier * 1000 |
| 37 | + for suite, tests in by_suite.items(): |
| 38 | + setup_ms = timing.setup_ms(suite) |
| 39 | + if suite in SERIAL_SUITES: |
| 40 | + batches.append( |
| 41 | + Batch(suite, tuple(tests), setup_ms, sum(timing.estimate_ms(n) for n in tests)) |
| 42 | + ) |
| 43 | + continue |
| 44 | + |
| 45 | + # Keep setup at most one quarter of the estimate, even when setup cost |
| 46 | + # pushes the minimum batch above the requested target. |
| 47 | + minimum_work_ms = 3 * setup_ms |
| 48 | + target_work_ms = max(target_ms - setup_ms, minimum_work_ms) |
| 49 | + current: list[str] = [] |
| 50 | + work_ms = 0.0 |
| 51 | + for name in tests: |
| 52 | + test_ms = timing.estimate_ms(name) |
| 53 | + if current and work_ms >= target_work_ms: |
| 54 | + batches.append(Batch(suite, tuple(current), setup_ms, work_ms)) |
| 55 | + current = [] |
| 56 | + work_ms = 0.0 |
| 57 | + current.append(name) |
| 58 | + work_ms += test_ms |
| 59 | + if current: |
| 60 | + if batches and batches[-1].suite == suite and work_ms < minimum_work_ms: |
| 61 | + previous = batches.pop() |
| 62 | + current = list(previous.tests) + current |
| 63 | + work_ms += previous.work_ms |
| 64 | + batches.append(Batch(suite, tuple(current), setup_ms, work_ms)) |
| 65 | + return batches |
| 66 | + |
| 67 | + |
| 68 | +def ordered_batches(batches: list[Batch]) -> list[tuple[int, Batch]]: |
| 69 | + """Return stable longest-processing-time-first queue priorities.""" |
| 70 | + return sorted(enumerate(batches), key=lambda item: (-item[1].estimate_ms, item[0])) |
| 71 | + |
| 72 | + |
| 73 | +def print_plan(names: list[str], batches: list[Batch], jobs: int, mode: str) -> None: |
| 74 | + suites = {} |
| 75 | + for batch in batches: |
| 76 | + suites[batch.suite] = batch.setup_ms |
| 77 | + print(f"plan: {len(names)} tests, {len(batches)} batches, {jobs} workers, {mode}") |
| 78 | + for suite in sorted(suites): |
| 79 | + print(f"setup {suite}: {suites[suite]:.0f}ms") |
| 80 | + for index, batch in ordered_batches(batches): |
| 81 | + print( |
| 82 | + f"batch {index + 1}: {batch.suite}, {len(batch.tests)} tests, " |
| 83 | + f"setup {batch.setup_ms:.0f}ms, work {batch.work_ms:.0f}ms, " |
| 84 | + f"estimate {batch.estimate_ms:.0f}ms" |
| 85 | + ) |
0 commit comments