Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| from contextlib import asynccontextmanager | |
| import os | |
| from pathlib import Path | |
| import shlex | |
| import shutil | |
| import subprocess | |
| import time | |
| from urllib.error import URLError | |
| from urllib.request import Request as UrlRequest | |
| from urllib.request import urlopen | |
| import modal | |
| from fastapi import FastAPI, HTTPException, Request | |
| from fastapi.responses import JSONResponse, StreamingResponse | |
| import httpx | |
| from starlette.background import BackgroundTask | |
| try: | |
| from scripts.llamacpp_modal_config import ( | |
| DeploymentConfig, | |
| ModalAppConfig, | |
| ModelConfig, | |
| load_deployment_config, | |
| ) | |
| except ModuleNotFoundError: | |
| from llamacpp_modal_config import ( # type: ignore[no-redef] | |
| DeploymentConfig, | |
| ModalAppConfig, | |
| ModelConfig, | |
| load_deployment_config, | |
| ) | |
| MINUTES = 60 | |
| DEFAULT_CONFIG_PATH = Path("config/modal-models.yaml") | |
| SERVER_BIN = os.getenv("MODAL_LLAMA_CPP_SERVER_BIN", "/app/llama-server") | |
| def deployment_from_environment() -> DeploymentConfig: | |
| configured_path = os.getenv("MODAL_LLAMA_CPP_CONFIG", "").strip() | |
| if configured_path: | |
| return load_deployment_config(configured_path) | |
| if DEFAULT_CONFIG_PATH.is_file(): | |
| return load_deployment_config(DEFAULT_CONFIG_PATH) | |
| app_config = ModalAppConfig( | |
| name=os.getenv("MODAL_LLAMA_CPP_APP_NAME", ModalAppConfig.name), | |
| gpu=os.getenv("MODAL_LLAMA_CPP_GPU", ModalAppConfig.gpu), | |
| image=os.getenv("MODAL_LLAMA_CPP_IMAGE", ModalAppConfig.image), | |
| secret_name=os.getenv( | |
| "MODAL_LLAMA_CPP_SECRET_NAME", ModalAppConfig.secret_name | |
| ), | |
| timeout_minutes=int( | |
| os.getenv("MODAL_LLAMA_CPP_TIMEOUT_MINUTES", "15") | |
| ), | |
| scaledown_minutes=int( | |
| os.getenv("MODAL_LLAMA_CPP_SCALEDOWN_MINUTES", "15") | |
| ), | |
| startup_timeout_seconds=int( | |
| os.getenv("MODAL_LLAMA_CPP_STARTUP_TIMEOUT", "900") | |
| ), | |
| max_concurrent_inputs=int( | |
| os.getenv("MODAL_LLAMA_CPP_MAX_CONCURRENT_INPUTS", "12") | |
| ), | |
| ) | |
| model = ModelConfig( | |
| route="", | |
| role="text", | |
| model_ref=os.getenv( | |
| "MODAL_LLAMA_CPP_MODEL_REF", "openbmb/MiniCPM-V-4_5-gguf:Q4_K_M" | |
| ), | |
| port=8001, | |
| ctx_size=int(os.getenv("MODAL_LLAMA_CPP_CTX_SIZE", "8192")), | |
| gpu_layers=int(os.getenv("MODAL_LLAMA_CPP_GPU_LAYERS", "999")), | |
| threads=int(os.getenv("MODAL_LLAMA_CPP_THREADS", "8")), | |
| parallel=int(os.getenv("MODAL_LLAMA_CPP_PARALLEL", "2")), | |
| extra_args=tuple( | |
| shlex.split(os.getenv("MODAL_LLAMA_CPP_EXTRA_ARGS", "")) | |
| ), | |
| ) | |
| return DeploymentConfig( | |
| app=app_config, | |
| models=(model,), | |
| source="legacy environment variables", | |
| legacy=True, | |
| ) | |
| DEPLOYMENT = deployment_from_environment() | |
| APP_CONFIG = DEPLOYMENT.app | |
| RUNTIME_CONFIG_PATH = "/root/config/modal-models.yaml" | |
| app = modal.App(APP_CONFIG.name) | |
| hf_cache = modal.Volume.from_name("secret-student-hf-cache", create_if_missing=True) | |
| llama_cache = modal.Volume.from_name( | |
| "secret-student-llamacpp-cache", create_if_missing=True | |
| ) | |
| image = ( | |
| modal.Image.from_registry(APP_CONFIG.image, add_python="3.12") | |
| .entrypoint([]) | |
| .pip_install( | |
| "fastapi>=0.136.3", | |
| "httpx>=0.28.1", | |
| "uvicorn[standard]>=0.49.0", | |
| ) | |
| .add_local_python_source("scripts.llamacpp_modal_config", copy=True) | |
| .env( | |
| { | |
| "HF_XET_HIGH_PERFORMANCE": "1", | |
| "LLAMA_CACHE": "/root/.cache/llama.cpp", | |
| } | |
| ) | |
| ) | |
| if not DEPLOYMENT.legacy: | |
| image = image.add_local_file( | |
| DEPLOYMENT.source, | |
| RUNTIME_CONFIG_PATH, | |
| copy=True, | |
| ).env({"MODAL_LLAMA_CPP_CONFIG": RUNTIME_CONFIG_PATH}) | |
| def resolve_llama_server_bin() -> str: | |
| candidates = [ | |
| SERVER_BIN, | |
| shutil.which("llama-server") or "", | |
| "/app/llama-server", | |
| "/usr/local/bin/llama-server", | |
| ] | |
| for candidate in candidates: | |
| if candidate and os.path.isfile(candidate) and os.access(candidate, os.X_OK): | |
| return candidate | |
| raise FileNotFoundError(f"Could not find llama-server. Tried: {candidates}") | |
| def build_server_command( | |
| model: ModelConfig, server_bin: str, api_key: str | |
| ) -> list[str]: | |
| return [ | |
| server_bin, | |
| "--host", | |
| "127.0.0.1", | |
| "--port", | |
| str(model.port), | |
| "-hf", | |
| model.model_ref, | |
| "--ctx-size", | |
| str(model.ctx_size), | |
| "--n-gpu-layers", | |
| str(model.gpu_layers), | |
| "--threads", | |
| str(model.threads), | |
| "--parallel", | |
| str(model.parallel), | |
| "--api-key", | |
| api_key, | |
| *model.extra_args, | |
| ] | |
| def start_model_servers(api_key: str) -> dict[str, subprocess.Popen]: | |
| server_bin = resolve_llama_server_bin() | |
| processes: dict[str, subprocess.Popen] = {} | |
| try: | |
| for model in DEPLOYMENT.models: | |
| command = build_server_command(model, server_bin, api_key) | |
| label = model.route or "root" | |
| printable = ["<redacted>" if value == api_key else value for value in command] | |
| print(f"Starting llama.cpp model {label}: {' '.join(printable)}") | |
| processes[model.route] = subprocess.Popen( | |
| command, start_new_session=True | |
| ) | |
| wait_for_model_servers(processes, api_key) | |
| return processes | |
| except BaseException: | |
| stop_model_servers(processes) | |
| raise | |
| def wait_for_model_servers( | |
| processes: dict[str, subprocess.Popen], api_key: str | |
| ) -> None: | |
| pending = {model.route: model for model in DEPLOYMENT.models} | |
| deadline = time.monotonic() + APP_CONFIG.startup_timeout_seconds | |
| headers = {"Authorization": f"Bearer {api_key}"} | |
| while pending and time.monotonic() < deadline: | |
| for route, model in tuple(pending.items()): | |
| process = processes[route] | |
| return_code = process.poll() | |
| if return_code is not None: | |
| raise RuntimeError( | |
| f"llama-server for {route or 'root'} exited during startup " | |
| f"with code {return_code}." | |
| ) | |
| try: | |
| health_url = f"http://127.0.0.1:{model.port}/health" | |
| with urlopen(UrlRequest(health_url, headers=headers), timeout=2) as response: | |
| if response.status < 500: | |
| print(f"llama.cpp model {route or 'root'} is ready on port {model.port}") | |
| pending.pop(route) | |
| except (OSError, URLError): | |
| pass | |
| if pending: | |
| time.sleep(2) | |
| if pending: | |
| routes = ", ".join(route or "root" for route in pending) | |
| raise TimeoutError( | |
| f"llama-server models did not become ready within " | |
| f"{APP_CONFIG.startup_timeout_seconds}s: {routes}" | |
| ) | |
| def stop_model_servers(processes: dict[str, subprocess.Popen]) -> None: | |
| for process in processes.values(): | |
| if process.poll() is None: | |
| process.terminate() | |
| deadline = time.monotonic() + 15 | |
| for process in processes.values(): | |
| if process.poll() is not None: | |
| continue | |
| try: | |
| process.wait(timeout=max(0.1, deadline - time.monotonic())) | |
| except subprocess.TimeoutExpired: | |
| process.kill() | |
| def create_proxy_app(api_key: str): | |
| processes: dict[str, subprocess.Popen] = {} | |
| client: httpx.AsyncClient | None = None | |
| models_by_route = {model.route: model for model in DEPLOYMENT.models} | |
| async def lifespan(_app: FastAPI): | |
| nonlocal processes, client | |
| processes = start_model_servers(api_key) | |
| client = httpx.AsyncClient(timeout=None) | |
| try: | |
| yield | |
| finally: | |
| if client is not None: | |
| await client.aclose() | |
| stop_model_servers(processes) | |
| proxy = FastAPI(title="Secret Student llama.cpp router", lifespan=lifespan) | |
| async def aggregate_health(): | |
| statuses = { | |
| (route or "root"): { | |
| "model_ref": models_by_route[route].model_ref, | |
| "port": models_by_route[route].port, | |
| "running": process.poll() is None, | |
| "return_code": process.poll(), | |
| } | |
| for route, process in processes.items() | |
| } | |
| ready = len(statuses) == len(models_by_route) and all( | |
| status["running"] for status in statuses.values() | |
| ) | |
| return JSONResponse( | |
| status_code=200 if ready else 503, | |
| content={"ready": ready, "models": statuses}, | |
| ) | |
| async def forward(request: Request, route: str, path: str): | |
| if client is None: | |
| raise HTTPException(status_code=503, detail="Model router is starting.") | |
| model = models_by_route.get(route) | |
| if model is None: | |
| raise HTTPException(status_code=404, detail=f"Unknown model route: {route}") | |
| process = processes.get(route) | |
| if process is None or process.poll() is not None: | |
| raise HTTPException( | |
| status_code=503, detail=f"Model route {route or 'root'} is unavailable." | |
| ) | |
| target = f"http://127.0.0.1:{model.port}/{path}" | |
| if request.url.query: | |
| target = f"{target}?{request.url.query}" | |
| excluded_request_headers = {"host", "content-length", "connection"} | |
| headers = { | |
| key: value | |
| for key, value in request.headers.items() | |
| if key.lower() not in excluded_request_headers | |
| } | |
| upstream_request = client.build_request( | |
| request.method, | |
| target, | |
| headers=headers, | |
| content=await request.body(), | |
| ) | |
| upstream = await client.send(upstream_request, stream=True) | |
| excluded_response_headers = { | |
| "content-length", | |
| "connection", | |
| "keep-alive", | |
| "proxy-authenticate", | |
| "proxy-authorization", | |
| "te", | |
| "trailers", | |
| "transfer-encoding", | |
| "upgrade", | |
| } | |
| response_headers = { | |
| key: value | |
| for key, value in upstream.headers.items() | |
| if key.lower() not in excluded_response_headers | |
| } | |
| return StreamingResponse( | |
| upstream.aiter_raw(), | |
| status_code=upstream.status_code, | |
| headers=response_headers, | |
| background=BackgroundTask(upstream.aclose), | |
| ) | |
| if DEPLOYMENT.legacy: | |
| async def forward_legacy(request: Request, path: str): | |
| return await forward(request, "", path) | |
| else: | |
| async def forward_route_root(request: Request, route: str): | |
| return await forward(request, route, "") | |
| async def forward_route(request: Request, route: str, path: str): | |
| return await forward(request, route, path) | |
| return proxy | |
| def serve(): | |
| api_key = ( | |
| os.getenv("LLAMA_ARG_API_KEY", "").strip() | |
| or os.getenv("LLM_API_KEY", "").strip() | |
| ) | |
| if not api_key: | |
| raise RuntimeError( | |
| f"Modal secret {APP_CONFIG.secret_name!r} must contain " | |
| "LLAMA_ARG_API_KEY or LLM_API_KEY." | |
| ) | |
| return create_proxy_app(api_key) | |
| async def main() -> None: | |
| url = (await serve.get_web_url.aio()).rstrip("/") | |
| print(f"Modal config={DEPLOYMENT.source}") | |
| for model in DEPLOYMENT.models: | |
| base_url = url if DEPLOYMENT.legacy else f"{url}/{model.route}" | |
| if model.role == "text": | |
| print(f"LLM_BASE_URL={base_url}") | |
| print(f"LLM_MODEL={model.model_ref}") | |
| elif model.role == "vision": | |
| print(f"VISION_LLM_BASE_URL={base_url}") | |
| print(f"VISION_LLM_MODEL={model.model_ref}") | |
| else: | |
| variable = f"{model.route.upper().replace('-', '_')}_BASE_URL" | |
| print(f"{variable}={base_url}") | |
| print(f"MODEL_ROUTE {model.route or '/'} -> {model.model_ref}") | |
| print(f"Modal GPU={APP_CONFIG.gpu}; llama.cpp image={APP_CONFIG.image}") | |