From 1cb70a78a041137ea6c014a259534477dacf98fe Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Wed, 19 Aug 2026 15:12:48 +0000 Subject: [PATCH 01/12] fix(gateway): one config object per request; drop LM cache (#292) Consolidate split config into a single ITSRequestConfig snapshot that holds every knob governing a request (LM target + scaling parameters). - Drop the LM client cache (_lm_cache, _get_or_create_lm, _hash_api_key); LM clients are now created and closed per request - Wire temperature through to OpenAICompatibleLanguageModel - Drop provider/extra_args from ConfigRequest - Fix metadata alg being read after await: arun_chat_completion returns the resolved alg in its result dict, sourced from the pre-await snapshot - Move SUPPORTED_ALGORITHMS and VALID_TOOL_VOTE_OPTIONS to api/types.py and validate all config fields in ITSRequestConfig.__post_init__ - Gateway.configure() merges non-None fields over the current default; arun_chat_completion() merges per-request fields over the default - Expand __repr__ to dump all fields (api_key masked) - Add 4 natural regression tests (real HTTP server, minimal mocks) - Fix incorrect docs: LM cache claims, /health endpoint, alg values Signed-off-by: Harrison Stropkay --- docs/iaas-service.md | 30 +- its_hub/api/__init__.py | 11 +- its_hub/api/types.py | 71 +++- its_hub/core/algorithms/self_consistency.py | 13 +- its_hub/core/gateway.py | 232 +++++------- its_hub/core/lms/openai_lm.py | 2 +- its_hub/integration/iaas/app.py | 73 ++-- its_hub/integration/iaas/models.py | 11 +- tests/conftest.py | 19 +- tests/mocks/recording_llm.py | 78 ++++ tests/test_gateway.py | 393 ++++++++++++-------- tests/test_iaas.py | 156 +++++--- 12 files changed, 647 insertions(+), 442 deletions(-) create mode 100644 tests/mocks/recording_llm.py diff --git a/docs/iaas-service.md b/docs/iaas-service.md index 9fed5582..d2cd7314 100644 --- a/docs/iaas-service.md +++ b/docs/iaas-service.md @@ -34,8 +34,8 @@ used. | Request body | `budget` | Compute budget | 2 | | `/configure` | `budget`, `endpoint`, `api_key` | Service defaults | 3 (lowest) | -Priority chain: **header > body > service default**. Headers are intended for Envoy -ext_proc routing but are also accepted on the standalone IaaS endpoint. +Priority chain: **header > body > service default**. Headers are intended for +Envoy ext_proc routing but are also accepted on the standalone IaaS endpoint. ### Algorithm Selection @@ -69,11 +69,10 @@ API keys can enter the system through three paths: **Security properties:** - Keys are **never logged** — the gateway logs endpoint and model but not credentials -- Keys are **hashed** (SHA-256, truncated to 16 hex chars) in LM cache keys to prevent - credential cross-contamination between requests using different API keys - Keys are **not persisted** to disk — they exist only in memory for the lifetime of the - service process -- On shutdown, all cached LM clients (and their associated keys) are cleared + request (an LM client is created per request and closed when it completes) +- Keys are **never shared between requests** — each request builds its own LM client, so + credentials supplied via header or `/configure` cannot cross-contaminate ## Prerequisites @@ -181,7 +180,7 @@ All configurations support: - `endpoint`: OpenAI-compatible API endpoint URL - `api_key`: API key for the provider - `model`: Model identifier -- `alg`: Algorithm name - `"self-consistency"` or `"best-of-n"` +- `alg`: Algorithm name - `"self-consistency"`, `"adaptive-self-consistency"`, or `"beta-self-consistency"` ## Usage Examples @@ -429,15 +428,15 @@ arrives as a burst rather than incrementally. ## Restart and Scaling -- **State**: The service holds an in-memory LM client cache (LRU, default 64 entries), - a gateway instance, and the service config set via `/configure`. No state is persisted +- **State**: The service holds a gateway instance and the service config set via + `/configure`. LM clients are created and discarded per request. No state is persisted to disk. -- **Restart**: Restarting clears all cached LM clients and the service config. - `/configure` must be called again after restart. +- **Restart**: Restarting clears the service config. `/configure` must be called again + after restart. - **Horizontal scaling**: Multiple IaaS instances can run independently. Each maintains - its own LM client cache, config, and gateway. There is no shared state between - instances. A load balancer must route `/configure` to all instances or each instance - must be configured independently. + its own config and gateway. There is no shared state between instances. A load + balancer must route `/configure` to all instances or each instance must be configured + independently. ## API Endpoints @@ -450,7 +449,6 @@ arrives as a burst rather than incrementally. ### Health Check - `GET /docs` - API documentation -- `GET /health` - Service health (if available) ## Troubleshooting @@ -489,7 +487,7 @@ curl -X GET http://localhost:8109/docs **5. Slow Responses** - This is expected behavior for inference-time scaling - Reduce `budget` parameter for faster responses -- Best-of-N with budget=4 typically takes 30-60 seconds +- Self-consistency with budget=4 typically takes 30-60 seconds ### Log Files diff --git a/its_hub/api/__init__.py b/its_hub/api/__init__.py index 3bc1725b..057f7556 100644 --- a/its_hub/api/__init__.py +++ b/its_hub/api/__init__.py @@ -24,7 +24,14 @@ from .orchestrator import AbstractOrchestrator from .reward_models.orm import AbstractOutcomeRewardModel from .reward_models.prm import AbstractProcessRewardModel -from .types import ChatMessage, ChatMessages, GenerationUsage, ITSRequestConfig +from .types import ( + SUPPORTED_ALGORITHMS, + VALID_TOOL_VOTE_OPTIONS, + ChatMessage, + ChatMessages, + GenerationUsage, + ITSRequestConfig, +) __all__ = [ # noqa: RUF022 # Algorithm abstractions @@ -44,6 +51,8 @@ "ChatMessages", "GenerationUsage", "ITSRequestConfig", + "SUPPORTED_ALGORITHMS", + "VALID_TOOL_VOTE_OPTIONS", # Error types "APIError", "RateLimitError", diff --git a/its_hub/api/types.py b/its_hub/api/types.py index c0203208..1dde650e 100644 --- a/its_hub/api/types.py +++ b/its_hub/api/types.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import re from dataclasses import dataclass, fields from typing import Literal @@ -18,7 +19,7 @@ class ChatMessage: """ role: Literal["system", "user", "assistant", "tool"] - content: str | list[dict] | None + content: str | list[dict] | None = None tool_calls: list[dict] | None = None # Store as plain dicts tool_call_id: str | None = None @@ -164,26 +165,78 @@ def total_tokens(self) -> int: return self.prompt_tokens + self.completion_tokens +SELF_CONSISTENCY_ALGORITHMS = frozenset( + {"self-consistency", "adaptive-self-consistency", "beta-self-consistency"} +) + +# All currently supported algorithms are self-consistency variants. +SUPPORTED_ALGORITHMS = SELF_CONSISTENCY_ALGORITHMS + +# Valid values for ITSRequestConfig.tool_vote / SelfConsistency.tool_vote. +VALID_TOOL_VOTE_OPTIONS = frozenset( + {"tool_name", "tool_args", "tool_hierarchical", "tool_flat_all"} +) + + @dataclass class ITSRequestConfig: """Per-request configuration for ITS execution. - Can be constructed incrementally — headers provide budget/endpoint/api_key, - then model is set from the request body. + Holds every knob that governs a request: the LM target + (``api_endpoint``, ``model``, ``api_key``, ``temperature``) and the scaling + parameters (``budget``, ``alg``, ``regex_patterns``, ``tool_vote``, + ``exclude_tool_args``, ``threshold``, ``confidence_threshold``). """ - budget: int - api_endpoint: str + # LM target + api_endpoint: str | None = None model: str | None = None api_key: str | None = None + temperature: float | None = None + # Scaling parameters + budget: int | None = None + alg: str | None = None + regex_patterns: list[str] | None = None + tool_vote: str | None = None + exclude_tool_args: list[str] | None = None + threshold: float | None = None + confidence_threshold: float | None = None def __post_init__(self): - if self.budget < 1 or self.budget > 1000: + if self.budget is not None and not (1 <= self.budget <= 1000): raise ValueError("budget must be between 1 and 1000") + if self.alg is not None and self.alg not in SUPPORTED_ALGORITHMS: + raise ValueError( + f"Algorithm {self.alg!r} not supported. " + f"Choose from: {SUPPORTED_ALGORITHMS}" + ) + if self.regex_patterns is not None: + for p in self.regex_patterns: + try: + re.compile(p) + except re.error as e: + raise ValueError(f"Invalid regex pattern {p!r}: {e}") from e + if self.threshold is not None and not (0.5 < self.threshold <= 1.0): + raise ValueError(f"threshold must be in (0.5, 1.0], got: {self.threshold}") + if self.confidence_threshold is not None and not ( + 0.5 < self.confidence_threshold <= 1.0 + ): + raise ValueError( + f"confidence_threshold must be in (0.5, 1.0], " + f"got: {self.confidence_threshold}" + ) + if self.tool_vote is not None and self.tool_vote not in VALID_TOOL_VOTE_OPTIONS: + raise ValueError( + f"tool_vote must be one of {VALID_TOOL_VOTE_OPTIONS}, " + f"got: {self.tool_vote}" + ) def __repr__(self) -> str: return ( - f"ITSRequestConfig(budget={self.budget}, " - f"api_endpoint='{self.api_endpoint}', " - f"model={self.model!r}, api_key={'***' if self.api_key else None})" + "ITSRequestConfig(" + + ", ".join( + f"{f.name}={'***' if f.name == 'api_key' else getattr(self, f.name)!r}" + for f in fields(self) + ) + + ")" ) diff --git a/its_hub/core/algorithms/self_consistency.py b/its_hub/core/algorithms/self_consistency.py index 638884f5..a4fde46b 100644 --- a/its_hub/core/algorithms/self_consistency.py +++ b/its_hub/core/algorithms/self_consistency.py @@ -8,6 +8,7 @@ from dataclasses import dataclass from its_hub.api import ( + VALID_TOOL_VOTE_OPTIONS, AbstractLanguageModel, AbstractOrchestrator, AbstractScalingAlgorithm, @@ -170,16 +171,10 @@ def __init__( ValueError: If tool_vote is not one of the supported options. """ # Validate tool_vote parameter - only validation needed since typing handles the rest - valid_tool_vote_options = { - None, - "tool_name", - "tool_args", - "tool_hierarchical", - "tool_flat_all", - } - if tool_vote not in valid_tool_vote_options: + valid_options = VALID_TOOL_VOTE_OPTIONS | {None} + if tool_vote not in valid_options: raise ValueError( - f"tool_vote must be one of {valid_tool_vote_options}, got: {tool_vote}" + f"tool_vote must be one of {valid_options}, got: {tool_vote}" ) # Set default projection function if provided None self.consistency_space_projection_func = ( diff --git a/its_hub/core/gateway.py b/its_hub/core/gateway.py index 3dd8b045..3a14d09a 100644 --- a/its_hub/core/gateway.py +++ b/its_hub/core/gateway.py @@ -4,10 +4,9 @@ similar to how HTTP services reuse client instances across requests. """ -import hashlib import logging import time -from collections import OrderedDict +from dataclasses import fields, replace from typing import Any from its_hub.api import ( @@ -24,7 +23,6 @@ from its_hub.core.algorithms.self_consistency import ( SelfConsistency, create_regex_projection_function, - validate_regex_patterns, ) from its_hub.core.lms.openai_lm import OpenAICompatibleLanguageModel from its_hub.core.orchestrator import LMOrchestrator @@ -44,6 +42,10 @@ # All currently supported algorithms are self-consistency variants. SUPPORTED_ALGORITHMS = SELF_CONSISTENCY_ALGORITHMS +# System defaults seeded into every gateway's default config. +_DEFAULT_BUDGET = 4 +_DEFAULT_ALG = "self-consistency" + class ITSGateway(AbstractGateway): """Long-lived gateway for running ITS algorithms. @@ -51,10 +53,6 @@ class ITSGateway(AbstractGateway): Initialized once at service startup and reused across all requests. Per-request configuration is passed as arguments to arun_chat_completion(). - Maintains a cache of LM clients keyed by (endpoint, model, hashed_api_key) - to reuse HTTP connections across requests while preventing credential - cross-contamination. - Example: # At service startup gateway = ITSGateway() @@ -69,141 +67,100 @@ class ITSGateway(AbstractGateway): result = await gateway.arun_chat_completion(config, messages) """ - _DEFAULT_MAX_LM_CACHE_SIZE = 64 - def __init__( self, - algorithm: AbstractScalingAlgorithm | None = None, orchestrator: AbstractOrchestrator | None = None, - max_lm_cache_size: int = _DEFAULT_MAX_LM_CACHE_SIZE, + default_config: ITSRequestConfig | None = None, ): if orchestrator is None: orchestrator = LMOrchestrator() self._orchestrator = orchestrator + self._default_config = default_config or ITSRequestConfig( + budget=_DEFAULT_BUDGET, alg=_DEFAULT_ALG + ) + logger.info( + f"ITSGateway initialized with default alg={self._default_config.alg} and budget={self._default_config.budget}" + ) - if algorithm is None: - algorithm = SelfConsistency(orchestrator=orchestrator) - self._algorithm = algorithm + @property + def default_config(self) -> ITSRequestConfig: + """The service-default config set via :meth:`configure`.""" + return self._default_config - self._algorithm_name = type(algorithm).__name__ - self._max_lm_cache_size = max_lm_cache_size - self._lm_cache: OrderedDict[ - tuple[str, str, str], OpenAICompatibleLanguageModel - ] = OrderedDict() - logger.info("ITSGateway initialized with algorithm=%s", self._algorithm_name) + def configure(self, config: ITSRequestConfig) -> None: + """Merge ``config`` into the service default. - def configure( - self, - alg: str, - regex_patterns: list[str] | None = None, - tool_vote: str | None = None, - exclude_tool_args: list[str] | None = None, - threshold: float | None = None, - confidence_threshold: float | None = None, - ) -> None: - """Configure the gateway's scaling algorithm at runtime. - - ``threshold`` applies only to ``adaptive-self-consistency`` and - ``confidence_threshold`` only to ``beta-self-consistency``; both fall - back to the algorithm's own default when left as ``None``. - - Raises ValueError for unsupported algorithms or invalid options. + Validation (alg, regex, thresholds) runs automatically in + ``ITSRequestConfig.__post_init__`` via ``_merge``. """ - if alg not in SUPPORTED_ALGORITHMS: - raise ValueError( - f"Algorithm {alg!r} not supported. Choose from: {SUPPORTED_ALGORITHMS}" - ) + self._default_config = self._merge(self._default_config, config) + logger.info( + "ITSGateway reconfigured with default alg=%s", self._default_config.alg + ) + + @staticmethod + def _merge(base: ITSRequestConfig, overlay: ITSRequestConfig) -> ITSRequestConfig: + """Return a copy of ``base`` with non-``None`` fields from ``overlay``.""" + return replace( + base, + **{ + f.name: getattr(overlay, f.name) + for f in fields(overlay) + if getattr(overlay, f.name) is not None + }, + ) + + def _build_algorithm(self, config: ITSRequestConfig) -> AbstractScalingAlgorithm: + """Construct a fresh scaling algorithm from a config snapshot.""" + alg = config.alg - # The self-consistency family shares projection/tool-vote plumbing. projection_func = None - if regex_patterns: - validate_regex_patterns(regex_patterns) - projection_func = create_regex_projection_function(regex_patterns) + if config.regex_patterns: + projection_func = create_regex_projection_function(config.regex_patterns) common_kwargs = { "consistency_space_projection_func": projection_func, - "exclude_args": exclude_tool_args, + "exclude_args": config.exclude_tool_args, "orchestrator": self._orchestrator, } # Only override the algorithm's default tool_vote when one is explicitly # provided; otherwise let it fall back to DEFAULT_TOOL_VOTE so that # tool-calling responses still vote sensibly with no extra config. - if tool_vote is not None: - common_kwargs["tool_vote"] = tool_vote + if config.tool_vote is not None: + common_kwargs["tool_vote"] = config.tool_vote if alg == "adaptive-self-consistency": - extra = {} if threshold is None else {"threshold": threshold} - algorithm = AdaptiveSelfConsistency(**common_kwargs, **extra) + extra = {} if config.threshold is None else {"threshold": config.threshold} + return AdaptiveSelfConsistency(**common_kwargs, **extra) elif alg == "beta-self-consistency": extra = ( {} - if confidence_threshold is None - else {"confidence_threshold": confidence_threshold} + if config.confidence_threshold is None + else {"confidence_threshold": config.confidence_threshold} ) - algorithm = BetaSelfConsistency(**common_kwargs, **extra) + return BetaSelfConsistency(**common_kwargs, **extra) else: # self-consistency - algorithm = SelfConsistency(**common_kwargs) - - self._algorithm = algorithm - self._algorithm_name = type(algorithm).__name__ - logger.info("ITSGateway reconfigured with algorithm=%s", self._algorithm_name) + return SelfConsistency(**common_kwargs) - @staticmethod - def _hash_api_key(api_key: str | None) -> str: - return hashlib.sha256((api_key or "").encode()).hexdigest()[:16] - - async def _get_or_create_lm( + def _build_lm( self, - endpoint: str, - model: str, - api_key: str | None = None, + config: ITSRequestConfig, request_id: str | None = None, ) -> OpenAICompatibleLanguageModel: - """Get cached LM client or create new one. - - Clients are cached by (endpoint, model, hashed_api_key) to reuse - HTTP connections while isolating different credentials. - """ - cache_key = (endpoint, model, self._hash_api_key(api_key)) + """Construct a one-shot LM client from a config snapshot.""" log_prefix = f"[{request_id}] " if request_id else "" - - if cache_key in self._lm_cache: - self._lm_cache.move_to_end(cache_key) - logger.debug( - "%sReusing cached LM client: endpoint=%s, model=%s", - log_prefix, - endpoint, - model, - ) - return self._lm_cache[cache_key] - - evicted_lm = None - if len(self._lm_cache) >= self._max_lm_cache_size: - evicted_key, evicted_lm = self._lm_cache.popitem(last=False) - logger.info( - "%sEvicting LM client from cache: endpoint=%s, model=%s", - log_prefix, - evicted_key[0], - evicted_key[1], - ) - logger.info( - "%sCreating new LM client: endpoint=%s, model=%s", + "%sCreating LM client: endpoint=%s, model=%s", log_prefix, - endpoint, - model, + config.api_endpoint, + config.model, ) - lm = OpenAICompatibleLanguageModel( - endpoint=endpoint, - api_key=api_key or "", - model_name=model, + return OpenAICompatibleLanguageModel( + endpoint=config.api_endpoint, + api_key=config.api_key, + model_name=config.model, + temperature=config.temperature, ) - self._lm_cache[cache_key] = lm - - if evicted_lm is not None: - await evicted_lm.close() - - return lm async def arun_chat_completion( self, @@ -217,42 +174,44 @@ async def arun_chat_completion( request_id = kwargs.get("request_id") log_prefix = f"[{request_id}] " if request_id else "" - if not config.api_endpoint: - raise ValueError("api_endpoint must be specified in ITSRequestConfig") - if not config.model: - raise ValueError( - "Model must be specified in ITSRequestConfig before running" - ) + merged = self._merge(self._default_config, config) - lm = await self._get_or_create_lm( - endpoint=config.api_endpoint, - model=config.model, - api_key=config.api_key, - request_id=request_id, - ) + if not merged.api_endpoint: + raise ValueError("api_endpoint must be specified") + if not merged.model: + raise ValueError("Model must be specified") + + # Snapshot the algorithm + LM for this request only. Nothing on `self` + # is mutated, so a concurrent /configure cannot swap the algorithm or + # close this request's HTTP session mid-flight. + algorithm = self._build_algorithm(merged) + lm = self._build_lm(merged, request_id) chat_messages = ChatMessages([ChatMessage.from_dict(msg) for msg in messages]) logger.info( "%sRunning ITS: algorithm=%s, budget=%s, endpoint=%s, model=%s, messages=%s, tools=%s", log_prefix, - self._algorithm_name, - config.budget, - config.api_endpoint, - config.model, + type(algorithm).__name__, + merged.budget, + merged.api_endpoint, + merged.model, len(messages), "yes" if tools else "no", ) t0 = time.monotonic() - result = await self._algorithm.ainfer( - lm=lm, - prompt_or_messages=chat_messages, - budget=config.budget, - return_response_only=False, - tools=tools, - tool_choice=tool_choice, - ) + try: + result = await algorithm.ainfer( + lm=lm, + prompt_or_messages=chat_messages, + budget=merged.budget, + return_response_only=False, + tools=tools, + tool_choice=tool_choice, + ) + finally: + await lm.close() usage_dict = {} if isinstance(result.usage, GenerationUsage): @@ -272,7 +231,7 @@ async def arun_chat_completion( duration_s, usage_dict, ) - return {"message": result.the_one, "usage": usage_dict} + return {"message": result.the_one, "usage": usage_dict, "alg": merged.alg} logger.info( "%sITS completed in %.2fs. Usage: %s (selected_index=%s)", @@ -287,18 +246,13 @@ async def arun_chat_completion( "selected_index": result.selected_index, "the_one": result.the_one, "usage": usage_dict, + "alg": merged.alg, } async def ashutdown(self) -> None: """Cleanup resources on service shutdown. - Closes all cached LM clients (releasing aiohttp sessions) then - clears the cache. + LM clients are created and closed per request, so there is nothing + cached to release here. """ - logger.info( - "ITSGateway shutting down, closing %d LM clients", - len(self._lm_cache), - ) - for lm in self._lm_cache.values(): - await lm.close() - self._lm_cache.clear() + logger.info("ITSGateway shutting down") diff --git a/its_hub/core/lms/openai_lm.py b/its_hub/core/lms/openai_lm.py index 70c83d64..409b8e06 100644 --- a/its_hub/core/lms/openai_lm.py +++ b/its_hub/core/lms/openai_lm.py @@ -27,7 +27,7 @@ class OpenAICompatibleLanguageModel(AbstractLanguageModel): def __init__( self, endpoint: str, - api_key: str, + api_key: str | None, model_name: str, system_prompt: str | None = None, is_async: bool = False, # Deprecated: parameter is ignored (always async internally) diff --git a/its_hub/integration/iaas/app.py b/its_hub/integration/iaas/app.py index 577ec485..5afefd1b 100644 --- a/its_hub/integration/iaas/app.py +++ b/its_hub/integration/iaas/app.py @@ -10,13 +10,11 @@ import time import uuid from contextlib import asynccontextmanager -from dataclasses import dataclass from fastapi import FastAPI, Header, HTTPException, status from fastapi.responses import StreamingResponse from its_hub.api.types import ITSRequestConfig -from its_hub.core.algorithms.self_consistency import SelfConsistency from its_hub.core.gateway import ITSGateway from its_hub.integration.iaas.models import ( ChatCompletionChoice, @@ -29,28 +27,14 @@ logger = logging.getLogger(__name__) -@dataclass -class _ServiceConfig: - """Mutable service-level defaults set via /configure.""" - - endpoint: str = "" - model: str = "" - api_key: str | None = None - budget: int = 4 - temperature: float | None = None - alg: str = "self-consistency" - - class _ServiceState: """Encapsulates all mutable service state (replaces module-level globals).""" def __init__(self): - self.gateway: ITSGateway = ITSGateway(algorithm=SelfConsistency()) - self.config = _ServiceConfig() + self.gateway: ITSGateway = ITSGateway() def reset(self): - self.gateway = ITSGateway(algorithm=SelfConsistency()) - self.config = _ServiceConfig() + self.gateway = ITSGateway() _state = _ServiceState() @@ -76,25 +60,18 @@ def _build_its_config( its_endpoint: str | None = None, its_api_key: str | None = None, ) -> ITSRequestConfig: - """Build ITSRequestConfig merging headers, body, and service defaults. + """Build a per-request ITSRequestConfig overlay. - Priority: header > body > service default. + Only per-request fields (header > body) are populated; service-default + fields are left ``None`` for the gateway to merge. """ - if its_budget is not None: - budget = its_budget - elif request.budget is not None: - budget = request.budget - else: - budget = _state.config.budget - - api_endpoint = its_endpoint or _state.config.endpoint - api_key = its_api_key if its_api_key is not None else _state.config.api_key - + budget = its_budget if its_budget is not None else request.budget return ITSRequestConfig( budget=budget, - api_endpoint=api_endpoint, + api_endpoint=its_endpoint, + api_key=its_api_key, model=request.model, - api_key=api_key, + temperature=request.temperature, ) @@ -102,7 +79,12 @@ def _build_its_config( async def config_service(request: ConfigRequest) -> dict[str, str]: """Configure the IaaS service with language model and scaling algorithm.""" try: - _state.gateway.configure( + config = ITSRequestConfig( + budget=request.budget, + api_endpoint=request.endpoint, + api_key=request.api_key, + model=request.model, + temperature=request.temperature, alg=request.alg, regex_patterns=request.regex_patterns, tool_vote=request.tool_vote, @@ -110,20 +92,14 @@ async def config_service(request: ConfigRequest) -> dict[str, str]: threshold=request.threshold, confidence_threshold=request.confidence_threshold, ) - _state.config.endpoint = request.endpoint - _state.config.model = request.model - _state.config.api_key = request.api_key - _state.config.alg = request.alg - if request.budget is not None: - _state.config.budget = request.budget - if request.temperature is not None: - _state.config.temperature = request.temperature + _state.gateway.configure(config) + resolved = _state.gateway.default_config logger.info( "Configured IaaS: model=%s, alg=%s, budget=%s", - request.model, - request.alg, - _state.config.budget, + resolved.model, + resolved.alg, + resolved.budget, ) return { "status": "success", @@ -145,11 +121,12 @@ async def config_service(request: ConfigRequest) -> dict[str, str]: @app.get("/v1/models") async def list_models() -> dict[str, list[dict[str, str]]]: """List available models (OpenAI-compatible endpoint).""" - if _state.config.model: + model = _state.gateway.default_config.model + if model: return { "data": [ { - "id": _state.config.model, + "id": model, "object": "model", "owned_by": "its_hub", } @@ -202,7 +179,7 @@ async def chat_completions( metadata = None if not request.return_response_only: metadata = { - "algorithm": _state.config.alg, + "algorithm": result["alg"], "all_responses": result.get("responses"), "response_counts": result.get("response_counts"), "selected_index": result.get("selected_index"), @@ -253,7 +230,7 @@ async def _generate(): its_config = _build_its_config(request, its_budget, its_endpoint, its_api_key) - if not its_config.api_endpoint: + if not (its_config.api_endpoint or _state.gateway.default_config.api_endpoint): yield f"data: {json.dumps({'error': 'Service not configured'})}\n\n" yield "data: [DONE]\n\n" return diff --git a/its_hub/integration/iaas/models.py b/its_hub/integration/iaas/models.py index 1d453792..d5201391 100644 --- a/its_hub/integration/iaas/models.py +++ b/its_hub/integration/iaas/models.py @@ -4,21 +4,16 @@ from pydantic import BaseModel, Field, field_validator, model_validator -from its_hub.api.types import ChatMessage -from its_hub.core.gateway import SUPPORTED_ALGORITHMS +from its_hub.api.types import SUPPORTED_ALGORITHMS, ChatMessage class ConfigRequest(BaseModel): """Configuration request for setting up the IaaS service.""" - provider: str = Field("openai", description="LM provider: 'openai'") endpoint: str = Field(..., description="Language model endpoint URL") api_key: str | None = Field(None, description="API key for the language model") model: str = Field(..., description="Model name identifier") alg: str = Field(..., description="Scaling algorithm to use") - extra_args: dict[str, Any] | None = Field( - None, description="Additional provider-specific arguments" - ) regex_patterns: list[str] | None = Field( None, description="Regex patterns for self-consistency projection function" ) @@ -66,8 +61,8 @@ def validate_algorithm(cls, v): @model_validator(mode="after") def validate_config_requirements(self): - if self.provider == "openai" and not self.api_key: - raise ValueError("api_key is required when using openai provider") + if not self.api_key: + raise ValueError("api_key is required") return self diff --git a/tests/conftest.py b/tests/conftest.py index ca40e326..80fab997 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,11 +4,12 @@ import socket import threading import time -from http.server import BaseHTTPRequestHandler, HTTPServer +from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer import pytest from its_hub import AbstractLanguageModel, AbstractOutcomeRewardModel +from tests.mocks.recording_llm import RecordingLLMHandler def find_free_port() -> int: @@ -351,3 +352,19 @@ def mock_process_reward_model(): "ERROR_TRIGGER": "trigger_error", "VLLM_ERROR_TRIGGER": "error", } + + +@pytest.fixture +def llm_server(): + """Controllable upstream LLM server; see RecordingLLMHandler.""" + RecordingLLMHandler.reset() + port = find_free_port() + server = ThreadingHTTPServer(("localhost", port), RecordingLLMHandler) + thread = threading.Thread(target=server.serve_forever) + thread.daemon = True + thread.start() + time.sleep(0.1) + yield f"http://localhost:{port}" + server.shutdown() + thread.join() + RecordingLLMHandler.reset() diff --git a/tests/mocks/recording_llm.py b/tests/mocks/recording_llm.py new file mode 100644 index 00000000..420d88b2 --- /dev/null +++ b/tests/mocks/recording_llm.py @@ -0,0 +1,78 @@ +"""Controllable upstream LLM server for concurrency / regression tests. + +Lives outside conftest.py: pytest loads conftest.py as plugin module ``conftest`` +(distinct from ``tests.conftest``), so a class defined there would be duplicated. +""" + +import json +import threading +from http.server import BaseHTTPRequestHandler +from typing import ClassVar + + +class RecordingLLMHandler(BaseHTTPRequestHandler): + """Upstream LLM stand-in that records requests and can hold them.""" + + received_bodies: ClassVar[list[dict]] = [] + _hold: ClassVar[threading.Event] = threading.Event() + _hold.set() + _lock: ClassVar[threading.Lock] = threading.Lock() + + @classmethod + def reset(cls) -> None: + with cls._lock: + cls.received_bodies = [] + cls._hold.set() + + @classmethod + def hold(cls) -> None: + cls._hold.clear() + + @classmethod + def release(cls) -> None: + cls._hold.set() + + def do_POST(self) -> None: + if self.path != "/v1/chat/completions": + self.send_response(404) + self.end_headers() + return + + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + with type(self)._lock: + type(self).received_bodies.append(body) + + type(self)._hold.wait(timeout=10) + + payload = json.dumps( + { + "model": body.get("model", "unknown"), + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": f"answer from {body.get('model', 'unknown')}", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 10, + "total_tokens": 20, + }, + } + ).encode() + + try: + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + except (BrokenPipeError, ConnectionResetError): + pass + + def log_message(self, *args) -> None: + pass diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 4d334e00..11708b2f 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -1,16 +1,17 @@ """Tests for ITSGateway (core/gateway.py).""" +import asyncio from collections import Counter from unittest.mock import AsyncMock, MagicMock, patch import pytest -from its_hub.api.types import GenerationUsage, ITSRequestConfig -from its_hub.core.gateway import SUPPORTED_ALGORITHMS, ITSGateway +from its_hub.api.types import SUPPORTED_ALGORITHMS, GenerationUsage, ITSRequestConfig +from its_hub.core.gateway import ITSGateway +from tests.mocks.recording_llm import RecordingLLMHandler def _make_result(usage=None): - """Build a mock SelfConsistencyResult-like object.""" result = MagicMock() result.the_one = {"role": "assistant", "content": "answer"} result.responses = [ @@ -26,6 +27,7 @@ def _make_result(usage=None): def _make_config(**overrides): defaults = { "budget": 3, + "alg": "self-consistency", "api_endpoint": "http://llm/v1", "model": "gpt-4", "api_key": "sk-test", @@ -37,96 +39,118 @@ def _make_config(**overrides): MESSAGES = [{"role": "user", "content": "What is 2+2?"}] +def _mock_lm(): + lm = MagicMock() + lm.close = AsyncMock() + return lm + + +def _patch_build_lm(gw): + return patch.object(gw, "_build_lm", return_value=_mock_lm()) + + +def _algo(gw): + return gw._build_algorithm(gw.default_config) + + class TestGatewayConstruction: def test_default_init(self): - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel"): - gw = ITSGateway() - assert gw._algorithm is not None + gw = ITSGateway() assert gw._orchestrator is not None - assert gw._algorithm_name == "SelfConsistency" + assert gw.default_config.alg == "self-consistency" - def test_custom_algorithm(self): - algo = MagicMock() - type(algo).__name__ = "CustomAlgo" - gw = ITSGateway(algorithm=algo) - assert gw._algorithm is algo - assert gw._algorithm_name == "CustomAlgo" + def test_custom_default_config(self): + config = ITSRequestConfig( + alg="beta-self-consistency", + api_endpoint="http://x/v1", + model="m", + api_key="k", + ) + gw = ITSGateway(default_config=config) + assert gw.default_config is config + assert gw.default_config.alg == "beta-self-consistency" - def test_custom_orchestrator_passed_to_default_algorithm(self): + def test_custom_orchestrator_passed_to_algorithm(self): orch = MagicMock() - with patch("its_hub.core.gateway.SelfConsistency") as sc_cls: - ITSGateway(orchestrator=orch) - sc_cls.assert_called_once_with(orchestrator=orch) + gw = ITSGateway(orchestrator=orch) + assert gw._orchestrator is orch + algo = gw._build_algorithm(_make_config()) + assert algo.orchestrator is orch -class TestLMClientCaching: +class TestLMClientLifecycle: @pytest.mark.asyncio - async def test_creates_new_client(self): - gw = ITSGateway(algorithm=MagicMock()) - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - lm_cls.return_value = MagicMock() - lm = await gw._get_or_create_lm("http://a/v1", "m1", "key1") - assert lm is lm_cls.return_value - lm_cls.assert_called_once() - - @pytest.mark.asyncio - async def test_reuses_cached_client(self): - gw = ITSGateway(algorithm=MagicMock()) - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - lm_cls.return_value = MagicMock() - lm1 = await gw._get_or_create_lm("http://a/v1", "m1", "key1") - lm2 = await gw._get_or_create_lm("http://a/v1", "m1", "key1") - assert lm1 is lm2 - assert lm_cls.call_count == 1 + async def test_creates_new_lm_per_request(self): + gw = ITSGateway() + algo = MagicMock() + algo.ainfer = AsyncMock(return_value=_make_result()) + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + patch.object(gw, "_build_lm", return_value=_mock_lm()) as build_lm, + ): + await gw.arun_chat_completion(_make_config(), MESSAGES) + await gw.arun_chat_completion(_make_config(), MESSAGES) + assert build_lm.call_count == 2 @pytest.mark.asyncio - async def test_different_endpoint_creates_new_client(self): - gw = ITSGateway(algorithm=MagicMock()) - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - lm_cls.return_value = MagicMock() - await gw._get_or_create_lm("http://a/v1", "m1", "key1") - lm_cls.return_value = MagicMock() - await gw._get_or_create_lm("http://b/v1", "m1", "key1") - assert lm_cls.call_count == 2 + async def test_closes_lm_after_request(self): + gw = ITSGateway() + algo = MagicMock() + algo.ainfer = AsyncMock(return_value=_make_result()) + lm = _mock_lm() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + patch.object(gw, "_build_lm", return_value=lm), + ): + await gw.arun_chat_completion(_make_config(), MESSAGES) + lm.close.assert_awaited_once() @pytest.mark.asyncio - async def test_different_api_key_creates_new_client(self): - gw = ITSGateway(algorithm=MagicMock()) - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - lm_cls.return_value = MagicMock() - await gw._get_or_create_lm("http://a/v1", "m1", "key-A") - lm_cls.return_value = MagicMock() - await gw._get_or_create_lm("http://a/v1", "m1", "key-B") - assert lm_cls.call_count == 2 + async def test_closes_lm_even_when_algorithm_raises(self): + gw = ITSGateway() + algo = MagicMock() + algo.ainfer = AsyncMock(side_effect=RuntimeError("boom")) + lm = _mock_lm() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + patch.object(gw, "_build_lm", return_value=lm), + pytest.raises(RuntimeError, match="boom"), + ): + await gw.arun_chat_completion(_make_config(), MESSAGES) + lm.close.assert_awaited_once() @pytest.mark.asyncio - async def test_ashutdown_closes_clients(self): - gw = ITSGateway(algorithm=MagicMock()) - mock_lm = MagicMock() - mock_lm.close = AsyncMock() - with patch( - "its_hub.core.gateway.OpenAICompatibleLanguageModel", return_value=mock_lm - ): - await gw._get_or_create_lm("http://a/v1", "m1", "key1") - assert len(gw._lm_cache) == 1 + async def test_ashutdown_is_noop(self): + gw = ITSGateway() await gw.ashutdown() - mock_lm.close.assert_awaited_once() - assert len(gw._lm_cache) == 0 @pytest.mark.asyncio - async def test_evicts_oldest_when_cache_full(self): - gw = ITSGateway(algorithm=MagicMock(), max_lm_cache_size=2) - mocks = [] - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - for i in range(3): - mock_lm = MagicMock() - mock_lm.close = AsyncMock() - lm_cls.return_value = mock_lm - mocks.append(mock_lm) - await gw._get_or_create_lm(f"http://host{i}/v1", "m1", "key1") + async def test_concurrent_different_models_do_not_interfere(self, llm_server): + gw = ITSGateway() + # On main the gateway has an LM cache; shrink it to 1 so the second + # request evicts the first's client. No-op on the fixed branch. + if hasattr(gw, "_max_lm_cache_size"): + gw._max_lm_cache_size = 1 - assert len(gw._lm_cache) == 2 - mocks[0].close.assert_awaited_once() + RecordingLLMHandler.hold() + + config_a = _make_config(api_endpoint=f"{llm_server}/v1", model="model-A") + config_b = _make_config(api_endpoint=f"{llm_server}/v1", model="model-B") + + task_a = asyncio.create_task(gw.arun_chat_completion(config_a, MESSAGES)) + await asyncio.sleep(0.3) + + task_b = asyncio.create_task(gw.arun_chat_completion(config_b, MESSAGES)) + await asyncio.sleep(0.3) + + RecordingLLMHandler.release() + + result_a = await task_a + result_b = await task_b + await gw.ashutdown() + + assert result_a["message"]["content"] == "answer from model-A" + assert result_b["message"]["content"] == "answer from model-B" class TestRunChatCompletion: @@ -135,8 +159,11 @@ async def test_response_only(self): usage = GenerationUsage(prompt_tokens=10, completion_tokens=20, num_calls=3) algo = MagicMock() algo.ainfer = AsyncMock(return_value=_make_result(usage)) - gw = ITSGateway(algorithm=algo) - with patch.object(gw, "_get_or_create_lm", return_value=MagicMock()): + gw = ITSGateway() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), + ): result = await gw.arun_chat_completion(_make_config(), MESSAGES) assert result["message"] == {"role": "assistant", "content": "answer"} assert result["usage"]["prompt_tokens"] == 10 @@ -149,8 +176,11 @@ async def test_full_result(self): usage = GenerationUsage(prompt_tokens=5, completion_tokens=10, num_calls=2) algo = MagicMock() algo.ainfer = AsyncMock(return_value=_make_result(usage)) - gw = ITSGateway(algorithm=algo) - with patch.object(gw, "_get_or_create_lm", return_value=MagicMock()): + gw = ITSGateway() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), + ): result = await gw.arun_chat_completion( _make_config(), MESSAGES, return_response_only=False ) @@ -164,9 +194,12 @@ async def test_full_result(self): async def test_tools_forwarded(self): algo = MagicMock() algo.ainfer = AsyncMock(return_value=_make_result()) - gw = ITSGateway(algorithm=algo) + gw = ITSGateway() tools = [{"type": "function", "function": {"name": "f"}}] - with patch.object(gw, "_get_or_create_lm", return_value=MagicMock()): + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), + ): await gw.arun_chat_completion( _make_config(), MESSAGES, tools=tools, tool_choice="auto" ) @@ -176,7 +209,7 @@ async def test_tools_forwarded(self): @pytest.mark.asyncio async def test_missing_model_raises(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() config = _make_config(model=None) with pytest.raises(ValueError, match="Model must be specified"): await gw.arun_chat_completion(config, MESSAGES) @@ -185,9 +218,10 @@ async def test_missing_model_raises(self): async def test_algorithm_exception_propagates(self): algo = MagicMock() algo.ainfer = AsyncMock(side_effect=RuntimeError("boom")) - gw = ITSGateway(algorithm=algo) + gw = ITSGateway() with ( - patch.object(gw, "_get_or_create_lm", return_value=MagicMock()), + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), pytest.raises(RuntimeError, match="boom"), ): await gw.arun_chat_completion(_make_config(), MESSAGES) @@ -196,141 +230,182 @@ async def test_algorithm_exception_propagates(self): async def test_usage_empty_when_none(self): algo = MagicMock() algo.ainfer = AsyncMock(return_value=_make_result(usage=None)) - gw = ITSGateway(algorithm=algo) - with patch.object(gw, "_get_or_create_lm", return_value=MagicMock()): + gw = ITSGateway() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), + ): result = await gw.arun_chat_completion(_make_config(), MESSAGES) assert result["usage"] == {} + @pytest.mark.asyncio + async def test_temperature_forwarded_to_lm(self, llm_server): + gw = ITSGateway() + config = _make_config(api_endpoint=f"{llm_server}/v1", temperature=0.7) + await gw.arun_chat_completion(config, MESSAGES) + await gw.ashutdown() + assert RecordingLLMHandler.received_bodies[-1].get("temperature") == 0.7 -class TestHashApiKey: - def test_same_key_same_hash(self): - assert ITSGateway._hash_api_key("sk-abc") == ITSGateway._hash_api_key("sk-abc") + @pytest.mark.asyncio + async def test_reconfigure_does_not_swap_in_flight_algorithm(self, llm_server): + """Reconfigure mid-request must not swap the algorithm used by the + in-flight request.""" + gw = ITSGateway() + base = ITSRequestConfig( + api_endpoint=f"{llm_server}/v1", + model="m", + api_key="k", + budget=4, + alg="self-consistency", + ) + gw.configure(base) - def test_different_keys_different_hash(self): - assert ITSGateway._hash_api_key("sk-abc") != ITSGateway._hash_api_key("sk-xyz") + RecordingLLMHandler.hold() - def test_none_key(self): - h = ITSGateway._hash_api_key(None) - assert isinstance(h, str) - assert len(h) == 16 + task = asyncio.create_task(gw.arun_chat_completion(base, MESSAGES)) + await asyncio.sleep(0.12) - def test_hash_length(self): - assert len(ITSGateway._hash_api_key("sk-test")) == 16 + gw.configure( + ITSRequestConfig( + api_endpoint=f"{llm_server}/v1", + model="m", + api_key="k", + budget=4, + alg="beta-self-consistency", + confidence_threshold=0.51, + ) + ) + + RecordingLLMHandler.release() + result = await task + await gw.ashutdown() + + assert result["alg"] == "self-consistency" + assert gw.default_config.alg == "beta-self-consistency" class TestConfigure: def test_configure_self_consistency(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], + _make_config(alg="self-consistency", regex_patterns=[r"\\boxed{([^}]+)}"]), ) - assert gw._algorithm_name == "SelfConsistency" + assert gw.default_config.alg == "self-consistency" + assert type(_algo(gw)).__name__ == "SelfConsistency" def test_configure_defaults_tool_vote_when_omitted(self): - """With neither regex nor tool_vote, the algorithm's DEFAULT_TOOL_VOTE applies.""" - gw = ITSGateway(algorithm=MagicMock()) - gw.configure(alg="self-consistency") - assert gw._algorithm.tool_vote == "tool_hierarchical" + gw = ITSGateway() + gw.configure(_make_config(alg="self-consistency")) + assert _algo(gw).tool_vote == "tool_hierarchical" def test_configure_with_tool_vote(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - tool_vote="tool_name", - exclude_tool_args=["timestamp"], + _make_config( + alg="self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + tool_vote="tool_name", + exclude_tool_args=["timestamp"], + ), ) - assert gw._algorithm_name == "SelfConsistency" + algo = _algo(gw) + assert algo.tool_vote == "tool_name" + assert algo.exclude_args == ["timestamp"] def test_configure_unsupported_algorithm(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() with pytest.raises(ValueError, match="not supported"): - gw.configure(alg="beam-search") + gw.configure(_make_config(alg="beam-search")) def test_configure_invalid_regex(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() with pytest.raises(ValueError, match="Invalid regex pattern"): - gw.configure(alg="self-consistency", regex_patterns=["[invalid("]) + gw.configure( + _make_config(alg="self-consistency", regex_patterns=["[invalid("]), + ) def test_configure_preserves_orchestrator(self): orch = MagicMock() - gw = ITSGateway(algorithm=MagicMock(), orchestrator=orch) - with patch("its_hub.core.gateway.SelfConsistency") as sc_cls: - sc_cls.return_value = MagicMock() - type(sc_cls.return_value).__name__ = "SelfConsistency" - gw.configure( - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - ) - sc_cls.assert_called_once() - assert sc_cls.call_args.kwargs["orchestrator"] is orch + gw = ITSGateway(orchestrator=orch) + gw.configure( + _make_config(alg="self-consistency", regex_patterns=[r"\\boxed{([^}]+)}"]), + ) + assert _algo(gw).orchestrator is orch def test_configure_invalid_tool_vote(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() with pytest.raises(ValueError, match="tool_vote must be one of"): gw.configure( - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - tool_vote="invalid", + _make_config( + alg="self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + tool_vote="invalid", + ), ) def test_configure_adaptive_self_consistency(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="adaptive-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], + _make_config( + alg="adaptive-self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + ), ) - assert gw._algorithm_name == "AdaptiveSelfConsistency" - # Unset threshold falls back to the class default. - assert gw._algorithm.threshold == pytest.approx(0.75) + algo = _algo(gw) + assert type(algo).__name__ == "AdaptiveSelfConsistency" + assert algo.threshold == pytest.approx(0.75) def test_configure_adaptive_threshold_plumbed(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="adaptive-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - threshold=0.9, + _make_config( + alg="adaptive-self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + threshold=0.9, + ), ) - assert gw._algorithm.threshold == pytest.approx(0.9) + assert _algo(gw).threshold == pytest.approx(0.9) def test_configure_beta_self_consistency(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="beta-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], + _make_config( + alg="beta-self-consistency", regex_patterns=[r"\\boxed{([^}]+)}"] + ), ) - assert gw._algorithm_name == "BetaSelfConsistency" - # Unset confidence_threshold falls back to the class default. - assert gw._algorithm.confidence_threshold == pytest.approx(0.95) + algo = _algo(gw) + assert type(algo).__name__ == "BetaSelfConsistency" + assert algo.confidence_threshold == pytest.approx(0.95) def test_configure_beta_confidence_threshold_plumbed(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="beta-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - confidence_threshold=0.8, + _make_config( + alg="beta-self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + confidence_threshold=0.8, + ), ) - assert gw._algorithm.confidence_threshold == pytest.approx(0.8) + assert _algo(gw).confidence_threshold == pytest.approx(0.8) @pytest.mark.parametrize( "alg", ["adaptive-self-consistency", "beta-self-consistency"], ) def test_configure_family_shares_tool_vote(self, alg): - """Adaptive/beta inherit the tool-vote voting surface from SelfConsistency.""" - gw = ITSGateway(algorithm=MagicMock()) - gw.configure(alg=alg, tool_vote="tool_name") - assert gw._algorithm.tool_vote == "tool_name" + gw = ITSGateway() + gw.configure(_make_config(alg=alg, tool_vote="tool_name")) + assert _algo(gw).tool_vote == "tool_name" def test_configure_preserves_orchestrator_for_family(self): orch = MagicMock() - gw = ITSGateway(algorithm=MagicMock(), orchestrator=orch) + gw = ITSGateway(orchestrator=orch) gw.configure( - alg="beta-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], + _make_config( + alg="beta-self-consistency", regex_patterns=[r"\\boxed{([^}]+)}"] + ), ) - assert gw._algorithm.orchestrator is orch + assert _algo(gw).orchestrator is orch class TestSupportedAlgorithms: @@ -350,7 +425,7 @@ def test_is_frozenset(self): class TestEndpointValidation: @pytest.mark.asyncio async def test_missing_endpoint_raises(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() config = _make_config(api_endpoint="") with pytest.raises(ValueError, match="api_endpoint"): await gw.arun_chat_completion(config, MESSAGES) diff --git a/tests/test_iaas.py b/tests/test_iaas.py index 124a6763..83312435 100644 --- a/tests/test_iaas.py +++ b/tests/test_iaas.py @@ -1,19 +1,24 @@ """Tests for the Inference-as-a-Service (IaaS) integration.""" +import asyncio import logging from collections import Counter from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import pytest_asyncio from fastapi.testclient import TestClient +from httpx import ASGITransport -from its_hub.api.types import ChatMessage +from its_hub.api.types import ChatMessage, ITSRequestConfig from its_hub.integration.iaas.app import _state, app from its_hub.integration.iaas.models import ( ChatCompletionRequest, ConfigRequest, ) from tests.conftest import TEST_CONSTANTS +from tests.mocks.recording_llm import RecordingLLMHandler from tests.mocks.test_data import TestDataFactory @@ -25,12 +30,42 @@ def iaas_client(): _state.reset() +@pytest_asyncio.fixture +async def async_iaas_client(): + """Async client (allows concurrent /configure + /v1/...).""" + _state.reset() + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + yield client + _state.reset() + + @pytest.fixture(scope="session") def vllm_endpoint(vllm_server): - """Alias the vllm_server fixture for clarity in IaaS tests.""" return vllm_server +def _install_mock_gateway(mock_return=None, side_effect=None, default=None): + """Replace _state.gateway with a mock, preserving default_config for merging.""" + if default is None: + default = _state.gateway.default_config + mock_gw = MagicMock() + mock_gw.default_config = default + if side_effect: + mock_gw.arun_chat_completion = AsyncMock(side_effect=side_effect) + else: + # Mirror the real gateway: inject the resolved alg from the merged + # default so metadata tests see the configured value, not a hardcoded one. + patched = {**(mock_return or {}), "alg": default.alg} + + async def _arun(*args, **kwargs): + return patched + + mock_gw.arun_chat_completion = AsyncMock(side_effect=_arun) + _state.gateway = mock_gw + return mock_gw + + def _mock_gateway_result(content="answer", usage=None): """Build a dict matching ITSGateway.arun_chat_completion return (response_only=True).""" if usage is None: @@ -135,8 +170,10 @@ def test_configure_creates_gateway(self, iaas_client, vllm_endpoint): response = iaas_client.post("/configure", json=config) assert response.status_code == 200 assert _state.gateway is not None - assert _state.config.model == TEST_CONSTANTS["DEFAULT_MODEL_NAME"] - assert _state.config.endpoint == vllm_endpoint + assert ( + _state.gateway.default_config.model == TEST_CONSTANTS["DEFAULT_MODEL_NAME"] + ) + assert _state.gateway.default_config.api_endpoint == vllm_endpoint def test_models_endpoint_after_configure(self, iaas_client, vllm_endpoint): config = { @@ -226,7 +263,7 @@ def test_tool_vote_algorithm_usage_verification(self, iaas_client, vllm_endpoint } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - + _state.gateway._build_algorithm(_state.gateway.default_config) mock_sc.assert_called_once() call_args = mock_sc.call_args assert call_args.kwargs["tool_vote"] == "tool_hierarchical" @@ -251,7 +288,7 @@ def test_family_basic_configuration(self, iaas_client, vllm_endpoint, alg): response = iaas_client.post("/configure", json=config) assert response.status_code == 200 assert "success" in response.json()["status"] - assert _state.config.alg == alg + assert _state.gateway.default_config.alg == alg def test_adaptive_threshold_forwarded_to_algorithm( self, iaas_client, vllm_endpoint @@ -268,6 +305,7 @@ def test_adaptive_threshold_forwarded_to_algorithm( } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 + _state.gateway._build_algorithm(_state.gateway.default_config) mock_alg.assert_called_once() assert mock_alg.call_args.kwargs["threshold"] == 0.9 @@ -286,6 +324,7 @@ def test_beta_confidence_threshold_forwarded_to_algorithm( } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 + _state.gateway._build_algorithm(_state.gateway.default_config) mock_alg.assert_called_once() assert mock_alg.call_args.kwargs["confidence_threshold"] == 0.8 @@ -302,6 +341,7 @@ def test_threshold_omitted_uses_algorithm_default(self, iaas_client, vllm_endpoi } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 + _state.gateway._build_algorithm(_state.gateway.default_config) assert "threshold" not in mock_alg.call_args.kwargs @pytest.mark.parametrize( @@ -338,6 +378,7 @@ def test_family_tool_vote_forwarded(self, iaas_client, vllm_endpoint): } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 + _state.gateway._build_algorithm(_state.gateway.default_config) assert mock_alg.call_args.kwargs["tool_vote"] == "tool_hierarchical" assert mock_alg.call_args.kwargs["exclude_args"] == ["timestamp"] @@ -353,11 +394,9 @@ def test_chat_completion_with_gateway(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock( - return_value=_mock_gateway_result("Tool voting response") + mock_gw = _install_mock_gateway( + mock_return=_mock_gateway_result("Tool voting response") ) - _state.gateway = mock_gw request_data = TestDataFactory.create_chat_completion_request( user_content="What is 2+2?", budget=8 @@ -380,11 +419,7 @@ def test_chat_completion_full_result(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock( - return_value=_mock_gateway_full_result() - ) - _state.gateway = mock_gw + _install_mock_gateway(mock_return=_mock_gateway_full_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) request_data["return_response_only"] = False @@ -409,11 +444,7 @@ def test_chat_completion_metadata_reports_configured_algorithm( } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock( - return_value=_mock_gateway_full_result() - ) - _state.gateway = mock_gw + _install_mock_gateway(mock_return=_mock_gateway_full_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) request_data["return_response_only"] = False @@ -421,6 +452,44 @@ def test_chat_completion_metadata_reports_configured_algorithm( assert response.status_code == 200 assert response.json()["metadata"]["algorithm"] == "beta-self-consistency" + @pytest.mark.asyncio + async def test_metadata_reports_request_time_algorithm( + self, llm_server, async_iaas_client + ): + """A reconfigure during a request must not change the reported algorithm.""" + base = { + "endpoint": f"{llm_server}/v1", + "api_key": "test-key", + "model": "test-model", + "budget": 1, + } + await async_iaas_client.post( + "/configure", json={**base, "alg": "self-consistency"} + ) + + RecordingLLMHandler.hold() + + request_task = asyncio.create_task( + async_iaas_client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "return_response_only": False, + }, + ) + ) + await asyncio.sleep(0.3) + + await async_iaas_client.post( + "/configure", json={**base, "alg": "beta-self-consistency"} + ) + + RecordingLLMHandler.release() + resp = await request_task + assert resp.status_code == 200 + assert resp.json()["metadata"]["algorithm"] == "self-consistency" + @pytest.mark.parametrize( "invalid_request", [ @@ -457,16 +526,15 @@ def _parse_sse(self, response): def _configure_and_mock( self, iaas_client, endpoint, mock_return=None, side_effect=None ): - _state.config.endpoint = endpoint - _state.config.model = TEST_CONSTANTS["DEFAULT_MODEL_NAME"] - _state.config.api_key = TEST_CONSTANTS["DEFAULT_API_KEY"] - mock_gw = MagicMock() - if side_effect: - mock_gw.arun_chat_completion = AsyncMock(side_effect=side_effect) - else: - mock_gw.arun_chat_completion = AsyncMock(return_value=mock_return) - _state.gateway = mock_gw - return mock_gw + return _install_mock_gateway( + mock_return=mock_return, + side_effect=side_effect, + default=ITSRequestConfig( + api_endpoint=endpoint, + model=TEST_CONSTANTS["DEFAULT_MODEL_NAME"], + api_key=TEST_CONSTANTS["DEFAULT_API_KEY"], + ), + ) def test_stream_content_response(self, iaas_client, vllm_endpoint): self._configure_and_mock( @@ -643,9 +711,7 @@ def test_budget_from_header_overrides_body(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + mock_gw = _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -670,9 +736,7 @@ def test_endpoint_from_header_overrides_service_default( } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + mock_gw = _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -695,9 +759,7 @@ def test_api_key_from_header(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + mock_gw = _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -712,9 +774,7 @@ def test_api_key_from_header(self, iaas_client, vllm_endpoint): def test_headers_without_service_config(self, iaas_client): """Headers alone can configure a request — no /configure needed.""" - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + mock_gw = _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -773,9 +833,7 @@ def test_negative_budget_header_rejected(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -796,9 +854,7 @@ def test_over_max_budget_header_rejected(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -835,11 +891,9 @@ def test_generation_error_is_sanitized(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock( + _install_mock_gateway( side_effect=RuntimeError("Connection to http://internal:8100 refused") ) - _state.gateway = mock_gw request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post("/v1/chat/completions", json=request_data) From 314bc652a63a9ff9c495b5315f0de29f8414a45f Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Wed, 19 Aug 2026 15:14:08 +0000 Subject: [PATCH 02/12] refactor: drop no-op shutdown hooks and simplify test mock - Delete ashutdown (was a no-op after LM cache removal) from AbstractGateway, ITSGateway, IaaS lifespan, and ext_proc processor - Delete processor.shutdown() and its call site in server.py - Delete the empty _lifespan context manager from IaaS app - Simplify RecordingLLMHandler: drop ClassVar annotations, log_message, and return-type hints - Drop test_ashutdown_is_noop and stray ashutdown calls in natural tests - Document config merge semantics: None means trickle down the hierarchy Signed-off-by: Harrison Stropkay --- docs/ext-proc-gateway.md | 3 ++- docs/iaas-service.md | 9 ++++++++- its_hub/api/gateway.py | 5 ----- its_hub/core/gateway.py | 8 -------- its_hub/integration/ext_proc/processor.py | 5 ----- its_hub/integration/ext_proc/server.py | 1 - its_hub/integration/iaas/app.py | 8 -------- tests/mocks/recording_llm.py | 20 ++++++++------------ tests/test_gateway.py | 8 -------- 9 files changed, 18 insertions(+), 49 deletions(-) diff --git a/docs/ext-proc-gateway.md b/docs/ext-proc-gateway.md index 89f0ce43..879d5baf 100644 --- a/docs/ext-proc-gateway.md +++ b/docs/ext-proc-gateway.md @@ -168,6 +168,7 @@ and the pass-through fallback — is identical. - **Per-request configuration.** In Approach 1, all ITS parameters travel in `X-ITS-*` headers; no prior configuration is required. In Approach 2, the upstream LLM is configured once via `POST /configure` and each request supplies only its `budget` in the request body; `X-ITS-*` headers are also accepted - as per-request overrides (header > body > `/configure` default). + as per-request overrides. A field is applied if it is not `None`; otherwise the next tier down + (header > body > `/configure` default) supplies the value. - **Ports.** The values shown are defaults and may be changed: `:8108` (Envoy), `:50051` (ext_proc gRPC), `:8109` (IaaS service), and `:8100` (the upstream LLM / `llm_upstream`). diff --git a/docs/iaas-service.md b/docs/iaas-service.md index d2cd7314..5b3f6e01 100644 --- a/docs/iaas-service.md +++ b/docs/iaas-service.md @@ -34,7 +34,8 @@ used. | Request body | `budget` | Compute budget | 2 | | `/configure` | `budget`, `endpoint`, `api_key` | Service defaults | 3 (lowest) | -Priority chain: **header > body > service default**. Headers are intended for +Priority chain: **header > body > service default**. A field is applied if it is +not `None`; otherwise the next tier down supplies the value. Headers are intended for Envoy ext_proc routing but are also accepted on the standalone IaaS endpoint. ### Algorithm Selection @@ -57,6 +58,12 @@ tool calls using the `tool_hierarchical` strategy, and text responses fall back exact-content matching. Supply `regex_patterns` to vote on extracted text answers, or set `tool_vote` to pick a different tool-voting strategy. +> **Note on clearing optional fields.** Because `None` means "use the tier below," a +> `/configure` call cannot explicitly reset `tool_vote` back to its default once set. +> To switch voting strategies, supply the new value; to revert to the built-in default +> (`tool_hierarchical`), pass `"tool_hierarchical"` explicitly. `regex_patterns` and +> `exclude_tool_args` can be cleared by passing an empty list (`[]`). + ## API Key Handling API keys can enter the system through three paths: diff --git a/its_hub/api/gateway.py b/its_hub/api/gateway.py index 0582bc8d..42de51cb 100644 --- a/its_hub/api/gateway.py +++ b/its_hub/api/gateway.py @@ -52,8 +52,3 @@ def run_chat_completion( ) -> dict[str, Any]: """Synchronous wrapper for arun_chat_completion.""" return asyncio.run(self.arun_chat_completion(config, messages, **kwargs)) - - @abstractmethod - async def ashutdown(self) -> None: - """Cleanup resources on service shutdown.""" - pass diff --git a/its_hub/core/gateway.py b/its_hub/core/gateway.py index 3a14d09a..e5b4906d 100644 --- a/its_hub/core/gateway.py +++ b/its_hub/core/gateway.py @@ -248,11 +248,3 @@ async def arun_chat_completion( "usage": usage_dict, "alg": merged.alg, } - - async def ashutdown(self) -> None: - """Cleanup resources on service shutdown. - - LM clients are created and closed per request, so there is nothing - cached to release here. - """ - logger.info("ITSGateway shutting down") diff --git a/its_hub/integration/ext_proc/processor.py b/its_hub/integration/ext_proc/processor.py index 8233f71a..964acdc9 100644 --- a/its_hub/integration/ext_proc/processor.py +++ b/its_hub/integration/ext_proc/processor.py @@ -349,8 +349,3 @@ def _parse_its_headers(self, headers: dict[str, str]) -> ITSRequestConfig | None except (ValueError, TypeError) as e: logger.error("Failed to parse ITS headers: %s", e) return None - - async def shutdown(self): - """Cleanup resources on service shutdown.""" - logger.info("External Processor shutting down") - await self.gateway.ashutdown() diff --git a/its_hub/integration/ext_proc/server.py b/its_hub/integration/ext_proc/server.py index 2f2eca88..a0716889 100644 --- a/its_hub/integration/ext_proc/server.py +++ b/its_hub/integration/ext_proc/server.py @@ -82,7 +82,6 @@ async def serve(port: int = 50051): await server.wait_for_termination() except (KeyboardInterrupt, asyncio.CancelledError): logger.info("Received shutdown signal") - await processor.shutdown() await server.stop(grace=5) diff --git a/its_hub/integration/iaas/app.py b/its_hub/integration/iaas/app.py index 5afefd1b..846153ec 100644 --- a/its_hub/integration/iaas/app.py +++ b/its_hub/integration/iaas/app.py @@ -9,7 +9,6 @@ import logging import time import uuid -from contextlib import asynccontextmanager from fastapi import FastAPI, Header, HTTPException, status from fastapi.responses import StreamingResponse @@ -40,17 +39,10 @@ def reset(self): _state = _ServiceState() -@asynccontextmanager -async def _lifespan(application: FastAPI): - yield - await _state.gateway.ashutdown() - - app = FastAPI( title="its_hub Inference-as-a-Service", description="OpenAI-compatible API for inference-time scaling algorithms", version="0.1.0-alpha", - lifespan=_lifespan, ) diff --git a/tests/mocks/recording_llm.py b/tests/mocks/recording_llm.py index 420d88b2..32e74146 100644 --- a/tests/mocks/recording_llm.py +++ b/tests/mocks/recording_llm.py @@ -7,32 +7,31 @@ import json import threading from http.server import BaseHTTPRequestHandler -from typing import ClassVar class RecordingLLMHandler(BaseHTTPRequestHandler): """Upstream LLM stand-in that records requests and can hold them.""" - received_bodies: ClassVar[list[dict]] = [] - _hold: ClassVar[threading.Event] = threading.Event() + received_bodies: list[dict] = [] # noqa: RUF012 - shared across per-request instances + _hold = threading.Event() _hold.set() - _lock: ClassVar[threading.Lock] = threading.Lock() + _lock = threading.Lock() @classmethod - def reset(cls) -> None: + def reset(cls): with cls._lock: - cls.received_bodies = [] + cls.received_bodies.clear() cls._hold.set() @classmethod - def hold(cls) -> None: + def hold(cls): cls._hold.clear() @classmethod - def release(cls) -> None: + def release(cls): cls._hold.set() - def do_POST(self) -> None: + def do_POST(self): if self.path != "/v1/chat/completions": self.send_response(404) self.end_headers() @@ -73,6 +72,3 @@ def do_POST(self) -> None: self.wfile.write(payload) except (BrokenPipeError, ConnectionResetError): pass - - def log_message(self, *args) -> None: - pass diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 11708b2f..983d94ba 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -119,11 +119,6 @@ async def test_closes_lm_even_when_algorithm_raises(self): await gw.arun_chat_completion(_make_config(), MESSAGES) lm.close.assert_awaited_once() - @pytest.mark.asyncio - async def test_ashutdown_is_noop(self): - gw = ITSGateway() - await gw.ashutdown() - @pytest.mark.asyncio async def test_concurrent_different_models_do_not_interfere(self, llm_server): gw = ITSGateway() @@ -147,7 +142,6 @@ async def test_concurrent_different_models_do_not_interfere(self, llm_server): result_a = await task_a result_b = await task_b - await gw.ashutdown() assert result_a["message"]["content"] == "answer from model-A" assert result_b["message"]["content"] == "answer from model-B" @@ -243,7 +237,6 @@ async def test_temperature_forwarded_to_lm(self, llm_server): gw = ITSGateway() config = _make_config(api_endpoint=f"{llm_server}/v1", temperature=0.7) await gw.arun_chat_completion(config, MESSAGES) - await gw.ashutdown() assert RecordingLLMHandler.received_bodies[-1].get("temperature") == 0.7 @pytest.mark.asyncio @@ -278,7 +271,6 @@ async def test_reconfigure_does_not_swap_in_flight_algorithm(self, llm_server): RecordingLLMHandler.release() result = await task - await gw.ashutdown() assert result["alg"] == "self-consistency" assert gw.default_config.alg == "beta-self-consistency" From 3f84ef96f3fee6be163794a27bfcb646ea6ed767 Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Wed, 19 Aug 2026 18:26:00 +0000 Subject: [PATCH 03/12] break up ITSRequestConfig into ITSRequestConfig and ITSRequestConfigUpdate --- its_hub/api/__init__.py | 2 + its_hub/api/gateway.py | 9 +- its_hub/api/types.py | 190 +++++++++++++++++----- its_hub/core/gateway.py | 86 ++++------ its_hub/core/lms/openai_lm.py | 2 +- its_hub/integration/ext_proc/processor.py | 8 +- its_hub/integration/iaas/app.py | 10 +- tests/test_gateway.py | 51 ++++-- tests/test_iaas.py | 4 +- 9 files changed, 235 insertions(+), 127 deletions(-) diff --git a/its_hub/api/__init__.py b/its_hub/api/__init__.py index 057f7556..b7ee20c6 100644 --- a/its_hub/api/__init__.py +++ b/its_hub/api/__init__.py @@ -31,6 +31,7 @@ ChatMessages, GenerationUsage, ITSRequestConfig, + ITSRequestConfigUpdate, ) __all__ = [ # noqa: RUF022 @@ -51,6 +52,7 @@ "ChatMessages", "GenerationUsage", "ITSRequestConfig", + "ITSRequestConfigUpdate", "SUPPORTED_ALGORITHMS", "VALID_TOOL_VOTE_OPTIONS", # Error types diff --git a/its_hub/api/gateway.py b/its_hub/api/gateway.py index 42de51cb..2ed7db66 100644 --- a/its_hub/api/gateway.py +++ b/its_hub/api/gateway.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from typing import Any -from its_hub.api.types import ITSRequestConfig +from its_hub.api.types import ITSRequestConfigUpdate class AbstractGateway(ABC): @@ -21,7 +21,7 @@ class AbstractGateway(ABC): @abstractmethod async def arun_chat_completion( self, - config: ITSRequestConfig, + config: ITSRequestConfigUpdate, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, @@ -31,7 +31,8 @@ async def arun_chat_completion( """Run chat completion with ITS algorithm asynchronously. Args: - config: Per-request ITS configuration + config: Per-request ITS configuration overlay (merged over the + gateway's service default at run time) messages: OpenAI-format conversation messages tools: Optional tool definitions for function calling tool_choice: Optional tool choice strategy @@ -46,7 +47,7 @@ async def arun_chat_completion( def run_chat_completion( self, - config: ITSRequestConfig, + config: ITSRequestConfigUpdate, messages: list[dict[str, Any]], **kwargs, ) -> dict[str, Any]: diff --git a/its_hub/api/types.py b/its_hub/api/types.py index 1dde650e..a119b074 100644 --- a/its_hub/api/types.py +++ b/its_hub/api/types.py @@ -4,7 +4,7 @@ import logging import re -from dataclasses import dataclass, fields +from dataclasses import dataclass, fields, replace from typing import Literal @@ -178,14 +178,110 @@ def total_tokens(self) -> int: ) -@dataclass +def _validate_optional_fields( + budget: int | None, + alg: str | None, + regex_patterns: list[str] | None, + tool_vote: str | None, + threshold: float | None, + confidence_threshold: float | None, +) -> None: + """Validate the optional scaling fields. Shared by ITSRequestConfig and + ITSRequestConfigUpdate so the two cannot drift. Each check is guarded so + that ``None`` ("not supplied") is always accepted.""" + if budget is not None and not (1 <= budget <= 1000): + raise ValueError("budget must be between 1 and 1000") + if alg is not None and alg not in SUPPORTED_ALGORITHMS: + raise ValueError( + f"Algorithm {alg!r} not supported. " + f"Choose from: {SUPPORTED_ALGORITHMS}" + ) + if regex_patterns is not None: + for p in regex_patterns: + try: + re.compile(p) + except re.error as e: + raise ValueError(f"Invalid regex pattern {p!r}: {e}") from e + if threshold is not None and not (0.5 < threshold <= 1.0): + raise ValueError(f"threshold must be in (0.5, 1.0], got: {threshold}") + if confidence_threshold is not None and not (0.5 < confidence_threshold <= 1.0): + raise ValueError( + f"confidence_threshold must be in (0.5, 1.0], " + f"got: {confidence_threshold}" + ) + if tool_vote is not None and tool_vote not in VALID_TOOL_VOTE_OPTIONS: + raise ValueError( + f"tool_vote must be one of {VALID_TOOL_VOTE_OPTIONS}, got: {tool_vote}" + ) + + +def _config_repr(obj: ITSRequestConfig | ITSRequestConfigUpdate) -> str: + return ( + type(obj).__name__ + + "(" + + ", ".join( + f"{f.name}={'***' if f.name == 'api_key' else getattr(obj, f.name)!r}" + for f in fields(obj) + ) + + ")" + ) + + +@dataclass(frozen=True) class ITSRequestConfig: - """Per-request configuration for ITS execution. + """A fully-resolved configuration snapshot, ready to drive one request. - Holds every knob that governs a request: the LM target - (``api_endpoint``, ``model``, ``api_key``, ``temperature``) and the scaling - parameters (``budget``, ``alg``, ``regex_patterns``, ``tool_vote``, - ``exclude_tool_args``, ``threshold``, ``confidence_threshold``). + ``api_endpoint`` and ``model`` are mandatory — constructing without them + raises. The remaining fields carry system defaults (``budget``, ``alg``, + ``api_key``) or are genuinely optional (``None`` is a meaningful "not set", + e.g. ``temperature=None`` means "don't send it upstream"). + """ + + # LM target + api_endpoint: str + model: str + api_key: str | None = None + temperature: float | None = None + # Scaling parameters + budget: int = 4 + alg: str = "self-consistency" + regex_patterns: list[str] | None = None + tool_vote: str | None = None + exclude_tool_args: list[str] | None = None + threshold: float | None = None + confidence_threshold: float | None = None + + def __post_init__(self): + if not self.api_endpoint: + raise ValueError("api_endpoint must be specified") + if not self.model: + raise ValueError("model must be specified") + _validate_optional_fields( + self.budget, + self.alg, + self.regex_patterns, + self.tool_vote, + self.threshold, + self.confidence_threshold, + ) + + def __repr__(self) -> str: + return _config_repr(self) + + +@dataclass +class ITSRequestConfigUpdate: + """A partial configuration overlay. + + Every field is ``Optional``: ``None`` means "not supplied — keep the value + below in the merge hierarchy". Because of this, ``tool_vote`` cannot be + cleared back to its default once set (pass the new value, or + ``"tool_hierarchical"`` explicitly); ``regex_patterns`` and + ``exclude_tool_args`` can be cleared by passing ``[]``. + + Use :meth:`merge` to layer one update over another, and :meth:`resolve` to + materialize a complete :class:`ITSRequestConfig` (the only point at which + mandatory-field presence is checked). """ # LM target @@ -203,40 +299,52 @@ class ITSRequestConfig: confidence_threshold: float | None = None def __post_init__(self): - if self.budget is not None and not (1 <= self.budget <= 1000): - raise ValueError("budget must be between 1 and 1000") - if self.alg is not None and self.alg not in SUPPORTED_ALGORITHMS: - raise ValueError( - f"Algorithm {self.alg!r} not supported. " - f"Choose from: {SUPPORTED_ALGORITHMS}" - ) - if self.regex_patterns is not None: - for p in self.regex_patterns: - try: - re.compile(p) - except re.error as e: - raise ValueError(f"Invalid regex pattern {p!r}: {e}") from e - if self.threshold is not None and not (0.5 < self.threshold <= 1.0): - raise ValueError(f"threshold must be in (0.5, 1.0], got: {self.threshold}") - if self.confidence_threshold is not None and not ( - 0.5 < self.confidence_threshold <= 1.0 - ): - raise ValueError( - f"confidence_threshold must be in (0.5, 1.0], " - f"got: {self.confidence_threshold}" - ) - if self.tool_vote is not None and self.tool_vote not in VALID_TOOL_VOTE_OPTIONS: - raise ValueError( - f"tool_vote must be one of {VALID_TOOL_VOTE_OPTIONS}, " - f"got: {self.tool_vote}" - ) + _validate_optional_fields( + self.budget, + self.alg, + self.regex_patterns, + self.tool_vote, + self.threshold, + self.confidence_threshold, + ) def __repr__(self) -> str: - return ( - "ITSRequestConfig(" - + ", ".join( - f"{f.name}={'***' if f.name == 'api_key' else getattr(self, f.name)!r}" - for f in fields(self) - ) - + ")" + return _config_repr(self) + + def merge(self, overlay: ITSRequestConfigUpdate) -> ITSRequestConfigUpdate: + """Return a copy of ``self`` with non-``None`` fields of ``overlay`` + applied. ``replace()`` re-fires ``__post_init__``, so the result is + format-validated.""" + return replace( + self, + **{ + f.name: getattr(overlay, f.name) + for f in fields(overlay) + if getattr(overlay, f.name) is not None + }, + ) + + def resolve(self) -> ITSRequestConfig: + """Materialize a complete :class:`ITSRequestConfig`. + + Raises ``ValueError`` if ``api_endpoint`` or ``model`` are absent; + this is the single point at which mandatory-field presence is + checked. Optional fields are forwarded only when non-``None``; + otherwise the resolved config's defaults (``budget=4``, + ``alg="self-consistency"``, ``api_key=None``) apply. + """ + if not self.api_endpoint: + raise ValueError("api_endpoint must be specified") + if not self.model: + raise ValueError("model must be specified") + optionals = { + f.name: getattr(self, f.name) + for f in fields(self) + if f.name not in {"api_endpoint", "model"} + and getattr(self, f.name) is not None + } + return ITSRequestConfig( + api_endpoint=self.api_endpoint, + model=self.model, + **optionals, ) diff --git a/its_hub/core/gateway.py b/its_hub/core/gateway.py index e5b4906d..bcb3b6fd 100644 --- a/its_hub/core/gateway.py +++ b/its_hub/core/gateway.py @@ -6,7 +6,6 @@ import logging import time -from dataclasses import fields, replace from typing import Any from its_hub.api import ( @@ -17,6 +16,7 @@ ChatMessages, GenerationUsage, ITSRequestConfig, + ITSRequestConfigUpdate, ) from its_hub.core.algorithms.adaptive_self_consistency import AdaptiveSelfConsistency from its_hub.core.algorithms.beta_self_consistency import BetaSelfConsistency @@ -42,10 +42,6 @@ # All currently supported algorithms are self-consistency variants. SUPPORTED_ALGORITHMS = SELF_CONSISTENCY_ALGORITHMS -# System defaults seeded into every gateway's default config. -_DEFAULT_BUDGET = 4 -_DEFAULT_ALG = "self-consistency" - class ITSGateway(AbstractGateway): """Long-lived gateway for running ITS algorithms. @@ -58,7 +54,7 @@ class ITSGateway(AbstractGateway): gateway = ITSGateway() # Per request - config = ITSRequestConfig( + config = ITSRequestConfigUpdate( budget=10, api_endpoint="http://envoy-cluster/v1", model="gpt-4", @@ -70,44 +66,33 @@ class ITSGateway(AbstractGateway): def __init__( self, orchestrator: AbstractOrchestrator | None = None, - default_config: ITSRequestConfig | None = None, + default_config: ITSRequestConfigUpdate | None = None, ): if orchestrator is None: orchestrator = LMOrchestrator() self._orchestrator = orchestrator - self._default_config = default_config or ITSRequestConfig( - budget=_DEFAULT_BUDGET, alg=_DEFAULT_ALG - ) + # Stored as a partial update (format-validated, never completeness- + # validated): endpoint/model may be absent until a request supplies + # them. System defaults (budget=4, alg="self-consistency") are NOT + # stored here — they live on ITSRequestConfig and are injected at + # resolve() time. + self.default_config = default_config or ITSRequestConfigUpdate() logger.info( - f"ITSGateway initialized with default alg={self._default_config.alg} and budget={self._default_config.budget}" + "ITSGateway initialized with default alg=%s and budget=%s", + self.default_config.alg, + self.default_config.budget, ) - @property - def default_config(self) -> ITSRequestConfig: - """The service-default config set via :meth:`configure`.""" - return self._default_config + def configure(self, update: ITSRequestConfigUpdate) -> None: + """Merge ``update`` into the service default. - def configure(self, config: ITSRequestConfig) -> None: - """Merge ``config`` into the service default. - - Validation (alg, regex, thresholds) runs automatically in - ``ITSRequestConfig.__post_init__`` via ``_merge``. + Format-validated (via ``ITSRequestConfigUpdate.__post_init__`` re-fired + by ``merge``); not completeness-validated — ``api_endpoint``/``model`` + may remain absent until a request supplies them. """ - self._default_config = self._merge(self._default_config, config) + self.default_config = self.default_config.merge(update) logger.info( - "ITSGateway reconfigured with default alg=%s", self._default_config.alg - ) - - @staticmethod - def _merge(base: ITSRequestConfig, overlay: ITSRequestConfig) -> ITSRequestConfig: - """Return a copy of ``base`` with non-``None`` fields from ``overlay``.""" - return replace( - base, - **{ - f.name: getattr(overlay, f.name) - for f in fields(overlay) - if getattr(overlay, f.name) is not None - }, + "ITSGateway reconfigured with default alg=%s", self.default_config.alg ) def _build_algorithm(self, config: ITSRequestConfig) -> AbstractScalingAlgorithm: @@ -164,7 +149,7 @@ def _build_lm( async def arun_chat_completion( self, - config: ITSRequestConfig, + config: ITSRequestConfigUpdate, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, @@ -174,18 +159,15 @@ async def arun_chat_completion( request_id = kwargs.get("request_id") log_prefix = f"[{request_id}] " if request_id else "" - merged = self._merge(self._default_config, config) - - if not merged.api_endpoint: - raise ValueError("api_endpoint must be specified") - if not merged.model: - raise ValueError("Model must be specified") + # Merge the per-request overlay over the service default, then resolve + # to a complete snapshot. resolve() is the single completeness check + # (raises ValueError if api_endpoint/model are still absent). Nothing + # on `self` is mutated, so a concurrent /configure cannot swap the + # algorithm or close this request's HTTP session mid-flight. + resolved = self.default_config.merge(config).resolve() - # Snapshot the algorithm + LM for this request only. Nothing on `self` - # is mutated, so a concurrent /configure cannot swap the algorithm or - # close this request's HTTP session mid-flight. - algorithm = self._build_algorithm(merged) - lm = self._build_lm(merged, request_id) + algorithm = self._build_algorithm(resolved) + lm = self._build_lm(resolved, request_id) chat_messages = ChatMessages([ChatMessage.from_dict(msg) for msg in messages]) @@ -193,9 +175,9 @@ async def arun_chat_completion( "%sRunning ITS: algorithm=%s, budget=%s, endpoint=%s, model=%s, messages=%s, tools=%s", log_prefix, type(algorithm).__name__, - merged.budget, - merged.api_endpoint, - merged.model, + resolved.budget, + resolved.api_endpoint, + resolved.model, len(messages), "yes" if tools else "no", ) @@ -205,7 +187,7 @@ async def arun_chat_completion( result = await algorithm.ainfer( lm=lm, prompt_or_messages=chat_messages, - budget=merged.budget, + budget=resolved.budget, return_response_only=False, tools=tools, tool_choice=tool_choice, @@ -231,7 +213,7 @@ async def arun_chat_completion( duration_s, usage_dict, ) - return {"message": result.the_one, "usage": usage_dict, "alg": merged.alg} + return {"message": result.the_one, "usage": usage_dict, "alg": resolved.alg} logger.info( "%sITS completed in %.2fs. Usage: %s (selected_index=%s)", @@ -246,5 +228,5 @@ async def arun_chat_completion( "selected_index": result.selected_index, "the_one": result.the_one, "usage": usage_dict, - "alg": merged.alg, + "alg": resolved.alg, } diff --git a/its_hub/core/lms/openai_lm.py b/its_hub/core/lms/openai_lm.py index 409b8e06..af6552d3 100644 --- a/its_hub/core/lms/openai_lm.py +++ b/its_hub/core/lms/openai_lm.py @@ -93,7 +93,7 @@ def __init__( # set up headers for API requests self.headers = { "Content-Type": "application/json", - "Authorization": f"Bearer {self.api_key}", + "Authorization": f"Bearer {self.api_key or ''}", } # raw response preservation diff --git a/its_hub/integration/ext_proc/processor.py b/its_hub/integration/ext_proc/processor.py index 964acdc9..a8e26cd4 100644 --- a/its_hub/integration/ext_proc/processor.py +++ b/its_hub/integration/ext_proc/processor.py @@ -18,7 +18,7 @@ from envoy.service.ext_proc.v3 import external_processor_pb2_grpc as ext_proc_grpc from envoy.type.v3 import http_status_pb2 -from its_hub.api import ChatMessage, ITSRequestConfig +from its_hub.api import ChatMessage, ITSRequestConfigUpdate from its_hub.core.gateway import ITSGateway _ITS_HEADER_PREFIX = "x-its-" @@ -185,7 +185,7 @@ async def Process( # noqa: N802 — gRPC servicer interface requires this name async def _apply_its( self, - its_config: ITSRequestConfig, + its_config: ITSRequestConfigUpdate, body: bytes, request_id: str, ) -> ext_proc_pb2.ProcessingResponse: @@ -324,7 +324,7 @@ def _headers_continue( ) ) - def _parse_its_headers(self, headers: dict[str, str]) -> ITSRequestConfig | None: + def _parse_its_headers(self, headers: dict[str, str]) -> ITSRequestConfigUpdate | None: """Parse ITS configuration from request headers. Model is NOT extracted from headers - it will be set from request body later. @@ -339,7 +339,7 @@ def _parse_its_headers(self, headers: dict[str, str]) -> ITSRequestConfig | None budget = int(budget_str) api_key = headers.get("x-its-api-key") - return ITSRequestConfig( + return ITSRequestConfigUpdate( budget=budget, api_endpoint=endpoint, api_key=api_key, diff --git a/its_hub/integration/iaas/app.py b/its_hub/integration/iaas/app.py index 846153ec..11cededb 100644 --- a/its_hub/integration/iaas/app.py +++ b/its_hub/integration/iaas/app.py @@ -13,7 +13,7 @@ from fastapi import FastAPI, Header, HTTPException, status from fastapi.responses import StreamingResponse -from its_hub.api.types import ITSRequestConfig +from its_hub.api.types import ITSRequestConfigUpdate from its_hub.core.gateway import ITSGateway from its_hub.integration.iaas.models import ( ChatCompletionChoice, @@ -51,14 +51,14 @@ def _build_its_config( its_budget: int | None = None, its_endpoint: str | None = None, its_api_key: str | None = None, -) -> ITSRequestConfig: - """Build a per-request ITSRequestConfig overlay. +) -> ITSRequestConfigUpdate: + """Build a per-request ITSRequestConfigUpdate overlay. Only per-request fields (header > body) are populated; service-default fields are left ``None`` for the gateway to merge. """ budget = its_budget if its_budget is not None else request.budget - return ITSRequestConfig( + return ITSRequestConfigUpdate( budget=budget, api_endpoint=its_endpoint, api_key=its_api_key, @@ -71,7 +71,7 @@ def _build_its_config( async def config_service(request: ConfigRequest) -> dict[str, str]: """Configure the IaaS service with language model and scaling algorithm.""" try: - config = ITSRequestConfig( + config = ITSRequestConfigUpdate( budget=request.budget, api_endpoint=request.endpoint, api_key=request.api_key, diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 983d94ba..b29e3f2a 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -6,7 +6,11 @@ import pytest -from its_hub.api.types import SUPPORTED_ALGORITHMS, GenerationUsage, ITSRequestConfig +from its_hub.api.types import ( + SUPPORTED_ALGORITHMS, + GenerationUsage, + ITSRequestConfigUpdate, +) from its_hub.core.gateway import ITSGateway from tests.mocks.recording_llm import RecordingLLMHandler @@ -33,7 +37,7 @@ def _make_config(**overrides): "api_key": "sk-test", } defaults.update(overrides) - return ITSRequestConfig(**defaults) + return ITSRequestConfigUpdate(**defaults) MESSAGES = [{"role": "user", "content": "What is 2+2?"}] @@ -50,17 +54,23 @@ def _patch_build_lm(gw): def _algo(gw): - return gw._build_algorithm(gw.default_config) + return gw._build_algorithm(gw.default_config.resolve()) class TestGatewayConstruction: def test_default_init(self): gw = ITSGateway() assert gw._orchestrator is not None - assert gw.default_config.alg == "self-consistency" + # Stored default is partial; system defaults (alg, budget) materialize + # at resolve() time on ITSRequestConfig. + resolved = gw.default_config.merge( + ITSRequestConfigUpdate(api_endpoint="http://x/v1", model="m") + ).resolve() + assert resolved.alg == "self-consistency" + assert resolved.budget == 4 def test_custom_default_config(self): - config = ITSRequestConfig( + config = ITSRequestConfigUpdate( alg="beta-self-consistency", api_endpoint="http://x/v1", model="m", @@ -74,7 +84,7 @@ def test_custom_orchestrator_passed_to_algorithm(self): orch = MagicMock() gw = ITSGateway(orchestrator=orch) assert gw._orchestrator is orch - algo = gw._build_algorithm(_make_config()) + algo = gw._build_algorithm(_make_config().resolve()) assert algo.orchestrator is orch @@ -205,7 +215,7 @@ async def test_tools_forwarded(self): async def test_missing_model_raises(self): gw = ITSGateway() config = _make_config(model=None) - with pytest.raises(ValueError, match="Model must be specified"): + with pytest.raises(ValueError, match="model must be specified"): await gw.arun_chat_completion(config, MESSAGES) @pytest.mark.asyncio @@ -244,28 +254,25 @@ async def test_reconfigure_does_not_swap_in_flight_algorithm(self, llm_server): """Reconfigure mid-request must not swap the algorithm used by the in-flight request.""" gw = ITSGateway() - base = ITSRequestConfig( + # `alg` is deliberately omitted from the per-request overlay so the + # algorithm is sourced from the (reconfigurable) service default — + # only then does a mid-flight /configure actually exercise the + # snapshot. With `alg` in the overlay, `result["alg"]` would be + # "self-consistency" regardless of any reconfigure. + base = ITSRequestConfigUpdate( api_endpoint=f"{llm_server}/v1", model="m", - api_key="k", - budget=4, - alg="self-consistency", ) gw.configure(base) RecordingLLMHandler.hold() task = asyncio.create_task(gw.arun_chat_completion(base, MESSAGES)) - await asyncio.sleep(0.12) + await asyncio.sleep(0.1) gw.configure( - ITSRequestConfig( - api_endpoint=f"{llm_server}/v1", - model="m", - api_key="k", - budget=4, + ITSRequestConfigUpdate( alg="beta-self-consistency", - confidence_threshold=0.51, ) ) @@ -304,6 +311,14 @@ def test_configure_with_tool_vote(self): assert algo.tool_vote == "tool_name" assert algo.exclude_args == ["timestamp"] + def test_configure_tool_vote_persists_when_omitted(self): + """`None` means trickle down, so a /configure omitting tool_vote + cannot clear a previously-set value.""" + gw = ITSGateway() + gw.configure(_make_config(alg="self-consistency", tool_vote="tool_name")) + gw.configure(_make_config(alg="self-consistency")) # no tool_vote + assert _algo(gw).tool_vote == "tool_name" + def test_configure_unsupported_algorithm(self): gw = ITSGateway() with pytest.raises(ValueError, match="not supported"): diff --git a/tests/test_iaas.py b/tests/test_iaas.py index 83312435..2d409d35 100644 --- a/tests/test_iaas.py +++ b/tests/test_iaas.py @@ -11,7 +11,7 @@ from fastapi.testclient import TestClient from httpx import ASGITransport -from its_hub.api.types import ChatMessage, ITSRequestConfig +from its_hub.api.types import ChatMessage, ITSRequestConfigUpdate from its_hub.integration.iaas.app import _state, app from its_hub.integration.iaas.models import ( ChatCompletionRequest, @@ -529,7 +529,7 @@ def _configure_and_mock( return _install_mock_gateway( mock_return=mock_return, side_effect=side_effect, - default=ITSRequestConfig( + default=ITSRequestConfigUpdate( api_endpoint=endpoint, model=TEST_CONSTANTS["DEFAULT_MODEL_NAME"], api_key=TEST_CONSTANTS["DEFAULT_API_KEY"], From 491e2205406f3b70a687fe41913295538c38a205 Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Wed, 19 Aug 2026 19:49:39 +0000 Subject: [PATCH 04/12] style: apply ruff format to gateway config types and ext_proc Signed-off-by: Harrison Stropkay --- its_hub/api/types.py | 6 ++---- its_hub/integration/ext_proc/processor.py | 4 +++- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/its_hub/api/types.py b/its_hub/api/types.py index a119b074..bbfa550b 100644 --- a/its_hub/api/types.py +++ b/its_hub/api/types.py @@ -193,8 +193,7 @@ def _validate_optional_fields( raise ValueError("budget must be between 1 and 1000") if alg is not None and alg not in SUPPORTED_ALGORITHMS: raise ValueError( - f"Algorithm {alg!r} not supported. " - f"Choose from: {SUPPORTED_ALGORITHMS}" + f"Algorithm {alg!r} not supported. Choose from: {SUPPORTED_ALGORITHMS}" ) if regex_patterns is not None: for p in regex_patterns: @@ -206,8 +205,7 @@ def _validate_optional_fields( raise ValueError(f"threshold must be in (0.5, 1.0], got: {threshold}") if confidence_threshold is not None and not (0.5 < confidence_threshold <= 1.0): raise ValueError( - f"confidence_threshold must be in (0.5, 1.0], " - f"got: {confidence_threshold}" + f"confidence_threshold must be in (0.5, 1.0], got: {confidence_threshold}" ) if tool_vote is not None and tool_vote not in VALID_TOOL_VOTE_OPTIONS: raise ValueError( diff --git a/its_hub/integration/ext_proc/processor.py b/its_hub/integration/ext_proc/processor.py index a8e26cd4..5ebdf587 100644 --- a/its_hub/integration/ext_proc/processor.py +++ b/its_hub/integration/ext_proc/processor.py @@ -324,7 +324,9 @@ def _headers_continue( ) ) - def _parse_its_headers(self, headers: dict[str, str]) -> ITSRequestConfigUpdate | None: + def _parse_its_headers( + self, headers: dict[str, str] + ) -> ITSRequestConfigUpdate | None: """Parse ITS configuration from request headers. Model is NOT extracted from headers - it will be set from request body later. From 1473b4abc141b29b8c29992c0a96018bb239e65c Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Thu, 20 Aug 2026 15:22:51 +0000 Subject: [PATCH 05/12] add connection cache; reduce CA bundle reads --- its_hub/core/gateway.py | 20 ++++++++++++++++++++ its_hub/core/lms/openai_lm.py | 19 +++++++++++++------ 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/its_hub/core/gateway.py b/its_hub/core/gateway.py index bcb3b6fd..89bececf 100644 --- a/its_hub/core/gateway.py +++ b/its_hub/core/gateway.py @@ -5,8 +5,13 @@ """ import logging +import ssl import time from typing import Any +from urllib.parse import urlparse + +import aiohttp +import certifi from its_hub.api import ( AbstractGateway, @@ -77,6 +82,9 @@ def __init__( # stored here — they live on ITSRequestConfig and are injected at # resolve() time. self.default_config = default_config or ITSRequestConfigUpdate() + self._ssl_context = ssl.create_default_context(cafile=certifi.where()) + # One connection pool per upstream endpoint + self._connector_pool: dict[str, aiohttp.TCPConnector] = {} logger.info( "ITSGateway initialized with default alg=%s and budget=%s", self.default_config.alg, @@ -127,6 +135,16 @@ def _build_algorithm(self, config: ITSRequestConfig) -> AbstractScalingAlgorithm else: # self-consistency return SelfConsistency(**common_kwargs) + def _get_connector(self, endpoint: str) -> aiohttp.TCPConnector: + """Get or create the pooled connector.""" + parsed = urlparse(endpoint) + origin = f"{parsed.scheme}://{parsed.netloc}" + conn = self._connector_pool.get(origin) + if conn is None: + conn = aiohttp.TCPConnector(ssl=self._ssl_context) + self._connector_pool[origin] = conn + return conn + def _build_lm( self, config: ITSRequestConfig, @@ -145,6 +163,8 @@ def _build_lm( api_key=config.api_key, model_name=config.model, temperature=config.temperature, + ssl_context=self._ssl_context, + connector=self._get_connector(config.api_endpoint), ) async def arun_chat_completion( diff --git a/its_hub/core/lms/openai_lm.py b/its_hub/core/lms/openai_lm.py index af6552d3..217eeef6 100644 --- a/its_hub/core/lms/openai_lm.py +++ b/its_hub/core/lms/openai_lm.py @@ -44,6 +44,8 @@ def __init__( ssl_context: ssl.SSLContext | None = None, # Raw response preservation include_raw_choices: bool = False, + # Borrowed connection pool (e.g. from ITSGateway); not closed by this LM + connector: aiohttp.BaseConnector | None = None, ): assert max_concurrency == -1 or max_concurrency > 0, ( "max_concurrency must be -1 (unlimited concurrency) or a positive integer" @@ -86,8 +88,6 @@ def __init__( self.ssl_context.check_hostname = False self.ssl_context.verify_mode = ssl.CERT_NONE else: - # For async requests, create SSL context using the same CA bundle as requests - # This ensures aiohttp uses the same certificates as requests library self.ssl_context = ssl.create_default_context(cafile=certifi.where()) # set up headers for API requests @@ -104,6 +104,7 @@ def __init__( # Session cache: one session per event loop, entry is auto-cleaned via weak references. # Session(s) need to be closed via close or close_session function calls. + self._borrowed_connector = connector self._sessions: weakref.WeakKeyDictionary[ asyncio.AbstractEventLoop, aiohttp.ClientSession ] = weakref.WeakKeyDictionary() @@ -124,10 +125,16 @@ def _get_session(self, loop: asyncio.AbstractEventLoop) -> aiohttp.ClientSession if session is not None and not session.closed: return session - # Create new session for the event loop. - # trust_env=True so aiohttp honors HTTP(S)_PROXY/NO_PROXY/.netrc. - connector = aiohttp.TCPConnector(ssl=self.ssl_context) - session = aiohttp.ClientSession(connector=connector, trust_env=True) + if self._borrowed_connector is not None: + # Borrowed pool: the session uses it but must not close it. + session = aiohttp.ClientSession( + connector=self._borrowed_connector, + connector_owner=False, + trust_env=True, + ) + else: + connector = aiohttp.TCPConnector(ssl=self.ssl_context) + session = aiohttp.ClientSession(connector=connector, trust_env=True) self._sessions[loop] = session return session From 4f796088423dc117dcfcb75342301385e574adbd Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Thu, 20 Aug 2026 15:32:25 +0000 Subject: [PATCH 06/12] address review: api_key optional, auth header, dead constants, docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - iaas/models: drop the `api_key is required` model validator so a no-auth local endpoint can be configured via POST /configure; add regression tests (model-level ConfigRequest + HTTP-level /configure) for the no-api_key path - openai_lm: omit the Authorization header entirely when no api_key is set, instead of sending an empty Bearer token - gateway: remove the dead duplicate SELF_CONSISTENCY_ALGORITHMS / SUPPORTED_ALGORITHMS (the shared constants in its_hub.api.types are the ones actually consumed); simplify the init log - docs/iaas-service: correct the API-key lifetime — header keys are request-scoped, /configure keys persist as service defaults Signed-off-by: Harrison Stropkay --- docs/iaas-service.md | 6 ++++-- its_hub/core/gateway.py | 19 +------------------ its_hub/core/lms/openai_lm.py | 10 +++++----- its_hub/integration/iaas/models.py | 8 +------- tests/test_iaas.py | 29 +++++++++++++++++++++-------- 5 files changed, 32 insertions(+), 40 deletions(-) diff --git a/docs/iaas-service.md b/docs/iaas-service.md index 5b3f6e01..0abd0e31 100644 --- a/docs/iaas-service.md +++ b/docs/iaas-service.md @@ -76,8 +76,10 @@ API keys can enter the system through three paths: **Security properties:** - Keys are **never logged** — the gateway logs endpoint and model but not credentials -- Keys are **not persisted** to disk — they exist only in memory for the lifetime of the - request (an LM client is created per request and closed when it completes) +- Keys are **not persisted** to disk. A key supplied via the `X-ITS-API-Key` header is + request-scoped: it lives only for that request's LM client, which is closed when the + request completes. A key set via `/configure` remains in memory as a service default + (`ITSGateway.default_config`) until replaced or the process restarts. - Keys are **never shared between requests** — each request builds its own LM client, so credentials supplied via header or `/configure` cannot cross-contaminate diff --git a/its_hub/core/gateway.py b/its_hub/core/gateway.py index 89bececf..33fe443c 100644 --- a/its_hub/core/gateway.py +++ b/its_hub/core/gateway.py @@ -34,19 +34,6 @@ logger = logging.getLogger(__name__) -# The self-consistency family shares the same voting surface (regex/tool-vote -# projection) and only differs in how it decides when to stop sampling. -SELF_CONSISTENCY_ALGORITHMS = frozenset( - { - "self-consistency", - "adaptive-self-consistency", - "beta-self-consistency", - } -) - -# All currently supported algorithms are self-consistency variants. -SUPPORTED_ALGORITHMS = SELF_CONSISTENCY_ALGORITHMS - class ITSGateway(AbstractGateway): """Long-lived gateway for running ITS algorithms. @@ -85,11 +72,7 @@ def __init__( self._ssl_context = ssl.create_default_context(cafile=certifi.where()) # One connection pool per upstream endpoint self._connector_pool: dict[str, aiohttp.TCPConnector] = {} - logger.info( - "ITSGateway initialized with default alg=%s and budget=%s", - self.default_config.alg, - self.default_config.budget, - ) + logger.info("ITSGateway initialized") def configure(self, update: ITSRequestConfigUpdate) -> None: """Merge ``update`` into the service default. diff --git a/its_hub/core/lms/openai_lm.py b/its_hub/core/lms/openai_lm.py index 217eeef6..5125b9b0 100644 --- a/its_hub/core/lms/openai_lm.py +++ b/its_hub/core/lms/openai_lm.py @@ -90,11 +90,11 @@ def __init__( else: self.ssl_context = ssl.create_default_context(cafile=certifi.where()) - # set up headers for API requests - self.headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {self.api_key or ''}", - } + # set up headers for API requests; omit Authorization when no API key is + # configured (e.g. a local no-auth vLLM) rather than sending an empty token. + self.headers = {"Content-Type": "application/json"} + if self.api_key: + self.headers["Authorization"] = f"Bearer {self.api_key}" # raw response preservation self.include_raw_choices = include_raw_choices diff --git a/its_hub/integration/iaas/models.py b/its_hub/integration/iaas/models.py index d5201391..348a556b 100644 --- a/its_hub/integration/iaas/models.py +++ b/its_hub/integration/iaas/models.py @@ -2,7 +2,7 @@ from typing import Any -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, Field, field_validator from its_hub.api.types import SUPPORTED_ALGORITHMS, ChatMessage @@ -59,12 +59,6 @@ def validate_algorithm(cls, v): ) return v - @model_validator(mode="after") - def validate_config_requirements(self): - if not self.api_key: - raise ValueError("api_key is required") - return self - class ChatCompletionRequest(BaseModel): """Chat completion request with inference-time scaling support.""" diff --git a/tests/test_iaas.py b/tests/test_iaas.py index 2d409d35..718b8fbc 100644 --- a/tests/test_iaas.py +++ b/tests/test_iaas.py @@ -175,6 +175,18 @@ def test_configure_creates_gateway(self, iaas_client, vllm_endpoint): ) assert _state.gateway.default_config.api_endpoint == vllm_endpoint + def test_configure_without_api_key(self, iaas_client, vllm_endpoint): + """A no-auth local endpoint can be configured without an api_key.""" + config = { + "endpoint": vllm_endpoint, + "model": TEST_CONSTANTS["DEFAULT_MODEL_NAME"], + "alg": "self-consistency", + } + response = iaas_client.post("/configure", json=config) + assert response.status_code == 200 + assert _state.gateway is not None + assert _state.gateway.default_config.api_key is None + def test_models_endpoint_after_configure(self, iaas_client, vllm_endpoint): config = { "endpoint": vllm_endpoint, @@ -937,14 +949,15 @@ def test_self_consistency_accepts_tool_vote_without_regex(self): assert config.regex_patterns is None assert config.tool_vote == "tool_hierarchical" - def test_openai_requires_api_key_when_omitted(self): - with pytest.raises(ValueError, match="api_key is required"): - ConfigRequest( - endpoint="http://example.com:8000", - model="test-model", - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - ) + def test_api_key_is_optional_when_omitted(self): + """A no-auth endpoint can be configured without an api_key.""" + config = ConfigRequest( + endpoint="http://example.com:8000", + model="test-model", + alg="self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + ) + assert config.api_key is None class TestExtProcessor: From 6a9258262ff4009ca6964263e15d98a3dd75396e Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Thu, 20 Aug 2026 17:06:54 +0000 Subject: [PATCH 07/12] address review: test determinism, fd leak, temperature validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - recording_llm: add wait_for_bodies(n) bounded-poll helper so tests wait deterministically for the expected upstream request count instead of fixed asyncio.sleep; replaces sleep-based coordination in test_gateway (budget=3 → counts 3/6; reconfigure test → count 1) and test_iaas (budget=1 → count 1) - conftest: call server.server_close() in teardown of vllm_server, openai_server, and llm_server fixtures to close the listening socket and stop leaking one fd per test - test_gateway: remove the dead _max_lm_cache_size hasattr shim left over from the removed LM cache; rename test_configure_{unsupported_algorithm, invalid_regex,invalid_tool_vote} → test_overlay_* since they validate ITSRequestConfigUpdate in _make_config, never reaching ITSGateway.configure - types: validate temperature ∈ [0.0, 2.0] in _validate_optional_fields so direct ITSRequestConfig / ITSRequestConfigUpdate construction enforces the same range as ConfigRequest and ChatCompletionRequest; add test_overlay_invalid_temperature regression Signed-off-by: Harrison Stropkay --- its_hub/api/types.py | 5 +++++ tests/conftest.py | 3 +++ tests/mocks/recording_llm.py | 18 ++++++++++++++++++ tests/test_gateway.py | 21 +++++++++++---------- tests/test_iaas.py | 2 +- 5 files changed, 38 insertions(+), 11 deletions(-) diff --git a/its_hub/api/types.py b/its_hub/api/types.py index bbfa550b..65908b11 100644 --- a/its_hub/api/types.py +++ b/its_hub/api/types.py @@ -185,6 +185,7 @@ def _validate_optional_fields( tool_vote: str | None, threshold: float | None, confidence_threshold: float | None, + temperature: float | None, ) -> None: """Validate the optional scaling fields. Shared by ITSRequestConfig and ITSRequestConfigUpdate so the two cannot drift. Each check is guarded so @@ -207,6 +208,8 @@ def _validate_optional_fields( raise ValueError( f"confidence_threshold must be in (0.5, 1.0], got: {confidence_threshold}" ) + if temperature is not None and not (0.0 <= temperature <= 2.0): + raise ValueError(f"temperature must be in [0.0, 2.0], got: {temperature}") if tool_vote is not None and tool_vote not in VALID_TOOL_VOTE_OPTIONS: raise ValueError( f"tool_vote must be one of {VALID_TOOL_VOTE_OPTIONS}, got: {tool_vote}" @@ -261,6 +264,7 @@ def __post_init__(self): self.tool_vote, self.threshold, self.confidence_threshold, + self.temperature, ) def __repr__(self) -> str: @@ -304,6 +308,7 @@ def __post_init__(self): self.tool_vote, self.threshold, self.confidence_threshold, + self.temperature, ) def __repr__(self) -> str: diff --git a/tests/conftest.py b/tests/conftest.py index 80fab997..4ee6e1e8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -304,6 +304,7 @@ def vllm_server(): server.shutdown() server_thread.join() + server.server_close() @pytest.fixture(scope="session") @@ -322,6 +323,7 @@ def openai_server(): server.shutdown() server_thread.join() + server.server_close() @pytest.fixture @@ -367,4 +369,5 @@ def llm_server(): yield f"http://localhost:{port}" server.shutdown() thread.join() + server.server_close() RecordingLLMHandler.reset() diff --git a/tests/mocks/recording_llm.py b/tests/mocks/recording_llm.py index 32e74146..8f301b89 100644 --- a/tests/mocks/recording_llm.py +++ b/tests/mocks/recording_llm.py @@ -4,6 +4,7 @@ (distinct from ``tests.conftest``), so a class defined there would be duplicated. """ +import asyncio import json import threading from http.server import BaseHTTPRequestHandler @@ -31,6 +32,23 @@ def hold(cls): def release(cls): cls._hold.set() + @classmethod + async def wait_for_bodies(cls, n: int, timeout: float = 5.0) -> None: + """Poll until at least ``n`` request bodies have been recorded. + + Replaces fixed ``asyncio.sleep`` coordination so tests fail fast and + deterministically when expected upstream traffic never arrives. + """ + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while len(cls.received_bodies) < n: + if loop.time() >= deadline: + raise AssertionError( + f"timed out after {timeout}s waiting for {n} recorded " + f"bodies; got {len(cls.received_bodies)}" + ) + await asyncio.sleep(0.01) + def do_POST(self): if self.path != "/v1/chat/completions": self.send_response(404) diff --git a/tests/test_gateway.py b/tests/test_gateway.py index b29e3f2a..77474114 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -132,10 +132,6 @@ async def test_closes_lm_even_when_algorithm_raises(self): @pytest.mark.asyncio async def test_concurrent_different_models_do_not_interfere(self, llm_server): gw = ITSGateway() - # On main the gateway has an LM cache; shrink it to 1 so the second - # request evicts the first's client. No-op on the fixed branch. - if hasattr(gw, "_max_lm_cache_size"): - gw._max_lm_cache_size = 1 RecordingLLMHandler.hold() @@ -143,10 +139,10 @@ async def test_concurrent_different_models_do_not_interfere(self, llm_server): config_b = _make_config(api_endpoint=f"{llm_server}/v1", model="model-B") task_a = asyncio.create_task(gw.arun_chat_completion(config_a, MESSAGES)) - await asyncio.sleep(0.3) + await RecordingLLMHandler.wait_for_bodies(3) task_b = asyncio.create_task(gw.arun_chat_completion(config_b, MESSAGES)) - await asyncio.sleep(0.3) + await RecordingLLMHandler.wait_for_bodies(6) RecordingLLMHandler.release() @@ -268,7 +264,7 @@ async def test_reconfigure_does_not_swap_in_flight_algorithm(self, llm_server): RecordingLLMHandler.hold() task = asyncio.create_task(gw.arun_chat_completion(base, MESSAGES)) - await asyncio.sleep(0.1) + await RecordingLLMHandler.wait_for_bodies(1) gw.configure( ITSRequestConfigUpdate( @@ -319,12 +315,12 @@ def test_configure_tool_vote_persists_when_omitted(self): gw.configure(_make_config(alg="self-consistency")) # no tool_vote assert _algo(gw).tool_vote == "tool_name" - def test_configure_unsupported_algorithm(self): + def test_overlay_unsupported_algorithm(self): gw = ITSGateway() with pytest.raises(ValueError, match="not supported"): gw.configure(_make_config(alg="beam-search")) - def test_configure_invalid_regex(self): + def test_overlay_invalid_regex(self): gw = ITSGateway() with pytest.raises(ValueError, match="Invalid regex pattern"): gw.configure( @@ -339,7 +335,7 @@ def test_configure_preserves_orchestrator(self): ) assert _algo(gw).orchestrator is orch - def test_configure_invalid_tool_vote(self): + def test_overlay_invalid_tool_vote(self): gw = ITSGateway() with pytest.raises(ValueError, match="tool_vote must be one of"): gw.configure( @@ -350,6 +346,11 @@ def test_configure_invalid_tool_vote(self): ), ) + def test_overlay_invalid_temperature(self): + gw = ITSGateway() + with pytest.raises(ValueError, match="temperature must be in"): + gw.configure(_make_config(temperature=3.0)) + def test_configure_adaptive_self_consistency(self): gw = ITSGateway() gw.configure( diff --git a/tests/test_iaas.py b/tests/test_iaas.py index 718b8fbc..f203e8bd 100644 --- a/tests/test_iaas.py +++ b/tests/test_iaas.py @@ -491,7 +491,7 @@ async def test_metadata_reports_request_time_algorithm( }, ) ) - await asyncio.sleep(0.3) + await RecordingLLMHandler.wait_for_bodies(1) await async_iaas_client.post( "/configure", json={**base, "alg": "beta-self-consistency"} From 72ffb737e94c77d56a8e248ca45d9e7f93b1d0fd Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Mon, 24 Aug 2026 18:15:51 +0000 Subject: [PATCH 08/12] update types.py --- its_hub/api/types.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/its_hub/api/types.py b/its_hub/api/types.py index 65908b11..2390adc0 100644 --- a/its_hub/api/types.py +++ b/its_hub/api/types.py @@ -282,8 +282,8 @@ class ITSRequestConfigUpdate: ``exclude_tool_args`` can be cleared by passing ``[]``. Use :meth:`merge` to layer one update over another, and :meth:`resolve` to - materialize a complete :class:`ITSRequestConfig` (the only point at which - mandatory-field presence is checked). + materialize a complete :class:`ITSRequestConfig`; mandatory-field presence + is enforced by :class:`ITSRequestConfig` at construction time. """ # LM target @@ -330,16 +330,12 @@ def merge(self, overlay: ITSRequestConfigUpdate) -> ITSRequestConfigUpdate: def resolve(self) -> ITSRequestConfig: """Materialize a complete :class:`ITSRequestConfig`. - Raises ``ValueError`` if ``api_endpoint`` or ``model`` are absent; - this is the single point at which mandatory-field presence is - checked. Optional fields are forwarded only when non-``None``; - otherwise the resolved config's defaults (``budget=4``, - ``alg="self-consistency"``, ``api_key=None``) apply. + Optional fields are forwarded only when non-``None``; otherwise + the resolved config's defaults (``budget=4``, + ``alg="self-consistency"``, ``api_key=None``) apply. Mandatory-field + presence (``api_endpoint``, ``model``) is enforced by + :class:`ITSRequestConfig` itself at construction time. """ - if not self.api_endpoint: - raise ValueError("api_endpoint must be specified") - if not self.model: - raise ValueError("model must be specified") optionals = { f.name: getattr(self, f.name) for f in fields(self) From 6120d1ed858b6e966d764ba8ec08505ed009e4e7 Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Mon, 24 Aug 2026 19:43:05 +0000 Subject: [PATCH 09/12] re-add gateway shutdown --- its_hub/core/gateway.py | 11 ++++++++ its_hub/integration/ext_proc/server.py | 2 ++ its_hub/integration/iaas/app.py | 8 ++++++ tests/test_gateway.py | 37 ++++++++++++++++++++++++++ 4 files changed, 58 insertions(+) diff --git a/its_hub/core/gateway.py b/its_hub/core/gateway.py index 33fe443c..e1216031 100644 --- a/its_hub/core/gateway.py +++ b/its_hub/core/gateway.py @@ -233,3 +233,14 @@ async def arun_chat_completion( "usage": usage_dict, "alg": resolved.alg, } + + async def aclose(self) -> None: + """Close all pooled connectors. + + Call on service shutdown to release keep-alive connections. + Idempotent: closing an already-closed connector is a no-op. + """ + for conn in self._connector_pool.values(): + if not conn.closed: + await conn.close() + self._connector_pool.clear() diff --git a/its_hub/integration/ext_proc/server.py b/its_hub/integration/ext_proc/server.py index a0716889..82bfc12f 100644 --- a/its_hub/integration/ext_proc/server.py +++ b/its_hub/integration/ext_proc/server.py @@ -83,6 +83,8 @@ async def serve(port: int = 50051): except (KeyboardInterrupt, asyncio.CancelledError): logger.info("Received shutdown signal") await server.stop(grace=5) + finally: + await processor.gateway.aclose() def _print_config() -> None: diff --git a/its_hub/integration/iaas/app.py b/its_hub/integration/iaas/app.py index 11cededb..0befc222 100644 --- a/its_hub/integration/iaas/app.py +++ b/its_hub/integration/iaas/app.py @@ -9,6 +9,7 @@ import logging import time import uuid +from contextlib import asynccontextmanager from fastapi import FastAPI, Header, HTTPException, status from fastapi.responses import StreamingResponse @@ -39,10 +40,17 @@ def reset(self): _state = _ServiceState() +@asynccontextmanager +async def _lifespan(application: FastAPI): + yield + await _state.gateway.aclose() + + app = FastAPI( title="its_hub Inference-as-a-Service", description="OpenAI-compatible API for inference-time scaling algorithms", version="0.1.0-alpha", + lifespan=_lifespan, ) diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 77474114..3d639a93 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -153,6 +153,43 @@ async def test_concurrent_different_models_do_not_interfere(self, llm_server): assert result_b["message"]["content"] == "answer from model-B" +class TestGatewayShutdown: + @pytest.mark.asyncio + async def test_aclose_closes_pooled_connectors(self): + gw = ITSGateway() + c1 = gw._get_connector("http://upstream-a/v1") + c2 = gw._get_connector("http://upstream-b/v1") + assert not c1.closed + assert not c2.closed + assert len(gw._connector_pool) == 2 + + await gw.aclose() + + assert c1.closed + assert c2.closed + assert len(gw._connector_pool) == 0 + + @pytest.mark.asyncio + async def test_aclose_is_idempotent(self): + gw = ITSGateway() + gw._get_connector("http://upstream/v1") + await gw.aclose() + await gw.aclose() # second call must not raise + + @pytest.mark.asyncio + async def test_get_connector_rebuilds_after_close(self): + """A request arriving after aclose() gets a fresh connector.""" + gw = ITSGateway() + c1 = gw._get_connector("http://upstream/v1") + await gw.aclose() + assert c1.closed + + c2 = gw._get_connector("http://upstream/v1") + assert not c2.closed + assert c2 is not c1 + await gw.aclose() + + class TestRunChatCompletion: @pytest.mark.asyncio async def test_response_only(self): From c718e071f684599ef4b59231b76d655e47678f2d Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Mon, 24 Aug 2026 19:59:07 +0000 Subject: [PATCH 10/12] add underscore to default config --- docs/iaas-service.md | 2 +- its_hub/core/gateway.py | 8 ++++---- its_hub/integration/iaas/app.py | 6 +++--- tests/test_gateway.py | 12 ++++++------ tests/test_iaas.py | 24 ++++++++++++------------ 5 files changed, 26 insertions(+), 26 deletions(-) diff --git a/docs/iaas-service.md b/docs/iaas-service.md index 0abd0e31..8d1fbf3c 100644 --- a/docs/iaas-service.md +++ b/docs/iaas-service.md @@ -79,7 +79,7 @@ API keys can enter the system through three paths: - Keys are **not persisted** to disk. A key supplied via the `X-ITS-API-Key` header is request-scoped: it lives only for that request's LM client, which is closed when the request completes. A key set via `/configure` remains in memory as a service default - (`ITSGateway.default_config`) until replaced or the process restarts. + (`ITSGateway._default_config`) until replaced or the process restarts. - Keys are **never shared between requests** — each request builds its own LM client, so credentials supplied via header or `/configure` cannot cross-contaminate diff --git a/its_hub/core/gateway.py b/its_hub/core/gateway.py index e1216031..5184819e 100644 --- a/its_hub/core/gateway.py +++ b/its_hub/core/gateway.py @@ -68,7 +68,7 @@ def __init__( # them. System defaults (budget=4, alg="self-consistency") are NOT # stored here — they live on ITSRequestConfig and are injected at # resolve() time. - self.default_config = default_config or ITSRequestConfigUpdate() + self._default_config = default_config or ITSRequestConfigUpdate() self._ssl_context = ssl.create_default_context(cafile=certifi.where()) # One connection pool per upstream endpoint self._connector_pool: dict[str, aiohttp.TCPConnector] = {} @@ -81,9 +81,9 @@ def configure(self, update: ITSRequestConfigUpdate) -> None: by ``merge``); not completeness-validated — ``api_endpoint``/``model`` may remain absent until a request supplies them. """ - self.default_config = self.default_config.merge(update) + self._default_config = self._default_config.merge(update) logger.info( - "ITSGateway reconfigured with default alg=%s", self.default_config.alg + "ITSGateway reconfigured with default alg=%s", self._default_config.alg ) def _build_algorithm(self, config: ITSRequestConfig) -> AbstractScalingAlgorithm: @@ -167,7 +167,7 @@ async def arun_chat_completion( # (raises ValueError if api_endpoint/model are still absent). Nothing # on `self` is mutated, so a concurrent /configure cannot swap the # algorithm or close this request's HTTP session mid-flight. - resolved = self.default_config.merge(config).resolve() + resolved = self._default_config.merge(config).resolve() algorithm = self._build_algorithm(resolved) lm = self._build_lm(resolved, request_id) diff --git a/its_hub/integration/iaas/app.py b/its_hub/integration/iaas/app.py index 0befc222..b007efaa 100644 --- a/its_hub/integration/iaas/app.py +++ b/its_hub/integration/iaas/app.py @@ -94,7 +94,7 @@ async def config_service(request: ConfigRequest) -> dict[str, str]: ) _state.gateway.configure(config) - resolved = _state.gateway.default_config + resolved = _state.gateway._default_config logger.info( "Configured IaaS: model=%s, alg=%s, budget=%s", resolved.model, @@ -121,7 +121,7 @@ async def config_service(request: ConfigRequest) -> dict[str, str]: @app.get("/v1/models") async def list_models() -> dict[str, list[dict[str, str]]]: """List available models (OpenAI-compatible endpoint).""" - model = _state.gateway.default_config.model + model = _state.gateway._default_config.model if model: return { "data": [ @@ -230,7 +230,7 @@ async def _generate(): its_config = _build_its_config(request, its_budget, its_endpoint, its_api_key) - if not (its_config.api_endpoint or _state.gateway.default_config.api_endpoint): + if not (its_config.api_endpoint or _state.gateway._default_config.api_endpoint): yield f"data: {json.dumps({'error': 'Service not configured'})}\n\n" yield "data: [DONE]\n\n" return diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 3d639a93..a5b06a85 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -54,7 +54,7 @@ def _patch_build_lm(gw): def _algo(gw): - return gw._build_algorithm(gw.default_config.resolve()) + return gw._build_algorithm(gw._default_config.resolve()) class TestGatewayConstruction: @@ -63,7 +63,7 @@ def test_default_init(self): assert gw._orchestrator is not None # Stored default is partial; system defaults (alg, budget) materialize # at resolve() time on ITSRequestConfig. - resolved = gw.default_config.merge( + resolved = gw._default_config.merge( ITSRequestConfigUpdate(api_endpoint="http://x/v1", model="m") ).resolve() assert resolved.alg == "self-consistency" @@ -77,8 +77,8 @@ def test_custom_default_config(self): api_key="k", ) gw = ITSGateway(default_config=config) - assert gw.default_config is config - assert gw.default_config.alg == "beta-self-consistency" + assert gw._default_config is config + assert gw._default_config.alg == "beta-self-consistency" def test_custom_orchestrator_passed_to_algorithm(self): orch = MagicMock() @@ -313,7 +313,7 @@ async def test_reconfigure_does_not_swap_in_flight_algorithm(self, llm_server): result = await task assert result["alg"] == "self-consistency" - assert gw.default_config.alg == "beta-self-consistency" + assert gw._default_config.alg == "beta-self-consistency" class TestConfigure: @@ -322,7 +322,7 @@ def test_configure_self_consistency(self): gw.configure( _make_config(alg="self-consistency", regex_patterns=[r"\\boxed{([^}]+)}"]), ) - assert gw.default_config.alg == "self-consistency" + assert gw._default_config.alg == "self-consistency" assert type(_algo(gw)).__name__ == "SelfConsistency" def test_configure_defaults_tool_vote_when_omitted(self): diff --git a/tests/test_iaas.py b/tests/test_iaas.py index f203e8bd..b9a3a6a8 100644 --- a/tests/test_iaas.py +++ b/tests/test_iaas.py @@ -46,11 +46,11 @@ def vllm_endpoint(vllm_server): def _install_mock_gateway(mock_return=None, side_effect=None, default=None): - """Replace _state.gateway with a mock, preserving default_config for merging.""" + """Replace _state.gateway with a mock, preserving _default_config for merging.""" if default is None: - default = _state.gateway.default_config + default = _state.gateway._default_config mock_gw = MagicMock() - mock_gw.default_config = default + mock_gw._default_config = default if side_effect: mock_gw.arun_chat_completion = AsyncMock(side_effect=side_effect) else: @@ -171,9 +171,9 @@ def test_configure_creates_gateway(self, iaas_client, vllm_endpoint): assert response.status_code == 200 assert _state.gateway is not None assert ( - _state.gateway.default_config.model == TEST_CONSTANTS["DEFAULT_MODEL_NAME"] + _state.gateway._default_config.model == TEST_CONSTANTS["DEFAULT_MODEL_NAME"] ) - assert _state.gateway.default_config.api_endpoint == vllm_endpoint + assert _state.gateway._default_config.api_endpoint == vllm_endpoint def test_configure_without_api_key(self, iaas_client, vllm_endpoint): """A no-auth local endpoint can be configured without an api_key.""" @@ -185,7 +185,7 @@ def test_configure_without_api_key(self, iaas_client, vllm_endpoint): response = iaas_client.post("/configure", json=config) assert response.status_code == 200 assert _state.gateway is not None - assert _state.gateway.default_config.api_key is None + assert _state.gateway._default_config.api_key is None def test_models_endpoint_after_configure(self, iaas_client, vllm_endpoint): config = { @@ -275,7 +275,7 @@ def test_tool_vote_algorithm_usage_verification(self, iaas_client, vllm_endpoint } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway.default_config) + _state.gateway._build_algorithm(_state.gateway._default_config) mock_sc.assert_called_once() call_args = mock_sc.call_args assert call_args.kwargs["tool_vote"] == "tool_hierarchical" @@ -300,7 +300,7 @@ def test_family_basic_configuration(self, iaas_client, vllm_endpoint, alg): response = iaas_client.post("/configure", json=config) assert response.status_code == 200 assert "success" in response.json()["status"] - assert _state.gateway.default_config.alg == alg + assert _state.gateway._default_config.alg == alg def test_adaptive_threshold_forwarded_to_algorithm( self, iaas_client, vllm_endpoint @@ -317,7 +317,7 @@ def test_adaptive_threshold_forwarded_to_algorithm( } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway.default_config) + _state.gateway._build_algorithm(_state.gateway._default_config) mock_alg.assert_called_once() assert mock_alg.call_args.kwargs["threshold"] == 0.9 @@ -336,7 +336,7 @@ def test_beta_confidence_threshold_forwarded_to_algorithm( } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway.default_config) + _state.gateway._build_algorithm(_state.gateway._default_config) mock_alg.assert_called_once() assert mock_alg.call_args.kwargs["confidence_threshold"] == 0.8 @@ -353,7 +353,7 @@ def test_threshold_omitted_uses_algorithm_default(self, iaas_client, vllm_endpoi } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway.default_config) + _state.gateway._build_algorithm(_state.gateway._default_config) assert "threshold" not in mock_alg.call_args.kwargs @pytest.mark.parametrize( @@ -390,7 +390,7 @@ def test_family_tool_vote_forwarded(self, iaas_client, vllm_endpoint): } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway.default_config) + _state.gateway._build_algorithm(_state.gateway._default_config) assert mock_alg.call_args.kwargs["tool_vote"] == "tool_hierarchical" assert mock_alg.call_args.kwargs["exclude_args"] == ["timestamp"] From b8914d910be182cadccbed8d5f9d016b7a5a5a38 Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Mon, 24 Aug 2026 20:04:54 +0000 Subject: [PATCH 11/12] use resolve() in test_iaas.py --- tests/test_iaas.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_iaas.py b/tests/test_iaas.py index b9a3a6a8..e99f92f1 100644 --- a/tests/test_iaas.py +++ b/tests/test_iaas.py @@ -275,7 +275,7 @@ def test_tool_vote_algorithm_usage_verification(self, iaas_client, vllm_endpoint } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway._default_config) + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) mock_sc.assert_called_once() call_args = mock_sc.call_args assert call_args.kwargs["tool_vote"] == "tool_hierarchical" @@ -317,7 +317,7 @@ def test_adaptive_threshold_forwarded_to_algorithm( } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway._default_config) + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) mock_alg.assert_called_once() assert mock_alg.call_args.kwargs["threshold"] == 0.9 @@ -336,7 +336,7 @@ def test_beta_confidence_threshold_forwarded_to_algorithm( } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway._default_config) + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) mock_alg.assert_called_once() assert mock_alg.call_args.kwargs["confidence_threshold"] == 0.8 @@ -353,7 +353,7 @@ def test_threshold_omitted_uses_algorithm_default(self, iaas_client, vllm_endpoi } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway._default_config) + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) assert "threshold" not in mock_alg.call_args.kwargs @pytest.mark.parametrize( @@ -390,7 +390,7 @@ def test_family_tool_vote_forwarded(self, iaas_client, vllm_endpoint): } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - _state.gateway._build_algorithm(_state.gateway._default_config) + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) assert mock_alg.call_args.kwargs["tool_vote"] == "tool_hierarchical" assert mock_alg.call_args.kwargs["exclude_args"] == ["timestamp"] From 424b783ec76277f2053cc81d26e04e0f5243623d Mon Sep 17 00:00:00 2001 From: Harrison Stropkay Date: Mon, 24 Aug 2026 20:41:22 +0000 Subject: [PATCH 12/12] freeze ITSRequestConfig list fields into tuples regex_patterns and exclude_tool_args are now stored as immutable tuples on the resolved snapshot so a caller cannot mutate the source list (or the snapshot itself) after validation and affect an in-flight request. Also make the test mock gateway's aclose awaitable (lifespan shutdown awaits it) and report the resolved alg fallback. Signed-off-by: Harrison Stropkay --- its_hub/api/types.py | 18 ++++++++++++++++-- tests/test_gateway.py | 2 +- tests/test_iaas.py | 13 ++++++++----- 3 files changed, 25 insertions(+), 8 deletions(-) diff --git a/its_hub/api/types.py b/its_hub/api/types.py index 2390adc0..ef6b0980 100644 --- a/its_hub/api/types.py +++ b/its_hub/api/types.py @@ -216,6 +216,13 @@ def _validate_optional_fields( ) +def _freeze_list(value: list[str] | None) -> tuple[str, ...] | None: + """Copy a list into an immutable tuple; leave ``None`` unchanged.""" + if value is None: + return None + return tuple(value) + + def _config_repr(obj: ITSRequestConfig | ITSRequestConfigUpdate) -> str: return ( type(obj).__name__ @@ -246,9 +253,9 @@ class ITSRequestConfig: # Scaling parameters budget: int = 4 alg: str = "self-consistency" - regex_patterns: list[str] | None = None + regex_patterns: tuple[str, ...] | None = None tool_vote: str | None = None - exclude_tool_args: list[str] | None = None + exclude_tool_args: tuple[str, ...] | None = None threshold: float | None = None confidence_threshold: float | None = None @@ -266,6 +273,13 @@ def __post_init__(self): self.confidence_threshold, self.temperature, ) + # frozen=True blocks attribute rebinding, not mutation of a referenced + # list. Copy these into tuples so neither the source list nor the + # snapshot can be mutated after validation to affect an in-flight request. + object.__setattr__(self, "regex_patterns", _freeze_list(self.regex_patterns)) + object.__setattr__( + self, "exclude_tool_args", _freeze_list(self.exclude_tool_args) + ) def __repr__(self) -> str: return _config_repr(self) diff --git a/tests/test_gateway.py b/tests/test_gateway.py index a5b06a85..251050a2 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -342,7 +342,7 @@ def test_configure_with_tool_vote(self): ) algo = _algo(gw) assert algo.tool_vote == "tool_name" - assert algo.exclude_args == ["timestamp"] + assert algo.exclude_args == ("timestamp",) def test_configure_tool_vote_persists_when_omitted(self): """`None` means trickle down, so a /configure omitting tool_vote diff --git a/tests/test_iaas.py b/tests/test_iaas.py index e99f92f1..c1371417 100644 --- a/tests/test_iaas.py +++ b/tests/test_iaas.py @@ -51,12 +51,15 @@ def _install_mock_gateway(mock_return=None, side_effect=None, default=None): default = _state.gateway._default_config mock_gw = MagicMock() mock_gw._default_config = default + # _lifespan awaits aclose() on shutdown; make the mock awaitable. + mock_gw.aclose = AsyncMock() if side_effect: mock_gw.arun_chat_completion = AsyncMock(side_effect=side_effect) else: - # Mirror the real gateway: inject the resolved alg from the merged - # default so metadata tests see the configured value, not a hardcoded one. - patched = {**(mock_return or {}), "alg": default.alg} + # Mirror resolve(): alg defaults to "self-consistency" when the update + # omits it. We can't call resolve() directly because the header-only + # path has no endpoint/model yet. + patched = {**(mock_return or {}), "alg": default.alg or "self-consistency"} async def _arun(*args, **kwargs): return patched @@ -279,7 +282,7 @@ def test_tool_vote_algorithm_usage_verification(self, iaas_client, vllm_endpoint mock_sc.assert_called_once() call_args = mock_sc.call_args assert call_args.kwargs["tool_vote"] == "tool_hierarchical" - assert call_args.kwargs["exclude_args"] == ["timestamp", "id"] + assert call_args.kwargs["exclude_args"] == ("timestamp", "id") class TestAdaptiveAndBetaSelfConsistency: @@ -392,7 +395,7 @@ def test_family_tool_vote_forwarded(self, iaas_client, vllm_endpoint): assert response.status_code == 200 _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) assert mock_alg.call_args.kwargs["tool_vote"] == "tool_hierarchical" - assert mock_alg.call_args.kwargs["exclude_args"] == ["timestamp"] + assert mock_alg.call_args.kwargs["exclude_args"] == ("timestamp",) class TestChatCompletions: