diff --git a/.github/workflows/backend.yml b/.github/workflows/backend.yml index 50cc9b180bff..53ea01db91c1 100644 --- a/.github/workflows/backend.yml +++ b/.github/workflows/backend.yml @@ -79,6 +79,19 @@ jobs: dockerfile: "./backend/Dockerfile.python" context: "./" ubuntu-version: '2404' + - build-type: '' + cuda-major-version: "" + cuda-minor-version: "" + platforms: 'linux/amd64' + tag-latest: 'auto' + tag-suffix: '-cpu-unsloth' + runs-on: 'ubuntu-latest' + base-image: "ubuntu:24.04" + skip-drivers: 'true' + backend: "unsloth" + dockerfile: "./backend/Dockerfile.python" + context: "./" + ubuntu-version: '2404' - build-type: '' cuda-major-version: "" cuda-minor-version: "" @@ -444,6 +457,19 @@ jobs: dockerfile: "./backend/Dockerfile.python" context: "./" ubuntu-version: '2404' + - build-type: 'cublas' + cuda-major-version: "12" + cuda-minor-version: "8" + platforms: 'linux/amd64' + tag-latest: 'auto' + tag-suffix: '-gpu-nvidia-cuda-12-unsloth' + runs-on: 'ubuntu-latest' + base-image: "ubuntu:24.04" + skip-drivers: 'false' + backend: "unsloth" + dockerfile: "./backend/Dockerfile.python" + context: "./" + ubuntu-version: '2404' - build-type: 'cublas' cuda-major-version: "12" cuda-minor-version: "8" @@ -965,6 +991,19 @@ jobs: dockerfile: "./backend/Dockerfile.python" context: "./" ubuntu-version: '2404' + - build-type: 'cublas' + cuda-major-version: "13" + cuda-minor-version: "0" + platforms: 'linux/amd64' + tag-latest: 'auto' + tag-suffix: '-gpu-nvidia-cuda-13-unsloth' + runs-on: 'ubuntu-latest' + base-image: "ubuntu:24.04" + skip-drivers: 'false' + backend: "unsloth" + dockerfile: "./backend/Dockerfile.python" + context: "./" + ubuntu-version: '2404' - build-type: 'cublas' cuda-major-version: "13" cuda-minor-version: "0" diff --git a/Makefile b/Makefile index c429097b6f31..c79f9190d1e1 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ # Disable parallel execution for backend builds -.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/outetts backends/piper backends/stablediffusion-ggml backends/whisper backends/faster-whisper backends/silero-vad backends/local-store backends/huggingface backends/rfdetr backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus +.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/outetts backends/piper backends/stablediffusion-ggml backends/whisper backends/faster-whisper backends/silero-vad backends/local-store backends/huggingface backends/rfdetr backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/unsloth GOCMD=go GOTEST=$(GOCMD) test @@ -421,6 +421,7 @@ prepare-test-extra: protogen-python $(MAKE) -C backend/python/voxcpm $(MAKE) -C backend/python/whisperx $(MAKE) -C backend/python/ace-step + $(MAKE) -C backend/python/unsloth test-extra: prepare-test-extra $(MAKE) -C backend/python/transformers test @@ -440,6 +441,7 @@ test-extra: prepare-test-extra $(MAKE) -C backend/python/voxcpm test $(MAKE) -C backend/python/whisperx test $(MAKE) -C backend/python/ace-step test + $(MAKE) -C backend/python/unsloth test DOCKER_IMAGE?=local-ai IMAGE_TYPE?=core @@ -571,6 +573,7 @@ BACKEND_NEMO = nemo|python|.|false|true BACKEND_VOXCPM = voxcpm|python|.|false|true BACKEND_WHISPERX = whisperx|python|.|false|true BACKEND_ACE_STEP = ace-step|python|.|false|true +BACKEND_UNSLOTH = unsloth|python|.|false|true BACKEND_MLX_DISTRIBUTED = mlx-distributed|python|./|false|true # Helper function to build docker image for a backend @@ -627,6 +630,7 @@ $(eval $(call generate-docker-build-target,$(BACKEND_NEMO))) $(eval $(call generate-docker-build-target,$(BACKEND_VOXCPM))) $(eval $(call generate-docker-build-target,$(BACKEND_WHISPERX))) $(eval $(call generate-docker-build-target,$(BACKEND_ACE_STEP))) +$(eval $(call generate-docker-build-target,$(BACKEND_UNSLOTH))) $(eval $(call generate-docker-build-target,$(BACKEND_ACESTEP_CPP))) $(eval $(call generate-docker-build-target,$(BACKEND_MLX_DISTRIBUTED))) @@ -634,7 +638,7 @@ $(eval $(call generate-docker-build-target,$(BACKEND_MLX_DISTRIBUTED))) docker-save-%: backend-images docker save local-ai-backend:$* -o backend-images/$*.tar -docker-build-backends: docker-build-llama-cpp docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed +docker-build-backends: docker-build-llama-cpp docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-unsloth ######################################################## ### Mock Backend for E2E Tests diff --git a/backend/backend.proto b/backend/backend.proto index 8e4224789b6e..dd624e9491d8 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -39,6 +39,8 @@ service Backend { rpc AudioDecode(AudioDecodeRequest) returns (AudioDecodeResult) {} rpc ModelMetadata(ModelOptions) returns (ModelMetadataResponse) {} + + rpc TrainStream(TrainRequest) returns (stream TrainResponse) {} } // Define the empty request @@ -528,3 +530,49 @@ message ModelMetadataResponse { string rendered_template = 2; // The rendered chat template with enable_thinking=true (empty if not applicable) ToolFormatMarkers tool_format = 3; // Auto-detected tool format markers from differential template analysis } + +message TrainRequest { + string model = 1; // Base model name or HuggingFace model ID + string dataset = 2; // Path or HuggingFace dataset ID + string output_dir = 3; // Output directory for the fine-tuned model + + // Hyperparameters + int32 epochs = 4; // Number of training epochs + int32 batch_size = 5; // Training batch size + float learning_rate = 6; // Learning rate + int32 max_seq_length = 7; // Maximum sequence length + + // LoRA parameters + int32 lora_rank = 8; // LoRA rank (r) + int32 lora_alpha = 9; // LoRA alpha scaling factor + float lora_dropout = 10; // LoRA dropout rate + repeated string target_modules = 11; // LoRA target modules (e.g., "q_proj", "v_proj") + + // Quantization + string quantization = 12; // Quantization method (e.g., "4bit", "8bit", "none") + + // Additional options + map options = 13; // Generic key-value options for backend-specific settings +} + +message TrainResponse { + enum TrainStatus { + UNKNOWN = 0; + STARTING = 1; + RUNNING = 2; + COMPLETED = 3; + FAILED = 4; + } + + TrainStatus status = 1; // Current training status + float progress = 2; // Overall progress percentage (0-100) + int32 current_epoch = 3; // Current epoch number + int32 total_epochs = 4; // Total number of epochs + int32 current_step = 5; // Current training step + int32 total_steps = 6; // Total number of training steps + float loss = 7; // Current training loss + float learning_rate_current = 8; // Current learning rate (may change with schedulers) + string message = 9; // Human-readable status message + string error = 10; // Error message if status is FAILED + string output_path = 11; // Path to the output model (set on completion) +} diff --git a/backend/index.yaml b/backend/index.yaml index e5a4d3bbe5e1..d4d49259eb9a 100644 --- a/backend/index.yaml +++ b/backend/index.yaml @@ -530,6 +530,25 @@ nvidia-cuda-12: "cuda12-chatterbox" nvidia-l4t-cuda-12: "nvidia-l4t-arm64-chatterbox" nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-chatterbox" +- &unsloth + urls: + - https://github.com/unslothai/unsloth + description: | + Fine-tuning backend powered by Unsloth. Supports QLoRA, LoRA, and full fine-tuning + of LLMs with 2x faster training and 70% less memory. Includes GGUF export support. + tags: + - fine-tuning + - training + - lora + - qlora + license: apache-2.0 + name: "unsloth" + alias: "unsloth" + capabilities: + nvidia: "cuda12-unsloth" + default: "cpu-unsloth" + nvidia-cuda-13: "cuda13-unsloth" + nvidia-cuda-12: "cuda12-unsloth" - &vibevoice urls: - https://github.com/microsoft/VibeVoice @@ -3029,3 +3048,41 @@ uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-voxtral" mirrors: - localai/localai-backends:master-metal-darwin-arm64-voxtral +## unsloth +- !!merge <<: *unsloth + name: "unsloth-development" + capabilities: + nvidia: "cuda12-unsloth-development" + default: "cpu-unsloth-development" + nvidia-cuda-13: "cuda13-unsloth-development" + nvidia-cuda-12: "cuda12-unsloth-development" +- !!merge <<: *unsloth + name: "cpu-unsloth" + uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-unsloth" + mirrors: + - localai/localai-backends:latest-cpu-unsloth +- !!merge <<: *unsloth + name: "cpu-unsloth-development" + uri: "quay.io/go-skynet/local-ai-backends:master-cpu-unsloth" + mirrors: + - localai/localai-backends:master-cpu-unsloth +- !!merge <<: *unsloth + name: "cuda12-unsloth" + uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-unsloth" + mirrors: + - localai/localai-backends:latest-gpu-nvidia-cuda-12-unsloth +- !!merge <<: *unsloth + name: "cuda12-unsloth-development" + uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-unsloth" + mirrors: + - localai/localai-backends:master-gpu-nvidia-cuda-12-unsloth +- !!merge <<: *unsloth + name: "cuda13-unsloth" + uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-13-unsloth" + mirrors: + - localai/localai-backends:latest-gpu-nvidia-cuda-13-unsloth +- !!merge <<: *unsloth + name: "cuda13-unsloth-development" + uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-13-unsloth" + mirrors: + - localai/localai-backends:master-gpu-nvidia-cuda-13-unsloth diff --git a/backend/python/unsloth/Makefile b/backend/python/unsloth/Makefile new file mode 100644 index 000000000000..875d93e7e830 --- /dev/null +++ b/backend/python/unsloth/Makefile @@ -0,0 +1,23 @@ +.PHONY: unsloth +unsloth: + bash install.sh + +.PHONY: run +run: unsloth + @echo "Running unsloth..." + bash run.sh + @echo "unsloth run." + +.PHONY: test +test: unsloth + @echo "Testing unsloth..." + bash test.sh + @echo "unsloth tested." + +.PHONY: protogen-clean +protogen-clean: + $(RM) backend_pb2_grpc.py backend_pb2.py + +.PHONY: clean +clean: protogen-clean + rm -rf venv __pycache__ diff --git a/backend/python/unsloth/backend.py b/backend/python/unsloth/backend.py new file mode 100644 index 000000000000..134c81e2ddca --- /dev/null +++ b/backend/python/unsloth/backend.py @@ -0,0 +1,449 @@ +#!/usr/bin/env python3 +""" +Unsloth fine-tuning backend for LocalAI. + +Implements the TrainStream gRPC RPC to fine-tune LLMs using Unsloth +with QLoRA, LoRA, or full fine-tuning. Streams progress updates back +to the caller. +""" +from concurrent import futures + +import argparse +import signal +import sys +import os +import json +import traceback +import asyncio + +import backend_pb2 +import backend_pb2_grpc + +import grpc + +_ONE_DAY_IN_SECONDS = 60 * 60 * 24 + +MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1')) + + +def detect_hardware(): + """Detect available hardware and return device info.""" + import torch + + if torch.cuda.is_available(): + device_name = torch.cuda.get_device_name(0) + vram_gb = torch.cuda.get_device_properties(0).total_mem / (1024**3) + return { + "device": "cuda", + "device_name": device_name, + "vram_gb": round(vram_gb, 1), + "cuda_version": torch.version.cuda, + } + + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return {"device": "mps", "device_name": "Apple Silicon (MPS)"} + + return {"device": "cpu", "device_name": "CPU"} + + +def load_dataset_for_training(dataset_path, options): + """Load a dataset from a local path or HuggingFace hub.""" + from datasets import load_dataset + + dataset_format = options.get("dataset_format", "auto") + dataset_split = options.get("dataset_split", "train") + dataset_text_field = options.get("dataset_text_field", "text") + + # Local file + if os.path.exists(dataset_path): + ext = os.path.splitext(dataset_path)[1].lower() + if ext == ".jsonl" or ext == ".json": + ds = load_dataset("json", data_files=dataset_path, split=dataset_split) + elif ext == ".csv": + ds = load_dataset("csv", data_files=dataset_path, split=dataset_split) + elif ext == ".parquet": + ds = load_dataset("parquet", data_files=dataset_path, split=dataset_split) + else: + ds = load_dataset("json", data_files=dataset_path, split=dataset_split) + else: + # HuggingFace dataset + ds = load_dataset(dataset_path, split=dataset_split) + + # If dataset has chat format (conversations/messages column), handle it + if dataset_format == "chatml" or "conversations" in ds.column_names: + return ds, "chatml" + if "messages" in ds.column_names: + return ds, "messages" + + # If we have a text field, use it directly + if dataset_text_field in ds.column_names: + return ds, "text" + + # For instruction datasets with input/output/instruction columns + if "instruction" in ds.column_names: + return ds, "alpaca" + + return ds, "text" + + +def format_dataset(ds, fmt, tokenizer, dataset_text_field="text", max_seq_length=2048): + """Format a dataset for SFT training based on its format type.""" + from unsloth.chat_templates import get_chat_template, standardize_sharegpt + + if fmt == "chatml" or fmt == "messages": + tokenizer = get_chat_template( + tokenizer, + chat_template="chatml", + ) + col = "conversations" if "conversations" in ds.column_names else "messages" + ds = standardize_sharegpt(ds) + + def apply_template(examples): + convos = examples[col] + texts = [ + tokenizer.apply_chat_template( + convo, tokenize=False, add_generation_prompt=False + ) + for convo in convos + ] + return {"text": texts} + + ds = ds.map(apply_template, batched=True) + return ds, tokenizer + + if fmt == "alpaca": + from unsloth.chat_templates import get_chat_template + + tokenizer = get_chat_template(tokenizer, chat_template="alpaca") + + def format_alpaca(examples): + instructions = examples.get("instruction", []) + inputs = examples.get("input", [""] * len(instructions)) + outputs = examples.get("output", []) + texts = [] + for inst, inp, out in zip(instructions, inputs, outputs): + if inp: + text = f"### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n{out}" + else: + text = f"### Instruction:\n{inst}\n\n### Response:\n{out}" + texts.append(text) + return {"text": texts} + + ds = ds.map(format_alpaca, batched=True) + return ds, tokenizer + + # Plain text format — already has the text field + return ds, tokenizer + + +class BackendServicer(backend_pb2_grpc.BackendServicer): + """gRPC servicer for Unsloth fine-tuning.""" + + def Health(self, request, context): + return backend_pb2.Reply(message=bytes("OK", "utf-8")) + + def LoadModel(self, request, context): + # Not used for fine-tuning; training is self-contained via TrainStream + return backend_pb2.Result( + message="Unsloth backend: use TrainStream for fine-tuning", + success=True, + ) + + def TrainStream(self, request, context): + """ + Run a fine-tuning job and stream progress updates. + + Expects a TrainRequest with model, dataset, output_dir, and + hyperparameters. Yields TrainResponse messages with status, + progress, loss, etc. + """ + try: + yield from self._run_training(request, context) + except Exception as e: + traceback.print_exc(file=sys.stderr) + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.FAILED, + error=str(e), + message=f"Training failed: {e}", + ) + + def _run_training(self, request, context): + # -- Parse options -- + options = dict(request.options) if request.options else {} + + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.STARTING, + message="Detecting hardware...", + ) + + hw = detect_hardware() + print(f"[unsloth] Hardware: {hw}", file=sys.stderr) + + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.STARTING, + message=f"Hardware: {hw['device_name']}. Loading model...", + ) + + # -- Determine parameters -- + model_name = request.model + max_seq_length = request.max_seq_length if request.max_seq_length > 0 else 2048 + epochs = request.epochs if request.epochs > 0 else 1 + batch_size = request.batch_size if request.batch_size > 0 else 4 + learning_rate = request.learning_rate if request.learning_rate > 0 else 2e-4 + lora_rank = request.lora_rank if request.lora_rank > 0 else 16 + lora_alpha = request.lora_alpha if request.lora_alpha > 0 else 16 + lora_dropout = request.lora_dropout if request.lora_dropout > 0 else 0.0 + quantization = request.quantization or "4bit" + output_dir = request.output_dir or "/tmp/unsloth-output" + + gradient_accumulation_steps = int(options.get("gradient_accumulation_steps", "4")) + warmup_steps = int(options.get("warmup_steps", "5")) + weight_decay = float(options.get("weight_decay", "0.01")) + lr_scheduler = options.get("lr_scheduler_type", "linear") + save_steps = int(options.get("save_steps", "0")) + logging_steps = int(options.get("logging_steps", "1")) + dataset_text_field = options.get("dataset_text_field", "text") + training_type = options.get("training_type", "lora") # lora, qlora, full + export_gguf = options.get("export_gguf", "false").lower() == "true" + gguf_quantization = options.get("gguf_quantization", "q4_k_m") + + # Determine load_in_4bit based on quantization and training type + load_in_4bit = quantization == "4bit" or training_type == "qlora" + + # Target modules for LoRA + if request.target_modules: + target_modules = list(request.target_modules) + else: + target_modules = [ + "q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj", + ] + + # -- Load model with Unsloth -- + from unsloth import FastLanguageModel + + model, tokenizer = FastLanguageModel.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + dtype=None, # Auto-detect + load_in_4bit=load_in_4bit, + ) + + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.STARTING, + message="Model loaded. Applying LoRA adapters...", + ) + + # -- Apply LoRA / QLoRA -- + if training_type in ("lora", "qlora"): + model = FastLanguageModel.get_peft_model( + model, + r=lora_rank, + target_modules=target_modules, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias="none", + use_gradient_checkpointing="unsloth", + random_state=3407, + ) + + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.STARTING, + message="LoRA applied. Loading dataset...", + ) + + # -- Load and format dataset -- + ds, ds_format = load_dataset_for_training(request.dataset, options) + ds, tokenizer = format_dataset( + ds, ds_format, tokenizer, + dataset_text_field=dataset_text_field, + max_seq_length=max_seq_length, + ) + + print(f"[unsloth] Dataset loaded: {len(ds)} samples, format={ds_format}", file=sys.stderr) + + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.STARTING, + message=f"Dataset loaded ({len(ds)} samples). Starting training...", + ) + + # -- Setup trainer -- + from trl import SFTTrainer + from transformers import TrainingArguments, TrainerCallback + + # Estimate total steps + steps_per_epoch = max(1, len(ds) // (batch_size * gradient_accumulation_steps)) + total_steps = steps_per_epoch * epochs + + # Custom callback to stream progress via gRPC + progress_state = {"responses": []} + + class ProgressCallback(TrainerCallback): + def on_log(self, args, state, control, logs=None, **kwargs): + if logs is None: + return + current_step = state.global_step + current_epoch = int(state.epoch) if state.epoch else 0 + loss = logs.get("loss", 0.0) + lr = logs.get("learning_rate", 0.0) + progress = (current_step / total_steps * 100) if total_steps > 0 else 0 + + resp = backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.RUNNING, + progress=min(progress, 99.9), + current_epoch=current_epoch, + total_epochs=epochs, + current_step=current_step, + total_steps=total_steps, + loss=loss, + learning_rate_current=lr, + message=f"Step {current_step}/{total_steps} | Loss: {loss:.4f}", + ) + progress_state["responses"].append(resp) + + training_args = TrainingArguments( + per_device_train_batch_size=batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + warmup_steps=warmup_steps, + num_train_epochs=epochs, + learning_rate=learning_rate, + fp16=hw["device"] == "cuda", + bf16=False, + logging_steps=logging_steps, + optim="adamw_8bit", + weight_decay=weight_decay, + lr_scheduler_type=lr_scheduler, + seed=3407, + output_dir=output_dir, + save_steps=save_steps if save_steps > 0 else total_steps, + report_to="none", + ) + + trainer = SFTTrainer( + model=model, + tokenizer=tokenizer, + train_dataset=ds, + dataset_text_field="text", + max_seq_length=max_seq_length, + dataset_num_proc=2, + args=training_args, + callbacks=[ProgressCallback()], + ) + + # -- Run training in a thread so we can yield progress -- + import threading + + train_result = {"error": None, "stats": None} + + def do_train(): + try: + stats = trainer.train() + train_result["stats"] = stats + except Exception as e: + train_result["error"] = e + + train_thread = threading.Thread(target=do_train, daemon=True) + train_thread.start() + + # Yield progress responses as they arrive + import time + last_idx = 0 + while train_thread.is_alive(): + time.sleep(0.5) + responses = progress_state["responses"] + while last_idx < len(responses): + yield responses[last_idx] + last_idx += 1 + + # Drain remaining responses + responses = progress_state["responses"] + while last_idx < len(responses): + yield responses[last_idx] + last_idx += 1 + + train_thread.join() + + if train_result["error"] is not None: + raise train_result["error"] + + # -- Save model -- + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.RUNNING, + progress=95.0, + message="Training complete. Saving model...", + ) + + model.save_pretrained(output_dir) + tokenizer.save_pretrained(output_dir) + + final_output = output_dir + + # -- Optional GGUF export -- + if export_gguf: + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.RUNNING, + progress=97.0, + message=f"Exporting to GGUF ({gguf_quantization})...", + ) + + gguf_output = os.path.join(output_dir, "gguf") + try: + model.save_pretrained_gguf( + gguf_output, + tokenizer, + quantization_method=gguf_quantization, + ) + final_output = gguf_output + print(f"[unsloth] GGUF exported to {gguf_output}", file=sys.stderr) + except Exception as e: + print(f"[unsloth] GGUF export failed: {e}", file=sys.stderr) + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.RUNNING, + progress=98.0, + message=f"GGUF export failed (model still saved as safetensors): {e}", + ) + + # -- Done -- + yield backend_pb2.TrainResponse( + status=backend_pb2.TrainResponse.COMPLETED, + progress=100.0, + current_epoch=epochs, + total_epochs=epochs, + current_step=total_steps, + total_steps=total_steps, + message="Fine-tuning completed successfully", + output_path=final_output, + ) + + print(f"[unsloth] Training complete. Output: {final_output}", file=sys.stderr) + + +async def serve(address): + server = grpc.aio.server( + migration_thread_pool=futures.ThreadPoolExecutor(max_workers=MAX_WORKERS), + options=[ + ("grpc.max_message_length", 50 * 1024 * 1024), + ("grpc.max_send_message_length", 50 * 1024 * 1024), + ("grpc.max_receive_message_length", 50 * 1024 * 1024), + ], + ) + backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server) + server.add_insecure_port(address) + + loop = asyncio.get_event_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, lambda: asyncio.ensure_future(server.stop(5))) + + await server.start() + print("Server started. Listening on: " + address, file=sys.stderr) + await server.wait_for_termination() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run the gRPC server.") + parser.add_argument( + "--addr", default="localhost:50051", help="The address to bind the server to." + ) + args = parser.parse_args() + + asyncio.run(serve(args.addr)) diff --git a/backend/python/unsloth/install.sh b/backend/python/unsloth/install.sh new file mode 100755 index 000000000000..4136d8765589 --- /dev/null +++ b/backend/python/unsloth/install.sh @@ -0,0 +1,11 @@ +#!/bin/bash +set -e + +backend_dir=$(dirname $0) +if [ -d $backend_dir/common ]; then + source $backend_dir/common/libbackend.sh +else + source $backend_dir/../common/libbackend.sh +fi + +installRequirements diff --git a/backend/python/unsloth/requirements-cpu.txt b/backend/python/unsloth/requirements-cpu.txt new file mode 100644 index 000000000000..ea40b8fc77ed --- /dev/null +++ b/backend/python/unsloth/requirements-cpu.txt @@ -0,0 +1,9 @@ +--extra-index-url https://download.pytorch.org/whl/cpu +unsloth +torch +transformers +datasets +trl +peft +accelerate +sentencepiece diff --git a/backend/python/unsloth/requirements-cublas12.txt b/backend/python/unsloth/requirements-cublas12.txt new file mode 100644 index 000000000000..a46842d67baf --- /dev/null +++ b/backend/python/unsloth/requirements-cublas12.txt @@ -0,0 +1,11 @@ +--extra-index-url https://download.pytorch.org/whl/cu121 +unsloth[cu121-ampere-torch251] +torch +transformers +datasets +trl +peft +accelerate +bitsandbytes +sentencepiece +packaging diff --git a/backend/python/unsloth/requirements-cublas13.txt b/backend/python/unsloth/requirements-cublas13.txt new file mode 100644 index 000000000000..9421da5fd21e --- /dev/null +++ b/backend/python/unsloth/requirements-cublas13.txt @@ -0,0 +1,11 @@ +--extra-index-url https://download.pytorch.org/whl/cu126 +unsloth[cu124-ampere-torch251] +torch +transformers +datasets +trl +peft +accelerate +bitsandbytes +sentencepiece +packaging diff --git a/backend/python/unsloth/requirements-hipblas.txt b/backend/python/unsloth/requirements-hipblas.txt new file mode 100644 index 000000000000..3cbff7f39b52 --- /dev/null +++ b/backend/python/unsloth/requirements-hipblas.txt @@ -0,0 +1,10 @@ +--extra-index-url https://download.pytorch.org/whl/rocm6.3 +unsloth +torch +transformers +datasets +trl +peft +accelerate +bitsandbytes +sentencepiece diff --git a/backend/python/unsloth/requirements-mps.txt b/backend/python/unsloth/requirements-mps.txt new file mode 100644 index 000000000000..dcc5fdb567ed --- /dev/null +++ b/backend/python/unsloth/requirements-mps.txt @@ -0,0 +1,8 @@ +unsloth +torch +transformers +datasets +trl +peft +accelerate +sentencepiece diff --git a/backend/python/unsloth/requirements.txt b/backend/python/unsloth/requirements.txt new file mode 100644 index 000000000000..2a43ebece955 --- /dev/null +++ b/backend/python/unsloth/requirements.txt @@ -0,0 +1,4 @@ +grpcio==1.78.1 +protobuf==6.33.5 +certifi +setuptools diff --git a/backend/python/unsloth/run.sh b/backend/python/unsloth/run.sh new file mode 100755 index 000000000000..eae121f37b0b --- /dev/null +++ b/backend/python/unsloth/run.sh @@ -0,0 +1,9 @@ +#!/bin/bash +backend_dir=$(dirname $0) +if [ -d $backend_dir/common ]; then + source $backend_dir/common/libbackend.sh +else + source $backend_dir/../common/libbackend.sh +fi + +startBackend $@ diff --git a/backend/python/unsloth/test.py b/backend/python/unsloth/test.py new file mode 100644 index 000000000000..d85910edf19a --- /dev/null +++ b/backend/python/unsloth/test.py @@ -0,0 +1,32 @@ +"""Unit tests for the Unsloth fine-tuning backend.""" +import unittest +import sys +import os + +import backend_pb2 +import backend_pb2_grpc +from backend import BackendServicer, detect_hardware + + +class TestBackendServicer(unittest.TestCase): + def setUp(self): + self.servicer = BackendServicer() + + def test_health(self): + request = backend_pb2.HealthMessage() + resp = self.servicer.Health(request, None) + self.assertEqual(resp.message, b"OK") + + def test_load_model(self): + request = backend_pb2.ModelOptions() + resp = self.servicer.LoadModel(request, None) + self.assertTrue(resp.success) + + def test_detect_hardware(self): + hw = detect_hardware() + self.assertIn("device", hw) + self.assertIn(hw["device"], ("cuda", "mps", "cpu")) + + +if __name__ == "__main__": + unittest.main() diff --git a/backend/python/unsloth/test.sh b/backend/python/unsloth/test.sh new file mode 100755 index 000000000000..eb59f2aaf3f3 --- /dev/null +++ b/backend/python/unsloth/test.sh @@ -0,0 +1,11 @@ +#!/bin/bash +set -e + +backend_dir=$(dirname $0) +if [ -d $backend_dir/common ]; then + source $backend_dir/common/libbackend.sh +else + source $backend_dir/../common/libbackend.sh +fi + +runUnittests diff --git a/core/application/application.go b/core/application/application.go index f1adc71449ed..20cfe075e558 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -14,20 +14,21 @@ import ( ) type Application struct { - backendLoader *config.ModelConfigLoader - modelLoader *model.ModelLoader - applicationConfig *config.ApplicationConfig - startupConfig *config.ApplicationConfig // Stores original config from env vars (before file loading) - templatesEvaluator *templates.Evaluator - galleryService *services.GalleryService - agentJobService *services.AgentJobService - agentPoolService atomic.Pointer[services.AgentPoolService] - watchdogMutex sync.Mutex - watchdogStop chan bool - p2pMutex sync.Mutex - p2pCtx context.Context - p2pCancel context.CancelFunc - agentJobMutex sync.Mutex + backendLoader *config.ModelConfigLoader + modelLoader *model.ModelLoader + applicationConfig *config.ApplicationConfig + startupConfig *config.ApplicationConfig // Stores original config from env vars (before file loading) + templatesEvaluator *templates.Evaluator + galleryService *services.GalleryService + agentJobService *services.AgentJobService + trainingJobService *services.TrainingJobService + agentPoolService atomic.Pointer[services.AgentPoolService] + watchdogMutex sync.Mutex + watchdogStop chan bool + p2pMutex sync.Mutex + p2pCtx context.Context + p2pCancel context.CancelFunc + agentJobMutex sync.Mutex } func newApplication(appConfig *config.ApplicationConfig) *Application { @@ -70,6 +71,10 @@ func (a *Application) AgentJobService() *services.AgentJobService { return a.agentJobService } +func (a *Application) TrainingJobService() *services.TrainingJobService { + return a.trainingJobService +} + func (a *Application) AgentPoolService() *services.AgentPoolService { return a.agentPoolService.Load() } @@ -103,6 +108,20 @@ func (a *Application) start() error { a.agentJobService = agentJobService + // Initialize training job service + trainingJobService := services.NewTrainingJobService( + a.ApplicationConfig(), + a.ModelLoader(), + a.ModelConfigLoader(), + ) + + err = trainingJobService.Start(a.ApplicationConfig().Context) + if err != nil { + return err + } + + a.trainingJobService = trainingJobService + return nil } diff --git a/core/backend/training.go b/core/backend/training.go new file mode 100644 index 000000000000..b03fd2a432ff --- /dev/null +++ b/core/backend/training.go @@ -0,0 +1,32 @@ +package backend + +import ( + "context" + "fmt" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/model" +) + +// TrainModel initiates a training job on a backend and streams progress via the callback. +func TrainModel( + ctx context.Context, + trainReq *proto.TrainRequest, + loader *model.ModelLoader, + appConfig *config.ApplicationConfig, + modelConfig config.ModelConfig, + progressCallback func(*proto.TrainResponse), +) error { + opts := ModelOptions(modelConfig, appConfig) + backend, err := loader.Load(opts...) + if err != nil { + return fmt.Errorf("loading backend for training: %w", err) + } + + if backend == nil { + return fmt.Errorf("could not load backend %q for training", modelConfig.Backend) + } + + return backend.TrainStream(ctx, trainReq, progressCallback) +} diff --git a/core/http/endpoints/localai/training.go b/core/http/endpoints/localai/training.go new file mode 100644 index 000000000000..7b7d409dca21 --- /dev/null +++ b/core/http/endpoints/localai/training.go @@ -0,0 +1,136 @@ +package localai + +import ( + "net/http" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/application" + "github.com/mudler/LocalAI/core/schema" +) + +// StartTrainingEndpoint creates a new fine-tuning job +// @Summary Start a fine-tuning job +// @Description Start a new fine-tuning job with the specified model and dataset +// @Tags training +// @Accept json +// @Produce json +// @Param request body schema.TrainingJobRequest true "Training job configuration" +// @Success 201 {object} schema.TrainingJob "Training job created" +// @Failure 400 {object} map[string]string "Invalid request" +// @Failure 500 {object} map[string]string "Internal server error" +// @Router /api/training/jobs [post] +func StartTrainingEndpoint(app *application.Application) echo.HandlerFunc { + return func(c echo.Context) error { + svc := app.TrainingJobService() + if svc == nil { + return c.JSON(http.StatusServiceUnavailable, map[string]string{"error": "Training service not available"}) + } + + var req schema.TrainingJobRequest + if err := c.Bind(&req); err != nil { + return c.JSON(http.StatusBadRequest, map[string]string{"error": "Invalid request body: " + err.Error()}) + } + + job, err := svc.CreateJob(req) + if err != nil { + return c.JSON(http.StatusBadRequest, map[string]string{"error": err.Error()}) + } + + return c.JSON(http.StatusCreated, job) + } +} + +// GetTrainingJobEndpoint returns the status of a training job +// @Summary Get training job status +// @Description Get the current status and progress of a training job +// @Tags training +// @Produce json +// @Param id path string true "Job ID" +// @Success 200 {object} schema.TrainingJob "Training job status" +// @Failure 404 {object} map[string]string "Job not found" +// @Router /api/training/jobs/{id} [get] +func GetTrainingJobEndpoint(app *application.Application) echo.HandlerFunc { + return func(c echo.Context) error { + svc := app.TrainingJobService() + if svc == nil { + return c.JSON(http.StatusServiceUnavailable, map[string]string{"error": "Training service not available"}) + } + + id := c.Param("id") + job, ok := svc.GetJob(id) + if !ok { + return c.JSON(http.StatusNotFound, map[string]string{"error": "Job not found: " + id}) + } + + return c.JSON(http.StatusOK, job) + } +} + +// ListTrainingJobsEndpoint lists all training jobs +// @Summary List training jobs +// @Description List all fine-tuning jobs +// @Tags training +// @Produce json +// @Success 200 {array} schema.TrainingJob "List of training jobs" +// @Router /api/training/jobs [get] +func ListTrainingJobsEndpoint(app *application.Application) echo.HandlerFunc { + return func(c echo.Context) error { + svc := app.TrainingJobService() + if svc == nil { + return c.JSON(http.StatusServiceUnavailable, map[string]string{"error": "Training service not available"}) + } + + return c.JSON(http.StatusOK, svc.ListJobs()) + } +} + +// CancelTrainingEndpoint cancels a running training job +// @Summary Cancel a training job +// @Description Cancel a running or pending training job +// @Tags training +// @Produce json +// @Param id path string true "Job ID" +// @Success 200 {object} map[string]string "Job cancelled" +// @Failure 400 {object} map[string]string "Job cannot be cancelled" +// @Failure 404 {object} map[string]string "Job not found" +// @Router /api/training/jobs/{id}/cancel [post] +func CancelTrainingEndpoint(app *application.Application) echo.HandlerFunc { + return func(c echo.Context) error { + svc := app.TrainingJobService() + if svc == nil { + return c.JSON(http.StatusServiceUnavailable, map[string]string{"error": "Training service not available"}) + } + + id := c.Param("id") + if err := svc.CancelJob(id); err != nil { + return c.JSON(http.StatusBadRequest, map[string]string{"error": err.Error()}) + } + + return c.JSON(http.StatusOK, map[string]string{"message": "Job cancelled"}) + } +} + +// DeleteTrainingEndpoint deletes a training job +// @Summary Delete a training job +// @Description Delete a training job record +// @Tags training +// @Produce json +// @Param id path string true "Job ID" +// @Success 200 {object} map[string]string "Job deleted" +// @Failure 404 {object} map[string]string "Job not found" +// @Router /api/training/jobs/{id} [delete] +func DeleteTrainingEndpoint(app *application.Application) echo.HandlerFunc { + return func(c echo.Context) error { + svc := app.TrainingJobService() + if svc == nil { + return c.JSON(http.StatusServiceUnavailable, map[string]string{"error": "Training service not available"}) + } + + id := c.Param("id") + if err := svc.DeleteJob(id); err != nil { + return c.JSON(http.StatusNotFound, map[string]string{"error": err.Error()}) + } + + return c.JSON(http.StatusOK, map[string]string{"message": "Job deleted"}) + } +} diff --git a/core/http/routes/localai.go b/core/http/routes/localai.go index 50bd8218ef0f..f40e2666d429 100644 --- a/core/http/routes/localai.go +++ b/core/http/routes/localai.go @@ -182,6 +182,15 @@ func RegisterLocalAIRoutes(router *echo.Echo, router.OPTIONS("/api/cors-proxy", localai.CORSProxyOptionsEndpoint()) } + // Training job routes + if app != nil && app.TrainingJobService() != nil { + router.POST("/api/training/jobs", localai.StartTrainingEndpoint(app)) + router.GET("/api/training/jobs/:id", localai.GetTrainingJobEndpoint(app)) + router.GET("/api/training/jobs", localai.ListTrainingJobsEndpoint(app)) + router.POST("/api/training/jobs/:id/cancel", localai.CancelTrainingEndpoint(app)) + router.DELETE("/api/training/jobs/:id", localai.DeleteTrainingEndpoint(app)) + } + // Agent job routes (MCP CI Jobs — requires MCP to be enabled) if app != nil && app.AgentJobService() != nil && !appConfig.DisableMCP { router.POST("/api/agent/tasks", localai.CreateTaskEndpoint(app)) diff --git a/core/schema/training.go b/core/schema/training.go new file mode 100644 index 000000000000..b96932cf5b01 --- /dev/null +++ b/core/schema/training.go @@ -0,0 +1,71 @@ +package schema + +import "time" + +// TrainingJobStatus represents the status of a training job +type TrainingJobStatus string + +const ( + TrainingJobStatusPending TrainingJobStatus = "pending" + TrainingJobStatusRunning TrainingJobStatus = "running" + TrainingJobStatusCompleted TrainingJobStatus = "completed" + TrainingJobStatusFailed TrainingJobStatus = "failed" + TrainingJobStatusCancelled TrainingJobStatus = "cancelled" +) + +// TrainingJob represents a fine-tuning job +type TrainingJob struct { + ID string `json:"id"` + Model string `json:"model"` + Backend string `json:"backend,omitempty"` + Dataset string `json:"dataset"` + OutputDir string `json:"output_dir"` + Status TrainingJobStatus `json:"status"` + Progress float32 `json:"progress"` + CurrentEpoch int32 `json:"current_epoch"` + TotalEpochs int32 `json:"total_epochs"` + CurrentStep int32 `json:"current_step"` + TotalSteps int32 `json:"total_steps"` + Loss float32 `json:"loss"` + LearningRate float32 `json:"learning_rate_current"` + Message string `json:"message"` + Error string `json:"error,omitempty"` + OutputPath string `json:"output_path,omitempty"` + CreatedAt time.Time `json:"created_at"` + StartedAt *time.Time `json:"started_at,omitempty"` + CompletedAt *time.Time `json:"completed_at,omitempty"` + Parameters TrainingParams `json:"parameters"` +} + +// TrainingParams holds the training hyperparameters +type TrainingParams struct { + Epochs int32 `json:"epochs,omitempty"` + BatchSize int32 `json:"batch_size,omitempty"` + LearningRate float32 `json:"learning_rate,omitempty"` + MaxSeqLength int32 `json:"max_seq_length,omitempty"` + LoraRank int32 `json:"lora_rank,omitempty"` + LoraAlpha int32 `json:"lora_alpha,omitempty"` + LoraDropout float32 `json:"lora_dropout,omitempty"` + TargetModules []string `json:"target_modules,omitempty"` + Quantization string `json:"quantization,omitempty"` + + // Extra options passed through to the backend + Options map[string]string `json:"options,omitempty"` +} + +// TrainingJobRequest is the HTTP request body for creating a training job +type TrainingJobRequest struct { + Model string `json:"model"` + Backend string `json:"backend,omitempty"` + Dataset string `json:"dataset"` + OutputDir string `json:"output_dir,omitempty"` + Epochs int32 `json:"epochs,omitempty"` + BatchSize int32 `json:"batch_size,omitempty"` + LearningRate float32 `json:"learning_rate,omitempty"` + MaxSeqLength int32 `json:"max_seq_length,omitempty"` + LoraRank int32 `json:"lora_rank,omitempty"` + LoraAlpha int32 `json:"lora_alpha,omitempty"` + LoraDropout float32 `json:"lora_dropout,omitempty"` + Quantization string `json:"quantization,omitempty"` + Options map[string]string `json:"options,omitempty"` +} diff --git a/core/services/training.go b/core/services/training.go new file mode 100644 index 000000000000..ac17c8ffd86e --- /dev/null +++ b/core/services/training.go @@ -0,0 +1,395 @@ +package services + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "sync" + "time" + + "github.com/google/uuid" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/xsync" + "github.com/mudler/xlog" +) + +// TrainingJobService manages fine-tuning job lifecycle +type TrainingJobService struct { + appConfig *config.ApplicationConfig + modelLoader *model.ModelLoader + configLoader *config.ModelConfigLoader + + jobs *xsync.SyncedMap[string, schema.TrainingJob] + cancellations *xsync.SyncedMap[string, context.CancelFunc] + jobsFile string + + ctx context.Context + cancel context.CancelFunc + + fileMutex sync.Mutex +} + +// NewTrainingJobService creates a new TrainingJobService +func NewTrainingJobService( + appConfig *config.ApplicationConfig, + modelLoader *model.ModelLoader, + configLoader *config.ModelConfigLoader, +) *TrainingJobService { + dataDir := appConfig.DataPath + if dataDir == "" { + dataDir = appConfig.DynamicConfigsDir + } + + jobsFile := "" + if dataDir != "" { + jobsFile = filepath.Join(dataDir, "training_jobs.json") + } + + return &TrainingJobService{ + appConfig: appConfig, + modelLoader: modelLoader, + configLoader: configLoader, + jobs: xsync.NewSyncedMap[string, schema.TrainingJob](), + cancellations: xsync.NewSyncedMap[string, context.CancelFunc](), + jobsFile: jobsFile, + } +} + +// Start initializes the service and loads persisted jobs +func (s *TrainingJobService) Start(ctx context.Context) error { + s.ctx, s.cancel = context.WithCancel(ctx) + + if err := s.loadJobs(); err != nil { + xlog.Warn("Failed to load training jobs from disk", "error", err) + } + + return nil +} + +// CreateJob creates a new training job and starts it asynchronously +func (s *TrainingJobService) CreateJob(req schema.TrainingJobRequest) (schema.TrainingJob, error) { + if req.Model == "" { + return schema.TrainingJob{}, fmt.Errorf("model is required") + } + if req.Dataset == "" { + return schema.TrainingJob{}, fmt.Errorf("dataset is required") + } + + id := uuid.New().String() + now := time.Now() + + outputDir := req.OutputDir + if outputDir == "" { + outputDir = filepath.Join(s.appConfig.GeneratedContentDir, "training", id) + } + + job := schema.TrainingJob{ + ID: id, + Model: req.Model, + Backend: req.Backend, + Dataset: req.Dataset, + OutputDir: outputDir, + Status: schema.TrainingJobStatusPending, + CreatedAt: now, + Parameters: schema.TrainingParams{ + Epochs: req.Epochs, + BatchSize: req.BatchSize, + LearningRate: req.LearningRate, + MaxSeqLength: req.MaxSeqLength, + LoraRank: req.LoraRank, + LoraAlpha: req.LoraAlpha, + LoraDropout: req.LoraDropout, + Quantization: req.Quantization, + Options: req.Options, + }, + } + + s.jobs.Set(id, job) + s.saveJobs() + + // Start training asynchronously + go s.executeJob(id) + + return job, nil +} + +// GetJob returns a training job by ID +func (s *TrainingJobService) GetJob(id string) (schema.TrainingJob, bool) { + return s.jobs.Get(id) +} + +// ListJobs returns all training jobs, sorted by creation time (newest first) +func (s *TrainingJobService) ListJobs() []schema.TrainingJob { + var jobs []schema.TrainingJob + s.jobs.Iterate(func(_ string, job schema.TrainingJob) bool { + jobs = append(jobs, job) + return true + }) + sort.Slice(jobs, func(i, j int) bool { + return jobs[i].CreatedAt.After(jobs[j].CreatedAt) + }) + return jobs +} + +// CancelJob cancels a running training job +func (s *TrainingJobService) CancelJob(id string) error { + job, ok := s.jobs.Get(id) + if !ok { + return fmt.Errorf("job not found: %s", id) + } + + if job.Status != schema.TrainingJobStatusPending && job.Status != schema.TrainingJobStatusRunning { + return fmt.Errorf("job %s is not cancellable (status: %s)", id, job.Status) + } + + cancelFn, ok := s.cancellations.Get(id) + if ok { + cancelFn() + } + + now := time.Now() + job.Status = schema.TrainingJobStatusCancelled + job.CompletedAt = &now + job.Message = "Job cancelled by user" + s.jobs.Set(id, job) + s.saveJobs() + + return nil +} + +// DeleteJob removes a job from the store +func (s *TrainingJobService) DeleteJob(id string) error { + job, ok := s.jobs.Get(id) + if !ok { + return fmt.Errorf("job not found: %s", id) + } + + // Cancel if still running + if job.Status == schema.TrainingJobStatusRunning || job.Status == schema.TrainingJobStatusPending { + if cancelFn, ok := s.cancellations.Get(id); ok { + cancelFn() + } + } + + s.jobs.Delete(id) + s.cancellations.Delete(id) + s.saveJobs() + return nil +} + +// executeJob runs the training job via the backend gRPC call +func (s *TrainingJobService) executeJob(id string) { + job, ok := s.jobs.Get(id) + if !ok { + return + } + + ctx, cancel := context.WithCancel(s.ctx) + s.cancellations.Set(id, cancel) + defer func() { + cancel() + s.cancellations.Delete(id) + }() + + now := time.Now() + job.Status = schema.TrainingJobStatusRunning + job.StartedAt = &now + job.Message = "Starting training..." + s.jobs.Set(id, job) + s.saveJobs() + + // Build gRPC TrainRequest + trainReq := &proto.TrainRequest{ + Model: job.Model, + Dataset: job.Dataset, + OutputDir: job.OutputDir, + Epochs: job.Parameters.Epochs, + BatchSize: job.Parameters.BatchSize, + LearningRate: job.Parameters.LearningRate, + MaxSeqLength: job.Parameters.MaxSeqLength, + LoraRank: job.Parameters.LoraRank, + LoraAlpha: job.Parameters.LoraAlpha, + LoraDropout: job.Parameters.LoraDropout, + TargetModules: job.Parameters.TargetModules, + Quantization: job.Parameters.Quantization, + Options: job.Parameters.Options, + } + + // Resolve model config — use the backend field from the request, defaulting to "unsloth" + backendName := job.Backend + if backendName == "" { + backendName = "unsloth" + } + + modelConfig := config.ModelConfig{ + Name: job.Model, + Model: job.Model, + Backend: backendName, + } + + // Attempt to load from existing model configs + if cfgs := s.configLoader.ListConfigs(); len(cfgs) > 0 { + for _, name := range cfgs { + if name == job.Model { + if cfg, exists := s.configLoader.GetConfig(name); exists { + modelConfig = cfg + if backendName != "" { + modelConfig.Backend = backendName + } + } + break + } + } + } + + // Load backend and call TrainStream directly (avoids import cycle with core/backend) + opts := []model.Option{ + model.WithBackendString(modelConfig.Backend), + model.WithModel(modelConfig.Model), + model.WithContext(s.appConfig.Context), + model.WithModelID(modelConfig.Name), + } + grpcBackend, err := s.modelLoader.Load(opts...) + if err != nil { + err = fmt.Errorf("loading backend for training: %w", err) + } else if grpcBackend == nil { + err = fmt.Errorf("could not load backend %q for training", modelConfig.Backend) + } else { + err = grpcBackend.TrainStream(ctx, trainReq, func(resp *proto.TrainResponse) { + s.updateJobFromResponse(id, resp) + }) + } + + // Final status update + job, _ = s.jobs.Get(id) + completedAt := time.Now() + + if err != nil { + if ctx.Err() != nil { + // Context was cancelled + job.Status = schema.TrainingJobStatusCancelled + job.Message = "Job cancelled" + } else { + job.Status = schema.TrainingJobStatusFailed + job.Error = err.Error() + job.Message = fmt.Sprintf("Training failed: %s", err.Error()) + } + } else if job.Status != schema.TrainingJobStatusCompleted { + // Mark completed if the backend didn't already + job.Status = schema.TrainingJobStatusCompleted + job.Progress = 100.0 + job.Message = "Training completed" + } + + job.CompletedAt = &completedAt + s.jobs.Set(id, job) + s.saveJobs() + + xlog.Info("Training job finished", "id", id, "status", job.Status, "output", job.OutputPath) +} + +// updateJobFromResponse updates job state from a streaming TrainResponse +func (s *TrainingJobService) updateJobFromResponse(id string, resp *proto.TrainResponse) { + job, ok := s.jobs.Get(id) + if !ok { + return + } + + job.Progress = resp.Progress + job.CurrentEpoch = resp.CurrentEpoch + job.TotalEpochs = resp.TotalEpochs + job.CurrentStep = resp.CurrentStep + job.TotalSteps = resp.TotalSteps + job.Loss = resp.Loss + job.LearningRate = resp.LearningRateCurrent + job.Message = resp.Message + + if resp.OutputPath != "" { + job.OutputPath = resp.OutputPath + } + if resp.Error != "" { + job.Error = resp.Error + } + + switch resp.Status { + case proto.TrainResponse_COMPLETED: + job.Status = schema.TrainingJobStatusCompleted + now := time.Now() + job.CompletedAt = &now + case proto.TrainResponse_FAILED: + job.Status = schema.TrainingJobStatusFailed + now := time.Now() + job.CompletedAt = &now + case proto.TrainResponse_RUNNING: + job.Status = schema.TrainingJobStatusRunning + } + + s.jobs.Set(id, job) +} + +// loadJobs reads persisted jobs from disk +func (s *TrainingJobService) loadJobs() error { + if s.jobsFile == "" { + return nil + } + + data, err := os.ReadFile(s.jobsFile) + if err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + + var jobs []schema.TrainingJob + if err := json.Unmarshal(data, &jobs); err != nil { + return err + } + + for _, job := range jobs { + // Mark any previously-running jobs as failed (server restarted) + if job.Status == schema.TrainingJobStatusRunning || job.Status == schema.TrainingJobStatusPending { + job.Status = schema.TrainingJobStatusFailed + job.Error = "server restarted during training" + job.Message = "Job interrupted by server restart" + now := time.Now() + job.CompletedAt = &now + } + s.jobs.Set(job.ID, job) + } + + return nil +} + +// saveJobs persists all jobs to disk +func (s *TrainingJobService) saveJobs() { + if s.jobsFile == "" { + return + } + + s.fileMutex.Lock() + defer s.fileMutex.Unlock() + + jobs := s.ListJobs() + data, err := json.MarshalIndent(jobs, "", " ") + if err != nil { + xlog.Error("Failed to marshal training jobs", "error", err) + return + } + + dir := filepath.Dir(s.jobsFile) + if err := os.MkdirAll(dir, 0750); err != nil { + xlog.Error("Failed to create training jobs directory", "error", err) + return + } + + if err := os.WriteFile(s.jobsFile, data, 0600); err != nil { + xlog.Error("Failed to save training jobs", "error", err) + } +} diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index 348463808172..1d071ce4d40a 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -63,4 +63,6 @@ type Backend interface { AudioDecode(ctx context.Context, in *pb.AudioDecodeRequest, opts ...grpc.CallOption) (*pb.AudioDecodeResult, error) ModelMetadata(ctx context.Context, in *pb.ModelOptions, opts ...grpc.CallOption) (*pb.ModelMetadataResponse, error) + + TrainStream(ctx context.Context, in *pb.TrainRequest, f func(resp *pb.TrainResponse), opts ...grpc.CallOption) error } diff --git a/pkg/grpc/base/base.go b/pkg/grpc/base/base.go index da94c10ea8f2..34469eecf79f 100644 --- a/pkg/grpc/base/base.go +++ b/pkg/grpc/base/base.go @@ -120,6 +120,10 @@ func (llm *Base) AudioDecode(*pb.AudioDecodeRequest) (*pb.AudioDecodeResult, err return nil, fmt.Errorf("unimplemented") } +func (llm *Base) TrainStream(*pb.TrainRequest, chan *pb.TrainResponse) error { + return fmt.Errorf("unimplemented") +} + func memoryUsage() *pb.MemoryUsageData { mud := pb.MemoryUsageData{ Breakdown: make(map[string]uint64), diff --git a/pkg/grpc/client.go b/pkg/grpc/client.go index eedd464c7e21..20932f343212 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -632,6 +632,54 @@ func (c *Client) AudioDecode(ctx context.Context, in *pb.AudioDecodeRequest, opt return client.AudioDecode(ctx, in, opts...) } +func (c *Client) TrainStream(ctx context.Context, in *pb.TrainRequest, f func(resp *pb.TrainResponse), opts ...grpc.CallOption) error { + if !c.parallel { + c.opMutex.Lock() + defer c.opMutex.Unlock() + } + c.setBusy(true) + defer c.setBusy(false) + c.wdMark() + defer c.wdUnMark() + conn, err := grpc.Dial(c.address, grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithDefaultCallOptions( + grpc.MaxCallRecvMsgSize(50*1024*1024), // 50MB + grpc.MaxCallSendMsgSize(50*1024*1024), // 50MB + )) + if err != nil { + return err + } + defer conn.Close() + client := pb.NewBackendClient(conn) + + stream, err := client.TrainStream(ctx, in, opts...) + if err != nil { + return err + } + + for { + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + + resp, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + if ctx.Err() != nil { + return ctx.Err() + } + return err + } + f(resp) + } + + return nil +} + func (c *Client) ModelMetadata(ctx context.Context, in *pb.ModelOptions, opts ...grpc.CallOption) (*pb.ModelMetadataResponse, error) { if !c.parallel { c.opMutex.Lock() diff --git a/pkg/grpc/embed.go b/pkg/grpc/embed.go index 905290c2ac86..40e6434b2f4f 100644 --- a/pkg/grpc/embed.go +++ b/pkg/grpc/embed.go @@ -10,6 +10,7 @@ import ( var _ Backend = new(embedBackend) var _ pb.Backend_PredictStreamServer = new(embedBackendServerStream) +var _ pb.Backend_TrainStreamServer = new(embedBackendTrainStream) type embedBackend struct { s *server @@ -119,6 +120,14 @@ func (e *embedBackend) ModelMetadata(ctx context.Context, in *pb.ModelOptions, o return e.s.ModelMetadata(ctx, in) } +func (e *embedBackend) TrainStream(ctx context.Context, in *pb.TrainRequest, f func(resp *pb.TrainResponse), opts ...grpc.CallOption) error { + bs := &embedBackendTrainStream{ + ctx: ctx, + fn: f, + } + return e.s.TrainStream(in, bs) +} + func (e *embedBackend) GetTokenMetrics(ctx context.Context, in *pb.MetricsRequest, opts ...grpc.CallOption) (*pb.MetricsResponse, error) { return e.s.GetMetrics(ctx, in) } @@ -158,3 +167,39 @@ func (e *embedBackendServerStream) SendMsg(m any) error { func (e *embedBackendServerStream) RecvMsg(m any) error { return nil } + +type embedBackendTrainStream struct { + ctx context.Context + fn func(resp *pb.TrainResponse) +} + +func (e *embedBackendTrainStream) Send(resp *pb.TrainResponse) error { + e.fn(resp) + return nil +} + +func (e *embedBackendTrainStream) SetHeader(md metadata.MD) error { + return nil +} + +func (e *embedBackendTrainStream) SendHeader(md metadata.MD) error { + return nil +} + +func (e *embedBackendTrainStream) SetTrailer(md metadata.MD) { +} + +func (e *embedBackendTrainStream) Context() context.Context { + return e.ctx +} + +func (e *embedBackendTrainStream) SendMsg(m any) error { + if x, ok := m.(*pb.TrainResponse); ok { + return e.Send(x) + } + return nil +} + +func (e *embedBackendTrainStream) RecvMsg(m any) error { + return nil +} diff --git a/pkg/grpc/interface.go b/pkg/grpc/interface.go index 6b1dd9a46e5d..f451e4173366 100644 --- a/pkg/grpc/interface.go +++ b/pkg/grpc/interface.go @@ -35,6 +35,8 @@ type AIModel interface { AudioDecode(*pb.AudioDecodeRequest) (*pb.AudioDecodeResult, error) ModelMetadata(*pb.ModelOptions) (*pb.ModelMetadataResponse, error) + + TrainStream(*pb.TrainRequest, chan *pb.TrainResponse) error } func newReply(s string) *pb.Reply { diff --git a/pkg/grpc/server.go b/pkg/grpc/server.go index 90d494f583a5..52a838a369ce 100644 --- a/pkg/grpc/server.go +++ b/pkg/grpc/server.go @@ -320,6 +320,27 @@ func (s *server) ModelMetadata(ctx context.Context, in *pb.ModelOptions) (*pb.Mo return res, nil } +func (s *server) TrainStream(in *pb.TrainRequest, stream pb.Backend_TrainStreamServer) error { + if s.llm.Locking() { + s.llm.Lock() + defer s.llm.Unlock() + } + resultChan := make(chan *pb.TrainResponse) + + done := make(chan bool) + go func() { + for result := range resultChan { + stream.Send(result) + } + done <- true + }() + + err := s.llm.TrainStream(in, resultChan) + <-done + + return err +} + func StartServer(address string, model AIModel) error { lis, err := net.Listen("tcp", address) if err != nil {