Coverage for core / src / sensorkit / auto / scheduler.py: 76%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4from collections.abc import Iterable, Sequence 

5from dataclasses import dataclass 

6from datetime import UTC, datetime, timedelta 

7from typing import Any, cast 

8 

9from intervaltree import Interval, IntervalTree 

10from pydantic import BaseModel, Field 

11 

12from sensorkit.auto.mode import Mode 

13from sensorkit.common.interval import ( 

14 combine_interval_trees, 

15 print_interval_tree, 

16 stack_interval_trees, 

17 stamp_interval_trees, 

18) 

19from sensorkit.common.logging import limited_logger 

20from sensorkit.core.program import OfferInterval 

21from sensorkit.core.task import TaskContextOverlay 

22 

23 

24class ProgramConfig(BaseModel): 

25 """Scheduling configuration for a single Program within a Controller.""" 

26 program: str 

27 priority: int = 5 

28 interrupt: bool = False 

29 contexts: TaskContextOverlay = Field(default_factory=TaskContextOverlay) 

30 

31 

32class ScheduleInterval(Interval): 

33 """A typed interval carrying a ScheduleIntent as its data payload.""" 

34 begin: datetime 

35 end: datetime 

36 data: ScheduleIntent | None 

37 

38 

39@dataclass 

40class ScheduleIntent: 

41 """The intended operating mode and ordered program candidates for a schedule interval.""" 

42 mode: str 

43 programs: tuple[str, ...] = tuple() 

44 

45 

46type Schedule = Sequence[ScheduleInterval] 

47 

48 

49class Scheduler: 

50 """Combines user configuration and Program tasking availability into an operating schedule.""" 

51 

52 def __init__( 

53 self, 

54 modes: Iterable[Mode], 

55 program_configs: list[ProgramConfig], 

56 ): 

57 # Sort the input modes in priority order. 

58 self.modes = sorted(modes, reverse=True) 

59 self.program_configs = {config.program: config for config in program_configs} 

60 self.combined_offers = IntervalTree() 

61 self.intervals = IntervalTree() 

62 self._offer_history = IntervalTree() 

63 

64 def get_schedule(self) -> Schedule: 

65 """Returns a serializable copy of the current schedule.""" 

66 return tuple(self.intervals) 

67 

68 def get_intent(self, dt: datetime): 

69 """Retrieves the scheduled intent at the given time, if any.""" 

70 ivs = self.intervals.at(dt) 

71 

72 if not ivs: 

73 return None 

74 

75 assert len(ivs) == 1 

76 return cast(ScheduleIntent, next(iter(ivs)).data) 

77 

78 def update( 

79 self, 

80 *, 

81 offers_dict: dict[str, Sequence[OfferInterval]], 

82 enabled_programs: set[str], 

83 mode_context: dict[str, Any] | None = None 

84 ): 

85 """Updates the schedule.""" 

86 mode_context = mode_context.copy() if mode_context else {} 

87 now = datetime.now(UTC) 

88 

89 # Snapshot the previous combined offers before rebuilding so we can 

90 # detect offers that were removed mid-window (e.g. task finished early). 

91 prev_combined = self.combined_offers 

92 

93 mode_context["time_ref"] = now 

94 mode_context["scheduler_previous_combined"] = prev_combined 

95 mode_context["scheduler_offer_history"] = self._offer_history 

96 

97 # Evaluate observing mode criteria. This gets us an IntervalTree for each configured 

98 # observing mode. 

99 evaluated_modes = [] 

100 

101 for mode in self.modes: 

102 try: 

103 intervals = mode.evaluate(mode_context) 

104 except Exception as e: 

105 exc_name = type(e).__name__ 

106 limited_logger(f"{mode.name}-{exc_name}", interval=300).warning( 

107 f"{exc_name} while evaluating mode {mode.name}: {e}" 

108 ) 

109 intervals = IntervalTree() 

110 

111 evaluated_modes.append(intervals) 

112 

113 # Get all the program offers for this Controller and sort them by configured priority. 

114 offers = sorted( 

115 (it for it in offers_dict.items() if it[0] in enabled_programs), 

116 key=lambda item: self.program_configs[item[0]].priority 

117 ) 

118 

119 # Combine into our updated schedule. 

120 self.combined_offers, self.intervals = create_schedule( 

121 self.modes, 

122 evaluated_modes, 

123 (it[0] for it in offers), 

124 (it[1] for it in offers), 

125 ) 

126 

127 self._support_after_activity(now, prev_combined) 

128 

129 def _support_after_activity(self, now: datetime, prev_combined: IntervalTree): 

130 # FIXME: Below only exists to support the `after_activity` mode criterion, which should 

131 # instead use closed-loop feedback about task execution (including manual). 

132 # Detect offers that disappeared while still in-progress, and inform schedule 

133 for iv in prev_combined: 

134 if iv.begin < now < iv.end and not self.combined_offers.overlaps(iv.begin, iv.end): 

135 truncated_end = now 

136 if not self._offer_history.containsi(iv.begin, truncated_end, iv.data): 

137 self._offer_history.addi(iv.begin, truncated_end, iv.data) 

138 

139 # Merge current offers into the persistent history 

140 for iv in self.combined_offers: 

141 if not self._offer_history.containsi(iv.begin, iv.end, iv.data): 

142 self._offer_history.addi(iv.begin, iv.end, iv.data) 

143 

144 # Prune historical entries 

145 cutoff = now - timedelta(hours=1) 

146 stale = [iv for iv in self._offer_history if iv.end <= cutoff] 

147 for iv in stale: 

148 self._offer_history.discard(iv) 

149 

150 

151def debug_print_schedule(arg: Schedule | Scheduler, *, lookahead_mins=30, print_func=print): 

152 """Print the schedule interval tree for debugging, showing the next *lookahead_mins* minutes.""" 

153 schedule = arg.intervals if isinstance(arg, Scheduler) else arg 

154 now = datetime.now(UTC) 

155 print_interval_tree( 

156 schedule, 

157 start_from=now, 

158 end_at=now + timedelta(minutes=lookahead_mins), 

159 key_func=lambda iv: iv.data.programs[0] if iv.data.programs else None, 

160 print_func=print_func, 

161 legend_show_count=True, 

162 legend_show_keys=True, 

163 ) 

164 

165 

166def create_schedule( 

167 modes: Iterable[Mode], 

168 evaluated_modes: Iterable[IntervalTree], 

169 programs: Iterable[str], 

170 offers: Iterable[Iterable[OfferInterval]], 

171): 

172 """Build a combined schedule from evaluated mode intervals and program offer windows. 

173 

174 Returns `(combined_offers, schedule)` where *schedule* is an IntervalTree whose 

175 intervals carry `ScheduleIntent` objects with the active mode and candidate programs. 

176 """ 

177 # Combine each set of intervals by stacking them in order, removing overlap regions. 

178 # The result is that operate modes take precedence over standby modes. We also set the 

179 # associated data with each interval to a new `ScheduleIntent` object. 

180 schedule = stack_interval_trees( 

181 *evaluated_modes, 

182 data=(ScheduleIntent(mode.name) for mode in modes), 

183 ) 

184 

185 # Combine the tasking schedules for all associated Programs such that each interval contains 

186 # a priority-ordered list of Program candidates. 

187 combined_offers = combine_interval_trees( 

188 *offers, 

189 data=programs, 

190 ) 

191 

192 # Merge program offer intervals into the mode intervals. This is done by splitting 

193 # the mode intervals on offer interval boundaries. The data fields (ScheduleIntent) are 

194 # updated to reflect which program offer, if any, overlaps each interval. 

195 def merge_func(target_data: ScheduleIntent, stamp_data: tuple[str, ...]): 

196 return ScheduleIntent(target_data.mode, stamp_data) 

197 

198 stamp_interval_trees(schedule, stamp=combined_offers, merge_func=merge_func) 

199 

200 return combined_offers, schedule