File size: 4,675 Bytes
9fa0cbc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """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()
|