"""LM Studio server lifecycle and exclusive model-residency management. The `lms` CLI is the sole server lifecycle controller. Model inspection, loading, and unloading use LM Studio's native v1 REST API so every applied configuration is captured in experiment telemetry. """ from __future__ import annotations from dataclasses import asdict, dataclass import json import os from pathlib import Path import shutil import subprocess import time from typing import Any from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen class LMStudioManagementError(RuntimeError): """Raised when server lifecycle or exclusive residency cannot be verified.""" @dataclass(frozen=True, slots=True) class ResidencyTransition: target_model: str | None target_context_length: int | None before_instances: tuple[str, ...] unloaded_instances: tuple[str, ...] loaded_instance: str | None after_instances: tuple[str, ...] load_response: dict[str, Any] | None elapsed_seconds: float reused: bool def to_dict(self) -> dict[str, Any]: return asdict(self) class LMStudioServer: """Control only the local LM Studio API server via `lms server`.""" def __init__( self, port: int = 1234, cli_path: Path | None = None, command_timeout_seconds: float = 30.0, ): discovered = shutil.which("lms") default = Path.home() / ".lmstudio" / "bin" / "lms" self.cli_path = Path(cli_path or discovered or default).resolve() self.port = port self.command_timeout_seconds = command_timeout_seconds def _run(self, *arguments: str) -> subprocess.CompletedProcess[str]: if not self.cli_path.exists(): raise LMStudioManagementError(f"LM Studio CLI is unavailable: {self.cli_path}") try: return subprocess.run( [str(self.cli_path), *arguments], text=True, capture_output=True, check=False, timeout=self.command_timeout_seconds, ) except (OSError, subprocess.SubprocessError) as exc: raise LMStudioManagementError(f"lms {' '.join(arguments)} failed: {exc}") from exc def status(self) -> dict[str, Any]: result = self._run("server", "status") # Current lms releases render human-facing status on stderr even when # the command succeeds. Treat both captured streams as the CLI record. rendered = result.stdout + "\n" + result.stderr normalized = rendered.lower() return { "running": ( result.returncode == 0 and "server is running" in normalized and "server is not running" not in normalized ), "returncode": result.returncode, "stdout": result.stdout, "stderr": result.stderr, "command": [str(self.cli_path), "server", "status"], } def _api_ready(self) -> bool: request = Request( f"http://127.0.0.1:{self.port}/api/v1/models", method="GET", headers={"Content-Type": "application/json"}, ) try: with urlopen(request, timeout=min(self.command_timeout_seconds, 2.0)): return True except HTTPError as exc: # Authentication failures still prove that the REST listener is ready. return exc.code in {401, 403} except (URLError, OSError): return False def _wait_until_ready(self) -> dict[str, Any]: deadline = time.monotonic() + self.command_timeout_seconds last_status: dict[str, Any] | None = None while time.monotonic() < deadline: last_status = self.status() if last_status["running"] and self._api_ready(): return last_status time.sleep(0.25) raise LMStudioManagementError( "lms reported startup activity, but the official REST API never became ready: " + repr(last_status) ) def ensure_running(self) -> dict[str, Any]: status = self.status() if status["running"]: return { "action": "already_running", "status": self._wait_until_ready(), } started = self._run("server", "start", "--port", str(self.port)) if started.returncode != 0: raise LMStudioManagementError( "lms server start failed: " + (started.stderr.strip() or started.stdout.strip()) ) return { "action": "started", "start_stdout": started.stdout, "start_stderr": started.stderr, "status": self._wait_until_ready(), } def stop(self) -> dict[str, Any]: result = self._run("server", "stop") if result.returncode != 0: raise LMStudioManagementError( "lms server stop failed: " + (result.stderr.strip() or result.stdout.strip()) ) return {"stdout": result.stdout, "stderr": result.stderr, "returncode": result.returncode} class LMStudioResidencyManager: """Use native v1 REST endpoints to enforce exactly one or zero loaded models.""" def __init__(self, base_url: str, api_token_env: str, timeout_seconds: float = 120.0): self.base_url = base_url.rstrip("/") self.api_token_env = api_token_env self.timeout_seconds = timeout_seconds def _headers(self) -> dict[str, str]: headers = {"Content-Type": "application/json"} token = os.environ.get(self.api_token_env, "").strip() if token: headers["Authorization"] = f"Bearer {token}" return headers def _request( self, method: str, endpoint: str, payload: dict[str, Any] | None = None ) -> dict[str, Any]: request = Request( self.base_url + endpoint, data=None if payload is None else json.dumps(payload).encode("utf-8"), method=method, headers=self._headers(), ) try: with urlopen(request, timeout=self.timeout_seconds) as response: body = response.read().decode("utf-8") except HTTPError as exc: detail = exc.read().decode("utf-8", errors="replace") raise LMStudioManagementError( f"LM Studio returned HTTP {exc.code} for {endpoint}: {detail}" ) from exc except URLError as exc: raise LMStudioManagementError( f"Cannot reach LM Studio management API at {self.base_url}: {exc.reason}" ) from exc try: value = json.loads(body) except json.JSONDecodeError as exc: raise LMStudioManagementError(f"LM Studio returned non-JSON data for {endpoint}") from exc if not isinstance(value, dict): raise LMStudioManagementError(f"Unexpected LM Studio response for {endpoint}") if isinstance(value.get("error"), dict): raise LMStudioManagementError(f"LM Studio management error for {endpoint}: {value['error']}") return value def models(self) -> tuple[dict[str, Any], ...]: value = self._request("GET", "/api/v1/models") models = value.get("models") if not isinstance(models, list): raise LMStudioManagementError("/api/v1/models response has no models array") return tuple(item for item in models if isinstance(item, dict)) def loaded_instances(self) -> tuple[dict[str, Any], ...]: result: list[dict[str, Any]] = [] for model in self.models(): instances = model.get("loaded_instances", []) if not isinstance(instances, list): continue for instance in instances: if isinstance(instance, dict) and instance.get("id"): result.append( { "model_key": str(model.get("key", "")), "type": str(model.get("type", "")), "instance_id": str(instance["id"]), "config": dict(instance.get("config", {})), } ) return tuple(result) def _unload(self, instance_id: str) -> dict[str, Any]: value = self._request( "POST", "/api/v1/models/unload", {"instance_id": instance_id} ) if value.get("instance_id") != instance_id: raise LMStudioManagementError( f"unload acknowledgement mismatch for {instance_id}: {value}" ) return value def unload_all(self) -> ResidencyTransition: started = time.monotonic() before = self.loaded_instances() unloaded: list[str] = [] for instance in before: self._unload(instance["instance_id"]) unloaded.append(instance["instance_id"]) after = self.loaded_instances() if after: raise LMStudioManagementError(f"models remained loaded after unload-all: {after}") return ResidencyTransition( target_model=None, target_context_length=None, before_instances=tuple(item["instance_id"] for item in before), unloaded_instances=tuple(unloaded), loaded_instance=None, after_instances=(), load_response=None, elapsed_seconds=time.monotonic() - started, reused=not before, ) def ensure_exclusive(self, model_key: str, context_length: int) -> ResidencyTransition: started = time.monotonic() before = self.loaded_instances() if len(before) == 1: current = before[0] current_context = current.get("config", {}).get("context_length") if current["model_key"] == model_key and current_context == context_length: return ResidencyTransition( target_model=model_key, target_context_length=context_length, before_instances=(current["instance_id"],), unloaded_instances=(), loaded_instance=current["instance_id"], after_instances=(current["instance_id"],), load_response=None, elapsed_seconds=time.monotonic() - started, reused=True, ) unloaded: list[str] = [] for instance in before: self._unload(instance["instance_id"]) unloaded.append(instance["instance_id"]) response = self._request( "POST", "/api/v1/models/load", { "model": model_key, "context_length": context_length, "echo_load_config": True, }, ) instance_id = response.get("instance_id") if response.get("status") != "loaded" or not isinstance(instance_id, str): raise LMStudioManagementError(f"model load did not succeed for {model_key}: {response}") load_config = response.get("load_config", {}) if not isinstance(load_config, dict) or load_config.get("context_length") != context_length: raise LMStudioManagementError( f"LM Studio did not apply context_length={context_length}: {response}" ) after = self.loaded_instances() if len(after) != 1: raise LMStudioManagementError(f"exclusive residency failed for {model_key}: {after}") only = after[0] if ( only["model_key"] != model_key or only["instance_id"] != instance_id or only.get("config", {}).get("context_length") != context_length ): raise LMStudioManagementError( f"loaded runtime does not match requested model/context: {after}" ) return ResidencyTransition( target_model=model_key, target_context_length=context_length, before_instances=tuple(item["instance_id"] for item in before), unloaded_instances=tuple(unloaded), loaded_instance=instance_id, after_instances=(instance_id,), load_response=response, elapsed_seconds=time.monotonic() - started, reused=False, )