Coverage for astrocyte/pipeline/compile_trigger.py: 91%
105 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""M8 W4: Threshold trigger + async compile queue.
3CompileQueue monitors per-bank memory growth and staleness, enqueueing
4background compile jobs when either threshold is exceeded — without ever
5blocking the retain path.
7Design reference: docs/_design/llm-wiki-compile.md §3.1, §8.
8"""
10from __future__ import annotations
12import asyncio
13import logging
14from dataclasses import dataclass
15from datetime import UTC, datetime
16from typing import TYPE_CHECKING
18from astrocyte._log_safety import safe as _safe_log
20if TYPE_CHECKING:
21 from astrocyte.pipeline.compile import CompileEngine
23_logger = logging.getLogger("astrocyte.compile_trigger")
25# Sentinel pushed into the queue to signal the worker to stop.
26_STOP: str | None = None
29# ---------------------------------------------------------------------------
30# Configuration
31# ---------------------------------------------------------------------------
34@dataclass
35class CompileTriggerConfig:
36 """Thresholds that govern automatic compile triggering.
38 Either condition is sufficient to enqueue a compile job:
40 - **Size**: ``memories_since_last_compile >= size_threshold``
41 - **Staleness**: last compile is older than ``staleness_days`` AND
42 ``memories_since_last_compile >= staleness_min_memories``
44 Use ``size_threshold=0`` to trigger on every retain (useful for tests or
45 high-frequency low-cost scenarios).
46 """
48 size_threshold: int = 50
49 """New memories since last compile required to trigger on size alone."""
51 staleness_days: float = 7.0
52 """Days since last compile required for the staleness trigger."""
54 staleness_min_memories: int = 10
55 """Minimum new memories for the staleness condition to fire."""
58# ---------------------------------------------------------------------------
59# Per-bank state
60# ---------------------------------------------------------------------------
63@dataclass
64class BankCompileState:
65 """Runtime state for one bank's compile trigger."""
67 bank_id: str
68 memories_since_last_compile: int = 0
69 last_compile_at: datetime | None = None
71 def should_trigger(self, config: CompileTriggerConfig) -> bool:
72 """Return True if either trigger condition is met."""
73 n = self.memories_since_last_compile
75 # Size trigger: enough new memories regardless of time
76 if n >= config.size_threshold:
77 return True
79 # Staleness trigger: old enough AND at least staleness_min_memories
80 if self.last_compile_at is not None and n >= config.staleness_min_memories:
81 age_days = (datetime.now(UTC) - self.last_compile_at).total_seconds() / 86400.0
82 if age_days >= config.staleness_days:
83 return True
85 return False
88# ---------------------------------------------------------------------------
89# Compile queue
90# ---------------------------------------------------------------------------
93class CompileQueue:
94 """Async background queue for compile jobs.
96 Usage::
98 engine = CompileEngine(vs, llm, ws)
99 queue = CompileQueue(engine, CompileTriggerConfig(size_threshold=10))
100 await queue.start()
102 # In retain path (sync-safe — notify_retain is non-blocking):
103 queue.notify_retain("user-alice")
105 # Graceful shutdown:
106 await queue.stop()
108 Call :meth:`drain` in tests to wait for all pending jobs to finish before
109 making assertions::
111 queue.notify_retain("bank")
112 await queue.drain()
113 # CompileEngine.run("bank") has now completed
114 """
116 def __init__(
117 self,
118 engine: CompileEngine,
119 config: CompileTriggerConfig | None = None,
120 *,
121 max_queue_size: int = 100,
122 ) -> None:
123 self._engine = engine
124 self._config = config or CompileTriggerConfig()
125 self._queue: asyncio.Queue[str | None] = asyncio.Queue(maxsize=max_queue_size)
126 self._pending: set[str] = set() # banks queued but not yet running
127 self._state: dict[str, BankCompileState] = {}
128 self._task: asyncio.Task[None] | None = None
129 self._started = False
131 # ------------------------------------------------------------------
132 # Public API
133 # ------------------------------------------------------------------
135 def notify_retain(self, bank_id: str) -> None:
136 """Increment the memory counter for *bank_id* and enqueue if threshold met.
138 Non-blocking and sync-safe — designed to be called from the hot retain
139 path without holding any lock or awaiting anything. Duplicate banks
140 already in the queue are silently skipped so the queue never fills with
141 the same bank repeated N times.
142 """
143 state = self._state.setdefault(bank_id, BankCompileState(bank_id=bank_id))
144 state.memories_since_last_compile += 1
146 if not state.should_trigger(self._config):
147 return
149 if bank_id in self._pending:
150 # Already queued — the in-flight compile will cover these memories.
151 return
153 try:
154 self._queue.put_nowait(bank_id)
155 self._pending.add(bank_id)
156 _logger.debug("Compile job enqueued for bank %r", _safe_log(bank_id))
157 except asyncio.QueueFull:
158 _logger.warning(
159 "Compile queue full (%d slots); skipping bank %r. "
160 "Consider increasing max_queue_size or lowering thresholds.",
161 self._queue.maxsize,
162 _safe_log(bank_id),
163 )
165 async def start(self) -> None:
166 """Start the background worker task. Idempotent."""
167 if self._started:
168 return
169 self._started = True
170 self._task = asyncio.create_task(self._worker(), name="astrocyte.compile_worker")
171 _logger.debug("Compile worker started")
173 async def stop(self, *, timeout: float = 5.0) -> None:
174 """Signal the worker to stop and wait up to *timeout* seconds."""
175 if not self._started or self._task is None:
176 return
177 self._started = False
178 try:
179 self._queue.put_nowait(_STOP)
180 except asyncio.QueueFull:
181 self._task.cancel()
182 try:
183 await asyncio.wait_for(self._task, timeout=timeout)
184 except (asyncio.TimeoutError, asyncio.CancelledError):
185 _logger.warning("Compile worker did not stop within %.1fs; cancelling.", timeout)
186 self._task.cancel()
187 self._task = None
188 _logger.debug("Compile worker stopped")
190 async def drain(self, *, timeout: float = 5.0) -> None:
191 """Wait for all currently queued jobs to finish.
193 Raises :exc:`asyncio.TimeoutError` if jobs do not complete within
194 *timeout* seconds. Safe to call when the worker is not started
195 (returns immediately if the queue is empty).
196 """
197 if self._queue.empty() and not self._pending:
198 return
199 await asyncio.wait_for(self._queue.join(), timeout=timeout)
201 @property
202 def pending_banks(self) -> set[str]:
203 """Banks currently in the queue awaiting compilation (snapshot)."""
204 return set(self._pending)
206 @property
207 def state(self) -> dict[str, BankCompileState]:
208 """Per-bank compile state, keyed by bank_id (snapshot)."""
209 return dict(self._state)
211 # ------------------------------------------------------------------
212 # Background worker
213 # ------------------------------------------------------------------
215 async def _worker(self) -> None:
216 """Process compile jobs from the queue until a stop sentinel is received."""
217 _logger.debug("Compile worker loop started")
218 while True:
219 bank_id = await self._queue.get()
220 try:
221 if bank_id is _STOP:
222 _logger.debug("Compile worker received stop signal")
223 break
225 self._pending.discard(bank_id)
226 _logger.info("Running background compile for bank %r", bank_id)
227 try:
228 result = await self._engine.run(bank_id)
229 _logger.info(
230 "Background compile done for bank %r: %d created, %d updated, %d noise, %d tokens, %dms",
231 bank_id,
232 result.pages_created,
233 result.pages_updated,
234 result.noise_memories,
235 result.tokens_used,
236 result.elapsed_ms,
237 )
238 if result.error:
239 _logger.error(
240 "Background compile for bank %r reported error: %s",
241 bank_id,
242 result.error,
243 )
244 # Reset counter; record compile time
245 state = self._state.setdefault(bank_id, BankCompileState(bank_id=bank_id))
246 state.memories_since_last_compile = 0
247 state.last_compile_at = datetime.now(UTC)
249 except Exception:
250 _logger.exception(
251 "Background compile for bank %r raised unexpectedly; skipping.",
252 bank_id,
253 )
254 finally:
255 self._queue.task_done()