Coverage for core / src / sensorkit / auto / agent.py: 63%
116 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
5from typing import Annotated, Any, Literal
7from loguru import logger
8from pydantic import AfterValidator, BaseModel, Field
10import sensorkit.api as sk
11from sensorkit.auto.operator import ControllerConfig, VirtualOperator
12from sensorkit.auto.scheduler import Schedule
13from sensorkit.core.task import TaskContextOverlay
16@sk.declare_keyword
17class ElectionVotes(BaseModel):
18 """Published election keyword: per-source, per-controller vote state for observability."""
19 votes: dict[str, dict[str, bool | None]] = Field(default_factory=dict)
22def _set_controller_config_names(configs: dict[str, ControllerConfig]):
23 for name, config in configs.items():
24 config.name = name
26 return configs
29ControllerConfigMap = Annotated[
30 dict[str, ControllerConfig],
31 AfterValidator(_set_controller_config_names),
32]
35class FirstRunConfig(BaseModel):
36 """Initial state to assume on the first run before any persisted agent state exists."""
37 operate_all: bool = False
38 enable_scheduling: bool = True
41class AgentConfig(BaseModel):
42 """Top-level agent service configuration."""
43 first_run: FirstRunConfig = Field(default_factory=FirstRunConfig)
44 controllers: ControllerConfigMap = Field(default_factory=dict)
45 contexts: TaskContextOverlay = Field(default_factory=TaskContextOverlay)
47 def model_post_init(self, __context: Any):
48 for config in self.controllers.values():
49 config.propagate_config(self.contexts)
52class Capabilities(BaseModel):
53 """Deprecated agent capability descriptor published to the KV store."""
54 type: Literal["agent"] = "agent"
55 controllers: ControllerConfigMap
58class AgentControllerInfo(BaseModel):
59 """Per-controller operating state: control gate, elected state, and demand override."""
60 control_enabled: bool
61 elected_state: bool | None
62 demand_override: bool | None
64 @classmethod
65 def default(cls):
66 """Return a default AgentControllerInfo with control enabled and no overrides."""
67 return cls(control_enabled=True, elected_state=None, demand_override=None)
70class AgentOperatingState(sk.Event):
71 """Event capturing the agent's global control gate and per-controller operating info."""
72 global_control_enabled: bool
73 controllers: dict[str, AgentControllerInfo] = Field(default_factory=dict)
75 def derive_for_request(self, req: AgentConfigureRequest):
76 """Return a new AgentOperatingState reflecting the changes in *req*."""
77 return AgentOperatingState(
78 global_control_enabled=(
79 req.global_control_enabled
80 if req.global_control_enabled is not None
81 else self.global_control_enabled
82 ),
83 controllers={
84 name: AgentControllerInfo(
85 control_enabled=req.controller_control_enabled.get(name, info.control_enabled),
86 elected_state=info.elected_state,
87 demand_override=req.controller_demand_override.get(name, info.demand_override),
88 )
89 for name, info in self.controllers.items()
90 },
91 )
93 def derive_for_status(self, *, votes: dict[str, bool]):
94 """Return a new AgentOperatingState updated with the latest election *votes*."""
95 return AgentOperatingState(
96 global_control_enabled=self.global_control_enabled,
97 controllers={
98 name: info.model_copy(update=dict(elected_state=votes.get(name)))
99 for name, info in self.controllers.items()
100 },
101 )
103 def derive_for_config_change(self, *, configs: ControllerConfigMap):
104 """Return an updated state reconciled against *configs*, or None if no change is needed."""
105 if set(self.controllers) == set(configs):
106 return None
108 controllers = {
109 name: info
110 for name, info in self.controllers.items()
111 if name in configs
112 }
114 for name in configs:
115 if name not in controllers:
116 controllers[name] = AgentControllerInfo.default()
118 return AgentOperatingState(
119 global_control_enabled=self.global_control_enabled,
120 controllers=controllers,
121 )
124class AgentSchedulerState(sk.Event):
125 """Event capturing the scheduler gate, program exclusions, and current schedule snapshots."""
126 scheduling_enabled: bool
127 excluded_programs: set[str] = Field(default_factory=set)
128 schedule: dict[str, Schedule] = Field(default_factory=dict)
130 def derive_for_request(self, req: AgentConfigureRequest):
131 """Return a new AgentSchedulerState reflecting the scheduler changes in *req*."""
132 return AgentSchedulerState(
133 scheduling_enabled=(
134 req.scheduling_enabled
135 if req.scheduling_enabled is not None
136 else self.scheduling_enabled
137 ),
138 excluded_programs=(
139 self.excluded_programs.union(req.add_program_exclusions).difference(
140 req.remove_program_exclusions
141 )
142 ),
143 schedule=self.schedule,
144 )
146 def derive_for_status(self, *, schedule: dict[str, Schedule]):
147 """Return a copy of this state updated with the latest *schedule* snapshots."""
148 return self.model_copy(update=dict(schedule=schedule))
151class AgentState(sk.EventSourcedState):
152 """Persisted agent state combining operating and scheduler sub-states."""
153 operating_state: AgentOperatingState
154 scheduler_state: AgentSchedulerState
156 async def apply_to_operator(self, operator: VirtualOperator, configs: ControllerConfigMap):
157 """Push this state into the VirtualOperator: programs, overrides, and lifecycle gates.
159 Reads only from this instance, so call it on a consistent snapshot (see update's
160 return_snapshot) when concurrent updates are possible.
161 """
162 if self.scheduler_state.scheduling_enabled:
163 await operator.programs.global_enable()
164 else:
165 await operator.programs.global_disable()
167 for program in operator.programs.all_programs:
168 if program in self.scheduler_state.excluded_programs:
169 await operator.programs.disable(program)
170 else:
171 await operator.programs.enable(program)
173 for controller in configs:
174 info = self.operating_state.controllers.get(controller)
176 if info:
177 operator.election.vote(
178 source="override", subject=controller, vote=info.demand_override
179 )
181 lifecycle = operator.drivers[controller].lifecycle
183 if self.operating_state.global_control_enabled:
184 if info is None or info.control_enabled:
185 lifecycle.enable()
186 continue
188 lifecycle.disable()
191class AgentConfigureRequest(BaseModel):
192 """Request payload for the agent configure RPC, specifying partial state changes."""
193 global_control_enabled: bool | None = None
194 controller_control_enabled: dict[str, bool] = Field(default_factory=dict)
195 controller_demand_override: dict[str, bool | None] = Field(default_factory=dict)
196 scheduling_enabled: bool | None = None
197 add_program_exclusions: set[str] = Field(default_factory=set)
198 remove_program_exclusions: set[str] = Field(default_factory=set)
201agent_configure_request = sk.Request.define(
202 "configure",
203 payload=AgentConfigureRequest,
204 response=AgentState,
205)
208# TODO: Move this to a generic entity impl. Want a core feature to collapse single-entity services
209# into a single entity first.
210@sk.service_entrypoint(version=sk.VERSION)
211async def agent_service(service: sk.Service):
212 # Register the service.
213 await service.register()
215 # Read config.
216 config = await service.context.kv_get_model(AgentConfig)
218 # Recover state.
219 state = await AgentState.recover_or_init(
220 service.context,
221 operating_state=AgentOperatingState(global_control_enabled=config.first_run.operate_all),
222 scheduler_state=AgentSchedulerState(scheduling_enabled=config.first_run.enable_scheduling),
223 )
225 # Make sure our state structure is in agreement with configuration.
226 if update := state.operating_state.derive_for_config_change(configs=config.controllers):
227 logger.debug("controllers manifest was updated due to config change")
228 await state.update(service.context, update)
230 # Create the virtual operator and apply our initial state to it.
231 operator = VirtualOperator(config.controllers.values())
232 await state.apply_to_operator(operator, config.controllers)
234 # Start the virtual operator.
235 with service.context.enter_context():
236 await operator.start(client=service.client, task_group=service.context.task_group)
238 # Handle configuration requests.
239 # FIXME: asyncio.Lock does not guarantee FIFO wakeup order, so this usage could result in
240 # scrambled concurrent request calls. Also, the number of waiters is unbounded. The
241 # request API should provide a "serialized" mode that internally utilizes a bounded
242 # Queue to guarantee FIFO ordering and put an upper bound on the number of waiting
243 # requests.
244 _configure_lock = asyncio.Lock()
246 async def _agent_configure_request(request: AgentConfigureRequest):
247 async with _configure_lock:
248 # Apply the configuration into the current state, taking a consistent snapshot to
249 # configure the operator against without racing the publish loop's updates.
250 snapshot = await state.update(
251 service.context,
252 state.scheduler_state.derive_for_request(request),
253 state.operating_state.derive_for_request(request),
254 return_snapshot=True,
255 )
257 # Configure the virtual operator to reflect the new state.
258 await snapshot.apply_to_operator(operator, config.controllers)
260 return state
262 await service.context.handle_request(agent_configure_request, _agent_configure_request)
264 async def publish():
265 while True:
266 await service.context.publish(
267 ElectionVotes(votes=operator.election._votes)
268 )
270 await state.update(
271 service.context,
272 state.scheduler_state.derive_for_status(
273 schedule={
274 name: driver.scheduler.get_schedule()
275 for name, driver in operator.drivers.items()
276 }
277 ),
278 state.operating_state.derive_for_status(votes=operator.election.evaluate()),
279 )
281 await asyncio.sleep(1)
283 _publish_task = asyncio.create_task(publish())
285 await service.context.kv_put_model(Capabilities(
286 controllers=config.controllers
287 ))
289 # Run the service.
290 await service.run()