feat(scheduler_gate): throttling + power-aware execution for local LLM (#1073) (#1252)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
sanil-23
2026-05-05 15:33:27 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent be9bc54fbb
commit 65032870f2
12 changed files with 781 additions and 48 deletions
+13
View File
@@ -245,6 +245,19 @@ pub async fn run_triage_with_resolved(
used_local = used_local,
"[triage::evaluator] dispatching {AGENT_RUN_TURN_METHOD}"
);
// Gate the LLM round-trip on the scheduler's global single slot
// ONLY when the resolved provider is local. Cloud routes already
// have their own (forthcoming) rate-limiter and are not the
// resource that the gate exists to protect — gating cloud here
// would just delay triage for no reason.
let _gate_permit = if used_local {
log::debug!("[triage::evaluator] local route — acquiring scheduler_gate llm permit");
crate::openhuman::scheduler_gate::wait_for_capacity().await
} else {
None
};
let response = match request_native_global::<AgentTurnRequest, AgentTurnResponse>(
AGENT_RUN_TURN_METHOD,
request,
+5 -1
View File
@@ -707,8 +707,12 @@ impl AutocompleteEngine {
v
};
// Interactive variant — bypasses the scheduler_gate's LLM permit
// so per-keystroke autocomplete doesn't queue behind a memory-tree
// backfill or a triage turn. See `inline_complete_interactive`
// docs in `local_ai/service/public_infer.rs`.
let generated = match service
.inline_complete(
.inline_complete_interactive(
&config,
&context,
&config.autocomplete.style_preset,
@@ -58,6 +58,26 @@ pub struct SchedulerGateConfig {
/// 60_000 (60s).
#[serde(default = "default_paused_poll_ms")]
pub paused_poll_ms: u64,
/// Hard CPU ceiling (recent global usage, 0..100). When the host CPU
/// climbs above this in `auto` mode, the gate flips to
/// `Paused { CpuPressure }` rather than just `Throttled` — every
/// background LLM call is held until the host calms down. Distinct
/// from `cpu_busy_threshold_pct`, which only triggers `Throttled`.
/// Default: 95.0.
#[serde(default = "default_cpu_severe_pct")]
pub cpu_severe_pct: f32,
/// When `true`, `auto` mode only runs background LLM work while the
/// laptop is on AC power. On battery the gate flips to
/// `Paused { OnBattery }` — no background inference at all,
/// regardless of charge level.
///
/// Default `false` to preserve the prior behavior (battery-floor
/// based throttling). Power-conscious users who never want
/// background inference on battery can flip this on.
#[serde(default)]
pub require_ac_power: bool,
}
fn default_battery_floor() -> f32 {
@@ -72,6 +92,9 @@ fn default_throttled_backoff_ms() -> u64 {
fn default_paused_poll_ms() -> u64 {
60_000
}
fn default_cpu_severe_pct() -> f32 {
95.0
}
impl Default for SchedulerGateConfig {
fn default() -> Self {
@@ -81,6 +104,8 @@ impl Default for SchedulerGateConfig {
cpu_busy_threshold_pct: default_cpu_busy_threshold(),
throttled_backoff_ms: default_throttled_backoff_ms(),
paused_poll_ms: default_paused_poll_ms(),
cpu_severe_pct: default_cpu_severe_pct(),
require_ac_power: false,
}
}
}
+9
View File
@@ -155,6 +155,15 @@ impl ReflectionHook {
async fn run_reflection(&self, prompt: &str) -> anyhow::Result<String> {
match self.config.reflection_source {
ReflectionSource::Local => {
// Local reflection acquires the scheduler_gate LLM
// permit transitively through `service.prompt` →
// `inference_with_temperature_internal`. Cloud
// reflection skips the gate (#1073 intentionally
// gates only local routes; cloud rate limiting is
// tracked separately).
log::debug!(
"[learning::reflection] local route — gate permit acquired via LocalAiService"
);
let service = crate::openhuman::local_ai::global(&self.full_config);
service
.prompt(&self.full_config, prompt, Some(512), true)
+112 -2
View File
@@ -53,6 +53,65 @@ impl LocalAiService {
style_instructions: Option<&str>,
style_examples: &[String],
max_tokens: Option<u32>,
) -> Result<String, String> {
self.inline_complete_internal(
config,
context,
style_preset,
style_instructions,
style_examples,
max_tokens,
/* gated = */ true,
)
.await
}
/// Latency-sensitive sibling of [`Self::inline_complete`] that
/// **bypasses the scheduler gate's LLM permit**.
///
/// Per-keystroke autocomplete must not block waiting for a
/// long-running memory-tree backfill or a triage turn to release
/// the global single slot. The user is at the keyboard; if the
/// background pipeline is busy we'd rather race the autocomplete
/// turn against it than show stale or empty completions for the
/// duration of the backfill.
///
/// This is the only path inside [`LocalAiService`] that opts out of
/// the gate. Every other entry point (`inference`, `prompt`,
/// `summarize`, `inline_complete`, `vision_prompt`, `embed`)
/// acquires before talking to Ollama.
pub async fn inline_complete_interactive(
&self,
config: &Config,
context: &str,
style_preset: &str,
style_instructions: Option<&str>,
style_examples: &[String],
max_tokens: Option<u32>,
) -> Result<String, String> {
log::trace!("[local_ai] inline_complete_interactive bypasses scheduler_gate permit");
self.inline_complete_internal(
config,
context,
style_preset,
style_instructions,
style_examples,
max_tokens,
/* gated = */ false,
)
.await
}
#[allow(clippy::too_many_arguments)]
async fn inline_complete_internal(
&self,
config: &Config,
context: &str,
style_preset: &str,
style_instructions: Option<&str>,
style_examples: &[String],
max_tokens: Option<u32>,
gated: bool,
) -> Result<String, String> {
if !config.local_ai.enabled {
return Ok(String::new());
@@ -102,6 +161,7 @@ impl LocalAiService {
max_tokens.or(Some(24)),
true,
0.05,
gated,
)
.await?;
Ok(sanitize_inline_completion(&raw, context))
@@ -128,6 +188,9 @@ impl LocalAiService {
return Err("messages must not be empty".to_string());
}
// Multi-turn local chat is background LLM-bound work — gate it.
let _gate_permit = crate::openhuman::scheduler_gate::wait_for_capacity().await;
tracing::debug!(
message_count = messages.len(),
model = %crate::openhuman::local_ai::model_ids::effective_chat_model_id(config),
@@ -223,6 +286,36 @@ impl LocalAiService {
.await
}
/// Latency-sensitive sibling of [`Self::inference`] that **bypasses
/// the scheduler gate's LLM permit**.
///
/// Used by user-arrival paths where the user is staring at the
/// output (push-to-talk dictation cleanup, in particular). If we
/// queue these behind a long-running memory backfill, the user
/// experiences a frozen UI; better to race the call against
/// background work and accept the contention than to silently
/// degrade interactivity.
///
/// Sibling to [`Self::inline_complete_interactive`] for autocomplete.
/// Every other entry point (`inference`, `prompt`, `summarize`,
/// `inline_complete`, `vision_prompt`, `embed`, `chat_with_history`)
/// remains gated.
pub(crate) async fn inference_interactive(
&self,
config: &Config,
system: &str,
prompt: &str,
max_tokens: Option<u32>,
no_think: bool,
) -> Result<String, String> {
log::trace!("[local_ai] inference_interactive bypasses scheduler_gate permit");
self.inference_with_temperature_internal(
config, system, prompt, max_tokens, no_think, 0.2, /* allow_empty = */ false,
/* gated = */ false,
)
.await
}
pub(crate) async fn inference_with_temperature(
&self,
config: &Config,
@@ -239,11 +332,13 @@ impl LocalAiService {
max_tokens,
no_think,
temperature,
false,
/* allow_empty = */ false,
/* gated = */ true,
)
.await
}
#[allow(clippy::too_many_arguments)]
async fn inference_with_temperature_allow_empty(
&self,
config: &Config,
@@ -252,6 +347,7 @@ impl LocalAiService {
max_tokens: Option<u32>,
no_think: bool,
temperature: f32,
gated: bool,
) -> Result<String, String> {
self.inference_with_temperature_internal(
config,
@@ -260,7 +356,8 @@ impl LocalAiService {
max_tokens,
no_think,
temperature,
true,
/* allow_empty = */ true,
gated,
)
.await
}
@@ -275,11 +372,24 @@ impl LocalAiService {
no_think: bool,
temperature: f32,
allow_empty: bool,
gated: bool,
) -> Result<String, String> {
if !matches!(self.status.lock().state.as_str(), "ready") {
self.bootstrap(config).await;
}
// Cooperative throttle + global single-slot acquisition for
// background LLM-bound work. Drop happens at end of scope so
// post-processing (status writes, logging) does NOT hold the
// permit any longer than necessary. Interactive autocomplete
// skips this via `gated = false` from
// `inline_complete_interactive`.
let _gate_permit = if gated {
crate::openhuman::scheduler_gate::wait_for_capacity().await
} else {
None
};
let started = std::time::Instant::now();
let model_id = model_ids::effective_chat_model_id(config);
@@ -162,3 +162,129 @@ async fn inline_complete_disabled_returns_empty_string() {
.unwrap();
assert!(out.is_empty());
}
#[tokio::test]
async fn inline_complete_interactive_disabled_returns_empty_string() {
// Interactive variant must match the gated variant on the
// disabled short-circuit so the autocomplete UX is identical.
let mut config = Config::default();
config.local_ai.enabled = false;
let service = LocalAiService::new(&config);
let out = service
.inline_complete_interactive(&config, "ctx", "casual", None, &[], None)
.await
.unwrap();
assert!(out.is_empty());
}
/// Interactive autocomplete (`inline_complete_interactive`) MUST NOT
/// block on a held LLM permit. Hold the global slot, race the
/// interactive variant against a tight deadline; if it queued behind
/// the permit it would deadlock or time out.
#[tokio::test]
async fn inline_complete_interactive_does_not_block_on_held_permit() {
let _guard = crate::openhuman::local_ai::LOCAL_AI_TEST_MUTEX
.lock()
.expect("local ai test mutex");
// Hold the global LLM permit for the duration of the test.
let _held = crate::openhuman::scheduler_gate::gate::try_acquire_llm_permit()
.expect("test must start with a free permit; previous test leaked one");
let app = Router::new().route(
"/api/generate",
post(|Json(_body): Json<serde_json::Value>| async move {
Json(json!({
"model": "test",
"response": "ip",
"done": true
}))
}),
);
let base = spawn_mock(app).await;
unsafe {
std::env::set_var("OPENHUMAN_OLLAMA_BASE_URL", &base);
}
let config = enabled_config();
let service = ready_service(&config);
// Tight 2s deadline — comfortably above mock RTT, well below any
// policy-paused-poll backoff. If the interactive call goes through
// the gate it'll never finish.
let result = tokio::time::timeout(
std::time::Duration::from_secs(2),
service.inline_complete_interactive(&config, "ctx", "casual", None, &[], Some(8)),
)
.await;
unsafe {
std::env::remove_var("OPENHUMAN_OLLAMA_BASE_URL");
}
let inner = result.expect("interactive variant must NOT block on held permit");
assert!(
inner.is_ok(),
"interactive call should have completed: {inner:?}"
);
}
/// Counterpart: the gated `inline_complete` (and `prompt`/`summarize`)
/// MUST queue behind a held permit. We assert this with a try-style
/// race: spawn the gated call, give it time to enter the wait, then
/// confirm it hasn't completed. We then drop the permit and verify
/// the call resolves.
#[tokio::test]
async fn gated_inline_complete_blocks_on_held_permit() {
let _guard = crate::openhuman::local_ai::LOCAL_AI_TEST_MUTEX
.lock()
.expect("local ai test mutex");
let held = crate::openhuman::scheduler_gate::gate::try_acquire_llm_permit()
.expect("test must start with a free permit");
let app = Router::new().route(
"/api/generate",
post(|Json(_body): Json<serde_json::Value>| async move {
Json(json!({
"model": "test",
"response": "x",
"done": true
}))
}),
);
let base = spawn_mock(app).await;
unsafe {
std::env::set_var("OPENHUMAN_OLLAMA_BASE_URL", &base);
}
let config = enabled_config();
let service = std::sync::Arc::new(ready_service(&config));
let svc = service.clone();
let cfg = config.clone();
let join = tokio::spawn(async move {
svc.inline_complete(&cfg, "ctx", "casual", None, &[], Some(8))
.await
});
// Give the spawned task a chance to enter `wait_for_capacity`.
tokio::time::sleep(std::time::Duration::from_millis(80)).await;
assert!(
!join.is_finished(),
"gated inline_complete must block while permit is held"
);
// Release the permit; the gated call should now resolve.
drop(held);
let resolved = tokio::time::timeout(std::time::Duration::from_secs(2), join)
.await
.expect("gated call must resolve once permit is released")
.expect("join")
.expect("ollama call");
assert!(!resolved.is_empty() || resolved.is_empty()); // sanity — value depends on sanitiser
unsafe {
std::env::remove_var("OPENHUMAN_OLLAMA_BASE_URL");
}
}
@@ -47,6 +47,10 @@ impl LocalAiService {
return Err("no valid image payloads were provided".to_string());
}
// Vision generation is background LLM-bound work; gate it through
// the scheduler's global LLM permit.
let _gate_permit = crate::openhuman::scheduler_gate::wait_for_capacity().await;
let body = OllamaGenerateRequest {
model: vision_model,
prompt: prompt.trim().to_string(),
@@ -147,6 +151,11 @@ impl LocalAiService {
self.ensure_ollama_model_available(&embedding_model, "embedding")
.await?;
// Embeds are bge-m3 calls (8K context, ~1.3 GB resident) — the
// single concurrent embed that has historically crashed the
// user's laptop when stacked with other Ollama work. Gate it.
let _gate_permit = crate::openhuman::scheduler_gate::wait_for_capacity().await;
let response = self
.http
.post(format!("{}/api/embed", ollama_base_url()))
+24 -23
View File
@@ -1,12 +1,18 @@
//! Worker pool: claims jobs from `mem_tree_jobs`, dispatches them through
//! [`handlers::handle_job`], and settles the row. A small global semaphore
//! caps concurrent LLM-bound work; non-LLM jobs run unrestricted.
//! [`handlers::handle_job`], and settles the row.
//!
//! Concurrency control for LLM-bound work is delegated to
//! [`crate::openhuman::scheduler_gate`] — its global single-slot
//! semaphore (`LlmPermit`) is the one source of truth across this
//! worker, voice cleanup, autocomplete, triage, and reflection. The
//! worker itself just calls `wait_for_capacity()`; non-LLM jobs
//! (`AppendBuffer`, `FlushStale`) run without acquiring a permit.
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use anyhow::Result;
use tokio::sync::{Notify, Semaphore};
use tokio::sync::Notify;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::jobs::handlers;
@@ -14,12 +20,6 @@ use crate::openhuman::memory::tree::jobs::store::{
claim_next, mark_done, mark_failed, recover_stale_locks, DEFAULT_LOCK_DURATION_MS,
};
// Held at 1 to keep concurrent bge-m3 embed calls (8K context, ~1.3 GB
// resident each) from saturating local RAM. The cloud-chat path itself
// would be fine at higher concurrency, but every worker also runs a
// bge-m3 embed step, and 3 concurrent embeds at 8K context have
// crashed the laptop in practice. See feedback memory
// `feedback_local_llm_load.md`.
const WORKER_COUNT: usize = 1;
const POLL_INTERVAL: Duration = Duration::from_secs(5);
@@ -47,18 +47,16 @@ pub fn start(config: Config) {
let notify = WORKER_NOTIFY
.get_or_init(|| Arc::new(Notify::new()))
.clone();
let llm_slots = Arc::new(Semaphore::new(1));
if let Err(err) = recover_stale_locks(&config) {
log::warn!("[memory_tree::jobs] recover_stale_locks failed at startup: {err:#}");
}
for idx in 0..WORKER_COUNT {
let notify = notify.clone();
let llm_slots = llm_slots.clone();
let cfg = config.clone();
tokio::spawn(async move {
loop {
match run_once_with_semaphore(&cfg, llm_slots.clone()).await {
match run_once(&cfg).await {
Ok(true) => continue,
Ok(false) => {
tokio::select! {
@@ -80,14 +78,8 @@ pub fn start(config: Config) {
}
/// Claim and run a single job. Returns `true` when work was processed,
/// `false` when no eligible row was available. Test entry point — the
/// production worker loop calls [`run_once_with_semaphore`] directly.
/// `false` when no eligible row was available.
pub async fn run_once(config: &Config) -> Result<bool> {
let llm_slots = Arc::new(Semaphore::new(1));
run_once_with_semaphore(config, llm_slots).await
}
async fn run_once_with_semaphore(config: &Config, llm_slots: Arc<Semaphore>) -> Result<bool> {
// Cooperative throttle BEFORE `claim_next()`. Holding the DB claim
// across an awaited `wait_for_capacity()` would let `Paused` mode
// sit on the row past `DEFAULT_LOCK_DURATION_MS`, after which
@@ -99,19 +91,28 @@ async fn run_once_with_semaphore(config: &Config, llm_slots: Arc<Semaphore>) ->
// is welcome; in Paused mode the user has explicitly asked us to
// stand down. Returns immediately in Aggressive/Normal so plugged-in
// desktops with headroom pay zero cost.
crate::openhuman::scheduler_gate::wait_for_capacity().await;
//
// For LLM-bound jobs the returned `LlmPermit` reserves the global
// single slot for the lifetime of `handle_job`. Non-LLM jobs
// (`AppendBuffer`, `FlushStale`) drop the permit before the
// handler runs so they don't block the slot.
let gate_permit = crate::openhuman::scheduler_gate::wait_for_capacity().await;
let Some(job) = claim_next(config, DEFAULT_LOCK_DURATION_MS)? else {
return Ok(false);
};
let permit = if job.kind.is_llm_bound() {
Some(llm_slots.acquire().await?)
let llm_permit = if job.kind.is_llm_bound() {
gate_permit
} else {
// Non-LLM jobs don't need the global slot; release it so an
// LLM-bound caller waiting elsewhere in the process can run.
drop(gate_permit);
None
};
let result = handlers::handle_job(config, &job).await;
drop(permit);
drop(llm_permit);
match result {
Ok(()) => {
+203 -15
View File
@@ -8,11 +8,52 @@ use std::sync::{Arc, OnceLock};
use std::time::Duration;
use parking_lot::RwLock;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use crate::openhuman::config::{Config, SchedulerGateConfig};
use crate::openhuman::scheduler_gate::policy::{decide, Policy};
use crate::openhuman::scheduler_gate::signals::Signals;
/// Process-wide ceiling on concurrent LLM-bound work.
///
/// Held at 1 to keep concurrent local-Ollama / bge-m3 calls (8K context,
/// ~1.3 GB resident each) from saturating local RAM. The cloud path
/// itself is bandwidth-bound but the gate also fronts triage and
/// reflection turns that fan out to `agent.run_turn`, where every local
/// route loads the same Ollama model — so a single global slot is the
/// safest contract until #1064 (per-toolkit triage toggle) and the
/// cloud-side rate limiter ship.
///
/// See `feedback_local_llm_load.md` — backfills with multiple
/// simultaneous Ollama requests have crashed the user's laptop twice.
const LLM_SLOTS: usize = 1;
static LLM_PERMITS: OnceLock<Arc<Semaphore>> = OnceLock::new();
fn llm_permits() -> &'static Arc<Semaphore> {
LLM_PERMITS.get_or_init(|| Arc::new(Semaphore::new(LLM_SLOTS)))
}
/// RAII guard returned by [`wait_for_capacity`] / [`acquire_llm_permit`].
///
/// While the caller holds an `LlmPermit`, no other LLM-bound caller in
/// the process can acquire one (the global semaphore has a single slot).
/// Drop the permit as soon as the LLM request returns — holding it past
/// post-processing serialises unrelated work for no reason.
///
/// This type is intentionally opaque: callers can't reach into the
/// underlying [`OwnedSemaphorePermit`] and risk forgetting to release it.
#[must_use = "drop the LlmPermit only after the LLM call returns"]
pub struct LlmPermit {
_permit: OwnedSemaphorePermit,
}
impl Drop for LlmPermit {
fn drop(&mut self) {
log::trace!("[scheduler_gate] llm permit released");
}
}
struct State {
cfg: SchedulerGateConfig,
signals: Signals,
@@ -110,35 +151,182 @@ pub fn current_signals() -> Signals {
})
}
/// Cooperatively block a worker until the host is ready for LLM-bound work.
/// Cooperatively block a caller until the host is ready for LLM-bound
/// work, then hand back an [`LlmPermit`] that holds a slot in the global
/// LLM semaphore.
///
/// * **Aggressive / Normal** — returns immediately.
/// * **Throttled** — sleeps `throttled_backoff_ms` so concurrent workers
/// serialise themselves and the host catches its breath between jobs.
/// * **Paused** — polls every `paused_poll_ms` until the policy changes.
/// Policy-driven backoff happens **before** semaphore acquisition so a
/// `Paused` mode doesn't pile up tasks queued for the slot — they sit
/// in the pause-poll loop, not in the semaphore wait queue.
///
/// Designed so existing semaphore-bounded worker pools can keep their pool
/// size and just gain a per-job throttle in front of the existing
/// `semaphore.acquire()` call.
pub async fn wait_for_capacity() {
/// * **Aggressive / Normal** — wait for the global slot; return immediately
/// once granted.
/// * **Throttled** — sleep `throttled_backoff_ms` first so concurrent
/// workers serialise themselves, then acquire the slot.
/// * **Paused** — poll every `paused_poll_ms` until the policy changes,
/// then acquire the slot.
///
/// Drop the returned [`LlmPermit`] as soon as the LLM call returns.
///
/// Returns `None` only if the global LLM semaphore has been closed
/// (never happens in production — the semaphore lives for the lifetime
/// of the process). Callers can safely treat `None` as "skip the
/// gate" rather than propagating an error.
pub async fn wait_for_capacity() -> Option<LlmPermit> {
loop {
let (policy, throttled_ms, paused_ms) = match STATE.get() {
Some(state) => {
let g = state.read();
(g.policy, g.cfg.throttled_backoff_ms, g.cfg.paused_poll_ms)
}
None => return,
None => {
// Gate not initialised (unit tests, early bootstrap).
// Acquire directly — no policy to consult.
return acquire_llm_permit_inner().await;
}
};
match policy {
Policy::Aggressive | Policy::Normal => return,
Policy::Throttled => {
tokio::time::sleep(Duration::from_millis(throttled_ms)).await;
return;
Policy::Aggressive | Policy::Normal => {
return acquire_llm_permit_inner().await;
}
Policy::Paused => {
Policy::Throttled => {
log::trace!(
"[scheduler_gate] throttled — sleeping {throttled_ms}ms before permit acquire"
);
tokio::time::sleep(Duration::from_millis(throttled_ms)).await;
return acquire_llm_permit_inner().await;
}
Policy::Paused { reason } => {
log::debug!(
"[scheduler_gate] paused ({}); polling every {paused_ms}ms",
reason.as_str()
);
tokio::time::sleep(Duration::from_millis(paused_ms)).await;
// re-evaluate; user may have toggled the gate back on.
}
}
}
}
async fn acquire_llm_permit_inner() -> Option<LlmPermit> {
let sem = llm_permits().clone();
match sem.acquire_owned().await {
Ok(permit) => {
log::trace!("[scheduler_gate] llm permit acquired");
Some(LlmPermit { _permit: permit })
}
Err(_) => {
// Semaphore closed — should never happen since we never
// close it. Log loudly and let the caller proceed without
// a permit so the pipeline doesn't deadlock.
log::warn!(
"[scheduler_gate] llm semaphore closed unexpectedly — proceeding without a permit"
);
None
}
}
}
/// Test/diagnostic hook: try to grab a permit without consulting the
/// gate policy. Returns `None` if no slots are free. **Do not** call
/// from production code — production callers should use
/// [`wait_for_capacity`] so the policy backoff applies.
#[cfg(test)]
pub fn try_acquire_llm_permit() -> Option<LlmPermit> {
let sem = llm_permits().clone();
sem.try_acquire_owned()
.ok()
.map(|p| LlmPermit { _permit: p })
}
/// Number of permits currently available. Test-only diagnostic.
#[cfg(test)]
pub fn available_llm_permits() -> usize {
llm_permits().available_permits()
}
#[cfg(test)]
mod tests {
//! These tests share the **process-wide** `LLM_PERMITS` semaphore
//! (which is intentional — that's what they're testing). They are
//! serialised via a module-local mutex so two test threads can't
//! both hold a permit at the same time and confuse each other's
//! `available_permits` reads.
use super::*;
use std::sync::Mutex;
use std::time::Instant;
use tokio::time::{timeout, Duration as TokioDuration};
static GATE_TEST_LOCK: Mutex<()> = Mutex::new(());
fn lock() -> std::sync::MutexGuard<'static, ()> {
// Tolerate poisoning so a panicking test doesn't block the rest.
GATE_TEST_LOCK.lock().unwrap_or_else(|p| p.into_inner())
}
#[tokio::test]
async fn wait_for_capacity_returns_permit_when_gate_uninit() {
let _g = lock();
let permit = wait_for_capacity().await;
assert!(
permit.is_some(),
"uninit gate must still hand back a permit"
);
assert_eq!(
available_llm_permits(),
0,
"permit must occupy the single LLM slot"
);
drop(permit);
assert_eq!(available_llm_permits(), 1, "drop must release the slot");
}
#[tokio::test]
async fn second_waiter_blocks_until_first_drops() {
let _g = lock();
let first = wait_for_capacity().await.expect("first permit");
assert_eq!(available_llm_permits(), 0);
// Spawn a second acquirer; it must block.
let handle = tokio::spawn(async move {
let started = Instant::now();
let p = wait_for_capacity().await;
(started.elapsed(), p)
});
// Give the second waiter a moment to start polling.
tokio::time::sleep(TokioDuration::from_millis(40)).await;
assert!(!handle.is_finished(), "second waiter must be blocked");
// Release the first permit; the second should resolve.
drop(first);
let (elapsed, second) = timeout(TokioDuration::from_secs(1), handle)
.await
.unwrap()
.unwrap();
assert!(
second.is_some(),
"second waiter must eventually get a permit"
);
assert!(
elapsed >= TokioDuration::from_millis(20),
"second waiter should have actually waited (got {elapsed:?})"
);
drop(second);
}
#[tokio::test]
async fn semaphore_size_is_one() {
let _g = lock();
let p1 = wait_for_capacity().await.expect("first permit");
// Try-acquire must fail while the slot is held.
assert!(
try_acquire_llm_permit().is_none(),
"semaphore must be size-1 — second try_acquire should fail"
);
drop(p1);
// Now another should succeed.
let p2 = try_acquire_llm_permit().expect("permit free after drop");
drop(p2);
}
}
+2 -2
View File
@@ -25,6 +25,6 @@ pub mod gate;
pub mod policy;
pub mod signals;
pub use gate::{current_policy, current_signals, init_global, wait_for_capacity};
pub use policy::Policy;
pub use gate::{current_policy, current_signals, init_global, wait_for_capacity, LlmPermit};
pub use policy::{PauseReason, Policy};
pub use signals::Signals;
+245 -4
View File
@@ -3,13 +3,48 @@
use crate::openhuman::config::SchedulerGateConfig;
use crate::openhuman::scheduler_gate::signals::Signals;
/// Why the gate is currently paused. Carried by [`Policy::Paused`] so
/// downstream consumers (UI, logging, observability) can surface a
/// specific user-facing reason instead of a generic "paused" label.
///
/// New variants will land alongside #1073's full power-aware work
/// (`OnBattery`, `CpuPressure`); `UserDisabled` covers the existing
/// `SchedulerGateMode::Off` path and `Unknown` is the safe fallback for
/// callers that don't have specific context yet.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PauseReason {
/// User explicitly turned the gate off in config.
UserDisabled,
/// Host on battery and gate's power-aware mode kicked in (#1073).
OnBattery,
/// CPU pressure exceeded the gate threshold (#1073).
CpuPressure,
/// Pause reason not yet classified — placeholder while #1073 is in flight.
Unknown,
}
impl PauseReason {
pub fn as_str(self) -> &'static str {
match self {
Self::UserDisabled => "user_disabled",
Self::OnBattery => "on_battery",
Self::CpuPressure => "cpu_pressure",
Self::Unknown => "unknown",
}
}
}
/// Background-AI scheduling tier. See module docs in `mod.rs` for semantics.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Policy {
Aggressive,
Normal,
Throttled,
Paused,
/// Gate paused. The `reason` is rendered to users in the memory-sync
/// status UI (#1136) and recorded in observability.
Paused {
reason: PauseReason,
},
}
impl Policy {
@@ -18,7 +53,17 @@ impl Policy {
Self::Aggressive => "aggressive",
Self::Normal => "normal",
Self::Throttled => "throttled",
Self::Paused => "paused",
Self::Paused { .. } => "paused",
}
}
/// `Some(reason)` when paused, `None` otherwise. Convenience for
/// callers that only need the reason and don't want to pattern-match
/// the whole enum (UI badges, log line construction).
pub fn pause_reason(self) -> Option<PauseReason> {
match self {
Self::Paused { reason } => Some(reason),
_ => None,
}
}
}
@@ -31,7 +76,11 @@ pub fn decide(signals: &Signals, cfg: &SchedulerGateConfig) -> Policy {
use crate::openhuman::config::SchedulerGateMode;
match cfg.mode {
SchedulerGateMode::Off => return Policy::Paused,
SchedulerGateMode::Off => {
return Policy::Paused {
reason: PauseReason::UserDisabled,
}
}
SchedulerGateMode::AlwaysOn => return Policy::Aggressive,
SchedulerGateMode::Auto => {}
}
@@ -46,6 +95,39 @@ pub fn decide(signals: &Signals, cfg: &SchedulerGateConfig) -> Policy {
// reject out-of-domain values for us.
let battery_floor = cfg.battery_floor.clamp(0.0, 1.0);
let cpu_threshold = cfg.cpu_busy_threshold_pct.clamp(0.0, 100.0);
let cpu_severe = cfg.cpu_severe_pct.clamp(0.0, 100.0);
// ── Pause checks come BEFORE the throttle gate — these are the
// "stand down completely" signals. Hierarchy:
// 1. user policy (`require_ac_power` on battery)
// 2. host on fire (CPU severely pegged)
// (1) Power-aware stand-down. Only consult `on_ac_power` when the
// user explicitly opts in — many desktops report `false` here
// because they have no battery + no AC sensor, and we don't
// want to silently disable background work for them.
if cfg.require_ac_power && !signals.on_ac_power {
log::debug!(
"[scheduler_gate] policy decision: paused on_battery (require_ac_power=true, on_ac={})",
signals.on_ac_power
);
return Policy::Paused {
reason: PauseReason::OnBattery,
};
}
// (2) Hard CPU ceiling — at >= cpu_severe_pct the host is unusable;
// a Throttled 30s backoff is not enough, hold every job.
if signals.cpu_usage_pct >= cpu_severe {
log::debug!(
"[scheduler_gate] policy decision: paused cpu_pressure (cpu={:.1}% >= severe={:.1}%)",
signals.cpu_usage_pct,
cpu_severe,
);
return Policy::Paused {
reason: PauseReason::CpuPressure,
};
}
let battery_ok = signals.on_ac_power
|| signals
@@ -74,6 +156,8 @@ mod tests {
cpu_busy_threshold_pct: 70.0,
throttled_backoff_ms: 30_000,
paused_poll_ms: 60_000,
cpu_severe_pct: 95.0,
require_ac_power: false,
}
}
@@ -92,7 +176,36 @@ mod tests {
&signals(true, None, 5.0, true),
&cfg(SchedulerGateMode::Off),
);
assert_eq!(p, Policy::Paused);
assert_eq!(
p,
Policy::Paused {
reason: PauseReason::UserDisabled
}
);
}
#[test]
fn pause_reason_helper_returns_user_disabled_for_off_mode() {
let p = decide(
&signals(true, None, 5.0, false),
&cfg(SchedulerGateMode::Off),
);
assert_eq!(p.pause_reason(), Some(PauseReason::UserDisabled));
}
#[test]
fn pause_reason_helper_returns_none_for_non_paused() {
assert_eq!(Policy::Aggressive.pause_reason(), None);
assert_eq!(Policy::Normal.pause_reason(), None);
assert_eq!(Policy::Throttled.pause_reason(), None);
}
#[test]
fn pause_reason_as_str_round_trips_each_variant() {
assert_eq!(PauseReason::UserDisabled.as_str(), "user_disabled");
assert_eq!(PauseReason::OnBattery.as_str(), "on_battery");
assert_eq!(PauseReason::CpuPressure.as_str(), "cpu_pressure");
assert_eq!(PauseReason::Unknown.as_str(), "unknown");
}
#[test]
@@ -166,8 +279,11 @@ mod tests {
#[test]
fn out_of_range_cpu_threshold_is_clamped() {
// 200.0 clamped to 100.0 — nothing above it, never throttles on CPU.
// Also push `cpu_severe_pct` to its max so the new pause-on-severe
// arm doesn't trip first.
let mut c = cfg(SchedulerGateMode::Auto);
c.cpu_busy_threshold_pct = 200.0;
c.cpu_severe_pct = 100.0;
let p = decide(&signals(true, None, 99.0, false), &c);
assert_eq!(p, Policy::Normal);
// -10.0 clamped to 0.0 — any positive CPU usage throttles.
@@ -185,4 +301,129 @@ mod tests {
);
assert_eq!(p, Policy::Normal);
}
// ── Power-aware require_ac_power gate (#1073) ─────────────────────
#[test]
fn require_ac_power_pauses_on_battery() {
let mut c = cfg(SchedulerGateMode::Auto);
c.require_ac_power = true;
// On battery, even with healthy charge + low CPU.
let p = decide(&signals(false, Some(0.95), 10.0, false), &c);
assert_eq!(
p,
Policy::Paused {
reason: PauseReason::OnBattery
}
);
}
#[test]
fn require_ac_power_normal_when_plugged_in() {
let mut c = cfg(SchedulerGateMode::Auto);
c.require_ac_power = true;
// Plugged in with headroom — should still run.
let p = decide(&signals(true, Some(0.90), 10.0, false), &c);
assert_eq!(p, Policy::Normal);
}
#[test]
fn require_ac_power_off_preserves_legacy_behavior_on_battery() {
// Default `require_ac_power = false` and a fresh battery means
// the legacy path runs: battery >= floor ⇒ Normal.
let mut c = cfg(SchedulerGateMode::Auto);
c.require_ac_power = false;
let p = decide(&signals(false, Some(0.95), 10.0, false), &c);
assert_eq!(p, Policy::Normal);
}
#[test]
fn require_ac_power_pause_resumes_when_back_on_ac() {
// Pause → re-evaluate after plugging in → Normal.
let mut c = cfg(SchedulerGateMode::Auto);
c.require_ac_power = true;
let s_battery = signals(false, Some(0.40), 5.0, false);
let s_ac = signals(true, Some(0.45), 5.0, false);
let p1 = decide(&s_battery, &c);
assert!(matches!(
p1,
Policy::Paused {
reason: PauseReason::OnBattery
}
));
let p2 = decide(&s_ac, &c);
assert_eq!(p2, Policy::Normal);
}
// ── Hard CPU ceiling (#1073) ──────────────────────────────────────
#[test]
fn cpu_severe_pauses_on_pressure() {
let mut c = cfg(SchedulerGateMode::Auto);
c.cpu_severe_pct = 90.0;
// CPU above severe ceiling, plugged in.
let p = decide(&signals(true, None, 96.0, false), &c);
assert_eq!(
p,
Policy::Paused {
reason: PauseReason::CpuPressure
}
);
}
#[test]
fn cpu_just_below_severe_throttles_not_pauses() {
let mut c = cfg(SchedulerGateMode::Auto);
c.cpu_busy_threshold_pct = 70.0;
c.cpu_severe_pct = 95.0;
// CPU above busy but below severe → Throttled, not Paused.
let p = decide(&signals(true, None, 80.0, false), &c);
assert_eq!(p, Policy::Throttled);
}
#[test]
fn cpu_severe_recovers_to_normal() {
let mut c = cfg(SchedulerGateMode::Auto);
c.cpu_severe_pct = 90.0;
let s_pegged = signals(true, None, 99.0, false);
let s_idle = signals(true, None, 5.0, false);
assert!(matches!(
decide(&s_pegged, &c),
Policy::Paused {
reason: PauseReason::CpuPressure
}
));
assert_eq!(decide(&s_idle, &c), Policy::Normal);
}
#[test]
fn out_of_range_cpu_severe_pct_is_clamped() {
// 200.0 clamped to 100.0 — only true 100% CPU triggers pause.
let mut c = cfg(SchedulerGateMode::Auto);
c.cpu_severe_pct = 200.0;
let p = decide(&signals(true, None, 99.9, false), &c);
// 99.9 < 100.0 (clamped), so we don't hit the pause arm and
// fall through to Throttled (cpu_busy_threshold=70).
assert_eq!(p, Policy::Throttled);
// Negative clamps to 0.0 — any positive CPU usage pauses.
c.cpu_severe_pct = -10.0;
let p = decide(&signals(true, None, 0.5, false), &c);
assert_eq!(
p,
Policy::Paused {
reason: PauseReason::CpuPressure
}
);
}
#[test]
fn server_mode_overrides_pause_signals() {
// Even on battery + CPU pegged, server mode stays Aggressive.
let mut c = cfg(SchedulerGateMode::Auto);
c.require_ac_power = true;
c.cpu_severe_pct = 50.0;
let p = decide(&signals(false, None, 99.0, true), &c);
assert_eq!(p, Policy::Aggressive);
}
}
+8 -1
View File
@@ -112,7 +112,14 @@ pub async fn cleanup_transcription(
// Hard timeout — dictation must feel instant. If the LLM doesn't
// respond within 3 seconds, fall back to the raw Whisper text.
let inference_fut = service.inference(config, CLEANUP_SYSTEM_PROMPT, &prompt, Some(512), true);
//
// Voice cleanup is a user-arrival path (mic press → STT → cleanup
// shown to user). It bypasses the scheduler_gate permit via
// `inference_interactive` so a long-running memory backfill does
// not push the cleanup past the 3s timeout. Mirrors the autocomplete
// bypass pattern (`inline_complete_interactive`).
let inference_fut =
service.inference_interactive(config, CLEANUP_SYSTEM_PROMPT, &prompt, Some(512), true);
let result: Result<String, String> =
match tokio::time::timeout(std::time::Duration::from_secs(3), inference_fut).await {
Ok(r) => r,