diff --git a/studio/backend/core/inference/llama_cpp_builder.py b/studio/backend/core/inference/llama_cpp_builder.py new file mode 100644 index 0000000000..268bddba8f --- /dev/null +++ b/studio/backend/core/inference/llama_cpp_builder.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Background llama.cpp compilation. + +On server startup, if the llama-server binary is not found, a background +thread clones and builds llama.cpp from source. Features that need +llama.cpp (GGUF chat, export, AI Assist) can await `wait_for_ready()` +which blocks until compilation finishes or fails. +""" + +import os +import shutil +import subprocess +import threading +from pathlib import Path +from typing import Optional + +import structlog +from loggers import get_logger + +logger = get_logger(__name__) + + +class LlamaCppBuilder: + """Manages background llama.cpp compilation.""" + + def __init__(self): + self._ready = threading.Event() + self._building = False + self._error: Optional[str] = None + self._thread: Optional[threading.Thread] = None + self._binary_path: Optional[str] = None + + @property + def is_ready(self) -> bool: + return self._ready.is_set() + + @property + def is_building(self) -> bool: + return self._building + + @property + def error(self) -> Optional[str]: + return self._error + + @property + def binary_path(self) -> Optional[str]: + return self._binary_path + + def check_and_build(self) -> None: + """Check if llama-server exists. If not, start background build.""" + from core.inference.llama_cpp import LlamaCppBackend + + binary = LlamaCppBackend._find_llama_server_binary() + if binary: + self._binary_path = binary + self._ready.set() + logger.info(f"llama-server binary found: {binary}") + return + + logger.info("llama-server binary not found, starting background build...") + self._building = True + self._thread = threading.Thread( + target=self._build, daemon=True, name="llama-cpp-build" + ) + self._thread.start() + + def wait_for_ready(self, timeout: Optional[float] = None) -> bool: + """Block until llama.cpp is ready. Returns True if ready, False on timeout.""" + return self._ready.wait(timeout=timeout) + + def _build(self) -> None: + """Clone and build llama.cpp in the background.""" + try: + import sys + + llama_dir = Path.home() / ".unsloth" / "llama.cpp" + binary_name = ( + "llama-server.exe" if sys.platform == "win32" else "llama-server" + ) + + # Don't rebuild if binary appeared while we were starting + from core.inference.llama_cpp import LlamaCppBackend + + binary = LlamaCppBackend._find_llama_server_binary() + if binary: + self._binary_path = binary + self._building = False + self._ready.set() + return + + # Check prerequisites + if not shutil.which("cmake"): + self._error = "cmake not found. Install cmake to enable GGUF inference." + self._building = False + logger.warning(self._error) + return + if not shutil.which("git"): + self._error = "git not found. Install git to enable GGUF inference." + self._building = False + logger.warning(self._error) + return + + # Clone + if llama_dir.exists(): + shutil.rmtree(llama_dir, ignore_errors=True) + + logger.info("Cloning llama.cpp...") + subprocess.run( + [ + "git", + "clone", + "--depth", + "1", + "https://github.com/ggml-org/llama.cpp.git", + str(llama_dir), + ], + check=True, + capture_output=True, + ) + + # Detect CUDA + cmake_args = [ + "-DBUILD_SHARED_LIBS=OFF", + ] + nvcc = shutil.which("nvcc") + if not nvcc: + for cuda_dir in sorted( + Path("/usr/local").glob("cuda-*/bin/nvcc"), reverse=True + ): + nvcc = str(cuda_dir) + break + if not nvcc and Path("/usr/local/cuda/bin/nvcc").is_file(): + nvcc = "/usr/local/cuda/bin/nvcc" + + if nvcc: + logger.info(f"CUDA detected: {nvcc}") + cmake_args.append("-DGGML_CUDA=ON") + # Detect compute capabilities + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=compute_cap", + "--format=csv,noheader", + ], + capture_output=True, + text=True, + timeout=10, + ) + if result.returncode == 0: + caps = set() + for line in result.stdout.strip().split("\n"): + cap = line.strip().replace(".", "") + if cap: + caps.add(cap) + if caps: + cmake_args.append( + f"-DCMAKE_CUDA_ARCHITECTURES={';'.join(sorted(caps))}" + ) + except Exception: + pass + cmake_args.append("-DCMAKE_CUDA_FLAGS=--threads=0") + else: + logger.info("No CUDA detected, building CPU-only llama.cpp") + + # Use Ninja if available + generator_args = [] + if shutil.which("ninja"): + generator_args = ["-G", "Ninja"] + + # Configure + build_dir = llama_dir / "build" + logger.info("Configuring llama.cpp build...") + subprocess.run( + [ + "cmake", + "-B", + str(build_dir), + "-S", + str(llama_dir), + ] + + generator_args + + cmake_args, + check=True, + capture_output=True, + ) + + # Build + logger.info("Building llama.cpp (this may take a few minutes)...") + subprocess.run( + [ + "cmake", + "--build", + str(build_dir), + "--config", + "Release", + "--target", + "llama-server", + "llama-quantize", + "-j", + ], + check=True, + capture_output=True, + ) + + # Symlink llama-quantize for unsloth-zoo's check_llama_cpp() + quantize_src = build_dir / "bin" / "llama-quantize" + quantize_link = llama_dir / "llama-quantize" + if quantize_src.is_file() and not quantize_link.exists(): + try: + quantize_link.symlink_to(quantize_src) + except Exception: + pass + + # Verify + expected = build_dir / "bin" / binary_name + if expected.is_file(): + self._binary_path = str(expected) + logger.info(f"llama.cpp build complete: {self._binary_path}") + else: + self._error = f"Build completed but binary not found at {expected}" + logger.error(self._error) + + except subprocess.CalledProcessError as e: + stderr = e.stderr.decode() if e.stderr else "" + self._error = f"llama.cpp build failed: {stderr[-500:]}" + logger.error(self._error) + except Exception as e: + self._error = f"llama.cpp build error: {e}" + logger.error(self._error) + finally: + self._building = False + self._ready.set() # Unblock waiters even on failure + + +# ── Singleton ──────────────────────────────────────────────────── + +_builder: Optional[LlamaCppBuilder] = None + + +def get_llama_cpp_builder() -> LlamaCppBuilder: + global _builder + if _builder is None: + _builder = LlamaCppBuilder() + return _builder diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index ecc11c30d5..af20146a7f 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -2352,11 +2352,12 @@ class UnslothTrainer: logger.info("Stopped during dataset loading\n") return None + n_rows = len(dataset) if hasattr(dataset, "__len__") else 0 self._update_progress( - status_message = f"Loaded dataset from HuggingFace: {dataset_source}" + status_message = f"Loaded dataset from HuggingFace: {dataset_source} ({n_rows:,} rows)" ) logger.info( - f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n" + f"Loaded dataset from Hugging Face: {dataset_source} ({n_rows:,} rows)\n" ) # Resolve eval split from a separate HF split (explicit or auto-detected) diff --git a/studio/backend/main.py b/studio/backend/main.py index 3ab846306e..4495c99e2f 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -75,6 +75,12 @@ async def lifespan(app: FastAPI): threading.Thread(target = _precache, daemon = True).start() + # Build llama.cpp in the background if the binary is missing. + # Studio starts immediately; GGUF features wait for the build to finish. + from core.inference.llama_cpp_builder import get_llama_cpp_builder + + get_llama_cpp_builder().check_and_build() + if storage.ensure_default_admin(): bootstrap_pw = storage.get_bootstrap_password() app.state.bootstrap_password = bootstrap_pw diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e93ba1d6d6..34a5f345cb 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -111,6 +111,19 @@ async def load_model( # ── GGUF path: load via llama-server ────────────────────── if config.is_gguf: + # Wait for llama.cpp to be compiled if a background build is in progress + from core.inference.llama_cpp_builder import get_llama_cpp_builder + + builder = get_llama_cpp_builder() + if not builder.is_ready: + logger.info("Waiting for llama.cpp build to finish...") + ready = await asyncio.to_thread(builder.wait_for_ready, timeout = 600) + if not ready or builder.error: + raise HTTPException( + status_code = 503, + detail = builder.error or "llama.cpp is still compiling. Please wait.", + ) + llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() @@ -538,6 +551,19 @@ async def get_status( raise HTTPException(status_code = 500, detail = f"Failed to get status: {str(e)}") +@router.get("/llama-cpp-status") +async def llama_cpp_status(): + """Check if llama.cpp is ready (built and available).""" + from core.inference.llama_cpp_builder import get_llama_cpp_builder + + builder = get_llama_cpp_builder() + return { + "ready": builder.is_ready and builder.error is None, + "building": builder.is_building, + "error": builder.error, + } + + # ===================================================================== # Audio (TTS) Generation (/audio/generate) # ===================================================================== diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 9d15b86ca1..1d2e4083ec 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -1092,6 +1092,11 @@ def format_and_template_dataset( # LLM FLOW (Existing code) else: # Step 1: Format the dataset + n_rows = len(dataset) if hasattr(dataset, "__len__") else None + if progress_callback and n_rows: + progress_callback( + status_message = f"Formatting dataset ({n_rows:,} rows)..." + ) dataset_info = format_dataset( dataset, format_type = format_type, @@ -1106,6 +1111,10 @@ def format_and_template_dataset( ) # Step 2: Apply chat template + if progress_callback and n_rows: + progress_callback( + status_message = f"Applying chat template ({n_rows:,} rows)..." + ) # Gemma emits a leading that must be stripped for text-only chatml/sharegpt. is_alpaca = format_type == "alpaca" or ( format_type == "auto" and dataset_info["detected_format"] == "alpaca" diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 01bc762d86..0f63e0d462 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -62,6 +62,15 @@ export async function getInferenceStatus(): Promise { return parseJsonOrThrow(response); } +export async function getLlamaCppStatus(): Promise<{ + ready: boolean; + building: boolean; + error: string | null; +}> { + const response = await authFetch("/api/inference/llama-cpp-status"); + return parseJsonOrThrow(response); +} + export async function loadModel( payload: LoadModelRequest, ): Promise { diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index b79c3073b2..e06418ce6d 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -9,6 +9,7 @@ import { getDownloadProgress, getGgufDownloadProgress, getInferenceStatus, + getLlamaCppStatus, listLoras, listModels, loadModel, @@ -318,6 +319,25 @@ export function useChatModelRuntime() { try { async function performLoad(): Promise { if (abortCtrl.signal.aborted) throw new Error("Cancelled"); + + // For GGUF models, check if llama.cpp is still compiling + if (ggufVariant) { + try { + const llamaStatus = await getLlamaCppStatus(); + if (llamaStatus.building) { + toast.info("Waiting for llama.cpp to compile...", { + description: "This is a one-time build. Studio is still usable for non-GGUF tasks.", + duration: 8000, + }); + } else if (llamaStatus.error) { + throw new Error(llamaStatus.error); + } + } catch (e) { + if (e instanceof Error && e.message.includes("not found")) throw e; + // Endpoint might not exist on older backends, ignore + } + } + let previousWasUnloaded = false; const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint;