diff --git a/setup.sh b/setup.sh index ab84c7a053..4277b55aa5 100755 --- a/setup.sh +++ b/setup.sh @@ -194,7 +194,7 @@ else VENV_T5_DIR="$SCRIPT_DIR/.venv_t5" mkdir -p "$VENV_T5_DIR" run_quiet "pip install transformers 5.x" pip install --target "$VENV_T5_DIR" --no-deps "transformers==5.1.0" - run_quiet "pip install huggingface_hub for t5" pip install --target "$VENV_T5_DIR" --no-deps "huggingface_hub>=1.3.0" + run_quiet "pip install huggingface_hub for t5" pip install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==0.36.0" echo "✅ Transformers 5.x pre-installed to .venv_t5/" # ── 7. WSL: pre-install GGUF build dependencies ── diff --git a/studio/backend/core/export/__init__.py b/studio/backend/core/export/__init__.py index 66154f48eb..0a883f5f3f 100644 --- a/studio/backend/core/export/__init__.py +++ b/studio/backend/core/export/__init__.py @@ -1,9 +1,17 @@ """ Export submodule - Model export operations + +The default get_export_backend() returns an ExportOrchestrator that +delegates to a subprocess. The original ExportBackend runs inside +the subprocess and can be imported directly from .export when needed. """ -from .export import ExportBackend, get_export_backend +from .orchestrator import ExportOrchestrator, get_export_backend + +# Expose ExportOrchestrator as ExportBackend for backward compat +ExportBackend = ExportOrchestrator __all__ = [ 'ExportBackend', + 'ExportOrchestrator', 'get_export_backend', ] diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py new file mode 100644 index 0000000000..99c61edc3f --- /dev/null +++ b/studio/backend/core/export/orchestrator.py @@ -0,0 +1,385 @@ +""" +Export orchestrator — subprocess-based. + +Provides the same API as ExportBackend, but delegates all ML work +to a persistent subprocess. The subprocess is spawned on first checkpoint +load and stays alive for subsequent export operations. + +When switching between checkpoints that need different transformers versions, +the old subprocess is killed and a new one is spawned with the correct version. + +Pattern follows core/inference/orchestrator.py. +""" +import atexit +import logging +import multiprocessing as mp +import queue +import threading +import time +from pathlib import Path +from typing import Any, List, Optional, Tuple + +logger = logging.getLogger(__name__) + +_CTX = mp.get_context("spawn") + + +class ExportOrchestrator: + """ + Export backend orchestrator — subprocess-based. + + Exposes the same API surface as ExportBackend so routes/export.py + needs minimal changes. Internally, all heavy ML operations happen in + a persistent subprocess. + """ + + def __init__(self): + # Subprocess state + self._proc: Optional[mp.Process] = None + self._cmd_queue: Any = None + self._resp_queue: Any = None + self._lock = threading.Lock() + + # Local state mirrors (updated from subprocess responses) + self.current_checkpoint: Optional[str] = None + self.is_vision: bool = False + self.is_peft: bool = False + + atexit.register(self._cleanup) + logger.info("ExportOrchestrator initialized (subprocess mode)") + + # ------------------------------------------------------------------ + # Subprocess lifecycle + # ------------------------------------------------------------------ + + def _spawn_subprocess(self, config: dict) -> None: + """Spawn a new export subprocess.""" + from .worker import run_export_process + + self._cmd_queue = _CTX.Queue() + self._resp_queue = _CTX.Queue() + + self._proc = _CTX.Process( + target=run_export_process, + kwargs={ + "cmd_queue": self._cmd_queue, + "resp_queue": self._resp_queue, + "config": config, + }, + daemon=True, + ) + self._proc.start() + logger.info("Export subprocess started (pid=%s)", self._proc.pid) + + def _shutdown_subprocess(self, timeout: float = 10.0) -> None: + """Gracefully shut down the export subprocess.""" + if self._proc is None or not self._proc.is_alive(): + self._proc = None + return + + # 1. Drain stale responses + self._drain_queue() + + # 2. Send shutdown command + try: + self._cmd_queue.put({"type": "shutdown"}) + except (OSError, ValueError): + pass + + # 3. Wait for graceful shutdown + try: + self._proc.join(timeout=timeout) + except Exception: + pass + + # 4. Force kill if still alive + if self._proc is not None and self._proc.is_alive(): + logger.warning("Export subprocess did not exit gracefully, terminating") + try: + self._proc.terminate() + self._proc.join(timeout=5) + except Exception: + pass + if self._proc is not None and self._proc.is_alive(): + logger.warning("Subprocess still alive after terminate, killing") + try: + self._proc.kill() + self._proc.join(timeout=3) + except Exception: + pass + + self._proc = None + self._cmd_queue = None + self._resp_queue = None + logger.info("Export subprocess shut down") + + def _cleanup(self): + """atexit handler.""" + self._shutdown_subprocess(timeout=5.0) + + def _ensure_subprocess_alive(self) -> bool: + """Check if subprocess is alive.""" + return self._proc is not None and self._proc.is_alive() + + # ------------------------------------------------------------------ + # Queue helpers + # ------------------------------------------------------------------ + + def _send_cmd(self, cmd: dict) -> None: + """Send a command to the subprocess.""" + if self._cmd_queue is None: + raise RuntimeError("No export subprocess running") + try: + self._cmd_queue.put(cmd) + except (OSError, ValueError) as exc: + raise RuntimeError(f"Failed to send command to subprocess: {exc}") + + def _read_resp(self, timeout: float = 1.0) -> Optional[dict]: + """Read a response from the subprocess (non-blocking with timeout).""" + if self._resp_queue is None: + return None + try: + return self._resp_queue.get(timeout=timeout) + except queue.Empty: + return None + except (EOFError, OSError, ValueError): + return None + + def _wait_response( + self, expected_type: str, timeout: float = 600.0 + ) -> dict: + """Block until a response of the expected type arrives. + + Export operations can take a long time (GGUF build, push to Hub), + so default timeout is 10 minutes. + """ + deadline = time.monotonic() + timeout + + while time.monotonic() < deadline: + remaining = max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout=min(remaining, 2.0)) + + if resp is None: + # Check subprocess health + if not self._ensure_subprocess_alive(): + raise RuntimeError("Export subprocess crashed during wait") + continue + + rtype = resp.get("type", "") + + if rtype == expected_type: + return resp + + if rtype == "error": + error_msg = resp.get("error", "Unknown error") + raise RuntimeError(f"Subprocess error: {error_msg}") + + if rtype == "status": + logger.info("Export subprocess status: %s", resp.get("message", "")) + continue + + # Other response types during wait — skip + logger.debug( + "Skipping response type '%s' while waiting for '%s'", + rtype, expected_type, + ) + + raise RuntimeError( + f"Timeout waiting for '{expected_type}' response after {timeout}s" + ) + + def _drain_queue(self) -> list: + """Drain all pending responses.""" + events = [] + if self._resp_queue is None: + return events + while True: + try: + events.append(self._resp_queue.get_nowait()) + except queue.Empty: + return events + except (EOFError, OSError, ValueError): + return events + + # ------------------------------------------------------------------ + # Public API — same interface as ExportBackend + # ------------------------------------------------------------------ + + def load_checkpoint( + self, + checkpoint_path: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + ) -> Tuple[bool, str]: + """Load a checkpoint for export. + + Always spawns a fresh subprocess to ensure a clean Python interpreter. + """ + project_root = str( + Path(__file__).resolve().parent.parent.parent.parent.parent + ) + + sub_config = { + "project_root": project_root, + "checkpoint_path": checkpoint_path, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + } + + # Always kill existing subprocess and spawn fresh. + if self._ensure_subprocess_alive(): + self._shutdown_subprocess() + elif self._proc is not None: + self._shutdown_subprocess(timeout=2) + + logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path) + self._spawn_subprocess(sub_config) + + try: + resp = self._wait_response("loaded", timeout=300) + except RuntimeError as exc: + self._shutdown_subprocess(timeout=5) + return False, str(exc) + + if resp.get("success"): + self.current_checkpoint = resp.get("checkpoint") + self.is_vision = resp.get("is_vision", False) + self.is_peft = resp.get("is_peft", False) + logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path) + return True, resp.get("message", "Loaded successfully") + else: + error = resp.get("message", "Failed to load checkpoint") + logger.error("Failed to load checkpoint: %s", error) + return False, error + + def export_merged_model( + self, + save_directory: str, + format_type: str = "16-bit (FP16)", + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + private: bool = False, + ) -> Tuple[bool, str]: + """Export merged PEFT model.""" + return self._run_export("merged", { + "save_directory": save_directory, + "format_type": format_type, + "push_to_hub": push_to_hub, + "repo_id": repo_id, + "hf_token": hf_token, + "private": private, + }) + + def export_base_model( + self, + save_directory: str, + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + private: bool = False, + base_model_id: Optional[str] = None, + ) -> Tuple[bool, str]: + """Export base model (non-PEFT).""" + return self._run_export("base", { + "save_directory": save_directory, + "push_to_hub": push_to_hub, + "repo_id": repo_id, + "hf_token": hf_token, + "private": private, + "base_model_id": base_model_id, + }) + + def export_gguf( + self, + save_directory: str, + quantization_method: str = "Q4_K_M", + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + ) -> Tuple[bool, str]: + """Export model in GGUF format.""" + return self._run_export("gguf", { + "save_directory": save_directory, + "quantization_method": quantization_method, + "push_to_hub": push_to_hub, + "repo_id": repo_id, + "hf_token": hf_token, + }) + + def export_lora_adapter( + self, + save_directory: str, + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + private: bool = False, + ) -> Tuple[bool, str]: + """Export LoRA adapter only.""" + return self._run_export("lora", { + "save_directory": save_directory, + "push_to_hub": push_to_hub, + "repo_id": repo_id, + "hf_token": hf_token, + "private": private, + }) + + def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str]: + """Send an export command to the subprocess and wait for result.""" + if not self._ensure_subprocess_alive(): + return False, "No export subprocess running. Load a checkpoint first." + + cmd = {"type": "export", "export_type": export_type, **params} + + try: + self._send_cmd(cmd) + resp = self._wait_response( + f"export_{export_type}_done", + timeout=600, # Export can take a long time + ) + return resp.get("success", False), resp.get("message", "") + except RuntimeError as exc: + return False, str(exc) + + def cleanup_memory(self) -> bool: + """Cleanup export-related models from memory.""" + if not self._ensure_subprocess_alive(): + # No subprocess — just clear local state + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return True + + try: + self._send_cmd({"type": "cleanup"}) + resp = self._wait_response("cleanup_done", timeout=30) + success = resp.get("success", False) + except RuntimeError: + success = False + + # Shut down subprocess after cleanup — no model loaded + self._shutdown_subprocess() + + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return success + + def scan_checkpoints( + self, outputs_dir: str = "./outputs" + ) -> List[Tuple[str, list]]: + """Scan for checkpoints — no ML imports needed, runs locally.""" + from utils.models.checkpoints import scan_checkpoints + return scan_checkpoints(outputs_dir=outputs_dir) + + +# ========== GLOBAL INSTANCE ========== +_export_backend = None + + +def get_export_backend() -> ExportOrchestrator: + """Get global export backend instance (orchestrator).""" + global _export_backend + if _export_backend is None: + _export_backend = ExportOrchestrator() + return _export_backend diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py new file mode 100644 index 0000000000..a184969c95 --- /dev/null +++ b/studio/backend/core/export/worker.py @@ -0,0 +1,330 @@ +""" +Export subprocess entry point. + +Each export session runs in a persistent subprocess (mp.get_context("spawn")). +This gives us a clean Python interpreter with no stale module state — +solving the transformers version-switching problem completely. + +The subprocess stays alive while a model is loaded, accepting commands +(load, export_merged, export_base, export_gguf, export_lora, cleanup, +shutdown) via mp.Queue. + +Pattern follows core/inference/worker.py and core/training/worker.py. +""" +from __future__ import annotations + +import logging +import os +import sys +import time +import traceback +from typing import Any + +logger = logging.getLogger(__name__) + + +def _activate_transformers_version(model_name: str, project_root: str) -> None: + """Activate the correct transformers version BEFORE any ML imports. + + If the model needs transformers 5.x, prepend the pre-installed .venv_t5/ + directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/). + """ + # Ensure backend is on path for utils imports + backend_path = os.path.join(project_root, "studio", "backend") + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.transformers_version import needs_transformers_5, _resolve_base_model + + resolved = _resolve_base_model(model_name) + if needs_transformers_5(resolved): + venv_t5 = os.path.join(project_root, ".venv_t5") + if os.path.isdir(venv_t5): + sys.path.insert(0, venv_t5) + logger.info("Activated transformers 5.x from %s", venv_t5) + else: + # Fallback: pip install at runtime (slower, ~10-15s) + logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5) + import subprocess as sp + os.makedirs(venv_t5, exist_ok=True) + sp.run( + [sys.executable, "-m", "pip", "install", "--target", venv_t5, + "--no-deps", "transformers==5.1.0"], + stdout=sp.PIPE, stderr=sp.STDOUT, + ) + sp.run( + [sys.executable, "-m", "pip", "install", "--target", venv_t5, + "--no-deps", "huggingface_hub==0.36.0"], + stdout=sp.PIPE, stderr=sp.STDOUT, + ) + if os.path.isdir(venv_t5): + sys.path.insert(0, venv_t5) + else: + logger.info("Using default transformers (4.57.x) for %s", model_name) + + +def _send_response(resp_queue: Any, response: dict) -> None: + """Send a response to the parent process.""" + try: + resp_queue.put(response) + except (OSError, ValueError) as exc: + logger.error("Failed to send response: %s", exc) + + +def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: + """Handle a load_checkpoint command.""" + checkpoint_path = cmd["checkpoint_path"] + max_seq_length = cmd.get("max_seq_length", 2048) + load_in_4bit = cmd.get("load_in_4bit", True) + + try: + _send_response(resp_queue, { + "type": "status", + "message": f"Loading checkpoint: {checkpoint_path}", + "ts": time.time(), + }) + + success, message = backend.load_checkpoint( + checkpoint_path=checkpoint_path, + max_seq_length=max_seq_length, + load_in_4bit=load_in_4bit, + ) + + _send_response(resp_queue, { + "type": "loaded", + "success": success, + "message": message, + "checkpoint": checkpoint_path if success else None, + "is_vision": backend.is_vision if success else False, + "is_peft": backend.is_peft if success else False, + "ts": time.time(), + }) + + except Exception as exc: + _send_response(resp_queue, { + "type": "loaded", + "success": False, + "message": str(exc), + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + + +def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: + """Handle any export command (merged, base, gguf, lora).""" + export_type = cmd["export_type"] # "merged", "base", "gguf", "lora" + response_type = f"export_{export_type}_done" + + try: + if export_type == "merged": + success, message = backend.export_merged_model( + save_directory=cmd.get("save_directory", ""), + format_type=cmd.get("format_type", "16-bit (FP16)"), + push_to_hub=cmd.get("push_to_hub", False), + repo_id=cmd.get("repo_id"), + hf_token=cmd.get("hf_token"), + private=cmd.get("private", False), + ) + elif export_type == "base": + success, message = backend.export_base_model( + save_directory=cmd.get("save_directory", ""), + push_to_hub=cmd.get("push_to_hub", False), + repo_id=cmd.get("repo_id"), + hf_token=cmd.get("hf_token"), + private=cmd.get("private", False), + base_model_id=cmd.get("base_model_id"), + ) + elif export_type == "gguf": + success, message = backend.export_gguf( + save_directory=cmd.get("save_directory", ""), + quantization_method=cmd.get("quantization_method", "Q4_K_M"), + push_to_hub=cmd.get("push_to_hub", False), + repo_id=cmd.get("repo_id"), + hf_token=cmd.get("hf_token"), + ) + elif export_type == "lora": + success, message = backend.export_lora_adapter( + save_directory=cmd.get("save_directory", ""), + push_to_hub=cmd.get("push_to_hub", False), + repo_id=cmd.get("repo_id"), + hf_token=cmd.get("hf_token"), + private=cmd.get("private", False), + ) + else: + success, message = False, f"Unknown export type: {export_type}" + + _send_response(resp_queue, { + "type": response_type, + "success": success, + "message": message, + "ts": time.time(), + }) + + except Exception as exc: + _send_response(resp_queue, { + "type": response_type, + "success": False, + "message": str(exc), + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + + +def _handle_cleanup(backend, resp_queue: Any) -> None: + """Handle a cleanup command.""" + try: + success = backend.cleanup_memory() + _send_response(resp_queue, { + "type": "cleanup_done", + "success": success, + "ts": time.time(), + }) + except Exception as exc: + _send_response(resp_queue, { + "type": "cleanup_done", + "success": False, + "message": str(exc), + "ts": time.time(), + }) + + +def run_export_process( + *, + cmd_queue: Any, + resp_queue: Any, + config: dict, +) -> None: + """Subprocess entrypoint. Persistent — runs command loop until shutdown. + + Args: + cmd_queue: mp.Queue for receiving commands from parent. + resp_queue: mp.Queue for sending responses to parent. + config: Initial configuration dict with checkpoint_path and project_root. + """ + import queue as _queue + + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + project_root = config["project_root"] + checkpoint_path = config["checkpoint_path"] + + # ── 1. Activate correct transformers version BEFORE any ML imports ── + try: + _activate_transformers_version(checkpoint_path, project_root) + except Exception as exc: + _send_response(resp_queue, { + "type": "error", + "error": f"Failed to activate transformers version: {exc}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + # ── 2. Import ML libraries (fresh in this clean process) ── + try: + _send_response(resp_queue, { + "type": "status", + "message": "Importing ML libraries...", + "ts": time.time(), + }) + + backend_path = os.path.join(project_root, "studio", "backend") + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from core.export.export import ExportBackend + + import transformers + logger.info("Export subprocess loaded transformers %s", transformers.__version__) + + except Exception as exc: + _send_response(resp_queue, { + "type": "error", + "error": f"Failed to import ML libraries: {exc}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + # ── 3. Create export backend and load initial checkpoint ── + try: + backend = ExportBackend() + + _handle_load(backend, config, resp_queue) + + except Exception as exc: + _send_response(resp_queue, { + "type": "error", + "error": f"Failed to initialize export backend: {exc}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + # ── 4. Command loop — process commands until shutdown ── + logger.info("Export subprocess ready, entering command loop") + + while True: + try: + cmd = cmd_queue.get(timeout=1.0) + except _queue.Empty: + continue + except (EOFError, OSError): + logger.info("Command queue closed, shutting down") + return + + if cmd is None: + continue + + cmd_type = cmd.get("type", "") + logger.info("Received command: %s", cmd_type) + + try: + if cmd_type == "load": + # Load a new checkpoint (reusing this subprocess) + backend.cleanup_memory() + _handle_load(backend, cmd, resp_queue) + + elif cmd_type == "export": + _handle_export(backend, cmd, resp_queue) + + elif cmd_type == "cleanup": + _handle_cleanup(backend, resp_queue) + + elif cmd_type == "status": + _send_response(resp_queue, { + "type": "status_response", + "checkpoint": backend.current_checkpoint, + "is_vision": backend.is_vision, + "is_peft": backend.is_peft, + "ts": time.time(), + }) + + elif cmd_type == "shutdown": + logger.info("Shutdown command received, cleaning up and exiting") + try: + backend.cleanup_memory() + except Exception: + pass + _send_response(resp_queue, { + "type": "shutdown_ack", + "ts": time.time(), + }) + return + + else: + logger.warning("Unknown command type: %s", cmd_type) + _send_response(resp_queue, { + "type": "error", + "error": f"Unknown command type: {cmd_type}", + "ts": time.time(), + }) + + except Exception as exc: + logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info=True) + _send_response(resp_queue, { + "type": "error", + "error": f"Command '{cmd_type}' failed: {exc}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 3e98fb0e3c..985ebc74a3 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -56,7 +56,7 @@ def _activate_transformers_version(model_name: str, project_root: str) -> None: ) sp.run( [sys.executable, "-m", "pip", "install", "--target", venv_t5, - "--no-deps", "huggingface_hub>=1.3.0"], + "--no-deps", "huggingface_hub==0.36.0"], stdout=sp.PIPE, stderr=sp.STDOUT, ) if os.path.isdir(venv_t5): diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index de7499f2a1..018345b0cf 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -51,7 +51,7 @@ def _activate_transformers_version(model_name: str, project_root: str) -> None: ) sp.run( [sys.executable, "-m", "pip", "install", "--target", venv_t5, - "--no-deps", "huggingface_hub>=1.3.0"], + "--no-deps", "huggingface_hub==0.36.0"], stdout=sp.PIPE, stderr=sp.STDOUT, ) if os.path.isdir(venv_t5): diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 11d6377480..1b66306842 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -61,10 +61,8 @@ async def load_checkpoint( Wraps ExportBackend.load_checkpoint. """ try: - # Ensure correct transformers version for this model architecture - from utils.transformers_version import ensure_transformers_version - ensure_transformers_version(request.checkpoint_path) - + # Version switching is handled automatically by the subprocess-based + # export backend — no need for ensure_transformers_version() here. backend = get_export_backend() success, message = backend.load_checkpoint( checkpoint_path=request.checkpoint_path, diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 5079021ffe..418105c145 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -171,7 +171,7 @@ def _ensure_venv_t5_exists() -> bool: logger.warning(".venv_t5 not found at %s — installing at runtime", _VENV_T5_DIR) os.makedirs(_VENV_T5_DIR, exist_ok=True) - for pkg in (f"transformers=={TRANSFORMERS_5_VERSION}", "huggingface_hub>=1.3.0"): + for pkg in (f"transformers=={TRANSFORMERS_5_VERSION}", "huggingface_hub==0.36.0"): cmd = [ sys.executable, "-m", "pip", "install", "--target", _VENV_T5_DIR,