wuenlp_tools.utils.local_ollama

  1"""Local Ollama lifecycle: per-job serve, model pull with flock, health checks."""
  2
  3from __future__ import annotations
  4
  5import atexit
  6import fcntl
  7import os
  8import re
  9import shutil
 10import socket
 11import subprocess
 12import time
 13from contextlib import contextmanager
 14from pathlib import Path
 15from typing import Iterator
 16from urllib.parse import urlparse
 17
 18import requests
 19from loguru import logger
 20
 21from wuenlp_tools.config import get_local_llm_config, get_pull_lock_dir
 22
 23_ollama_process: subprocess.Popen | None = None
 24_resolved_base_url: str | None = None
 25
 26_DEFAULT_PORT = 11434
 27_STARTUP_TIMEOUT_S = 120.0
 28_PULL_LOCK_TIMEOUT_S = 7200.0
 29_TRUTHY = frozenset({"1", "true", "yes", "on"})
 30_FALSY = frozenset({"0", "false", "no", "off"})
 31
 32
 33def _is_slurm() -> bool:
 34    return bool(os.environ.get("SLURM_JOB_ID"))
 35
 36
 37def _env_flag(name: str) -> bool | None:
 38    raw = os.environ.get(name, "").strip().lower()
 39    if raw in _TRUTHY:
 40        return True
 41    if raw in _FALSY:
 42        return False
 43    return None
 44
 45
 46def _resolve_manage_process(setting: str | bool) -> bool:
 47    override = _env_flag("WUENLP_LOCAL_LLM_MANAGE_PROCESS")
 48    if override is not None:
 49        return override
 50    if isinstance(setting, bool):
 51        return setting
 52    if setting == "auto":
 53        return _is_slurm()
 54    return bool(setting)
 55
 56
 57def _pick_free_port() -> int:
 58    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
 59        sock.bind(("127.0.0.1", 0))
 60        return int(sock.getsockname()[1])
 61
 62
 63def _resolve_port(port_setting: str | int) -> int:
 64    if isinstance(port_setting, int):
 65        return port_setting
 66    if port_setting == "auto":
 67        return _pick_free_port() if _is_slurm() else _DEFAULT_PORT
 68    return _DEFAULT_PORT
 69
 70
 71def _base_url_from_port(port: int) -> str:
 72    return f"http://127.0.0.1:{port}/v1/"
 73
 74
 75def _normalize_base_url(url: str) -> str:
 76    url = url.rstrip("/") + "/"
 77    if not url.endswith("/v1/"):
 78        if url.endswith("/v1"):
 79            url += "/"
 80        else:
 81            url = url.rstrip("/") + "/v1/"
 82    return url
 83
 84
 85def _ollama_host_from_base_url(base_url: str) -> str:
 86    parsed = urlparse(base_url)
 87    host = parsed.hostname or "127.0.0.1"
 88    port = parsed.port or _DEFAULT_PORT
 89    return f"{host}:{port}"
 90
 91
 92def _api_root(base_url: str) -> str:
 93    return _normalize_base_url(base_url).removesuffix("/v1/")
 94
 95
 96_OLLAMA_MISSING_MSG = (
 97    "Ollama is required for provider='local' but the `ollama` executable was not "
 98    "found on PATH. Install from https://ollama.com or use provider='ollama' for "
 99    "the remote team server."
100)
101
102
103def _require_ollama_on_path() -> str:
104    path = shutil.which("ollama")
105    if not path:
106        raise RuntimeError(_OLLAMA_MISSING_MSG)
107    return path
108
109
110def _ollama_binary() -> str:
111    return _require_ollama_on_path()
112
113
114def _wait_for_ollama(base_url: str, timeout: float = _STARTUP_TIMEOUT_S) -> None:
115    tags_url = f"{_api_root(base_url)}/api/tags"
116    deadline = time.monotonic() + timeout
117    while time.monotonic() < deadline:
118        try:
119            response = requests.get(tags_url, timeout=2.0)
120            if response.status_code == 200:
121                return
122        except requests.RequestException:
123            pass
124        if _ollama_process is not None and _ollama_process.poll() is not None:
125            raise RuntimeError(
126                f"ollama serve exited with code {_ollama_process.returncode}"
127            )
128        time.sleep(0.25)
129    raise TimeoutError(f"Ollama did not become ready at {base_url} within {timeout}s")
130
131
132def _start_ollama_serve(host: str) -> subprocess.Popen:
133    env = os.environ.copy()
134    env["OLLAMA_HOST"] = host
135    logger.info("Starting ollama serve (OLLAMA_HOST={})", host)
136    process = subprocess.Popen(
137        [_ollama_binary(), "serve"],
138        env=env,
139        stdout=subprocess.DEVNULL,
140        stderr=subprocess.DEVNULL,
141    )
142    return process
143
144
145def _stop_ollama_serve() -> None:
146    global _ollama_process
147    if _ollama_process is None:
148        return
149    if _ollama_process.poll() is None:
150        logger.debug("Stopping job-local ollama serve (pid={})", _ollama_process.pid)
151        _ollama_process.terminate()
152        try:
153            _ollama_process.wait(timeout=10)
154        except subprocess.TimeoutExpired:
155            _ollama_process.kill()
156            _ollama_process.wait(timeout=5)
157    _ollama_process = None
158
159
160def _ollama_is_reachable(base_url: str) -> bool:
161    try:
162        response = requests.get(f"{_api_root(base_url)}/api/tags", timeout=2.0)
163        return response.status_code == 200
164    except requests.RequestException:
165        return False
166
167
168def _is_loopback_host(hostname: str | None) -> bool:
169    return hostname in ("127.0.0.1", "localhost", "::1")
170
171
172def _assert_local_base_url(base_url: str) -> str:
173    parsed = urlparse(base_url)
174    if not _is_loopback_host(parsed.hostname):
175        raise ValueError(
176            f"provider='local' requires a loopback Ollama URL, got {base_url!r}. "
177            "Use provider='ollama' for remote servers."
178        )
179    return _normalize_base_url(base_url)
180
181
182def _unreachable_error(base_url: str, *, manage_hint: bool) -> RuntimeError:
183    if shutil.which("ollama") is None:
184        return RuntimeError(_OLLAMA_MISSING_MSG)
185    hint = (
186        "Start `ollama serve`, enable manage_process in local_llm config, "
187        "or set WUENLP_LOCAL_LLM_MANAGE_PROCESS=1."
188        if manage_hint
189        else "Start `ollama serve` on that host/port."
190    )
191    return RuntimeError(f"Local Ollama is not reachable at {base_url}. {hint}")
192
193
194def _ensure_reachable(base_url: str, *, manage: bool) -> None:
195    if _ollama_is_reachable(base_url):
196        return
197    if manage:
198        _require_ollama_on_path()
199        host = _ollama_host_from_base_url(base_url)
200        global _ollama_process
201        _ollama_process = _start_ollama_serve(host)
202        atexit.register(_stop_ollama_serve)
203        _wait_for_ollama(base_url)
204        return
205    raise _unreachable_error(base_url, manage_hint=True)
206
207
208def restart_local_ollama() -> str:
209    """Stop the job-local serve process and start a fresh one (after timeouts/hangs)."""
210    global _resolved_base_url
211    _stop_ollama_serve()
212    _resolved_base_url = None
213    return ensure_local_ollama()
214
215
216def ensure_local_ollama() -> str:
217    """Ensure a local Ollama server is reachable; return OpenAI-compatible base URL."""
218    global _ollama_process, _resolved_base_url
219
220    if _resolved_base_url is not None:
221        return _resolved_base_url
222
223    env_url = os.environ.get("WUENLP_LOCAL_LLM_BASE_URL")
224    if env_url:
225        _resolved_base_url = _assert_local_base_url(env_url)
226        manage = _resolve_manage_process(get_local_llm_config().manage_process)
227        _ensure_reachable(_resolved_base_url, manage=manage)
228        return _resolved_base_url
229
230    cfg = get_local_llm_config()
231    if cfg.base_url:
232        _resolved_base_url = _assert_local_base_url(cfg.base_url)
233        manage = _resolve_manage_process(cfg.manage_process)
234        _ensure_reachable(_resolved_base_url, manage=manage)
235        return _resolved_base_url
236
237    port = _resolve_port(cfg.port)
238    base_url = _base_url_from_port(port)
239    manage = _resolve_manage_process(cfg.manage_process)
240    if cfg.manage_process == "auto" and not _is_slurm():
241        manage = not _ollama_is_reachable(base_url)
242
243    if manage:
244        if _ollama_is_reachable(base_url):
245            logger.debug("Using existing Ollama at {}", base_url)
246        else:
247            _ensure_reachable(base_url, manage=True)
248    elif not _ollama_is_reachable(base_url):
249        raise _unreachable_error(base_url, manage_hint=True)
250
251    _resolved_base_url = _assert_local_base_url(base_url)
252    return _resolved_base_url
253
254
255def _safe_lock_name(model: str) -> str:
256    return re.sub(r"[^A-Za-z0-9._-]+", "_", model)
257
258
259@contextmanager
260def _pull_lock(model: str) -> Iterator[None]:
261    lock_dir = get_pull_lock_dir()
262    lock_dir.mkdir(parents=True, exist_ok=True)
263    lock_path = lock_dir / f"{_safe_lock_name(model)}.lock"
264    with open(lock_path, "a+", encoding="utf-8") as lock_file:
265        logger.debug("Waiting for pull lock: {}", lock_path)
266        deadline = time.monotonic() + _PULL_LOCK_TIMEOUT_S
267        while True:
268            try:
269                fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
270                break
271            except BlockingIOError:
272                if time.monotonic() > deadline:
273                    raise TimeoutError(
274                        f"Timed out waiting for pull lock {lock_path} "
275                        f"(remove the lock file if a job crashed mid-pull)"
276                    )
277                logger.info(
278                    "Waiting for another job to finish pulling {}…", model
279                )
280                time.sleep(2.0)
281        try:
282            yield
283        finally:
284            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
285
286
287def model_exists(base_url: str, model: str) -> bool:
288    try:
289        response = requests.get(f"{_api_root(base_url)}/api/tags", timeout=30.0)
290        response.raise_for_status()
291    except requests.RequestException as exc:
292        logger.debug("Could not list Ollama models: {}", exc)
293        return False
294
295    for entry in response.json().get("models", []):
296        name = entry.get("name", "")
297        if name == model or name.startswith(f"{model}:"):
298            return True
299    return False
300
301
302def _run_ollama_pull(base_url: str, model: str) -> None:
303    host = _ollama_host_from_base_url(base_url)
304    env = os.environ.copy()
305    env["OLLAMA_HOST"] = f"http://{host}"
306    logger.info("Pulling Ollama model {} (OLLAMA_HOST={})", model, env["OLLAMA_HOST"])
307    try:
308        subprocess.run(
309            [_ollama_binary(), "pull", model],
310            env=env,
311            check=True,
312        )
313    except subprocess.CalledProcessError as exc:
314        raise RuntimeError(
315            f"Failed to pull Ollama model {model!r} from {env['OLLAMA_HOST']}. "
316            "If this is a custom GGUF, import it with `ollama create` first, "
317            "then use the tag from `ollama list`."
318        ) from exc
319
320
321def ensure_model_pulled(base_url: str, model: str) -> None:
322    cfg = get_local_llm_config()
323    auto_pull = _env_flag("WUENLP_LOCAL_LLM_AUTO_PULL")
324    if auto_pull is None:
325        auto_pull = cfg.auto_pull
326    if not auto_pull:
327        return
328
329    if model_exists(base_url, model):
330        return
331
332    with _pull_lock(model):
333        if model_exists(base_url, model):
334            return
335        _require_ollama_on_path()
336        _run_ollama_pull(base_url, model)
337
338
339def _warmup_timeout_s() -> float:
340    raw = os.environ.get("WUENLP_LOCAL_LLM_WARMUP_TIMEOUT", "600").strip()
341    try:
342        return float(raw)
343    except ValueError:
344        return 600.0
345
346
347def _warmup_keep_alive() -> str:
348    return os.environ.get("WUENLP_LOCAL_LLM_KEEP_ALIVE", "-1").strip() or "-1"
349
350
351def _warmup_generate(base_url: str, model: str) -> None:
352    api = _api_root(base_url)
353    response = requests.post(
354        f"{api}/api/generate",
355        json={
356            "model": model,
357            "prompt": "OK",
358            "stream": False,
359            "keep_alive": _warmup_keep_alive(),
360            "options": {"num_predict": 4},
361        },
362        timeout=_warmup_timeout_s(),
363    )
364    response.raise_for_status()
365
366
367def warm_up_model(base_url: str, model: str) -> None:
368    """Load model weights with a tiny generate call (avoids hung first chat request)."""
369    if _env_flag("WUENLP_LOCAL_LLM_WARMUP") is False:
370        return
371    logger.info("Warming up Ollama model {} …", model)
372    last_exc: requests.RequestException | None = None
373    for attempt in range(2):
374        try:
375            _warmup_generate(base_url, model)
376            return
377        except requests.RequestException as exc:
378            last_exc = exc
379            if attempt == 0:
380                logger.warning(
381                    "Warm-up attempt 1 failed for {} ({}); restarting Ollama",
382                    model,
383                    exc,
384                )
385                base_url = restart_local_ollama()
386                continue
387    logger.warning(
388        "Ollama warm-up failed for {} after retry ({}); continuing without warm-up",
389        model,
390        last_exc,
391    )
392
393
394def is_model_not_found_error(exc: BaseException) -> bool:
395    message = str(exc).lower()
396    return any(
397        token in message
398        for token in ("not found", "404", "model ''", "does not exist")
399    )