mjolnir-cli 0.2.0__py3-none-any.whl

mjolnir_cli/litellm.py

@@ -0,0 +1,254 @@
+"""LiteLLM proxy stack: ``litellm`` + ``postgres`` + ``redis``.
+
+Independent of the vLLM server — ``mjolnir litellm down`` leaves the vLLM
+running, and ``mjolnir serve down`` leaves the proxy up. The proxy reaches
+the vLLM through ``host.docker.internal`` (the host-published vLLM port), so
+the two never share a docker network and neither knows the other's
+container name.
+
+The config is rendered on every ``mjolnir litellm up`` from the user's
+template (``$MJOLNIR_LITELLM_TEMPLATE`` or, by default, inside the litellm
+data dir: ``~/.local/share/mjolnir/litellm/litellm_config.template.yaml``)
+with the *currently selected* model/config (``mjolnir model``):
+``${MAIN_LLM_MODEL}`` = the served name, ``${MAIN_LLM_MODEL_NAME}`` = the raw
+``model:`` field, ``${RAW_MODEL_SUFFIXED}`` = ``<model>-<quant>``.
+Everything else resolves from the user's env file (``$MJOLNIR_LITELLM_ENV``,
+default ``~/.local/share/mjolnir/litellm/litellm.env``).
+
+ALL user-persisted litellm files live in one dir — ``$MJOLNIR_DATA/litellm``
+(default ``~/.local/share/mjolnir/litellm``): the template, the env file,
+the rendered config, and the stack data (postgres + redis).
+
+The compose machinery (env file, data dir, up/down/status/logs, rollback,
+health-wait) lives in :mod:`mjolnir.stack`; this module is the stack's glue:
+the spec, the template rendering, the model-variable injection.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import yaml
+
+from mjolnir_cli import stack
+from mjolnir_cli.config import Settings, config_field
+
+SPEC = stack.StackSpec(
+    name="litellm",
+    compose_rel=Path("docker/litellm/docker-compose.yaml"),
+    env_file_var="MJOLNIR_LITELLM_ENV",
+    data_dir_var="MJOLNIR_LITELLM_DATA",
+    data_dir_env="LITELLM_DATA_DIR",
+    data_subdirs=("db", "redis"),
+    port_env="LITELLM_PORT",
+    default_port=8000,
+    wait_timeout_s=600.0,
+)
+DEFAULT_PROXY_PORT = SPEC.default_port
+
+
+class LitellmError(stack.StackError):
+    pass
+
+
+# ── paths ──────────────────────────────────────────────────────────────────
+
+
+def template_path() -> Path:
+    """The user's template: ``$MJOLNIR_LITELLM_TEMPLATE`` or, by default,
+    in the litellm data dir:
+    ``~/.local/share/mjolnir/litellm/litellm_config.template.yaml``."""
+    env = os.environ.get("MJOLNIR_LITELLM_TEMPLATE")
+    if env:
+        return Path(env).expanduser()
+    return data_dir() / "litellm_config.template.yaml"
+
+
+def env_file_path() -> Path:
+    """The user's env file (optional): ``$MJOLNIR_LITELLM_ENV`` or, by
+    default, in the litellm data dir:
+    ``~/.local/share/mjolnir/litellm/litellm.env``."""
+    return SPEC.env_file_path()
+
+
+def data_dir() -> Path:
+    """Stack data root: ``$MJOLNIR_LITELLM_DATA`` or
+    ``$MJOLNIR_DATA/litellm`` (default ``~/.local/share/mjolnir/litellm``);
+    holds the template/env/rendered config + the postgres/redis state
+    dirs."""
+    return SPEC.data_dir()
+
+
+def compose_file() -> Path:
+    return SPEC.compose_file()
+
+
+# ── env loading + rendering ────────────────────────────────────────────────
+
+
+def _model_field(s: Settings) -> str:
+    """The raw ``model:`` field of the config yaml (falls back to the model
+    path) — the ``MAIN_LLM_MODEL_NAME`` the template's alias table wants."""
+    return config_field(s.config_path, "model") or s.model
+
+
+def injected_env(s: Settings, override_base_url: str | None = None) -> dict[str, str]:
+    """The variables mjolnir injects from the active model/config."""
+    name = _model_field(s)
+    return {
+        "MAIN_LLM_MODEL": s.served_model,
+        "MAIN_LLM_MODEL_NAME": name,
+        "RAW_MODEL_SUFFIXED": f"{name}-{s.quant}",
+        "LITELLM_MAIN_MODEL_BASE_URL": override_base_url or f"http://host.docker.internal:{s.port}/v1",
+    }
+
+
+def resolved_env(s: Settings, proxy_port: int | None = None) -> dict[str, str]:
+    """Merged env for rendering + compose interpolation:
+    user env file < injected model variables (the CLI selection wins) <
+    stack defaults (data dir, proxy port)."""
+    env = stack.load_env_file(env_file_path())
+    extra = injected_env(s, override_base_url=env.get("LITELLM_MAIN_MODEL_BASE_URL"))
+    # CLI flag > env file > default
+    if proxy_port:
+        extra["LITELLM_PORT"] = str(proxy_port)
+    return stack.resolve_env(SPEC, extra)
+
+
+def validate(rendered: str) -> list[str]:
+    """Problems with a rendered config (empty list = ok)."""
+    problems: list[str] = []
+    try:
+        doc = yaml.safe_load(rendered)
+    except yaml.YAMLError as e:
+        return [f"rendered config is not valid YAML: {e}".replace("\n", " ")]
+    ml = doc.get("model_list") if isinstance(doc, dict) else None
+    if not isinstance(ml, list):
+        problems.append("no model_list in the rendered config")
+    else:
+        for i, e in enumerate(ml):
+            name = (e or {}).get("model_name")
+            if not name or not str(name).strip():
+                problems.append(f"model_list[{i}] has an empty model_name — check the template variables")
+    return problems
+
+
+def render_config(s: Settings, proxy_port: int | None = None) -> Path:
+    """Render the template → ``data_dir()/litellm_config.yaml``; raises
+    ``LitellmError`` on a missing template or unresolved variables."""
+    tpl = template_path()
+    if not tpl.exists():
+        example = compose_file().parent / "litellm_config.template.example.yaml"
+        msg = (
+            f"template not found: {tpl}\n"
+            "  start from the example (or override the location with "
+            "$MJOLNIR_LITELLM_TEMPLATE):\n"
+            f"  cp {example} {tpl}"
+        )
+        legacy = Path.home() / "litellm_config.template.yaml"
+        if legacy.exists() and legacy != tpl:
+            msg += f"\n  found a legacy copy — move it:  mv {legacy} {tpl}"
+        raise LitellmError(msg)
+    env = resolved_env(s, proxy_port)
+    rendered, missing = stack.render(tpl.read_text(), env)
+    if missing:
+        src = env_file_path()
+        raise LitellmError(
+            "unresolved template variable(s): "
+            + ", ".join("${" + m + "}" for m in missing)
+            + f"\n  define them in {src} (see the examples in "
+            f"{compose_file().parent})"
+        )
+    problems = validate(rendered)
+    if problems:
+        raise LitellmError("rendered config invalid: " + "; ".join(problems))
+    out = data_dir() / "litellm_config.yaml"
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(rendered)
+    return out
+
+
+# ── stack glue (the compose driver is mjolnir.stack) ───────────────────────
+
+
+def _port_of(env: dict[str]) -> int:
+    try:
+        return int(env.get("LITELLM_PORT", DEFAULT_PROXY_PORT))
+    except ValueError:
+        return DEFAULT_PROXY_PORT
+
+
+def up(
+    s: Settings, proxy_port: int | None = None, wait: bool = True, wait_timeout_s: float = 600.0, dry_run: bool = False
+) -> int:
+    """Render the config and start the stack (proxy + postgres + redis)."""
+    cfg_path: Path | None = None
+
+    def prepare(s2: Settings) -> dict[str, str]:
+        nonlocal cfg_path
+        cfg_path = render_config(s2, proxy_port)
+        if dry_run:
+            print(f"(rendered config: {cfg_path})")
+        return {}
+
+    def health(env: dict[str]) -> str:
+        return f"http://127.0.0.1:{_port_of(env)}/health/liveliness"
+
+    def banner(env: dict[str]) -> list[str]:
+        p = _port_of(env)
+        return [
+            f"litellm proxy: http://127.0.0.1:{p} (master key: {env.get('LITELLM_MASTER_KEY', 'sk-1234')})",
+            f"  UI:  http://127.0.0.1:{p}/ui"
+            f"  ({env.get('LITELLM_UI_USERNAME', 'admin')}/"
+            f"{env.get('LITELLM_UI_PASSWORD', 'admin')})",
+            f"  vllm main model: {env['MAIN_LLM_MODEL']} @ {env['LITELLM_MAIN_MODEL_BASE_URL']}",
+        ]
+
+    try:
+        stack.up(
+            SPEC,
+            s,
+            prepare=prepare,
+            wait=wait,
+            wait_timeout_s=wait_timeout_s,
+            dry_run=dry_run,
+            health=health,
+            banner=banner,
+            compose_wait=True,
+        )
+    except stack.StackError as e:
+        env = resolved_env(s, proxy_port)
+        hint = (
+            f"\n  is host port {_port_of(env)} already taken? try: mjolnir litellm up --port <other>"
+            if "port is already allocated" in str(e)
+            else ""
+        )
+        raise LitellmError(str(e) + hint)
+    return 0
+
+
+def down(purge: bool = False, dry_run: bool = False) -> int:
+    """Stop the stack (data kept; ``purge`` wipes postgres + redis data)."""
+    return stack.down(SPEC, purge=purge, dry_run=dry_run, keep_note="postgres/redis state survives")
+
+
+def status(s: Settings, proxy_port: int | None = None) -> dict:
+    env = resolved_env(s, proxy_port)
+    port = _port_of(env)
+    url = f"http://127.0.0.1:{port}"
+    return stack.status(
+        SPEC,
+        extra={
+            "proxy": {
+                "url": url,
+                "liveliness": stack.probe(url + "/health/liveliness"),
+                "config": str(data_dir() / "litellm_config.yaml"),
+                "template": str(template_path()),
+            },
+        },
+    )
+
+
+def logs(follow: bool = False, tail: int = 100) -> int:
+    return stack.logs(SPEC, follow=follow, tail=tail)

mjolnir_cli/perses.py

@@ -0,0 +1,214 @@
+"""Perses + Prometheus observability stack for the vLLM server.
+
+Same contract as litellm: **independent** of the vLLM server — ``perses
+down`` leaves the server running, and ``serve down`` leaves the stack up.
+Prometheus scrapes the server's /metrics through
+``host.docker.internal:<port>`` (the resolved port, same as
+``mjolnir serve``), so the stack works whether the server is up or down.
+The two vLLM dashboards (vendored from vllm-project/vllm,
+``docker/perses/dashboards/``) are PROVISIONED into Perses from a read-only
+folder — no manual import.
+
+All user-persisted files live in one dir — ``$MJOLNIR_PERSER_DATA`` or
+``$MJOLNIR_DATA/perses`` (default ``~/.local/share/mjolnir/perses``): the
+env file, the rendered prometheus.yml, and the stack data (``db/`` = Perses
+file DB, ``prom/`` = Prometheus TSDB).
+
+The compose machinery (env file, data dir, up/down/status/logs, rollback,
+health-wait) lives in :mod:`mjolnir.stack`; this module is the stack's glue:
+the spec, the prometheus.yml rendering, the health URL, the banner.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+
+from mjolnir_cli import stack
+from mjolnir_cli.config import Settings, resource_path
+
+SPEC = stack.StackSpec(
+    name="perses",
+    compose_rel=Path("docker/perses/docker-compose.yaml"),
+    env_file_var="MJOLNIR_PERSER_ENV",
+    data_dir_var="MJOLNIR_PERSER_DATA",
+    data_dir_env="PERSER_DATA_DIR",
+    data_subdirs=("db", "prom"),
+    port_env="PERSER_PORT",
+    default_port=8080,
+    wait_timeout_s=300.0,
+)
+DEFAULT_PORT = SPEC.default_port
+PROMETHEUS_DEFAULT_PORT = 9090
+
+
+class PersesError(stack.StackError):
+    pass
+
+
+# ── prometheus.yml rendering ───────────────────────────────────────────────
+
+
+def vllm_metrics_target(s: Settings, override: str | None = None) -> str:
+    """The Prometheus scrape target for the VLLM server as ``host:port``
+    (prometheus target form — the path comes from ``metrics_path``): the
+    resolved port on host.docker.internal (the host-published port the
+    vLLM container listens on), or the ``PERSER_VLLM_TARGET`` env-file
+    override."""
+    return override or f"host.docker.internal:{s.port}"
+
+
+def render_prometheus(s: Settings, vllm_target: str | None = None) -> Path:
+    """Render ``prometheus.template.yaml`` → ``data_dir()/prometheus.yml``
+    with the vLLM scrape target of the active config."""
+    tpl = resource_path(SPEC.compose_rel.parent / "prometheus.template.yaml")
+    if not tpl.exists():
+        raise PersesError(f"prometheus template not found: {tpl}")
+    rendered, missing = stack.render(tpl.read_text(), {"VLLM_METRICS_TARGET": vllm_metrics_target(s, vllm_target)})
+    if missing:
+        raise PersesError(f"unresolved variable(s) in the prometheus template: {', '.join(missing)}")
+    try:
+        yaml.safe_load(rendered)
+    except yaml.YAMLError as e:
+        raise PersesError(f"rendered prometheus config invalid: {e}".replace("\n", " "))
+    out = SPEC.data_dir() / "prometheus.yml"
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(rendered)
+    return out
+
+
+# ── stack glue (the compose driver is mjolnir.stack) ───────────────────────
+
+
+def _port_of(env: dict[str]) -> int:
+    try:
+        return int(env.get("PERSER_PORT", DEFAULT_PORT))
+    except ValueError:
+        return DEFAULT_PORT
+
+
+def _prom_port_of(env: dict[str]) -> int:
+    try:
+        return int(env.get("PROMETHEUS_PORT", PROMETHEUS_DEFAULT_PORT))
+    except ValueError:
+        return PROMETHEUS_DEFAULT_PORT
+
+
+def up(
+    s: Settings, port: int | None = None, wait: bool = True, wait_timeout_s: float = 300.0, dry_run: bool = False
+) -> int:
+    """Start the stack (prometheus + perses) with the vLLM dashboards
+    provisioned. The vLLM scrape target = the active config's port (the
+    server may be up or down)."""
+    file_env = stack.load_env_file(SPEC.env_file_path())
+    target = vllm_metrics_target(s, file_env.get("PERSER_VLLM_TARGET"))
+    extra: dict[str, str] = {}
+    if port:
+        extra["PERSER_PORT"] = str(port)
+
+    cfg_path: Path | None = None
+
+    def prepare(s2: Settings) -> dict[str, str]:
+        nonlocal cfg_path
+        for sub in SPEC.data_subdirs:
+            d = SPEC.data_dir() / sub
+            d.mkdir(parents=True, exist_ok=True)
+            # the containers run nonroot (perses uid 65532, prometheus
+            # 65534) and need to write their state into these host dirs
+            d.chmod(0o777)
+        cfg_path = render_prometheus(s2, file_env.get("PERSER_VLLM_TARGET"))
+        if dry_run:
+            print(f"(rendered prometheus config: {cfg_path})")
+        return extra
+
+    def health(env: dict[str]) -> str:
+        return f"http://127.0.0.1:{_port_of(env)}/api/v1/projects"
+
+    def banner(env: dict[str]) -> list[str]:
+        p, pp = _port_of(env), _prom_port_of(env)
+        return [
+            f"perses:     http://127.0.0.1:{p}  (dashboards: Performance Statistics, Query Statistics)",
+            f"prometheus: http://127.0.0.1:{pp}  (scrapes {target}/metrics — targets: http://127.0.0.1:{pp}/targets)",
+        ]
+
+    try:
+        stack.up(
+            SPEC,
+            s,
+            prepare=prepare,
+            wait=wait,
+            wait_timeout_s=wait_timeout_s,
+            dry_run=dry_run,
+            health=health,
+            banner=banner,
+        )
+    except stack.StackError as e:
+        raise PersesError(str(e))
+    return 0
+
+
+def down(purge: bool = False, dry_run: bool = False) -> int:
+    """Stop the stack (the vLLM server keeps running; data kept by default,
+    ``purge`` wipes the perses db + prometheus TSDB)."""
+    return stack.down(SPEC, purge=purge, dry_run=dry_run, keep_note="prometheus TSDB + perses db survive")
+
+
+def _rendered_target() -> str | None:
+    """The scrape target the rendered ``prometheus.yml`` actually points at
+    (what the *running* stack is scraping) — the source of truth, which may
+    differ from the settings-resolved port (e.g. after an ``up`` with a
+    different port). None if not yet rendered / unreadable."""
+    yml = SPEC.data_dir() / "prometheus.yml"
+    if not yml.exists():
+        return None
+    try:
+        data = yaml.safe_load(yml.read_text())
+    except yaml.YAMLError:
+        return None
+    for job in (data or {}).get("scrape_configs", []):
+        for sc in job.get("static_configs", []):
+            targets = sc.get("targets") or []
+            if targets:
+                return targets[0]
+    return None
+
+
+def status(s: Settings, port: int | None = None) -> dict:
+    env = stack.resolve_env(SPEC, {"PERSER_PORT": str(port)} if port else None)
+    p, pp = _port_of(env), _prom_port_of(env)
+    target = _rendered_target() or vllm_metrics_target(
+        s, stack.load_env_file(SPEC.env_file_path()).get("PERSER_VLLM_TARGET")
+    )
+    services = stack.project_services(SPEC)
+    live = any(v.get("state") == "running" for v in services.values())
+    return stack.status(
+        SPEC,
+        extra={
+            "perses": {
+                "url": f"http://127.0.0.1:{p}",
+                "api": stack.probe(f"http://127.0.0.1:{p}/api/v1/projects") if live else "down",
+            },
+            "prometheus": {
+                "url": f"http://127.0.0.1:{pp}",
+                "health": stack.probe(f"http://127.0.0.1:{pp}/-/healthy") if live else "down",
+            },
+            "vllm_target": f"http://{target}/metrics",
+            "data_dir": str(SPEC.data_dir()),
+        },
+    )
+
+
+def logs(follow: bool = False, tail: int = 100, service: str | None = None) -> int:
+    """Tail a stack service's logs (default ``perses``; ``--service
+    prometheus`` for the scraper)."""
+    return stack.logs(SPEC, follow=follow, tail=tail, service=service)
+
+
+def running_port() -> int | None:
+    """The Perses UI port if the stack is up (the ``serve status`` hint),
+    else None."""
+    svc = stack.project_services(SPEC).get(SPEC.name) or {}
+    if svc.get("state") != "running":
+        return None
+    return _port_of(stack.resolve_env(SPEC))

mjolnir_cli/picker.py

@@ -0,0 +1,70 @@
+"""A minimal arrow-key selection menu on ``curses`` (stdlib — no new deps).
+
+Used by ``mjolnir model`` / ``mjolnir image``: pick the active model
+config or the vLLM image with the arrow keys and Enter. Esc/q cancels.
+"""
+
+from __future__ import annotations
+
+import curses
+import sys
+
+
+class Cancelled(Exception):
+    """The user cancelled the picker (Esc / q), or the terminal could not
+    run the curses UI."""
+
+
+_FOOTER = "↑/↓ move · enter select · esc cancel"
+
+
+def _run(stdscr, choices: list[str], start: int, title: str) -> str:
+    curses.curs_set(0)
+    cur = start
+    while True:
+        stdscr.erase()
+        cols = curses.COLS
+        lines = curses.LINES
+        stdscr.addstr(0, 0, title[: cols - 1])
+        room = max(lines - 3, 1)  # title row + footer row
+        top = max(0, min(cur - room + 1, len(choices) - room))
+        for i in range(top, min(top + room, len(choices))):
+            y = i - top + 1
+            line = " " + choices[i]
+            if i == cur:
+                stdscr.addstr(y, 0, line[: cols - 1], curses.A_REVERSE)
+            else:
+                stdscr.addstr(y, 0, line[: cols - 1])
+        stdscr.addstr(lines - 1, 0, _FOOTER[: cols - 1])
+        stdscr.refresh()
+
+        key = stdscr.getch()
+        if key in (curses.KEY_UP, ord("k")):
+            cur = max(0, cur - 1)
+        elif key in (curses.KEY_DOWN, ord("j")):
+            cur = min(len(choices) - 1, cur + 1)
+        elif key == curses.KEY_HOME:
+            cur = 0
+        elif key == curses.KEY_END:
+            cur = len(choices) - 1
+        elif key in (curses.KEY_ENTER, 10, 13, ord(" ")):
+            return choices[cur]
+        elif key in (27, ord("q")):
+            raise Cancelled
+
+
+def select(title: str, choices: list[str], default: str | None = None) -> str:
+    """Arrow-key menu over ``choices``; returns the chosen item.
+
+    Raises ``Cancelled`` on Esc/q or when the terminal can't run curses
+    (non-tty, TERM=dumb). Callers should check ``isatty`` first and offer
+    the explicit-argument path instead."""
+    if not choices:
+        raise ValueError("select() needs at least one choice")
+    start = choices.index(default) if default in choices else 0
+    if not (sys.stdin.isatty() and sys.stdout.isatty()):
+        raise Cancelled
+    try:
+        return curses.wrapper(_run, choices, start, title)
+    except (curses.error, OSError) as e:
+        raise Cancelled from e