Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| import logging | |
| import os | |
| import shlex | |
| import signal | |
| import subprocess | |
| import sys | |
| import time | |
| from urllib.error import URLError | |
| from urllib.request import Request, urlopen | |
| import uvicorn | |
| from app.config import Settings, get_settings | |
| logger = logging.getLogger(__name__) | |
| def _existing_model_path(path: str) -> str: | |
| return path if path and os.path.isfile(path) else "" | |
| def build_llama_cpp_command(settings: Settings) -> list[str]: | |
| model_path = _existing_model_path(settings.llama_cpp_model_path) | |
| command = [ | |
| settings.llama_cpp_server_bin, | |
| "--host", | |
| settings.llama_cpp_host, | |
| "--port", | |
| str(settings.llama_cpp_port), | |
| *( | |
| ["--model", model_path] | |
| if model_path | |
| else ["-hf", settings.llama_cpp_model_ref] | |
| ), | |
| "--ctx-size", | |
| str(settings.llama_cpp_ctx_size), | |
| "--n-gpu-layers", | |
| str(settings.llama_cpp_gpu_layers), | |
| "--threads", | |
| str(settings.llama_cpp_threads), | |
| "--parallel", | |
| str(settings.llama_cpp_parallel), | |
| ] | |
| if settings.llama_cpp_api_key: | |
| command.extend(["--api-key", settings.llama_cpp_api_key]) | |
| if settings.llama_cpp_extra_args: | |
| command.extend(shlex.split(settings.llama_cpp_extra_args)) | |
| return command | |
| def build_vision_llama_cpp_command(settings: Settings) -> list[str]: | |
| model_path = _existing_model_path(settings.vision_llama_cpp_model_path) | |
| mmproj_path = _existing_model_path(settings.vision_llama_cpp_mmproj_path) | |
| use_persistent_files = bool(model_path and mmproj_path) | |
| command = [ | |
| settings.llama_cpp_server_bin, | |
| "--host", | |
| settings.llama_cpp_host, | |
| "--port", | |
| str(settings.vision_llama_cpp_port), | |
| *( | |
| ["--model", model_path] | |
| if use_persistent_files | |
| else ["-hf", settings.vision_llama_cpp_model_ref] | |
| ), | |
| "--ctx-size", | |
| str(settings.vision_llama_cpp_ctx_size), | |
| "--n-gpu-layers", | |
| str(settings.vision_llama_cpp_gpu_layers), | |
| "--threads", | |
| str(settings.llama_cpp_threads), | |
| "--parallel", | |
| "1", | |
| ] | |
| if use_persistent_files: | |
| command.extend(["--mmproj", mmproj_path]) | |
| if settings.vision_llm_api_key: | |
| command.extend(["--api-key", settings.vision_llm_api_key]) | |
| if settings.vision_llama_cpp_extra_args: | |
| command.extend(shlex.split(settings.vision_llama_cpp_extra_args)) | |
| return command | |
| def wait_for_llama_cpp( | |
| settings: Settings, | |
| process: subprocess.Popen, | |
| base_url: str, | |
| api_key: str, | |
| ) -> None: | |
| deadline = time.monotonic() + settings.llama_cpp_startup_timeout | |
| health_url = f"{base_url}/health" | |
| headers = {} | |
| if api_key: | |
| headers["Authorization"] = f"Bearer {api_key}" | |
| while time.monotonic() < deadline: | |
| return_code = process.poll() | |
| if return_code is not None: | |
| raise RuntimeError(f"llama-server exited during startup with code {return_code}.") | |
| try: | |
| with urlopen(Request(health_url, headers=headers), timeout=5) as response: | |
| if response.status < 500: | |
| logger.info("Embedded llama.cpp is ready at %s", base_url) | |
| return | |
| except (OSError, URLError): | |
| time.sleep(2) | |
| raise TimeoutError( | |
| f"llama-server did not become ready within {settings.llama_cpp_startup_timeout}s." | |
| ) | |
| def start_embedded_llama_cpp(settings: Settings) -> subprocess.Popen: | |
| command = build_llama_cpp_command(settings) | |
| if settings.llama_cpp_model_path and "-hf" in command: | |
| logger.warning( | |
| "Text model file is missing at %s; downloading %s instead.", | |
| settings.llama_cpp_model_path, | |
| settings.llama_cpp_model_ref, | |
| ) | |
| printable = [ | |
| "<redacted>" if part == settings.llama_cpp_api_key else part | |
| for part in command | |
| ] | |
| logger.info("Starting embedded llama.cpp: %s", " ".join(printable)) | |
| process = subprocess.Popen(command, start_new_session=True) | |
| wait_for_llama_cpp( | |
| settings, process, settings.llm_base_url, settings.llama_cpp_api_key | |
| ) | |
| return process | |
| def start_embedded_vision_llama_cpp(settings: Settings) -> subprocess.Popen: | |
| command = build_vision_llama_cpp_command(settings) | |
| if settings.vision_llama_cpp_model_path and "-hf" in command: | |
| logger.warning( | |
| "Vision model or projector is missing at %s / %s; downloading %s instead.", | |
| settings.vision_llama_cpp_model_path, | |
| settings.vision_llama_cpp_mmproj_path, | |
| settings.vision_llama_cpp_model_ref, | |
| ) | |
| printable = [ | |
| "<redacted>" if part == settings.vision_llm_api_key else part | |
| for part in command | |
| ] | |
| logger.info("Starting embedded vision llama.cpp: %s", " ".join(printable)) | |
| process = subprocess.Popen(command, start_new_session=True) | |
| wait_for_llama_cpp( | |
| settings, | |
| process, | |
| settings.vision_llm_base_url, | |
| settings.vision_llm_api_key, | |
| ) | |
| return process | |
| def stop_process(process: subprocess.Popen | None) -> None: | |
| if process is None or process.poll() is not None: | |
| return | |
| try: | |
| os.killpg(process.pid, signal.SIGTERM) | |
| process.wait(timeout=15) | |
| except (ProcessLookupError, subprocess.TimeoutExpired): | |
| try: | |
| os.killpg(process.pid, signal.SIGKILL) | |
| except ProcessLookupError: | |
| pass | |
| def main() -> None: | |
| logging.basicConfig( | |
| level=os.getenv("LOG_LEVEL", "INFO").upper(), | |
| format="%(asctime)s %(levelname)s %(name)s: %(message)s", | |
| ) | |
| settings = get_settings() | |
| if settings.llm_runtime not in { | |
| "mock", | |
| "external", | |
| "embedded_llamacpp", | |
| "embedded_dual_llamacpp", | |
| }: | |
| raise ValueError( | |
| "LLM_RUNTIME must be mock, external, embedded_llamacpp, or " | |
| "embedded_dual_llamacpp; " | |
| f"received {settings.llm_runtime!r}." | |
| ) | |
| if settings.llm_runtime == "external" and not settings.llm_base_url: | |
| raise ValueError("LLM_BASE_URL is required when LLM_RUNTIME=external.") | |
| model_process: subprocess.Popen | None = None | |
| vision_process: subprocess.Popen | None = None | |
| try: | |
| if settings.llm_runtime in {"embedded_llamacpp", "embedded_dual_llamacpp"}: | |
| model_process = start_embedded_llama_cpp(settings) | |
| if settings.llm_runtime == "embedded_dual_llamacpp": | |
| vision_process = start_embedded_vision_llama_cpp(settings) | |
| uvicorn.run( | |
| "main:app", | |
| host=os.getenv("HOST", "0.0.0.0"), | |
| port=int(os.getenv("PORT", "7860")), | |
| log_level=os.getenv("LOG_LEVEL", "info").lower(), | |
| ) | |
| finally: | |
| stop_process(vision_process) | |
| stop_process(model_process) | |
| if __name__ == "__main__": | |
| try: | |
| main() | |
| except KeyboardInterrupt: | |
| sys.exit(130) | |