294 lines
12 KiB
Python
294 lines
12 KiB
Python
from __future__ import annotations
|
|
import os
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import date
|
|
from test_window import get_shift_sequence_with_capacity, next_window_start_date
|
|
from test_bundle import Test, TestBundle, build_test_bundles, bundle_pair_lookup
|
|
|
|
DUT = (os.getenv("DUT") or "DUT").strip()
|
|
REF = (os.getenv("REF") or "REF").strip()
|
|
|
|
@dataclass(frozen=True)
|
|
class ScheduleEntry:
|
|
test_id: str
|
|
device: str
|
|
scheduled_date: str
|
|
shift_index: int
|
|
sequence_in_shift: int
|
|
|
|
|
|
@dataclass
|
|
class ScheduleWindow:
|
|
index: int
|
|
shifts: list[tuple[date, int]]
|
|
capacity_minutes: int
|
|
remaining_minutes: int
|
|
assigned_device: str | None = None
|
|
assigned_config: tuple[str, ...] | None = None
|
|
|
|
|
|
class Scheduler:
|
|
def __init__(
|
|
self,
|
|
tests: list[Test],
|
|
top_priority_tests: set[tuple[str, str]],
|
|
start_date: str | None = None,
|
|
holiday_dates: set[str] = set(),
|
|
daytime_testing_today: bool = False,
|
|
priority_weight: int = 100,
|
|
):
|
|
self.top_priority_tests = top_priority_tests
|
|
self.tests = tests
|
|
self.active_dut: dict[str, Test] = {test.test_id: test for test in tests if test.device == DUT}
|
|
self.active_ref: dict[str, Test] = {test.test_id: test for test in tests if test.device == REF}
|
|
self.start_date = date.fromisoformat(start_date) if start_date else date.today()
|
|
self.holiday_dates = holiday_dates
|
|
self.daytime_testing_today = daytime_testing_today
|
|
self.priority_weight = priority_weight
|
|
self.pending_dut_mirror: list[TestBundle] = []
|
|
self.pending_ref_mirror: list[TestBundle] = []
|
|
self.schedule: list[ScheduleEntry] = []
|
|
self.scheduled_bundle_keys: set[tuple[int, str]] = set()
|
|
self.scheduled_test_ids: set[str] = set()
|
|
|
|
def compile_schedule(self) -> str | None:
|
|
bundles = build_test_bundles(self.active_dut, self.active_ref, self.top_priority_tests)
|
|
if not bundles:
|
|
return f"Error: No test bundles could be created from the provided tests."
|
|
|
|
active_dut_bundles = [b for b in bundles if b.device == DUT]
|
|
active_ref_bundles = [b for b in bundles if b.device == REF]
|
|
|
|
# Group by individual TC values (from config tuples)
|
|
all_tcs = set()
|
|
for bundle in bundles:
|
|
all_tcs.update(bundle.config if bundle.config else [None])
|
|
|
|
# Create dict: TC -> bundles supporting that TC
|
|
bundles_by_tc: dict = {}
|
|
# Sort with None last
|
|
sorted_tcs = sorted([tc for tc in all_tcs if tc is not None]) + ([None] if None in all_tcs else [])
|
|
for tc in sorted_tcs:
|
|
if tc is None:
|
|
bundles_by_tc[tc] = [b for b in bundles if not b.config]
|
|
else:
|
|
bundles_by_tc[tc] = [b for b in bundles if tc in b.config]
|
|
# Sort by priority, then index
|
|
bundles_by_tc[tc].sort(key=lambda b: (b.priority, b.index))
|
|
|
|
window_index = 0
|
|
cursor_date = self.start_date
|
|
|
|
# Sort TC keys by bundle count (largest first); tie-break by TC name.
|
|
sorted_tc_items = sorted(
|
|
[(tc, bundles_by_tc[tc]) for tc in bundles_by_tc if tc is not None],
|
|
key=lambda x: (-len(x[1]), x[0]),
|
|
)
|
|
if None in bundles_by_tc:
|
|
sorted_tc_items.append((None, bundles_by_tc[None]))
|
|
|
|
|
|
for tc, tc_bundles in sorted_tc_items:
|
|
window_device = DUT
|
|
dut_unscheduled = [b for b in tc_bundles if b.device == DUT]
|
|
ref_unscheduled = [b for b in tc_bundles if b.device == REF]
|
|
|
|
while dut_unscheduled or ref_unscheduled:
|
|
active_unscheduled = dut_unscheduled if window_device == DUT else ref_unscheduled
|
|
if len(active_unscheduled) == 0:
|
|
# If no unscheduled bundles for the current device, switch to the other device
|
|
window_device = REF if window_device == DUT else DUT
|
|
active_unscheduled = dut_unscheduled if window_device == DUT else ref_unscheduled
|
|
|
|
shifts, capacity = get_shift_sequence_with_capacity(cursor_date, self.holiday_dates, self.daytime_testing_today)
|
|
|
|
# Evaluate both devices without mutating queue state, then commit once.
|
|
mirrored_bundles, knapsack_bundles, remaining_time = self._select_bundles(
|
|
window_device,
|
|
capacity,
|
|
active_unscheduled,
|
|
tc,
|
|
mutate=True,
|
|
)
|
|
selected_bundles = mirrored_bundles + knapsack_bundles
|
|
if len(selected_bundles) == 0:
|
|
window_device = REF if window_device == DUT else DUT
|
|
mirrored_bundles, knapsack_bundles, remaining_time = self._select_bundles(
|
|
window_device,
|
|
capacity,
|
|
active_unscheduled,
|
|
tc,
|
|
mutate=True,
|
|
)
|
|
selected_bundles = mirrored_bundles + knapsack_bundles
|
|
if len(selected_bundles) == 0:
|
|
break
|
|
|
|
# Create a mirror of the other device for the next window
|
|
bundle_pairs = bundle_pair_lookup(knapsack_bundles, dut_unscheduled, ref_unscheduled)
|
|
|
|
if window_device == DUT:
|
|
pending_target = self.pending_ref_mirror
|
|
else:
|
|
pending_target = self.pending_dut_mirror
|
|
|
|
existing_pending = {(b.index, b.device) for b in pending_target}
|
|
for pair_bundle in bundle_pairs:
|
|
pair_key = (pair_bundle.index, pair_bundle.device)
|
|
if pair_key in self.scheduled_bundle_keys or pair_key in existing_pending:
|
|
continue
|
|
pending_target.append(pair_bundle)
|
|
existing_pending.add(pair_key)
|
|
|
|
|
|
window = ScheduleWindow(
|
|
index=window_index,
|
|
shifts=shifts,
|
|
capacity_minutes=capacity,
|
|
remaining_minutes=remaining_time,
|
|
assigned_config=tc,
|
|
assigned_device=window_device,
|
|
)
|
|
|
|
self._place_bundles_in_window(window, selected_bundles)
|
|
|
|
# Mark bundles as scheduled and remove from ALL TC buckets
|
|
selected_keys = {(b.index, b.device) for b in selected_bundles}
|
|
|
|
# Remove scheduled bundles from ALL TC buckets globally
|
|
for all_tc in bundles_by_tc:
|
|
bundles_by_tc[all_tc] = [b for b in bundles_by_tc[all_tc] if (b.index, b.device) not in selected_keys]
|
|
|
|
# Rebuild current TC unscheduled lists
|
|
dut_unscheduled = [b for b in bundles_by_tc[tc] if b.device == DUT]
|
|
ref_unscheduled = [b for b in bundles_by_tc[tc] if b.device == REF]
|
|
|
|
window_index += 1
|
|
window_device = REF if window_device == DUT else DUT
|
|
cursor_date = next_window_start_date(shifts)
|
|
|
|
return None
|
|
|
|
def _select_bundles(
|
|
self,
|
|
device: str,
|
|
capacity: int,
|
|
unscheduled: list[TestBundle],
|
|
tc,
|
|
mutate: bool,
|
|
) -> tuple[list[TestBundle], list[TestBundle], int]:
|
|
# Consume pending mirrored bundles only while there is room.
|
|
pending = self.pending_dut_mirror if device == DUT else self.pending_ref_mirror
|
|
mirrored_bundles: list[TestBundle] = []
|
|
still_pending: list[TestBundle] = [] # Bundles that couldn't fit in the remaining capacity
|
|
remaining_capacity = capacity
|
|
|
|
for bundle in pending:
|
|
bundle_key = (bundle.index, bundle.device)
|
|
if bundle_key in self.scheduled_bundle_keys:
|
|
continue
|
|
if bundle.device != device or (tc is not None and tc not in bundle.config) or (tc is None and bundle.config):
|
|
still_pending.append(bundle)
|
|
continue
|
|
if bundle.total_minutes <= remaining_capacity:
|
|
mirrored_bundles.append(bundle)
|
|
remaining_capacity -= bundle.total_minutes
|
|
else:
|
|
still_pending.append(bundle)
|
|
|
|
if mutate:
|
|
pending[:] = still_pending
|
|
|
|
# Run knapsack selection for remaining capacity
|
|
candidates = [
|
|
b
|
|
for b in unscheduled
|
|
if b not in mirrored_bundles
|
|
and b.device == device
|
|
and (b.index, b.device) not in self.scheduled_bundle_keys
|
|
]
|
|
knapsack_bundles = self._knapsack_select(remaining_capacity, candidates)
|
|
remaining_capacity -= sum(bundle.total_minutes for bundle in knapsack_bundles)
|
|
|
|
print(f"Selected {len(mirrored_bundles)} mirrored bundles and {len(knapsack_bundles)} knapsack bundles for device {device} with remaining capacity {remaining_capacity} minutes.")
|
|
return mirrored_bundles, knapsack_bundles, remaining_capacity
|
|
|
|
def get_schedule(self) -> list[ScheduleEntry]:
|
|
return self.schedule
|
|
|
|
def _knapsack_select(
|
|
self,
|
|
capacity: int,
|
|
candidates: list[TestBundle],
|
|
) -> list[TestBundle]:
|
|
if capacity <= 0 or not candidates:
|
|
return []
|
|
|
|
# Weights are total minutes of each bundle
|
|
weights = [bundle.total_minutes for bundle in candidates]
|
|
# Values are based on priority, lower priority number means higher value
|
|
values = [self.priority_weight - bundle.priority for bundle in candidates]
|
|
|
|
# Implement dynamic programming knapsack algorithm to select bundles
|
|
n = len(candidates)
|
|
dp = [[0] * (capacity + 1) for _ in range(n + 1)]
|
|
for i in range(1, n + 1):
|
|
for w in range(capacity + 1):
|
|
if weights[i - 1] <= w:
|
|
dp[i][w] = max(dp[i - 1][w], dp[i - 1][w - weights[i - 1]] + values[i - 1])
|
|
else:
|
|
dp[i][w] = dp[i - 1][w]
|
|
|
|
# Traceback to find selected bundles
|
|
w = capacity
|
|
selected_indices = []
|
|
for i in range(n, 0, -1):
|
|
if w <= 0:
|
|
break
|
|
if dp[i][w] != dp[i - 1][w]:
|
|
selected_indices.append(i - 1)
|
|
w -= weights[i - 1]
|
|
|
|
# Return the selected bundles in the order they were added
|
|
return [candidates[i] for i in reversed(selected_indices)]
|
|
|
|
def _place_bundles_in_window(
|
|
self,
|
|
window: ScheduleWindow,
|
|
bundles: list[TestBundle],
|
|
) -> None:
|
|
for bundle in bundles:
|
|
bundle_key = (bundle.index, bundle.device)
|
|
if bundle_key in self.scheduled_bundle_keys:
|
|
continue
|
|
self.scheduled_bundle_keys.add(bundle_key)
|
|
|
|
for idx, test in enumerate(bundle.tests):
|
|
schedule_key = f"{test}:{bundle.device}"
|
|
if schedule_key in self.scheduled_test_ids:
|
|
continue
|
|
self.scheduled_test_ids.add(schedule_key)
|
|
|
|
self.schedule.append(
|
|
ScheduleEntry(
|
|
test_id=test,
|
|
device=bundle.device,
|
|
scheduled_date=str(window.shifts[0][0]),
|
|
shift_index=window.shifts[0][1],
|
|
sequence_in_shift=idx + 1,
|
|
)
|
|
)
|
|
window.remaining_minutes -= self._get_test_minutes(test, bundle.device)
|
|
|
|
|
|
def _get_test_minutes(self, test_id: str, device: str) -> int:
|
|
if device == DUT and test_id in self.active_dut:
|
|
return self.active_dut[test_id].estimated_minutes
|
|
if device == REF and test_id in self.active_ref:
|
|
return self.active_ref[test_id].estimated_minutes
|
|
return 0
|
|
|
|
|
|
|