diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 4c88772327f26..7acdc4b1e1d7d 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -539,6 +539,38 @@ def _peek_pool_entry(provider: str) -> Optional[Any]: return None +# ---- SessionModelPool helper (auxiliary slot tracking) ---- + +_UNSET = object() # sentinel: distinguishes "never tried" from "pool is None/disabled" +_aux_pool_cache = _UNSET + + +def _get_session_model_pool(): + """Return the SessionModelPool singleton (or None if disabled/unavailable). + + Uses a sentinel so that a ``None`` result (pool disabled/not configured) + is not cached permanently — if the pool is created later in the process + lifetime, subsequent calls will find it. + + Import failures are also not cached, so transient import errors (e.g. + during startup) don't permanently disable pool tracking. + """ + global _aux_pool_cache + if _aux_pool_cache is not _UNSET: + return _aux_pool_cache + try: + from gateway.session_model_pool import get_session_model_pool + _pool = get_session_model_pool({}) + if _pool is not None: + # Cache only non-None (enabled) results. + _aux_pool_cache = _pool + # None (disabled/unconfigured) is NOT cached — next call retries. + except Exception: + # Import failure is NOT cached — next call retries. + pass + return _aux_pool_cache + + def _pool_runtime_api_key(entry: Any) -> str: if entry is None: return "" @@ -4875,7 +4907,10 @@ def call_llm( extra_body: Additional request body fields. Returns: - Response object with .choices[0].message.content + Response object with .choices[0].message.content. + Returns ``None`` when the SessionModelPool auxiliary slot is + blocked (pool enabled + model saturated). Callers should check + for ``None`` before accessing response attributes. Raises: RuntimeError: If no provider is configured. @@ -4968,6 +5003,27 @@ def call_llm( # Handle unsupported temperature, max_tokens vs max_completion_tokens retry, # then payment fallback. + # + # Session Model Pool: acquire an auxiliary slot before making the call + # so the pool can throttle concurrent auxiliary requests to the same model. + _pool_aux_acquired = False + _pool = _get_session_model_pool() + try: + if _pool and _pool.enabled: + _pool_aux_acquired = _pool.acquire_auxiliary_slot(final_model or "", resolved_provider or "") + except Exception: + pass + + if not _pool_aux_acquired and _pool and _pool.enabled: + logger.warning( + "Auxiliary %s: blocked by SessionModelPool for %s:%s — throttled", + task or "call", resolved_provider, final_model, + ) + raise RuntimeError( + f"Auxiliary call '{task or 'call'}' throttled by SessionModelPool: " + f"no auxiliary slots available for {resolved_provider}:{final_model}" + ) + try: return _validate_llm_response( client.chat.completions.create(**kwargs), task) @@ -5261,6 +5317,18 @@ def call_llm( logger.debug("Auxiliary: cache eviction after connection error failed", exc_info=True) raise + finally: + # Session Model Pool: release auxiliary slot after the call + # completes (success, error, or fallback). + if _pool_aux_acquired: + try: + if _pool and _pool.enabled: + _pool.release_auxiliary_slot(final_model or "", resolved_provider or "") + except Exception as _exc: + logger.error( + "SessionModelPool: FAILED to release auxiliary slot for %s:%s — " + "slot may be leaked: %s", resolved_provider, final_model, _exc, + ) def extract_content_or_reasoning(response) -> str: @@ -5337,6 +5405,9 @@ async def async_call_llm( """Centralized asynchronous LLM call. Same as call_llm() but async. See call_llm() for full documentation. + + Includes SessionModelPool auxiliary slot tracking (acquire before call, + release in finally block), mirroring the synchronous implementation. """ resolved_provider, resolved_model, resolved_base_url, resolved_api_key, resolved_api_mode = _resolve_task_provider_model( task, provider, model, base_url, api_key) @@ -5409,6 +5480,28 @@ async def async_call_llm( if _is_anthropic_compat_endpoint(resolved_provider, _client_base): kwargs["messages"] = _convert_openai_images_to_anthropic(kwargs["messages"]) + # Session Model Pool: acquire an auxiliary slot before making the async call + # so the pool can throttle concurrent auxiliary requests to the same model. + _async_pool_aux_acquired = False + _async_pool = _get_session_model_pool() + try: + if _async_pool and _async_pool.enabled: + _async_pool_aux_acquired = _async_pool.acquire_auxiliary_slot( + final_model or "", resolved_provider or "" + ) + except Exception: + pass + + if not _async_pool_aux_acquired and _async_pool and _async_pool.enabled: + logger.warning( + "Auxiliary %s (async): blocked by SessionModelPool for %s:%s — throttled", + task or "call", resolved_provider, final_model, + ) + raise RuntimeError( + f"Async auxiliary call '{task or 'call'}' throttled by SessionModelPool: " + f"no auxiliary slots available for {resolved_provider}:{final_model}" + ) + try: return _validate_llm_response( await client.chat.completions.create(**kwargs), task) @@ -5660,3 +5753,17 @@ async def async_call_llm( logger.debug("Auxiliary (async): cache eviction after connection error failed", exc_info=True) raise + finally: + # Session Model Pool: release auxiliary slot after the async call + # completes (success, error, or fallback). + if _async_pool_aux_acquired: + try: + if _async_pool and _async_pool.enabled: + _async_pool.release_auxiliary_slot( + final_model or "", resolved_provider or "" + ) + except Exception as _exc: + logger.error( + "SessionModelPool: FAILED to release async auxiliary slot for %s:%s — " + "slot may be leaked: %s", resolved_provider, final_model, _exc, + ) diff --git a/gateway/run.py b/gateway/run.py index f11686ccd360b..97cc7271f8f21 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -24,6 +24,7 @@ # means UTF-8 stdio setup is skipped on Windows; POSIX is unaffected. pass + import asyncio import dataclasses import inspect @@ -1876,6 +1877,12 @@ def __init__(self, config: Optional[GatewayConfig] = None): # Per-session reasoning effort overrides from /reasoning. # Key: session_key, Value: parsed reasoning config dict. self._session_reasoning_overrides: Dict[str, Dict[str, Any]] = {} + # Pool-assigned models (populated by SessionModelPool when enabled). + # Key: session_key, Value: dict with model/provider/context_length. + # These are weaker than manual /model overrides and are released + # when the session ends or when a manual override takes effect. + self._pool_assigned_models: Dict[str, Dict[str, Any]] = {} + self._pool_assigned_models_lock = threading.Lock() self._kanban_notifier_profile = self._active_profile_name() # Teams meeting pipeline runtime (bound later when msgraph_webhook adapter exists). self._teams_pipeline_runtime = None @@ -2513,6 +2520,41 @@ def _recover_telegram_topic_thread_id( return None return None + def _mark_pool_override(self, session_key: str) -> None: + """Mark a session as manually overridden in the pool. + + Called after ``_release_pool_slot`` so the pool knows not to + reassign a model on the next turn. + """ + try: + from gateway.session_model_pool import get_session_model_pool as _get_pool + _p = _get_pool({}) + if _p: + _p.mark_manual_override(session_key) + except Exception as _exc: + logger.debug("SessionModelPool: failed to mark override for %s: %s", session_key, _exc) + + def _release_pool_slot(self, session_key: str) -> None: + """Release a pool-assigned slot for a session (if one exists). + + Centralizes the release pattern used in 3 places: session reset, + /model override, in-place model switch, and any other override path. + Thread-safe: acquires ``_pool_assigned_models_lock`` internally. + """ + try: + with self._pool_assigned_models_lock: + _old_pool = self._pool_assigned_models.pop(session_key, None) + if _old_pool: + from gateway.session_model_pool import get_session_model_pool as _get_pool + # The singleton ignores config after first init; pass {} + # to avoid unnecessary disk I/O via _load_gateway_config(). + _p = _get_pool({}) + if _p: + _p.release_session_slot(session_key) + except Exception as _exc: + logger.debug("SessionModelPool: failed to release slot for %s: %s", session_key, _exc) + + def _resolve_session_agent_runtime( self, *, @@ -2535,6 +2577,42 @@ def _resolve_session_agent_runtime( model = _resolve_gateway_model(user_config) override = self._session_model_overrides.get(resolved_session_key) if resolved_session_key else None + # Will be set by pool integration below if a pool-assigned provider + # is available. Applied after runtime_kwargs is created. + _pool_provider_override = None + + # --- Session Model Pool integration --- + # If no manual override exists for this session, check whether the + # pool wants to assign a different model. Pool assignments are + # weaker than manual /model overrides and are released when the + # session ends or when a manual override takes effect. + if not override and resolved_session_key: + try: + from gateway.session_model_pool import get_session_model_pool as _get_pool + _cfg = user_config if user_config else _load_gateway_config() + _pool = _get_pool(_cfg) + if _pool and _pool.enabled: + # Always call acquire_session_slot — it is thread-safe + # internally and refreshes the session timestamp on every + # call. This prevents premature eviction of active sessions + # and avoids a TOCTOU race between the local cache check + # and the pool's own state. + _pool_assign = _pool.acquire_session_slot(resolved_session_key) + if _pool_assign: + with self._pool_assigned_models_lock: + self._pool_assigned_models[resolved_session_key] = _pool_assign + model = _pool_assign.get("model", model) + # Stash pool provider so it can be injected into + # runtime_kwargs after _resolve_runtime_agent_kwargs(). + _pool_provider_override = _pool_assign.get("provider") + logger.debug( + "SessionModelPool: session=%s using pool-assigned model=%s provider=%s", + resolved_session_key, model, _pool_assign.get("provider"), + ) + except Exception as _pool_exc: + logger.debug("SessionModelPool lookup failed: %s", _pool_exc) + # --- End Session Model Pool integration --- + if override: override_model = override.get("model", model) override_runtime = { @@ -2577,6 +2655,10 @@ def _resolve_session_agent_runtime( resolved_session_key, model, runtime_kwargs ) + # Apply pool-assigned provider (set during pool integration above). + if not override and _pool_provider_override: + runtime_kwargs["provider"] = _pool_provider_override + # When the config has no model.default but a provider was resolved # (e.g. user ran `hermes auth add openai-codex` without `hermes model`), # fall back to the provider's first catalog model so the API call @@ -8743,6 +8825,16 @@ async def _handle_message_with_agent(self, event, source, _quick_key: str, run_g self._set_session_reasoning_override(session_key, None) if hasattr(self, "_pending_model_notes"): self._pending_model_notes.pop(session_key, None) + # Release pool-assigned slot for the reset session. + self._release_pool_slot(session_key) + # Clear manual override so the pool can reassign on next turn. + try: + from gateway.session_model_pool import get_session_model_pool as _get_pool_rst + _p_rst = _get_pool_rst({}) + if _p_rst: + _p_rst.clear_manual_override(session_key) + except Exception as _exc: + logger.debug("SessionModelPool: failed to clear override for %s: %s", session_key, _exc) # Emit session:start for new or auto-reset sessions _is_new_session = ( @@ -10953,6 +11045,11 @@ async def _on_model_selected( "api_mode": result.api_mode, } + # Release pool-assigned slot for this session if one + # exists — the manual override takes precedence. + self._release_pool_slot(_session_key) + self._mark_pool_override(_session_key) + # Evict cached agent so the next turn creates a fresh # agent from the override rather than relying on the # stale cache signature to trigger a rebuild. @@ -11107,6 +11204,11 @@ async def _on_model_selected( "api_mode": result.api_mode, } + # Release pool-assigned slot for this session if one exists — + # the manual override takes precedence. + self._release_pool_slot(session_key) + self._mark_pool_override(session_key) + # Evict cached agent so the next turn creates a fresh agent from the # override rather than relying on cache signature mismatch detection. self._evict_cached_agent(session_key) diff --git a/gateway/session_model_pool.py b/gateway/session_model_pool.py new file mode 100644 index 0000000000000..632c2180546e8 --- /dev/null +++ b/gateway/session_model_pool.py @@ -0,0 +1,472 @@ +"""Session model pool — concurrency-aware auto-assignment for multi-session gateways. + +Distributes sessions and auxiliary calls across a configured pool of models, +respecting per-model concurrency limits and reserving slots for auxiliary tasks. + +**Important:** The module-level singleton (`get_session_model_pool`) caches the +pool instance for the lifetime of the gateway process. Runtime config changes +to ``session_model_pool`` require a gateway restart (or calling +``reset_session_model_pool()`` before the next access). +""" + +from __future__ import annotations + +import logging +import threading +import time +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Data classes +# --------------------------------------------------------------------------- + + +@dataclass +class PoolModelEntry: + """A single model in the pool with concurrency tracking.""" + + model: str + provider: str + max_concurrent: int = 1 + reserved_for_auxiliary: int = 0 + context_length: Optional[int] = None + priority: int = 5 + + # --- runtime state (not persisted) --- + session_slots: Dict[str, float] = field(default_factory=dict) # session_key -> last_activity + auxiliary_count: int = 0 + + @property + def session_count(self) -> int: + return len(self.session_slots) + + @property + def available_session_slots(self) -> int: + return max(0, self.max_concurrent - self.reserved_for_auxiliary - self.session_count) + + @property + def available_auxiliary_slots(self) -> int: + return max(0, self.reserved_for_auxiliary - self.auxiliary_count) + + @property + def is_saturated(self) -> bool: + """True when total usage (sessions + auxiliary) >= max_concurrent.""" + return (self.session_count + self.auxiliary_count) >= self.max_concurrent + + @property + def pool_key(self) -> str: + return f"{self.provider}:{self.model}" + + +# --------------------------------------------------------------------------- +# SessionModelPool +# --------------------------------------------------------------------------- + + +class SessionModelPool: + """Thread-safe pool that assigns models to sessions based on concurrency limits. + + Usage:: + + pool = SessionModelPool.from_config(config.get("session_model_pool", {})) + if pool and pool.enabled: + entry = pool.acquire_session_slot("discord:12345:67890") + if entry: + model = entry["model"] + # ... later ... + pool.release_session_slot("discord:12345:67890") + """ + + def __init__( + self, + *, + enabled: bool = False, + strategy: str = "round-robin", + inactive_timeout: int = 1800, + entries: Optional[List[PoolModelEntry]] = None, + ): + self.enabled = enabled + self.strategy = strategy + self.inactive_timeout = inactive_timeout # seconds + self._entries: List[PoolModelEntry] = entries or [] + self._entries_by_key: Dict[str, PoolModelEntry] = { + e.pool_key: e for e in self._entries + } + self._lock = threading.Lock() + # Condition variable for auxiliary slot waiting (replaces polling). + self._aux_condition = threading.Condition(self._lock) + # Track which session_key is currently assigned to which model (for release). + self._session_map: Dict[str, str] = {} # session_key -> pool_key + # Track manually overridden sessions (pool won't reassign). + self._manual_overrides: set = set() + + # ---- factory ---- + + @classmethod + def from_config(cls, config: dict) -> Optional["SessionModelPool"]: + """Build a pool from the ``session_model_pool`` config section. + + Returns ``None`` when the feature is disabled or config is empty/invalid. + """ + if not config or not config.get("enabled"): + return None + + raw_pool = config.get("pool", []) + if not raw_pool or not isinstance(raw_pool, list): + logger.warning("session_model_pool.pool is empty or not a list — disabled") + return None + + entries: List[PoolModelEntry] = [] + for raw in raw_pool: + if not isinstance(raw, dict): + logger.warning("Skipping non-dict entry in session_model_pool.pool") + continue + model = raw.get("model", "") + provider = raw.get("provider", "") + if not model or not provider: + logger.warning( + "Skipping pool entry missing model/provider: %s", raw + ) + continue + max_c = raw.get("max_concurrent", 1) + if not isinstance(max_c, int) or max_c < 1: + max_c = 1 + reserved = raw.get("reserved_for_auxiliary", 0) + if not isinstance(reserved, int) or reserved < 0: + reserved = 0 + if reserved > max_c: + logger.warning( + "reserved_for_auxiliary (%d) > max_concurrent (%d) for %s:%s — capping", + reserved, max_c, provider, model, + ) + reserved = max_c + ctx_len = raw.get("context_length") + priority = raw.get("priority", 5) + if not isinstance(priority, int): + priority = 5 + if priority < 1 or priority > 10: + logger.warning( + "priority (%d) out of 1-10 range for %s:%s — clamping", priority, provider, model, + ) + priority = max(1, min(priority, 10)) + entries.append( + PoolModelEntry( + model=model, + provider=provider, + max_concurrent=max_c, + reserved_for_auxiliary=reserved, + context_length=ctx_len if isinstance(ctx_len, int) else None, + priority=priority, + ) + ) + + if not entries: + logger.warning("session_model_pool has no valid entries — disabled") + return None + + # Deduplicate pool keys — if two entries share the same + # provider:model, only the LAST entry is kept. Earlier duplicates + # are logged and discarded to prevent inconsistent slot tracking. + _deduped: Dict[str, PoolModelEntry] = {} + for _entry in entries: + _key = _entry.pool_key + if _key in _deduped: + logger.warning( + "session_model_pool: duplicate pool_key '%s' (%s and %s). " + "Keeping the last entry; discarding the earlier one.", + _key, _deduped[_key].model, _entry.model, + ) + _deduped[_key] = _entry + entries = list(_deduped.values()) + + strategy = config.get("strategy", "round-robin") + if strategy not in ("round-robin", "least-loaded", "priority"): + logger.warning("Unknown strategy '%s' — defaulting to round-robin", strategy) + strategy = "round-robin" + + inactive_timeout = config.get("inactive_timeout", 1800) + if not isinstance(inactive_timeout, (int, float)) or inactive_timeout < 0: + inactive_timeout = 1800 + + logger.info( + "SessionModelPool enabled: strategy=%s, %d models, timeout=%ds", + strategy, len(entries), inactive_timeout, + ) + return cls( + enabled=True, + strategy=strategy, + inactive_timeout=inactive_timeout, + entries=entries, + ) + + # ---- session slot management ---- + + def acquire_session_slot(self, session_key: str) -> Optional[dict]: + """Find the best available model and assign the session to it. + + Returns ``{"model": str, "provider": str, "context_length": int|None}`` or + ``None`` if all models are saturated. + """ + if not self.enabled: + return None + + with self._lock: + # Don't reassign manually-overridden sessions. + if session_key in self._manual_overrides: + return None + + # Evict stale sessions based on per-session activity timestamps. + self._evict_inactive_sessions() + + # If session already has a slot, refresh its timestamp and return it. + existing = self._session_map.get(session_key) + if existing: + entry = self._entries_by_key.get(existing) + if entry: + entry.session_slots[session_key] = time.monotonic() + return { + "model": entry.model, + "provider": entry.provider, + "context_length": entry.context_length, + } + + # Find candidates with available session slots. + candidates = [ + e for e in self._entries + if e.available_session_slots > 0 + ] + + if not candidates: + logger.warning( + "SessionModelPool: all models saturated for session %s " + "(%d models, %d active sessions)", + session_key, len(self._entries), len(self._session_map), + ) + return None + + # Pick the best candidate based on strategy. + chosen = self._pick_candidate(candidates) + + chosen.session_slots[session_key] = time.monotonic() + self._session_map[session_key] = chosen.pool_key + + logger.debug( + "SessionModelPool: assigned session %s -> %s:%s " + "(session_slots=%d/%d, aux=%d/%d)", + session_key, chosen.provider, chosen.model, + chosen.session_count, chosen.max_concurrent, + chosen.auxiliary_count, chosen.reserved_for_auxiliary, + ) + return { + "model": chosen.model, + "provider": chosen.provider, + "context_length": chosen.context_length, + } + + def release_session_slot(self, session_key: str) -> None: + """Release a session's claimed slot.""" + if not self.enabled: + return + + with self._lock: + pool_key = self._session_map.pop(session_key, None) + if not pool_key: + return + entry = self._entries_by_key.get(pool_key) + if entry: + entry.session_slots.pop(session_key, None) + logger.debug( + "SessionModelPool: released session %s from %s:%s", + session_key, entry.provider, entry.model, + ) + + def mark_manual_override(self, session_key: str) -> None: + """Mark a session as manually overridden (e.g. via /model command). + + The pool will release any existing slot and won't auto-assign this session. + """ + if not self.enabled: + return + + with self._lock: + self._manual_overrides.add(session_key) + # Inline release (can't call release_session_slot which re-acquires lock). + pool_key = self._session_map.pop(session_key, None) + if pool_key: + entry = self._entries_by_key.get(pool_key) + if entry: + entry.session_slots.pop(session_key, None) + + def clear_manual_override(self, session_key: str) -> None: + """Remove manual override marker (e.g. on /new or /reset).""" + with self._lock: + self._manual_overrides.discard(session_key) + + # ---- auxiliary slot management ---- + + def acquire_auxiliary_slot( + self, model: str, provider: str, timeout: float = 5.0 + ) -> bool: + """Try to claim an auxiliary slot for a given model. + + Uses ``threading.Condition`` to wait efficiently (no busy-polling). + Returns ``True`` if the call is allowed to proceed. + """ + if not self.enabled: + return True # no pool = no restrictions + + pool_key = f"{provider}:{model}" + deadline = time.monotonic() + timeout + + with self._aux_condition: + while True: + entry = self._entries_by_key.get(pool_key) + if entry: + if entry.available_auxiliary_slots > 0: + entry.auxiliary_count += 1 + return True + # No auxiliary slot available — wait for release. + else: + # Model not in pool — allow without restriction. + return True + + remaining = deadline - time.monotonic() + if remaining <= 0: + logger.warning( + "SessionModelPool: auxiliary slot unavailable for %s:%s " + "after %.1fs (aux=%d, reserved=%d)", + provider, model, timeout, + self._aux_count(pool_key), + self._reserved_count(pool_key), + ) + return False + + # Wait for a slot to be released (signaled by release_auxiliary_slot). + self._aux_condition.wait(timeout=min(remaining, 1.0)) + + def release_auxiliary_slot(self, model: str, provider: str) -> None: + """Release an auxiliary slot.""" + if not self.enabled: + return + + pool_key = f"{provider}:{model}" + with self._aux_condition: + entry = self._entries_by_key.get(pool_key) + if entry and entry.auxiliary_count > 0: + entry.auxiliary_count -= 1 + # Wake up any threads waiting for auxiliary slots. + self._aux_condition.notify_all() + + # ---- strategy ---- + + def _pick_candidate(self, candidates: List[PoolModelEntry]) -> PoolModelEntry: + """Pick the best model from candidates based on configured strategy.""" + if self.strategy == "priority": + # Highest priority first; break ties by least loaded. + return max(candidates, key=lambda e: (e.priority, -e.session_count)) + elif self.strategy == "least-loaded": + # Fewest sessions first; break ties by priority. + return min(candidates, key=lambda e: (e.session_count, -e.priority)) + else: + # round-robin: pick the entry whose oldest session timestamp is earliest. + def _oldest_activity(e: PoolModelEntry) -> float: + if e.session_slots: + return min(e.session_slots.values()) + return 0.0 + return min(candidates, key=_oldest_activity) + + # ---- maintenance ---- + + def _evict_inactive_sessions(self) -> int: + """Remove sessions that have been inactive longer than the timeout. + + Must be called with ``self._lock`` held. + """ + now = time.monotonic() + evicted = 0 + for entry in self._entries: + stale_keys = [ + sk for sk, ts in entry.session_slots.items() + if (now - ts) > self.inactive_timeout + ] + for sk in stale_keys: + entry.session_slots.pop(sk, None) + self._session_map.pop(sk, None) + evicted += 1 + if evicted: + logger.info( + "SessionModelPool: evicted %d inactive sessions", evicted + ) + return evicted + + def get_pool_stats(self) -> Dict[str, Any]: + """Return current slot usage for logging/status.""" + with self._lock: + stats = { + "enabled": self.enabled, + "strategy": self.strategy, + "total_sessions": len(self._session_map), + "manual_overrides": len(self._manual_overrides), + "models": [], + } + for entry in self._entries: + stats["models"].append({ + "model": entry.model, + "provider": entry.provider, + "max_concurrent": entry.max_concurrent, + "reserved_for_auxiliary": entry.reserved_for_auxiliary, + "session_count": entry.session_count, + "available_sessions": entry.available_session_slots, + "auxiliary_count": entry.auxiliary_count, + "available_auxiliary": entry.available_auxiliary_slots, + "is_saturated": entry.is_saturated, + }) + return stats + + # ---- internal helpers ---- + + def _aux_count(self, pool_key: str) -> int: + entry = self._entries_by_key.get(pool_key) + return entry.auxiliary_count if entry else 0 + + def _reserved_count(self, pool_key: str) -> int: + entry = self._entries_by_key.get(pool_key) + return entry.reserved_for_auxiliary if entry else 0 + + +# --------------------------------------------------------------------------- +# Module-level singleton (lazy) +# --------------------------------------------------------------------------- + +_pool_instance: Optional[SessionModelPool] = None +_pool_lock = threading.Lock() + + +def get_session_model_pool(config: dict) -> Optional[SessionModelPool]: + """Return the global pool instance, creating it from *config* if needed. + + The pool is cached for the lifetime of the gateway process. + Runtime changes to ``session_model_pool`` in ``config.yaml`` require + either a gateway restart or a call to ``reset_session_model_pool()`` + before the next access. + """ + global _pool_instance + if _pool_instance is None: + with _pool_lock: + if _pool_instance is None: + pool_config = config.get("session_model_pool", {}) + _pool_instance = SessionModelPool.from_config(pool_config) + return _pool_instance + + +def reset_session_model_pool() -> None: + """Clear the singleton so it gets re-created from config on next access. + + Call this after runtime config changes to ``session_model_pool``. + """ + global _pool_instance + with _pool_lock: + _pool_instance = None diff --git a/hermes_cli/config.py b/hermes_cli/config.py index cec27809fdd0d..8320280715dbd 100644 --- a/hermes_cli/config.py +++ b/hermes_cli/config.py @@ -1264,7 +1264,18 @@ def _ensure_hermes_home_managed(home: Path): "extra_body": {}, }, }, - + + # Session Model Pool — concurrency-aware model assignment for gateway + # sessions. When multiple chat sessions run simultaneously, the pool + # assigns distinct models from this list so load is spread across + # providers. Disabled by default; set enabled: true to activate. + "session_model_pool": { + "enabled": False, + "strategy": "round-robin", # round-robin | least-loaded | priority + "inactive_timeout": 1800, # seconds — slot released after this idle time + "pool": [], # list of model entries (see docs) + }, + "display": { "compact": False, "personality": "", diff --git a/tests/gateway/test_session_model_pool.py b/tests/gateway/test_session_model_pool.py new file mode 100644 index 0000000000000..362e31c1814d7 --- /dev/null +++ b/tests/gateway/test_session_model_pool.py @@ -0,0 +1,572 @@ +"""Tests for gateway/session_model_pool.py""" + +import time +import threading +import pytest + +from gateway.session_model_pool import ( + PoolModelEntry, + SessionModelPool, + reset_session_model_pool, +) + + +# --------------------------------------------------------------------------- +# PoolModelEntry +# --------------------------------------------------------------------------- + + +class TestPoolModelEntry: + def test_available_session_slots(self): + e = PoolModelEntry(model="glm-5", provider="zai", max_concurrent=3, reserved_for_auxiliary=1) + e.session_slots = {"s1": time.monotonic()} + assert e.available_session_slots == 1 # 3 - 1(reserved) - 1(session) = 1 + + def test_available_auxiliary_slots(self): + e = PoolModelEntry(model="glm-5", provider="zai", max_concurrent=3, reserved_for_auxiliary=2) + e.auxiliary_count = 1 + assert e.available_auxiliary_slots == 1 + + def test_is_saturated(self): + e = PoolModelEntry(model="glm-5", provider="zai", max_concurrent=2, reserved_for_auxiliary=1) + e.session_slots = {"s1": time.monotonic()} + e.auxiliary_count = 1 + assert e.is_saturated + + def test_not_saturated(self): + e = PoolModelEntry(model="glm-5", provider="zai", max_concurrent=3, reserved_for_auxiliary=1) + e.session_slots = {"s1": time.monotonic()} + assert not e.is_saturated + + def test_pool_key(self): + e = PoolModelEntry(model="glm-5-turbo", provider="zai") + assert e.pool_key == "zai:glm-5-turbo" + + +# --------------------------------------------------------------------------- +# SessionModelPool.from_config +# --------------------------------------------------------------------------- + + +class TestFromConfig: + def test_disabled_returns_none(self): + assert SessionModelPool.from_config({"enabled": False}) is None + + def test_empty_config_returns_none(self): + assert SessionModelPool.from_config({}) is None + + def test_empty_pool_returns_none(self): + assert SessionModelPool.from_config({"enabled": True, "pool": []}) is None + + def test_missing_pool_returns_none(self): + assert SessionModelPool.from_config({"enabled": True}) is None + + def test_valid_config(self): + config = { + "enabled": True, + "strategy": "round-robin", + "inactive_timeout": 600, + "pool": [ + {"model": "glm-5-turbo", "provider": "zai", "max_concurrent": 1}, + {"model": "glm-4.7", "provider": "zai", "max_concurrent": 2, "reserved_for_auxiliary": 1}, + ], + } + pool = SessionModelPool.from_config(config) + assert pool is not None + assert pool.enabled + assert pool.strategy == "round-robin" + assert pool.inactive_timeout == 600 + assert len(pool._entries) == 2 + assert pool._entries[0].model == "glm-5-turbo" + assert pool._entries[1].reserved_for_auxiliary == 1 + + def test_invalid_entries_skipped(self): + config = { + "enabled": True, + "pool": [ + {"model": "", "provider": "zai"}, # no model + {"provider": "zai"}, # no model + {"model": "glm-5", "provider": "", "max_concurrent": 1}, # no provider + "not-a-dict", + {"model": "glm-5", "provider": "zai", "max_concurrent": 1}, # valid + ], + } + pool = SessionModelPool.from_config(config) + assert pool is not None + assert len(pool._entries) == 1 + + def test_unknown_strategy_defaults_to_round_robin(self): + config = { + "enabled": True, + "strategy": "invalid", + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 1}, + ], + } + pool = SessionModelPool.from_config(config) + assert pool.strategy == "round-robin" + + def test_reserved_capped_to_max_concurrent(self): + config = { + "enabled": True, + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 1, "reserved_for_auxiliary": 5}, + ], + } + pool = SessionModelPool.from_config(config) + assert pool._entries[0].reserved_for_auxiliary == 1 + + def test_priority_clamped_to_1_10(self): + config = { + "enabled": True, + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 1, "priority": 0}, + {"model": "glm-4", "provider": "zai", "max_concurrent": 1, "priority": 99}, + ], + } + pool = SessionModelPool.from_config(config) + assert pool._entries[0].priority == 1 + assert pool._entries[1].priority == 10 + + +# --------------------------------------------------------------------------- +# Session slot management +# --------------------------------------------------------------------------- + + +class TestSessionSlots: + def _make_pool(self, strategy="round-robin"): + config = { + "enabled": True, + "strategy": strategy, + "pool": [ + {"model": "glm-5-turbo", "provider": "zai", "max_concurrent": 1, "priority": 10}, + {"model": "glm-4.7", "provider": "zai", "max_concurrent": 2, "priority": 8}, + {"model": "glm-4.6", "provider": "zai", "max_concurrent": 3, "reserved_for_auxiliary": 1, "priority": 5}, + {"model": "glm-4.5", "provider": "zai", "max_concurrent": 10, "priority": 1}, + ], + } + return SessionModelPool.from_config(config) + + def test_acquire_session_slot(self): + pool = self._make_pool() + result = pool.acquire_session_slot("sess-1") + assert result is not None + assert result["model"] == "glm-5-turbo" # first entry, oldest activity → picked by round-robin + assert result["provider"] == "zai" + + def test_acquire_returns_existing(self): + pool = self._make_pool() + r1 = pool.acquire_session_slot("sess-1") + assert r1["model"] == "glm-5-turbo" + r2 = pool.acquire_session_slot("sess-1") + # Should return the same assignment + assert r2["model"] == "glm-5-turbo" + + def test_acquires_next_when_saturated(self): + pool = self._make_pool(strategy="priority") + # glm-5-turbo has 1 session slot, 0 reserved → 1 available + pool.acquire_session_slot("sess-1") # takes glm-5-turbo + # Next should get glm-4.7 (priority 8, 2 slots available) + result = pool.acquire_session_slot("sess-2") + assert result["model"] == "glm-4.7" + + def test_saturate_all_models(self): + pool = self._make_pool(strategy="round-robin") + # Available session slots: 1 + 2 + (3-1) + 10 = 15 + for i in range(15): + result = pool.acquire_session_slot(f"sess-{i}") + assert result is not None, f"Failed at session {i}" + # 16th should fail + result = pool.acquire_session_slot("sess-overflow") + assert result is None + + def test_release_session_slot(self): + pool = self._make_pool() + pool.acquire_session_slot("sess-1") + pool.release_session_slot("sess-1") + # After release, glm-5-turbo should be available again + stats = pool.get_pool_stats() + glm5 = next(m for m in stats["models"] if m["model"] == "glm-5-turbo") + assert glm5["session_count"] == 0 + + def test_release_nonexistent_is_noop(self): + pool = self._make_pool() + pool.release_session_slot("nonexistent") # should not raise + + def test_manual_override_prevents_assignment(self): + pool = self._make_pool() + pool.mark_manual_override("sess-manual") + result = pool.acquire_session_slot("sess-manual") + assert result is None # manual override = pool won't assign + + def test_mark_manual_override_releases_existing(self): + pool = self._make_pool() + pool.acquire_session_slot("sess-1") + pool.mark_manual_override("sess-1") + stats = pool.get_pool_stats() + glm5 = next(m for m in stats["models"] if m["model"] == "glm-5-turbo") + assert glm5["session_count"] == 0 + + def test_clear_manual_override_allows_reassignment(self): + pool = self._make_pool() + pool.mark_manual_override("sess-1") + pool.clear_manual_override("sess-1") + result = pool.acquire_session_slot("sess-1") + assert result is not None + + +# --------------------------------------------------------------------------- +# Strategy tests +# --------------------------------------------------------------------------- + + +class TestStrategies: + def _make_pool(self, strategy): + config = { + "enabled": True, + "strategy": strategy, + "pool": [ + {"model": "high", "provider": "zai", "max_concurrent": 2, "priority": 10}, + {"model": "low", "provider": "zai", "max_concurrent": 2, "priority": 1}, + ], + } + return SessionModelPool.from_config(config) + + def test_priority_strategy_prefers_highest(self): + pool = self._make_pool("priority") + result = pool.acquire_session_slot("s1") + assert result["model"] == "high" + + def test_least_loaded_strategy(self): + pool = self._make_pool("least-loaded") + r1 = pool.acquire_session_slot("s1") + r2 = pool.acquire_session_slot("s2") + # Both should go to different models (least loaded) + assert r1["model"] != r2["model"] + + def test_round_robin_strategy(self): + pool = self._make_pool("round-robin") + r1 = pool.acquire_session_slot("s1") + r2 = pool.acquire_session_slot("s2") + # round-robin uses oldest last_activity → should distribute + assert r1["model"] != r2["model"] + + def test_priority_tie_break_by_least_loaded(self): + """When two models have same priority, pick the least loaded.""" + config = { + "enabled": True, + "strategy": "priority", + "pool": [ + {"model": "a", "provider": "zai", "max_concurrent": 3, "priority": 5}, + {"model": "b", "provider": "zai", "max_concurrent": 3, "priority": 5}, + ], + } + pool = SessionModelPool.from_config(config) + r1 = pool.acquire_session_slot("s1") # fills model a (first) + r2 = pool.acquire_session_slot("s2") # tie-break → should pick b (0 sessions) + assert r1["model"] == "a" + assert r2["model"] == "b" + + def test_least_loaded_tie_break_by_priority(self): + """When two models have same load, pick the highest priority.""" + config = { + "enabled": True, + "strategy": "least-loaded", + "pool": [ + {"model": "lo-pri", "provider": "zai", "max_concurrent": 3, "priority": 1}, + {"model": "hi-pri", "provider": "zai", "max_concurrent": 3, "priority": 10}, + ], + } + pool = SessionModelPool.from_config(config) + r1 = pool.acquire_session_slot("s1") + r2 = pool.acquire_session_slot("s2") + # Both have 1 session after r1 is on first; tie-break → hi-pri + # Actually both start at 0 load → hi-pri (10) should be picked first + assert r1["model"] == "hi-pri" + # After s1 on hi-pri, lo-pri (0 load) wins over hi-pri (1 load) + assert r2["model"] == "lo-pri" + + +# --------------------------------------------------------------------------- +# Auxiliary slot management +# --------------------------------------------------------------------------- + + +class TestAuxiliarySlots: + def _make_pool(self): + config = { + "enabled": True, + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 2, "reserved_for_auxiliary": 1}, + ], + } + return SessionModelPool.from_config(config) + + def test_acquire_auxiliary_slot(self): + pool = self._make_pool() + assert pool.acquire_auxiliary_slot("glm-5", "zai") is True + + def test_auxiliary_blocked_when_full(self): + pool = self._make_pool() + pool.acquire_auxiliary_slot("glm-5", "zai") + # reserved=1, now aux_count=1 → no more auxiliary slots + assert pool.acquire_auxiliary_slot("glm-5", "zai", timeout=0.1) is False + + def test_release_auxiliary_slot(self): + pool = self._make_pool() + pool.acquire_auxiliary_slot("glm-5", "zai") + pool.release_auxiliary_slot("glm-5", "zai") + assert pool.acquire_auxiliary_slot("glm-5", "zai") is True # freed + + def test_auxiliary_model_not_in_pool_allowed(self): + pool = self._make_pool() + assert pool.acquire_auxiliary_slot("qwen-local", "local") is True + + def test_auxiliary_release_nonexistent_noop(self): + pool = self._make_pool() + pool.release_auxiliary_slot("nonexistent", "unknown") # no raise + + +# --------------------------------------------------------------------------- +# get_pool_stats +# --------------------------------------------------------------------------- + + +class TestPoolStats: + def test_stats_structure(self): + config = { + "enabled": True, + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 2, "reserved_for_auxiliary": 1}, + ], + } + pool = SessionModelPool.from_config(config) + stats = pool.get_pool_stats() + assert stats["enabled"] is True + assert stats["strategy"] == "round-robin" + assert len(stats["models"]) == 1 + assert stats["models"][0]["max_concurrent"] == 2 + + +# --------------------------------------------------------------------------- +# Singleton +# --------------------------------------------------------------------------- + + +class TestSingleton: + def test_get_session_model_pool(self): + reset_session_model_pool() + from gateway.session_model_pool import get_session_model_pool + pool = get_session_model_pool({"session_model_pool": {"enabled": False}}) + assert pool is None + reset_session_model_pool() + + def test_reset(self): + reset_session_model_pool() + from gateway.session_model_pool import _pool_instance + assert _pool_instance is None + + +# --------------------------------------------------------------------------- +# Thread safety +# --------------------------------------------------------------------------- + + +class TestThreadSafety: + def test_concurrent_acquire(self): + config = { + "enabled": True, + "pool": [ + {"model": "glm-4.5", "provider": "zai", "max_concurrent": 5}, + ], + } + pool = SessionModelPool.from_config(config) + results = [] + errors = [] + + def acquire(idx): + try: + r = pool.acquire_session_slot(f"sess-{idx}") + results.append(r) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=acquire, args=(i,)) for i in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(errors) == 0 + # 5 slots available, 10 threads → 5 successes + successes = [r for r in results if r is not None] + assert len(successes) == 5 + + +# --------------------------------------------------------------------------- +# Eviction (inactive session reaping) +# --------------------------------------------------------------------------- + + +class TestEviction: + def _make_pool(self, timeout=2.0): + config = { + "enabled": True, + "strategy": "priority", + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 2}, + ], + "inactive_timeout": timeout, + } + return SessionModelPool.from_config(config) + + def test_active_session_not_evicted(self): + pool = self._make_pool(timeout=10.0) + pool.acquire_session_slot("sess-1") + # Immediately after acquire, session should still be alive + stats = pool.get_pool_stats() + assert stats["total_sessions"] == 1 + + def test_inactive_session_evicted(self): + pool = self._make_pool(timeout=0.2) + pool.acquire_session_slot("sess-1") + # Wait longer than timeout so sess-1 becomes stale + time.sleep(0.5) + # Next acquire should trigger eviction of sess-1 and assign sess-2 + result = pool.acquire_session_slot("sess-2") + assert result is not None + # The old session should have been evicted, freeing the slot + stats = pool.get_pool_stats() + assert stats["total_sessions"] == 1 # only sess-2 remains + + def test_partial_eviction(self): + """Only stale sessions are evicted; active ones remain.""" + pool = self._make_pool(timeout=0.5) + pool.acquire_session_slot("sess-old") + # Refresh sess-new by re-acquiring (touches timestamp) + time.sleep(0.6) + pool.acquire_session_slot("sess-old") # re-acquire to refresh + pool.acquire_session_slot("sess-new") + stats = pool.get_pool_stats() + assert stats["total_sessions"] == 2 + + def test_eviction_frees_slots_for_new_sessions(self): + """After eviction, the freed slot can be used by a new session.""" + pool = self._make_pool(timeout=0.2) + # Fill both slots + pool.acquire_session_slot("sess-1") + pool.acquire_session_slot("sess-2") + # Pool is now saturated + assert pool.acquire_session_slot("sess-3") is None + # Wait for eviction + time.sleep(0.5) + # Acquire should trigger eviction and succeed + result = pool.acquire_session_slot("sess-3") + assert result is not None + + +# --------------------------------------------------------------------------- +# Duplicate pool_key validation +# --------------------------------------------------------------------------- + + +class TestDuplicatePoolKey: + def test_duplicate_pool_key_warns(self, caplog): + config = { + "enabled": True, + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 1}, + {"model": "glm-5", "provider": "zai", "max_concurrent": 2}, # duplicate pool_key + ], + } + pool = SessionModelPool.from_config(config) + assert pool is not None + # Should have warned about duplicate + assert any("duplicate pool_key" in r.message for r in caplog.records) + + def test_duplicate_pool_key_both_entries_retained(self): + """Duplicate pool_key is deduplicated — last entry wins, first discarded.""" + config = { + "enabled": True, + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 1}, + {"model": "glm-5", "provider": "zai", "max_concurrent": 3}, + ], + } + pool = SessionModelPool.from_config(config) + # Only the last entry is kept (deduplication) + assert len(pool._entries) == 1 + assert pool._entries[0].max_concurrent == 3 + + +# --------------------------------------------------------------------------- +# Auxiliary slot blocking edge cases +# --------------------------------------------------------------------------- + + +class TestAuxSlotBlocking: + def test_session_does_not_steal_aux_slot(self): + """Sessions can only use session slots, not auxiliary reserved slots.""" + config = { + "enabled": True, + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 3, "reserved_for_auxiliary": 2}, + ], + } + pool = SessionModelPool.from_config(config) + # Only 1 session slot (3 - 2 reserved = 1) + pool.acquire_session_slot("sess-1") + # Should be saturated for sessions + result = pool.acquire_session_slot("sess-2") + assert result is None + + def test_aux_slot_independent_of_session_slots(self): + """Auxiliary slots don't consume session capacity.""" + config = { + "enabled": True, + "pool": [ + {"model": "glm-5", "provider": "zai", "max_concurrent": 3, "reserved_for_auxiliary": 2}, + ], + } + pool = SessionModelPool.from_config(config) + # Use all session slots (3 - 2 reserved = 1) + pool.acquire_session_slot("sess-1") + # Auxiliary should still have 2 reserved slots available + assert pool.acquire_auxiliary_slot("glm-5", "zai") is True + assert pool.acquire_auxiliary_slot("glm-5", "zai") is True + # Third aux should fail (reserved=2) + assert pool.acquire_auxiliary_slot("glm-5", "zai", timeout=0.1) is False + + def test_concurrent_aux_acquire(self): + """Multiple auxiliary callers compete for limited slots.""" + config = { + "enabled": True, + "pool": [ + {"model": "glm-4.6", "provider": "zai", "max_concurrent": 4, "reserved_for_auxiliary": 2}, + ], + } + pool = SessionModelPool.from_config(config) + results = [] + errors = [] + + def acquire_aux(idx): + try: + ok = pool.acquire_auxiliary_slot("glm-4.6", "zai", timeout=0.5) + results.append(ok) + if ok: + time.sleep(0.2) + pool.release_auxiliary_slot("glm-4.6", "zai") + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=acquire_aux, args=(i,)) for i in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(errors) == 0 + # At least 2 should succeed (reserved=2), possibly more with timing + successes = [r for r in results if r is True] + assert len(successes) >= 2