from __future__ import annotations import json import os import shutil import site import stat import subprocess import tarfile import tempfile import zipfile from collections import deque from datetime import UTC, datetime from pathlib import Path from urllib.request import urlretrieve from huggingface_hub import hf_hub_download from src.config import settings from src.errors import ApiError _runtime_notes: list[str] = [] _events: deque[dict[str, object]] = deque(maxlen=200) def log_event(event: str, **fields: object) -> None: record: dict[str, object] = { "ts": datetime.now(UTC).isoformat(timespec="seconds"), "event": event, **fields, } _events.append(record) print(json.dumps(record, ensure_ascii=False, default=str), flush=True) def runtime_events(limit: int = 100) -> list[dict[str, object]]: limit = max(1, min(limit, 200)) return list(_events)[-limit:] def nvidia_library_paths() -> list[str]: paths: list[str] = [] seen: set[str] = set() roots = [*site.getsitepackages(), site.getusersitepackages()] for root in roots: if not root: continue base = Path(root) / "nvidia" if not base.exists(): continue for candidate in base.glob("*/lib"): if candidate.is_dir(): value = str(candidate) if value not in seen: seen.add(value) paths.append(value) return paths def ld_library_path_for(binary_path: Path | None = None) -> str: paths: list[str] = [] if binary_path: paths.append(str(binary_path.resolve().parent)) paths.extend(nvidia_library_paths()) existing = os.getenv("LD_LIBRARY_PATH", "") if existing: paths.append(existing) return ":".join(paths) def runtime_status() -> dict[str, object]: return { "bin": str(settings.llama_diffusion_bin), "bin_exists": settings.llama_diffusion_bin.exists(), "model_cache_dir": str(settings.model_cache_dir), "model_file": str(settings.model_cache_dir / settings.gguf_filename), "model_file_exists": (settings.model_cache_dir / settings.gguf_filename).exists(), "nvidia_library_paths": nvidia_library_paths(), "ld_library_path": ld_library_path_for(settings.llama_diffusion_bin), "notes": list(_runtime_notes[-20:]), } def _mark_executable(path: Path) -> None: mode = path.stat().st_mode path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) def _find_binary(root: Path) -> Path | None: for candidate in root.rglob("llama-diffusion-cli"): if candidate.is_file(): return candidate return None def _download_prebuilt_binary(url: str) -> None: settings.bin_dir.mkdir(parents=True, exist_ok=True) log_event("runner.download.start", url=url) with tempfile.TemporaryDirectory() as td: tempdir = Path(td) archive_path = tempdir / "llama-diffusion-download" urlretrieve(url, archive_path) extracted_dir = tempdir / "extract" extracted_dir.mkdir(parents=True, exist_ok=True) if tarfile.is_tarfile(archive_path): with tarfile.open(archive_path) as tf: tf.extractall(extracted_dir) binary = _find_binary(extracted_dir) elif zipfile.is_zipfile(archive_path): with zipfile.ZipFile(archive_path) as zf: zf.extractall(extracted_dir) binary = _find_binary(extracted_dir) else: binary = archive_path if not binary or not binary.exists(): raise ApiError("runtime_error", "Could not find llama-diffusion-cli in downloaded artifact", 500) # Keep sibling shared libraries with the binary. This mirrors the qwen36 # Space strategy: a CUDA runner package is not just the executable. for child in binary.parent.iterdir(): target = settings.bin_dir / child.name if child.is_file() or child.is_symlink(): shutil.copy2(child, target) elif child.is_dir() and child.name in {"lib", "lib64"}: if target.exists(): shutil.rmtree(target) shutil.copytree(child, target, symlinks=True) _mark_executable(settings.llama_diffusion_bin) _runtime_notes.append(f"Installed prebuilt binary from {url}") log_event("runner.download.finish", bin=str(settings.llama_diffusion_bin)) def build_llama_diffusion() -> None: script = Path(__file__).resolve().parent.parent / "scripts" / "build_llama_diffusion.sh" if not script.exists(): raise ApiError("runtime_error", f"Missing build script: {script}", 500) log_event("runner.build.start", script=str(script)) env = dict(os.environ) env.setdefault("LLAMA_SRC_DIR", str(settings.llama_src_dir)) env.setdefault("LLAMA_BIN_DIR", str(settings.bin_dir)) env.setdefault("LLAMA_DIFFUSION_BIN", str(settings.llama_diffusion_bin)) env.setdefault("LLAMA_BUILD_CUDA", "1" if settings.llama_build_cuda else "0") env.setdefault("LLAMA_CMAKE_EXTRA_ARGS", settings.llama_cmake_extra_args) env["LD_LIBRARY_PATH"] = ld_library_path_for(settings.llama_diffusion_bin) proc = subprocess.run( ["bash", str(script)], text=True, capture_output=True, env=env, timeout=int(os.getenv("LLAMA_BUILD_TIMEOUT_SECONDS", "1800")), ) if proc.returncode != 0: log_event("runner.build.failed", stderr_tail=proc.stderr[-1200:]) raise ApiError( "runtime_error", "Failed to build llama-diffusion-cli. stderr tail: " + proc.stderr[-4000:], 500, ) _runtime_notes.append("Built llama-diffusion-cli from llama.cpp DiffusionGemma PR") log_event("runner.build.finish", bin=str(settings.llama_diffusion_bin)) def ensure_runner_binary() -> Path: if settings.llama_diffusion_bin.exists(): _mark_executable(settings.llama_diffusion_bin) log_event("runner.cache_hit", bin=str(settings.llama_diffusion_bin)) return settings.llama_diffusion_bin if settings.llama_diffusion_bin_url: _download_prebuilt_binary(settings.llama_diffusion_bin_url) return settings.llama_diffusion_bin if settings.build_llama_diffusion: build_llama_diffusion() if settings.llama_diffusion_bin.exists(): _mark_executable(settings.llama_diffusion_bin) return settings.llama_diffusion_bin raise ApiError( "runtime_error", "llama-diffusion-cli not found. Set LLAMA_DIFFUSION_BIN, LLAMA_DIFFUSION_BIN_URL, or BUILD_LLAMA_DIFFUSION=1.", 500, ) def ensure_model_file(repo_id: str, filename: str) -> str: settings.model_cache_dir.mkdir(parents=True, exist_ok=True) log_event("model.download.start", repo_id=repo_id, filename=filename) path = hf_hub_download( repo_id=repo_id, filename=filename, local_dir=str(settings.model_cache_dir), token=os.getenv("HF_TOKEN") or None, ) _runtime_notes.append(f"Model ready: {repo_id}/{filename}") log_event("model.download.finish", path=str(path)) return path def prepare_runtime_if_requested() -> None: if settings.prepare_runtime_on_startup: ensure_runner_binary() if settings.download_model_on_startup: ensure_model_file(settings.gguf_repo_id, settings.gguf_filename)