scripts/run_colibri.py

"""Launcher for Colibri (https://github.com/JustVugg/colibri) serving OLMoE.

Colibri is a pure-C MoE engine built per model family: `make -C c olmoe` produces
a standalone `olmoe` binary that exposes an OpenAI-compatible API via `serve`.
Unlike the npm/pip Claw services, it needs two things this repo cannot bundle:

  1. the binary, compiled at image build time (guarded by `--build-arg COLIBRI_BUILD=1`)
  2. a Colibri-format weight container, which the user must provision

Colibri uses its own container format, and no pre-converted OLMoE container is
published on the Hub. If either piece is missing this launcher logs what is
needed and parks the process instead of exiting: supervisord runs every program
with `autorestart=true`, and a fast-exiting program turns into a restart loop
that has taken this Space down before (see the agent-os/odysseus notes in the
Dockerfile).
"""
import json
import os
import shutil
import sys
import time
from pathlib import Path

CONFIG_PATH = Path(os.environ.get("COLIBRI_CONFIG_PATH", "/app/colibri.json"))


def _load_config() -> dict:
    try:
        return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
    except Exception:
        return {}


def _log(message: str) -> None:
    print(f"[colibri] {message}", flush=True)


def _park(reason: str) -> None:
    """Stay alive without work so supervisord does not thrash on restart."""
    _log(reason)
    _log("service parked; set the missing piece and restart this program")
    while True:
        time.sleep(3600)


def _model_ready(model_dir: Path) -> bool:
    return model_dir.is_dir() and any(model_dir.iterdir())


def _download_model(repo: str, revision: str, model_dir: Path) -> bool:
    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        _log("huggingface_hub is not installed; cannot fetch the model container")
        return False

    _log(f"downloading Colibri container from {repo}@{revision} -> {model_dir}")
    try:
        model_dir.parent.mkdir(parents=True, exist_ok=True)
        snapshot_download(
            repo_id=repo,
            revision=revision,
            local_dir=str(model_dir),
            token=os.environ.get("HF_TOKEN") or None,
        )
    except Exception as exc:
        _log(f"download failed: {exc}")
        return False
    return _model_ready(model_dir)


def main() -> None:
    config = _load_config()
    gateway = config.get("gateway", {})
    model_cfg = config.get("model", {})
    runtime = config.get("runtime", {})

    host = os.environ.get("COLIBRI_HOST") or gateway.get("host") or "127.0.0.1"
    port = str(os.environ.get("COLIBRI_PORT") or gateway.get("port") or 18797)
    model_id = os.environ.get("COLIBRI_MODEL_ID") or model_cfg.get("model_id") or "olmoe"
    model_dir = Path(os.environ.get("COLIBRI_MODEL_DIR") or model_cfg.get("dir") or "/data/models/olmoe")
    repo = (os.environ.get("COLIBRI_MODEL_REPO") or model_cfg.get("repo") or "").strip()
    revision = os.environ.get("COLIBRI_MODEL_REVISION") or model_cfg.get("revision") or "main"

    binary = os.environ.get("COLIBRI_BIN", "coli-olmoe")
    resolved = shutil.which(binary)
    if not resolved:
        _park(
            f"engine binary {binary!r} not found - rebuild the image with "
            "`docker build --build-arg COLIBRI_BUILD=1`"
        )

    if not _model_ready(model_dir):
        if not repo:
            _park(
                f"no weights at {model_dir} and COLIBRI_MODEL_REPO is unset - convert OLMoE "
                "with c/tools/convert_olmoe_merged.py, push the container to a Hub repo, "
                "then set COLIBRI_MODEL_REPO"
            )
        if not _download_model(repo, revision, model_dir):
            _park(f"could not provision weights at {model_dir}")

    threads = str(os.environ.get("OMP_NUM_THREADS") or runtime.get("threads") or 2)
    env = dict(os.environ)
    env["OMP_NUM_THREADS"] = threads
    env["COLI_MODEL"] = str(model_dir)
    env.setdefault("COLI_MAX_QUEUE", str(runtime.get("max_queue", 4)))
    env.setdefault("COLI_QUEUE_TIMEOUT", str(runtime.get("queue_timeout_sec", 600)))
    env.setdefault("COLI_KV_SLOTS", str(runtime.get("kv_slots", 2)))

    extra_args = (os.environ.get("COLIBRI_SERVE_ARGS") or "").split()
    argv = [
        binary,
        "serve",
        "--host", host,
        "--port", port,
        "--model", str(model_dir),
        "--model-id", model_id,
        *extra_args,
    ]

    _log(f"starting: {' '.join(argv)} (OMP_NUM_THREADS={threads})")
    try:
        os.execvpe(resolved, argv, env)
    except OSError as exc:
        _park(f"exec failed: {exc}")


if __name__ == "__main__":
    main()