Coverage for core / src / sensorkit / auto / operator.py: 83%
248 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-02 00:03 +0000
1# SPDX-License-Identifier: Apache-2.0
2from __future__ import annotations
4import asyncio
5import contextlib
6import os
7from collections.abc import Iterable
8from datetime import UTC, datetime, timedelta
9from typing import Any
11from loguru import logger
12from pydantic import BaseModel, Field
14from sensorkit.astro.common import SitePosition
15from sensorkit.astro.observer import EarthObserver
16from sensorkit.auto.constraint import Constraint, ConstraintManager
17from sensorkit.auto.lifecycle import ControllerLifecycle
18from sensorkit.auto.mode import Mode, ModeList
19from sensorkit.auto.scheduler import ProgramConfig, Scheduler, debug_print_schedule
20from sensorkit.backend.base import KeyNotFound
21from sensorkit.common.aio import AsyncObserver
22from sensorkit.common.graph import StateElection
23from sensorkit.core.client import SensorKit
24from sensorkit.core.controller import InternalControllerState
25from sensorkit.core.program import ControllerOffers, ProgramClient, ProgramDiscovery, ProgramState
26from sensorkit.core.task import TaskContextMap, TaskContextOverlay
29class ControllerConfig(BaseModel):
30 """Configuration for a single Controller managed by the VirtualOperator."""
31 depends_on: list[str] = Field(default_factory=list)
32 estimated_startup_time: float = 0.0
33 estimated_shutdown_time: float = 0.0
34 modes: ModeList = Field(default_factory=list)
35 constraints: list[Constraint] = Field(default_factory=list)
36 tasking: list[ProgramConfig] = Field(default_factory=list)
37 contexts: TaskContextOverlay = Field(default_factory=TaskContextOverlay)
39 _name: str | None = None
40 _mode_index: dict[str, Mode]
41 _program_index: dict[str, ProgramConfig]
43 @property
44 def name(self):
45 """The name of this controller, set by the parent ControllerConfigMap validator."""
46 if self._name is None:
47 raise RuntimeError("name not set")
49 return self._name
51 @name.setter
52 def name(self, value: str):
53 """Set the name of this controller."""
54 self._name = value
56 def mode(self, name: str):
57 """Return the Mode configuration with the given name."""
58 return self._mode_index[name]
60 def program_config(self, program: str):
61 """Return the ProgramConfig for the given program entity name."""
62 return self._program_index[program]
64 def model_post_init(self, __context: Any):
65 self._mode_index = {mode.name: mode for mode in self.modes}
66 self._program_index = {config.program: config for config in self.tasking}
68 def propagate_config(self, contexts: TaskContextOverlay):
69 """Merge agent-level task contexts into this controller and its program configs."""
70 contexts.propagate(self.contexts)
72 for config in self.tasking:
73 self.contexts.propagate(config.contexts)
76class ControllerDriver:
77 """Operate a single Controller."""
79 def __init__(
80 self,
81 config: ControllerConfig,
82 program_discovery: ProgramDiscovery,
83 program_manager: ProgramStateManager,
84 ):
85 self.config = config
86 self.program_manager = program_manager
88 # Build task context maps. These are user-configured context keywords passed through
89 # to program tasking and direct task executions.
90 self._context_maps: dict[str | None, TaskContextMap] = {
91 None: config.contexts.build(),
92 **{cfg.program: cfg.contexts.build() for cfg in config.tasking},
93 }
95 # Create the Controller task offers monitor.
96 self.controller_offers = ControllerOffers(config.name, program_discovery)
98 # Create a scheduler.
99 self.scheduler = Scheduler(
100 modes=config.modes,
101 program_configs=config.tasking,
102 )
104 # Create constraint manager.
105 self.constraints = ConstraintManager(config.constraints)
107 # Create the lifecycle object.
108 self.lifecycle = ControllerLifecycle()
110 async def start(
111 self,
112 kit: SensorKit,
113 *,
114 task_group: Any = asyncio,
115 ):
116 """Start the lifecycle, constraints, offers monitor, and background tasks for this driver."""
117 self.kit = kit
118 controller = kit.controller(self.config.name)
119 logger.info(f"Driver for {self.config.name} starting")
121 # Start the lifecycle loop. This won't result in commanding until the lifecycle is enabled
122 # and a demand state is set.
123 self.lifecycle.start(controller, task_group=task_group)
125 # Start checking the configured constraints.
126 await self.constraints.start(task_group=task_group, ready_timeout=10.0, kit=kit)
128 # Start monitoring associated Programs for offers.
129 await self.controller_offers.start(task_group=task_group)
131 # Query controller site position.
132 while True:
133 try:
134 site = await controller.kv_get_model(SitePosition)
135 break
136 except KeyNotFound:
137 logger.debug(f"waiting for {self.config.name} SitePosition")
138 await asyncio.sleep(1)
140 self.observer = await EarthObserver.get(
141 lat_deg=site.latitude_degrees, lon_deg=site.longitude_degrees
142 )
144 logger.debug(f"driver for {self.config.name} ready")
146 async def stop(self):
147 """Stop the lifecycle and all background tasks for this driver."""
148 await self.lifecycle.stop()
150 def evaluate_state(self, election: StateElection):
151 """Update the schedule and cast election votes for schedule and constraints."""
152 # Build mode evaluation context.
153 mode_context = {
154 "time_range_parsers": {
155 "sunrise": self.observer.get_sunrise_time,
156 "sunset": self.observer.get_sunset_time,
157 },
158 }
160 # Update the schedule.
161 # FIXME: We do the full call twice here to ensure `tasking_available` and `after_activity`
162 # criteria pick up changes in a single iteration. This double-evaluation should be
163 # done more efficiently within the Scheduler itself.
164 for _ in range(2):
165 self.scheduler.update(
166 offers_dict=self.controller_offers.get_offers(),
167 enabled_programs=self.program_manager.enabled_programs(),
168 mode_context=mode_context,
169 )
171 if os.getenv("AGENT_DEBUG_SCHEDULE", None):
172 debug_print_schedule(self.scheduler, print_func=logger.debug)
174 # Check the schedule. Lookahead by startup time.
175 now = datetime.now(UTC)
176 intent = self.scheduler.get_intent(now)
177 if not intent and self.config.estimated_startup_time:
178 lookahead = now + timedelta(seconds=self.config.estimated_startup_time)
179 intent = self.scheduler.get_intent(lookahead)
181 election.vote(
182 "schedule",
183 subject=self.config.name,
184 vote=True if intent else None,
185 )
187 # Check constraints.
188 election.vote(
189 "constraint",
190 subject=self.config.name,
191 vote=False if self.constraints.is_constrained() else None,
192 )
194 def set_demand(self, operate: bool):
195 """Apply the elected operate/shutdown demand to the lifecycle, returning True if changed."""
196 state: InternalControllerState = InternalControllerState.SHUTDOWN
197 program: ProgramClient | None = None
199 if operate:
200 # Determine the target state (operate or standby) and which program to use, if any. If
201 # the controller's scheduler has something in mind, its intended state and program are
202 # used.
203 intent = self.scheduler.get_intent(datetime.now(UTC))
205 match intent and self.config.mode(intent.mode).state:
206 case "operate":
207 state = InternalControllerState.OPERATE
209 # Always use the highest-priority program given by the scheduler. Note that
210 # things like offline programs, temporary failure cooldowns, and other program
211 # availability constraints are already accounted for in the schedule. Any
212 # exclusions or re-prioritization here would not be visible in the published
213 # schedule, thus should be avoided. The full listing of candidate programs
214 # exists only for user informational purposes.
215 program = self.kit.program(intent.programs[0]) if intent.programs else None
216 case "standby":
217 state = InternalControllerState.STANDBY
218 case None:
219 # We are not scheduled to operate, so this must be a forced override scenario.
220 # Presently overrides alone cannot specify STANDBY.
221 state = InternalControllerState.OPERATE
222 case _ as it:
223 logger.error(
224 f"Invalid demand state {it!r} for mode {intent.mode}; "
225 "defaulting to SHUTDOWN"
226 )
227 state = InternalControllerState.SHUTDOWN
229 # Check constraints and determine appropriate contexts, whether we should interrupt. The
230 # contexts are a dict of dicts, mapping task types to the combined user configuration that
231 # applies for that task for this controller and program (if any).
232 constrained = self.constraints.is_constrained()
233 contexts = self._context_maps[None]
234 interrupt = False
236 if constrained:
237 if state != InternalControllerState.SHUTDOWN:
238 # Refuse to go up if a constraint is active. This can happen if a constraint goes
239 # high between schedule updates. Note this is not redundant to the wiring of
240 # constraints through the StateElection, since the latter informs multi-controller
241 # dependencies and general observability.
242 state = InternalControllerState.SHUTDOWN
244 interrupt = True
245 program = None
247 if program:
248 program_name = str(program.entity)
249 program_config = self.config.program_config(program_name)
251 # Use the configured contexts and interrupt policy for this program. Program-configured
252 # interrupts only apply in the OPERATE state, as program tasking must not interrupt
253 # lifecycle tasks.
254 contexts = self._context_maps[program_name]
255 interrupt |= (
256 self.lifecycle.belief_state == InternalControllerState.OPERATE
257 and program_config.interrupt
258 )
260 changed = self.lifecycle.set_demand_state(
261 state,
262 contexts=contexts,
263 program=program,
264 interrupt=interrupt,
265 )
267 if changed:
268 logger.debug(
269 f"{self.config.name} -> {str(state)} program={program.entity if program else None}"
270 )
272 return changed
275# FIXME: This is an AI implementation of this class; should be reviewed.
276class ProgramStateManager:
277 """Manages the desired enable state for programs and drives remote state."""
279 def __init__(self, programs: Iterable[str]):
280 self.all_programs = frozenset(programs)
281 self._desired_enable: dict[
282 str, str | None
283 ] = {} # program -> target_controller or None if disabled
284 self._global_enabled = True
285 self.on_change = AsyncObserver[str | None](None)
286 self._controller: dict[str, list[str]] = {}
288 def assign_controller(self, program: str, controller: str | None = None):
289 """Register a controller as a candidate for a program."""
290 if program not in self._controller:
291 self._controller[program] = []
292 if controller not in self._controller[program]:
293 self._controller[program].append(controller)
295 async def enable(self, program: str):
296 """Mark *program* as enabled for its first registered candidate controller."""
297 candidates = self._controller.get(program, [])
298 if not candidates:
299 return
300 target = candidates[0]
301 current = self._desired_enable.get(program)
303 if current and current != target:
304 raise RuntimeError(f"Program {program} already enabled for {current}")
306 if current == target:
307 return
309 self._desired_enable[program] = target
310 self.on_change.notify(program)
312 async def disable(self, program: str):
313 """Mark *program* as disabled, removing it from the enabled set."""
314 if self._desired_enable.get(program) is not None:
315 del self._desired_enable[program]
316 self.on_change.notify(program)
318 async def global_enable(self):
319 """Enable the global scheduling gate, allowing all individually-enabled programs to run."""
320 if not self._global_enabled:
321 self._global_enabled = True
322 self.on_change.notify(None)
324 async def global_disable(self):
325 """Disable the global scheduling gate, preventing all programs from running."""
326 if self._global_enabled:
327 self._global_enabled = False
328 self.on_change.notify(None)
330 def is_enabled(self, program: str) -> bool:
331 """Return True if *program* is individually enabled and the global gate is open."""
332 return self._global_enabled and self._desired_enable.get(program) is not None
334 def get_target_controller(self, program: str) -> str | None:
335 """Return the target controller entity name for *program*, or None if disabled."""
336 if not self._global_enabled:
337 return None
338 return self._desired_enable.get(program)
340 async def start(
341 self,
342 client: SensorKit,
343 discovery: ProgramDiscovery,
344 task_group: asyncio.TaskGroup,
345 ):
346 """Start the background task that drives remote program enable/disable state."""
347 self.client = client
348 self.discovery = discovery
349 self._retries = set()
350 self._program_locks: dict[str, asyncio.Lock] = {}
352 # Start monitoring tasks.
353 task_group.create_task(self._watch_discovery())
354 task_group.create_task(self._watch_changes())
356 async def _handle_program(self, program: str):
357 if program not in self._program_locks:
358 self._program_locks[program] = asyncio.Lock()
360 async with self._program_locks[program]:
361 program_client = self.client.program(program)
363 try:
364 state = await program_client.kv_get_model(ProgramState)
365 target = self.get_target_controller(program)
367 if target:
368 if not state.enable_state.enabled or state.enable_state.controller != target:
369 logger.info(f"Enabling program {program} for {target}")
370 await program_client.enable(target)
371 else:
372 if state.enable_state.enabled:
373 logger.info(f"Disabling program {program}")
374 await program_client.disable()
375 except Exception as e:
376 logger.error(f"Failed to drive program {program} state: {e}")
378 async def delayed_retry(delay: float = 10.0):
379 logger.debug(f"retrying in {delay} sec")
380 await asyncio.sleep(delay)
381 await self._handle_program(program)
383 self._retries.add(asyncio.create_task(delayed_retry()))
385 for task in list(self._retries):
386 if task.done():
387 self._retries.discard(task)
389 async def _watch_discovery(self):
390 previous_discovered = set()
392 async for programs in self.discovery.known_programs():
393 changed = programs.symmetric_difference(previous_discovered)
395 for p in changed:
396 if p in self.all_programs:
397 await self._handle_program(p)
399 previous_discovered = programs
401 async def _watch_changes(self):
402 async for program in self.on_change.consume(initial_value=True):
403 if program is None:
404 # Global change, re-evaluate all programs.
405 for p in self.all_programs:
406 await self._handle_program(p)
407 else:
408 await self._handle_program(program)
410 def enabled_programs(self):
411 """Return the set of currently enabled program entity names."""
412 return set(self._desired_enable.keys()) if self._global_enabled else set()
415class VirtualOperator:
416 """Orchestrate a set of Controllers."""
418 def __init__(self, controllers: Iterable[ControllerConfig]):
419 # Create the state election object. This election is the final say on whether a Controller
420 # operates or not. It unites operating schedule determined from mode configuration and
421 # program offerings, constraint evaluation, dependency relationships, and overrides.
422 self.election = StateElection({config.name: config.depends_on for config in controllers})
424 # Create a discovery object to find Programs and their target Controllers.
425 self.discovery = ProgramDiscovery()
427 # Create a manager to handle program enablement.
428 self.programs = ProgramStateManager(
429 tasking.program for config in controllers for tasking in config.tasking
430 )
432 for ctrl in controllers:
433 for prog in ctrl.tasking:
434 self.programs.assign_controller(prog.program, ctrl.name)
436 # Create objects to operate each of the configured Controllers.
437 self.drivers = {
438 config.name: ControllerDriver(config, self.discovery, self.programs)
439 for config in controllers
440 }
442 async def start(self, client: SensorKit, *, task_group: Any = asyncio):
443 """Start discovery, the program state manager, all drivers, and the main operator loop."""
444 # Start discovery.
445 await self.discovery.start(client)
447 # Start the program state manager.
448 await self.programs.start(client, self.discovery, task_group=task_group)
450 # Start all drivers.
451 # FIXME: One driver's failure to start should not affect others.
452 await asyncio.gather(
453 *(driver.start(client, task_group=task_group) for driver in self.drivers.values())
454 )
456 # Start the main loop.
457 logger.info("Starting virtual operator main loop")
458 self._loop_task = task_group.create_task(self._virtual_operator())
460 async def stop(self):
461 """Cancel the main loop, stop all controller drivers, and stop program discovery."""
462 self._loop_task.cancel()
464 await asyncio.gather(
465 *(driver.stop() for driver in self.drivers.values()),
466 return_exceptions=True,
467 )
469 with contextlib.suppress(asyncio.CancelledError):
470 await self._loop_task
472 await self.discovery.stop()
474 async def _virtual_operator(self):
475 while True:
476 # Tell each driver to update its state election votes.
477 for driver in self.drivers.values():
478 driver.evaluate_state(self.election)
480 # Evaluate the election to determine the up/down demand.
481 demand = self.election.evaluate()
483 # If there is no explicit demand for a given controller, we default to down.
484 for driver in self.drivers.values():
485 demand.setdefault(driver.config.name, False)
487 for name, state in demand.items():
488 if name not in self.drivers:
489 # FIXME: This should be a config error -- can't depend on a Controller that
490 # is not being orchestrated by the Agent.
491 continue
493 # Tell the driver to make it so.
494 self.drivers[name].set_demand(state)
496 await asyncio.sleep(2)