diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 966e045b13..d8f2e8fa37 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -310,7 +310,7 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """ Export merged model (for PEFT models). @@ -323,14 +323,21 @@ class ExportBackend: private: Whether to make the repo private Returns: - Tuple of (success: bool, message: str) + Tuple of (success, message, output_path). output_path is the + resolved absolute on-disk directory of the saved model when + ``save_directory`` was set, else None. """ if not self.current_model or not self.current_tokenizer: - return False, "No model loaded. Please select a checkpoint first." + return False, "No model loaded. Please select a checkpoint first.", None if not self.is_peft: - return False, "This is not a PEFT model. Use 'Export Base Model' instead." + return ( + False, + "This is not a PEFT model. Use 'Export Base Model' instead.", + None, + ) + output_path: Optional[str] = None try: # Determine save method if format_type == "4-bit (FP4)": @@ -354,6 +361,7 @@ class ExportBackend: # Write export metadata so the Chat page can identify the base model self._write_export_metadata(save_directory) logger.info(f"Model saved successfully to {save_directory}") + output_path = str(Path(save_directory).resolve()) # Push to hub if requested if push_to_hub: @@ -361,6 +369,7 @@ class ExportBackend: return ( False, "Repository ID and Hugging Face token required for Hub upload", + None, ) logger.info(f"Pushing merged model to Hub: {repo_id}") @@ -378,14 +387,14 @@ class ExportBackend: ) logger.info(f"Model pushed successfully to {repo_id}") - return True, "Model exported successfully" + return True, "Model exported successfully", output_path except Exception as e: logger.error(f"Error exporting merged model: {e}") import traceback logger.error(traceback.format_exc()) - return False, f"Export failed: {str(e)}" + return False, f"Export failed: {str(e)}", None def export_base_model( self, @@ -395,22 +404,26 @@ class ExportBackend: hf_token: Optional[str] = None, private: bool = False, base_model_id: Optional[str] = None, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """ Export base model (for non-PEFT models). Returns: - Tuple of (success: bool, message: str) + Tuple of (success, message, output_path). output_path is the + resolved absolute on-disk directory of the saved model when + ``save_directory`` was set, else None. """ if not self.current_model or not self.current_tokenizer: - return False, "No model loaded. Please select a checkpoint first." + return False, "No model loaded. Please select a checkpoint first.", None if self.is_peft: return ( False, "This is a PEFT model. Use 'Merged Model' export type instead.", + None, ) + output_path: Optional[str] = None try: # Save locally if requested if save_directory: @@ -424,6 +437,7 @@ class ExportBackend: # Write export metadata so the Chat page can identify the base model self._write_export_metadata(save_directory) logger.info(f"Model saved successfully to {save_directory}") + output_path = str(Path(save_directory).resolve()) # Push to hub if requested if push_to_hub: @@ -431,6 +445,7 @@ class ExportBackend: return ( False, "Repository ID and Hugging Face token required for Hub upload", + None, ) logger.info(f"Pushing base model to Hub: {repo_id}") @@ -472,16 +487,16 @@ class ExportBackend: ) logger.info(f"Model pushed successfully to {repo_id}") else: - return False, "Local save directory required for Hub upload" + return False, "Local save directory required for Hub upload", None - return True, "Model exported successfully" + return True, "Model exported successfully", output_path except Exception as e: logger.error(f"Error exporting base model: {e}") import traceback logger.error(traceback.format_exc()) - return False, f"Export failed: {str(e)}" + return False, f"Export failed: {str(e)}", None def export_gguf( self, @@ -490,7 +505,7 @@ class ExportBackend: push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """ Export model in GGUF format. @@ -502,11 +517,14 @@ class ExportBackend: hf_token: Hugging Face token Returns: - Tuple of (success: bool, message: str) + Tuple of (success, message, output_path). output_path is the + resolved absolute on-disk directory containing the .gguf + files when ``save_directory`` was set, else None. """ if not self.current_model or not self.current_tokenizer: - return False, "No model loaded. Please select a checkpoint first." + return False, "No model loaded. Please select a checkpoint first.", None + output_path: Optional[str] = None try: # Convert quantization method to lowercase for unsloth quant_method = quantization_method.lower() @@ -601,6 +619,7 @@ class ExportBackend: abs_save_dir, "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", ) + output_path = str(Path(abs_save_dir).resolve()) # Push to hub if requested if push_to_hub: @@ -608,6 +627,7 @@ class ExportBackend: return ( False, "Repository ID and Hugging Face token required for Hub upload", + None, ) logger.info(f"Pushing GGUF model to Hub: {repo_id}") @@ -620,14 +640,18 @@ class ExportBackend: ) logger.info(f"GGUF model pushed successfully to {repo_id}") - return True, f"GGUF model exported successfully ({quantization_method})" + return ( + True, + f"GGUF model exported successfully ({quantization_method})", + output_path, + ) except Exception as e: logger.error(f"Error exporting GGUF model: {e}") import traceback logger.error(traceback.format_exc()) - return False, f"GGUF export failed: {str(e)}" + return False, f"GGUF export failed: {str(e)}", None def export_lora_adapter( self, @@ -636,19 +660,22 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """ Export LoRA adapter only (not merged). Returns: - Tuple of (success: bool, message: str) + Tuple of (success, message, output_path). output_path is the + resolved absolute on-disk directory of the saved adapter + when ``save_directory`` was set, else None. """ if not self.current_model or not self.current_tokenizer: - return False, "No model loaded. Please select a checkpoint first." + return False, "No model loaded. Please select a checkpoint first.", None if not self.is_peft: - return False, "This is not a PEFT model. No adapter to export." + return False, "This is not a PEFT model. No adapter to export.", None + output_path: Optional[str] = None try: # Save locally if requested if save_directory: @@ -659,6 +686,7 @@ class ExportBackend: self.current_model.save_pretrained(save_directory) self.current_tokenizer.save_pretrained(save_directory) logger.info(f"Adapter saved successfully to {save_directory}") + output_path = str(Path(save_directory).resolve()) # Push to hub if requested if push_to_hub: @@ -666,6 +694,7 @@ class ExportBackend: return ( False, "Repository ID and Hugging Face token required for Hub upload", + None, ) logger.info(f"Pushing LoRA adapter to Hub: {repo_id}") @@ -676,14 +705,14 @@ class ExportBackend: ) logger.info(f"Adapter pushed successfully to {repo_id}") - return True, "LoRA adapter exported successfully" + return True, "LoRA adapter exported successfully", output_path except Exception as e: logger.error(f"Error exporting LoRA adapter: {e}") import traceback logger.error(traceback.format_exc()) - return False, f"Adapter export failed: {str(e)}" + return False, f"Adapter export failed: {str(e)}", None # Global export backend instance diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index e52aa8c3cd..206dbd6dbb 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -16,19 +16,25 @@ Pattern follows core/inference/orchestrator.py. import atexit import structlog +from collections import deque from loggers import get_logger import multiprocessing as mp import queue import threading import time from pathlib import Path -from typing import Any, List, Optional, Tuple +from typing import Any, Deque, Dict, List, Optional, Tuple from utils.paths import outputs_root logger = get_logger(__name__) _CTX = mp.get_context("spawn") +# Maximum number of captured log lines kept in memory per export +# orchestrator. Acts as scrollback for the live export log panel in the +# UI. 4000 lines is ~1 MB worst-case at 256 chars/line. +_LOG_BUFFER_MAXLEN = 4000 + class ExportOrchestrator: """ @@ -44,6 +50,9 @@ class ExportOrchestrator: self._proc: Optional[mp.Process] = None self._cmd_queue: Any = None self._resp_queue: Any = None + # Serializes export operations (load_checkpoint, export_*, + # cleanup) so concurrent HTTP requests can never interleave + # commands on the subprocess queue. Previously unused. self._lock = threading.Lock() # Local state mirrors (updated from subprocess responses) @@ -51,9 +60,103 @@ class ExportOrchestrator: self.is_vision: bool = False self.is_peft: bool = False + # ── Live log capture ───────────────────────────────────── + # Thread-safe ring buffer of log lines forwarded from the + # worker subprocess. Powers the GET /api/export/logs/stream + # SSE endpoint that the export dialog consumes. + self._log_buffer: Deque[Dict[str, Any]] = deque(maxlen = _LOG_BUFFER_MAXLEN) + self._log_lock = threading.Lock() + # Monotonically increasing sequence number. Never reset across + # operations, so SSE clients can use it as a stable cursor even + # if clear_logs() is called mid-session. + self._log_seq: int = 0 + # Snapshot of _log_seq captured at the start of the current run + # (updated by clear_logs()). The SSE endpoint defaults its + # cursor to this value so a client that connects AFTER the + # worker has already emitted its first lines still sees the + # full run. Every line appended during the current run has seq + # strictly greater than _run_start_seq, and every line from + # prior runs has seq less than or equal to it. + self._run_start_seq: int = 0 + # True while an export operation (load/export/cleanup) is + # running. The SSE endpoint ends the stream 1 second after + # this flips back to False to drain any trailing log lines. + self._export_active: bool = False + atexit.register(self._cleanup) logger.info("ExportOrchestrator initialized (subprocess mode)") + # ------------------------------------------------------------------ + # Live log capture helpers + # ------------------------------------------------------------------ + + def _append_log(self, entry: Dict[str, Any]) -> None: + """Append a log line from the worker subprocess to the buffer. + + Entries look like {"type": "log", "stream": "stdout"|"stderr", + "line": "...", "ts": ...}. Each is stamped with a monotonic + seq number before it lands in the buffer so SSE clients can + cursor through new lines. + """ + line = entry.get("line") + if not line: + return + with self._log_lock: + self._log_seq += 1 + self._log_buffer.append( + { + "seq": self._log_seq, + "stream": entry.get("stream", "stdout"), + "line": line, + "ts": entry.get("ts", time.time()), + } + ) + + def clear_logs(self) -> None: + """Drop any buffered log lines from a previous operation. + + Called at the start of each export op so the UI shows only the + output of the current run. The seq counter is NOT reset, so an + SSE client that captured the cursor before clear_logs() will + still see new lines (with strictly greater seq numbers). + + Also snapshots the current seq into ``_run_start_seq`` so the + SSE endpoint can anchor its default cursor at the start of + this run. Anything appended after this call has seq strictly + greater than the snapshot and is reachable via + ``get_logs_since(get_run_start_seq())``. + """ + with self._log_lock: + self._log_buffer.clear() + self._run_start_seq = self._log_seq + + def get_logs_since(self, cursor: int) -> Tuple[List[Dict[str, Any]], int]: + """Return log entries with seq > cursor, plus the new cursor.""" + with self._log_lock: + new_entries = [entry for entry in self._log_buffer if entry["seq"] > cursor] + if new_entries: + return new_entries, new_entries[-1]["seq"] + return [], cursor + + def get_current_log_seq(self) -> int: + """Return the current seq counter without reading any entries.""" + with self._log_lock: + return self._log_seq + + def get_run_start_seq(self) -> int: + """Return the seq value captured at the start of the current run. + + The SSE endpoint uses this as the default cursor so a client + that connects AFTER the worker has already started emitting + output still sees every line from the current run. + """ + with self._log_lock: + return self._run_start_seq + + def is_export_active(self) -> bool: + """True while an export / load / cleanup command is running.""" + return self._export_active + # ------------------------------------------------------------------ # Subprocess lifecycle # ------------------------------------------------------------------ @@ -179,8 +282,26 @@ class ExportOrchestrator: error_msg = resp.get("error", "Unknown error") raise RuntimeError(f"Subprocess error: {error_msg}") + if rtype == "log": + # Forwarded stdout/stderr line from the worker process. + self._append_log(resp) + continue + if rtype == "status": - logger.info("Export subprocess status: %s", resp.get("message", "")) + message = resp.get("message", "") + logger.info("Export subprocess status: %s", message) + # Surface status messages in the live log panel too so + # users see high level progress (e.g. "Importing + # Unsloth...", "Loading checkpoint: ...") alongside + # subprocess output. + if message: + self._append_log( + { + "stream": "status", + "line": message, + "ts": resp.get("ts", time.time()), + } + ) continue # Other response types during wait — skip @@ -231,37 +352,47 @@ class ExportOrchestrator: "hf_token": hf_token, } - # 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) + with self._lock: + # Start a fresh log buffer for this operation so the UI + # sees only the current run's output. + self.clear_logs() + self._export_active = True + try: + # 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) + logger.info( + "Spawning fresh export subprocess for '%s'", checkpoint_path + ) + self._spawn_subprocess(sub_config) - try: - resp = self._wait_response("loaded") - except RuntimeError as exc: - self._shutdown_subprocess(timeout = 5) - self.current_checkpoint = None - self.is_vision = False - self.is_peft = False - return False, str(exc) + try: + resp = self._wait_response("loaded") + except RuntimeError as exc: + self._shutdown_subprocess(timeout = 5) + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + 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) - self.current_checkpoint = None - self.is_vision = False - self.is_peft = False - return False, error + 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) + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return False, error + finally: + self._export_active = False def export_merged_model( self, @@ -271,7 +402,7 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """Export merged PEFT model.""" return self._run_export( "merged", @@ -293,7 +424,7 @@ class ExportOrchestrator: hf_token: Optional[str] = None, private: bool = False, base_model_id: Optional[str] = None, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """Export base model (non-PEFT).""" return self._run_export( "base", @@ -314,7 +445,7 @@ class ExportOrchestrator: push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """Export model in GGUF format.""" return self._run_export( "gguf", @@ -334,7 +465,7 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """Export LoRA adapter only.""" return self._run_export( "lora", @@ -347,46 +478,74 @@ class ExportOrchestrator: }, ) - 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." + def _run_export( + self, export_type: str, params: dict + ) -> Tuple[bool, str, Optional[str]]: + """Send an export command to the subprocess and wait for result. - cmd = {"type": "export", "export_type": export_type, **params} + Returns ``(success, message, output_path)``. ``output_path`` is the + resolved on-disk directory the worker actually wrote to (None when + the export only pushed to Hub or failed before any file was + written). Surfaced via the export route's ``details.output_path`` + so the dialog's success screen can show the user where the model + landed. + """ + with self._lock: + if not self._ensure_subprocess_alive(): + return ( + False, + "No export subprocess running. Load a checkpoint first.", + None, + ) - try: - self._send_cmd(cmd) - resp = self._wait_response( - f"export_{export_type}_done", - timeout = 3600, # GGUF for 30B+ models can take 30+ min - ) - return resp.get("success", False), resp.get("message", "") - except RuntimeError as exc: - return False, str(exc) + self.clear_logs() + self._export_active = True + try: + cmd = {"type": "export", "export_type": export_type, **params} + try: + self._send_cmd(cmd) + resp = self._wait_response( + f"export_{export_type}_done", + timeout = 3600, # GGUF for 30B+ models can take 30+ min + ) + return ( + resp.get("success", False), + resp.get("message", ""), + resp.get("output_path"), + ) + except RuntimeError as exc: + return False, str(exc), None + finally: + self._export_active = False 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 + with self._lock: + 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 + self._export_active = True + try: + 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() + # 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 + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return success + finally: + self._export_active = False def scan_checkpoints( self, outputs_dir: str = str(outputs_root()) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 3f3dc955fa..f77b1966c4 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -17,10 +17,12 @@ Pattern follows core/inference/worker.py and core/training/worker.py. from __future__ import annotations +import errno import structlog from loggers import get_logger import os import sys +import threading import time import traceback from pathlib import Path @@ -29,6 +31,154 @@ from typing import Any logger = get_logger(__name__) +# Gate that controls whether captured stdout/stderr lines are forwarded +# to the parent's resp_queue (and from there to the export-dialog SSE +# stream). Closed by default so the noisy bootstrap phase -- transformers +# venv activation, Unsloth/torch imports, base-model resolution, "Top +# GGUF/hub models" lists, vision detection, weight loading bars -- is +# suppressed in the UI. _handle_export() opens the gate at the start of +# the actual export work and leaves it open; the orchestrator always +# spawns a fresh subprocess for the next checkpoint load (see +# orchestrator._spawn_subprocess) which resets this state. +# +# Lines dropped while the gate is closed are still echoed to the saved +# original stdout/stderr fds so the server console / log file keeps the +# full output for debugging. +_log_forward_gate = threading.Event() + + +def _setup_log_capture(resp_queue: Any) -> None: + """Redirect fds 1 and 2 through pipes so every line printed by this + worker process and any child process it spawns is forwarded to the + parent process via resp_queue as {"type": "log", ...} messages. + + Must be called BEFORE LogConfig.setup_logging and BEFORE any ML + imports, otherwise library handlers may capture the original stderr + reference and bypass the pipe. + + Lines are also echoed back to the original stdout/stderr so the + server console keeps receiving the full subprocess output, even + while ``_log_forward_gate`` is closed. + """ + + try: + saved_out_fd = os.dup(1) + saved_err_fd = os.dup(2) + except OSError: + # dup failed (exotic platforms) - give up quietly, export still + # works, just no live log streaming. + return + + try: + r_out, w_out = os.pipe() + r_err, w_err = os.pipe() + except OSError: + os.close(saved_out_fd) + os.close(saved_err_fd) + return + + try: + os.dup2(w_out, 1) + os.dup2(w_err, 2) + except OSError: + for fd in (saved_out_fd, saved_err_fd, r_out, w_out, r_err, w_err): + try: + os.close(fd) + except OSError: + pass + return + + # Close the write ends we just dup2'd (fds 1 and 2 are the real + # write ends now). + os.close(w_out) + os.close(w_err) + + # Replace Python's sys.stdout/sys.stderr with line-buffered writers + # bound to the (now-redirected) fds 1 and 2. + try: + sys.stdout = os.fdopen(1, "w", buffering = 1, encoding = "utf-8", errors = "replace") + sys.stderr = os.fdopen(2, "w", buffering = 1, encoding = "utf-8", errors = "replace") + except Exception: + pass + + def _reader(read_fd: int, stream_name: str, echo_fd: int) -> None: + buf = bytearray() + while True: + try: + chunk = os.read(read_fd, 4096) + except OSError as exc: + if exc.errno == errno.EBADF: + break + continue + if not chunk: + break + # Echo to the original fd so the server console still sees + # the full output. + try: + os.write(echo_fd, chunk) + except OSError: + pass + buf.extend(chunk) + # Split on \n OR \r so tqdm-style progress bars update. + while True: + nl = -1 + for i, b in enumerate(buf): + if b == 0x0A or b == 0x0D: + nl = i + break + if nl < 0: + break + line = bytes(buf[:nl]).decode("utf-8", errors = "replace") + del buf[: nl + 1] + if not line: + continue + if not _log_forward_gate.is_set(): + # Gate closed (bootstrap phase) -- already echoed to + # the saved console fd above; drop the line so the + # export dialog doesn't see import / vendoring noise. + continue + try: + resp_queue.put_nowait( + { + "type": "log", + "stream": stream_name, + "line": line, + "ts": time.time(), + } + ) + except Exception: + # Queue put failed (full, closed, etc.) - drop the + # line rather than crash the reader thread. + pass + if buf and _log_forward_gate.is_set(): + try: + resp_queue.put_nowait( + { + "type": "log", + "stream": stream_name, + "line": bytes(buf).decode("utf-8", errors = "replace"), + "ts": time.time(), + } + ) + except Exception: + pass + + t_out = threading.Thread( + target = _reader, + args = (r_out, "stdout", saved_out_fd), + daemon = True, + name = "export-log-stdout", + ) + t_err = threading.Thread( + target = _reader, + args = (r_err, "stderr", saved_err_fd), + daemon = True, + name = "export-log-stderr", + ) + t_out.start() + t_err.start() + + def _activate_transformers_version(model_name: str) -> None: """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on path for utils imports @@ -117,9 +267,17 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: export_type = cmd["export_type"] # "merged", "base", "gguf", "lora" response_type = f"export_{export_type}_done" + # Open the log forwarding gate so the user sees the actual export + # progress (Unsloth merge bars, file copies, GGUF conversion, etc.) + # in the live log panel. The gate stays open for the rest of this + # subprocess's life; the orchestrator spawns a fresh subprocess for + # the next checkpoint load, which resets the gate to closed. + _log_forward_gate.set() + + output_path: Any = None try: if export_type == "merged": - success, message = backend.export_merged_model( + success, message, output_path = 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), @@ -128,7 +286,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: private = cmd.get("private", False), ) elif export_type == "base": - success, message = backend.export_base_model( + success, message, output_path = 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"), @@ -137,7 +295,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: base_model_id = cmd.get("base_model_id"), ) elif export_type == "gguf": - success, message = backend.export_gguf( + success, message, output_path = 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), @@ -145,7 +303,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: hf_token = cmd.get("hf_token"), ) elif export_type == "lora": - success, message = backend.export_lora_adapter( + success, message, output_path = 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"), @@ -161,6 +319,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: "type": response_type, "success": success, "message": message, + "output_path": output_path, "ts": time.time(), }, ) @@ -172,6 +331,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: "type": response_type, "success": False, "message": str(exc), + "output_path": None, "stack": traceback.format_exc(limit = 20), "ts": time.time(), }, @@ -217,10 +377,26 @@ def run_export_process( """ import queue as _queue + # Install fd-level stdout/stderr capture FIRST so every subsequent + # print and every child process inherits the redirected fds. This + # is what powers the live export log stream in the UI. + _setup_log_capture(resp_queue) + os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = ( "ignore" # Suppress warnings at C-level before imports ) + # Force unbuffered output from any child Python process (e.g. the + # GGUF converter) so their prints surface in the log stream as they + # happen rather than at the end. + os.environ["PYTHONUNBUFFERED"] = "1" + # tqdm defaults to a 10-second mininterval when stdout is not a tty + # (which it isn't here -- we redirected fd 1/2 to a pipe). That makes + # multi-step progress bars look frozen in the export log panel. Force + # frequent flushes so the user sees movement during merge / GGUF + # conversion. Has no effect on single-step bars (e.g. "Copying 1 + # files") which only emit start/end events regardless. + os.environ.setdefault("TQDM_MININTERVAL", "0.5") import warnings from loggers.config import LogConfig diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 3e60eaaf20..798859fc87 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -5,9 +5,15 @@ Export API routes: checkpoint discovery and model export operations. """ +import asyncio +import json import sys +import time from pathlib import Path -from fastapi import APIRouter, Depends, HTTPException, Query +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple + +from fastapi import APIRouter, Depends, HTTPException, Query, Request +from fastapi.responses import StreamingResponse import structlog from loggers import get_logger @@ -97,7 +103,11 @@ async def load_checkpoint( logger.warning("Could not stop training: %s", e) backend = get_export_backend() - success, message = backend.load_checkpoint( + # load_checkpoint spawns and waits on a subprocess and can take + # minutes. Run it in a worker thread so the event loop stays + # free to serve the live log SSE stream concurrently. + success, message = await asyncio.to_thread( + backend.load_checkpoint, checkpoint_path = request.checkpoint_path, max_seq_length = request.max_seq_length, load_in_4bit = request.load_in_4bit, @@ -129,7 +139,7 @@ async def cleanup_export_memory( """ try: backend = get_export_backend() - success = backend.cleanup_memory() + success = await asyncio.to_thread(backend.cleanup_memory) if not success: raise HTTPException( @@ -173,6 +183,17 @@ async def get_export_status( ) +def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]: + """Wrap the resolved on-disk export path into the details dict the + frontend reads to populate the Export Complete screen. Returns None + when the export had no local component (Hub-only push) so the + Pydantic field stays absent rather than ``{"output_path": null}``. + """ + if not output_path: + return None + return {"output_path": output_path} + + @router.post("/export/merged", response_model = ExportOperationResponse) async def export_merged_model( request: ExportMergedModelRequest, @@ -185,7 +206,8 @@ async def export_merged_model( """ try: backend = get_export_backend() - success, message = backend.export_merged_model( + success, message, output_path = await asyncio.to_thread( + backend.export_merged_model, save_directory = request.save_directory, format_type = request.format_type, push_to_hub = request.push_to_hub, @@ -197,7 +219,11 @@ async def export_merged_model( if not success: raise HTTPException(status_code = 400, detail = message) - return ExportOperationResponse(success = True, message = message) + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) except HTTPException: raise except Exception as e: @@ -220,7 +246,8 @@ async def export_base_model( """ try: backend = get_export_backend() - success, message = backend.export_base_model( + success, message, output_path = await asyncio.to_thread( + backend.export_base_model, save_directory = request.save_directory, push_to_hub = request.push_to_hub, repo_id = request.repo_id, @@ -232,7 +259,11 @@ async def export_base_model( if not success: raise HTTPException(status_code = 400, detail = message) - return ExportOperationResponse(success = True, message = message) + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) except HTTPException: raise except Exception as e: @@ -255,7 +286,8 @@ async def export_gguf( """ try: backend = get_export_backend() - success, message = backend.export_gguf( + success, message, output_path = await asyncio.to_thread( + backend.export_gguf, save_directory = request.save_directory, quantization_method = request.quantization_method, push_to_hub = request.push_to_hub, @@ -266,7 +298,11 @@ async def export_gguf( if not success: raise HTTPException(status_code = 400, detail = message) - return ExportOperationResponse(success = True, message = message) + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) except HTTPException: raise except Exception as e: @@ -289,7 +325,8 @@ async def export_lora_adapter( """ try: backend = get_export_backend() - success, message = backend.export_lora_adapter( + success, message, output_path = await asyncio.to_thread( + backend.export_lora_adapter, save_directory = request.save_directory, push_to_hub = request.push_to_hub, repo_id = request.repo_id, @@ -300,7 +337,11 @@ async def export_lora_adapter( if not success: raise HTTPException(status_code = 400, detail = message) - return ExportOperationResponse(success = True, message = message) + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) except HTTPException: raise except Exception as e: @@ -309,3 +350,155 @@ async def export_lora_adapter( status_code = 500, detail = f"Failed to export LoRA adapter: {str(e)}", ) + + +# ───────────────────────────────────────────────────────────────────── +# Live export log stream (Server-Sent Events) +# ───────────────────────────────────────────────────────────────────── +# +# The export worker subprocess redirects its stdout/stderr into a pipe +# that a reader thread forwards to the orchestrator as log entries (see +# core/export/worker.py::_setup_log_capture and +# core/export/orchestrator.py::_append_log). This endpoint streams +# those entries to the browser so the export dialog can show a live +# terminal-style output panel while load_checkpoint / export_merged / +# export_gguf / export_lora / export_base run. +# +# Shape follows the training progress SSE endpoint +# (routes/training.py::stream_training_progress): each event carries +# `id`, `event`, and `data` fields, the stream starts with a `retry:` +# directive, and `Last-Event-ID` is honored on reconnect. + + +def _format_sse(data: str, event: str, event_id: Optional[int] = None) -> str: + """Format a single SSE message with id/event/data fields.""" + lines = [] + if event_id is not None: + lines.append(f"id: {event_id}") + lines.append(f"event: {event}") + lines.append(f"data: {data}") + lines.append("") + lines.append("") + return "\n".join(lines) + + +@router.get("/logs/stream") +async def stream_export_logs( + request: Request, + since: Optional[int] = Query( + None, + description = "Return log entries with seq strictly greater than this cursor.", + ), + current_subject: str = Depends(get_current_subject), +): + """ + Stream live stdout/stderr output from the export worker subprocess + as Server-Sent Events. + + Events: + - `log` : a single log line (data: {"stream","line","ts"}) + - `heartbeat`: periodic keepalive when no new lines are available + - `complete` : emitted once the export worker is idle and no new + lines arrived for ~1 second. Clients should close. + - `error` : unrecoverable server-side error + + The `id:` field on each event is the log entry's monotonic seq + number so the browser can resume via `Last-Event-ID` on reconnect. + """ + backend = get_export_backend() + + # Determine starting cursor. Explicit `since` wins, then + # Last-Event-ID header on reconnect, otherwise start from the + # run-start snapshot captured by clear_logs() so the client sees + # every line emitted since the current run began -- even if the + # SSE connection opened after the POST that kicked off the export. + # Using get_current_log_seq() here would lose the early bootstrap + # lines that arrive in the gap between POST and SSE connect. + last_event_id = request.headers.get("last-event-id") + if since is None and last_event_id is not None: + try: + since = int(last_event_id) + except ValueError: + pass + + if since is None: + cursor = backend.get_run_start_seq() + else: + cursor = max(0, int(since)) + + async def event_generator() -> AsyncGenerator[str, None]: + nonlocal cursor + # Tell the browser to reconnect after 3 seconds if the + # connection drops mid-export. + yield "retry: 3000\n\n" + + last_yield = time.monotonic() + idle_since: Optional[float] = None + try: + while True: + if await request.is_disconnected(): + return + + entries, new_cursor = backend.get_logs_since(cursor) + if entries: + for entry in entries: + payload = json.dumps( + { + "stream": entry.get("stream", "stdout"), + "line": entry.get("line", ""), + "ts": entry.get("ts"), + } + ) + yield _format_sse( + payload, + event = "log", + event_id = int(entry.get("seq", 0)), + ) + cursor = new_cursor + last_yield = time.monotonic() + idle_since = None + else: + now = time.monotonic() + if now - last_yield > 10.0: + yield _format_sse("{}", event = "heartbeat") + last_yield = now + if not backend.is_export_active(): + # Give the reader thread a moment to drain any + # trailing lines the worker process printed + # just before signalling done. + if idle_since is None: + idle_since = now + elif now - idle_since > 1.0: + yield _format_sse( + "{}", + event = "complete", + event_id = cursor, + ) + return + else: + idle_since = None + + await asyncio.sleep(0.1) + except asyncio.CancelledError: + # Client disconnected mid-yield. Don't re-raise, just end + # the generator cleanly so StreamingResponse finalizes. + return + except Exception as exc: + logger.error("Export log stream failed: %s", exc, exc_info = True) + try: + yield _format_sse( + json.dumps({"error": str(exc)}), + event = "error", + ) + except Exception: + pass + + return StreamingResponse( + event_generator(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) diff --git a/studio/backend/tests/test_export_log_cursor.py b/studio/backend/tests/test_export_log_cursor.py new file mode 100644 index 0000000000..734ca522c9 --- /dev/null +++ b/studio/backend/tests/test_export_log_cursor.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Regression tests for the export log ring-buffer cursor semantics. + +Context: the live export log SSE stream has a race where the frontend +opens the SSE connection AFTER the POST that starts the export. Any +lines the worker subprocess emits during the gap between POST and SSE +connect get buffered with seqs 1..k, and then the SSE default cursor +`get_current_log_seq()` returns k -- so lines 1..k are forever +unreachable to that client. + +Fix: `clear_logs()` snapshots the pre-run seq into `_run_start_seq` +(exposed via `get_run_start_seq()`), and `routes/export.py` defaults +the SSE cursor to that snapshot instead of the current seq. Every line +appended during the current run has seq strictly greater than the +snapshot, so the client sees the full run regardless of when it +connects. + +These tests exercise the orchestrator-side contract only (no +subprocess, no FastAPI, no frontend). The routes-level integration +with get_run_start_seq() is a one-line edit covered by manual testing +and the frontend build. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +import pytest + + +# Backend root on sys.path so `from core.export.orchestrator import ...` +# and friends resolve without the studio app bootstrap. +_BACKEND_DIR = Path(__file__).resolve().parent.parent +if str(_BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(_BACKEND_DIR)) + +# ExportOrchestrator imports structlog and a few heavy modules at the +# top of orchestrator.py. Stub the ones we don't need in these unit +# tests so the import succeeds on machines without the full studio +# venv. +_loggers_stub = types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +# structlog is only used for a module-level import; a bare stub is +# enough because we never call into it in these tests. +sys.modules.setdefault("structlog", types.ModuleType("structlog")) + +# utils.paths.outputs_root is only called inside scan_checkpoints which +# we don't hit in these tests. Provide a stub module so the top-level +# import in orchestrator.py resolves. +_utils_pkg = types.ModuleType("utils") +_utils_pkg.__path__ = [] # mark as package +_utils_paths_stub = types.ModuleType("utils.paths") +_utils_paths_stub.outputs_root = lambda: Path("/tmp") +sys.modules.setdefault("utils", _utils_pkg) +sys.modules.setdefault("utils.paths", _utils_paths_stub) + + +@pytest.fixture +def orchestrator(): + """Fresh ExportOrchestrator with only the log-buffer state exercised.""" + from core.export.orchestrator import ExportOrchestrator + + return ExportOrchestrator() + + +def _append(orch, line: str, stream: str = "stdout") -> None: + """Shortcut for simulating a worker log message.""" + orch._append_log({"type": "log", "stream": stream, "line": line, "ts": 0.0}) + + +# --------------------------------------------------------------------------- +# clear_logs() semantics +# --------------------------------------------------------------------------- + + +def test_run_start_seq_is_zero_before_any_logs(orchestrator) -> None: + """A brand-new orchestrator must report run_start_seq == 0 so a + first SSE connection picks up every line from seq 1 onward.""" + assert orchestrator.get_run_start_seq() == 0 + + +def test_clear_logs_snapshots_current_seq(orchestrator) -> None: + """clear_logs() must capture _log_seq BEFORE clearing the buffer, + so subsequent runs can anchor their SSE cursor at the snapshot.""" + _append(orchestrator, "old run line 1") + _append(orchestrator, "old run line 2") + _append(orchestrator, "old run line 3") + assert orchestrator.get_current_log_seq() == 3 + + orchestrator.clear_logs() + + assert orchestrator.get_run_start_seq() == 3 + assert orchestrator.get_current_log_seq() == 3 # seq counter preserved + + +# --------------------------------------------------------------------------- +# Race regression: SSE connects AFTER lines have been emitted +# --------------------------------------------------------------------------- + + +def test_sse_default_cursor_catches_all_current_run_lines(orchestrator) -> None: + """Simulate the POST-then-SSE race: worker starts emitting lines + immediately after clear_logs(), SSE connects several lines later. + Using get_run_start_seq() as the default cursor MUST return every + line emitted since clear_logs() ran. + + Pre-fix, the SSE defaulted to get_current_log_seq() at connect + time, which would return the last-seen seq and miss lines N+1..M. + """ + # Previous run leaves some buffered lines. + _append(orchestrator, "previous run line A") + _append(orchestrator, "previous run line B") + + # New run starts: orchestrator clears the buffer and snapshots seq. + orchestrator.clear_logs() + run_start = orchestrator.get_run_start_seq() + + # Worker emits early lines BEFORE the SSE connects. + _append(orchestrator, "Importing Unsloth...") + _append(orchestrator, "Loading checkpoint: /foo/bar") + _append(orchestrator, "Starting export...") + + # SSE connects now and asks "give me everything after the run + # start cursor". + entries, new_cursor = orchestrator.get_logs_since(run_start) + + # All three early lines must be present. Pre-fix this was []. + lines = [e["line"] for e in entries] + assert lines == [ + "Importing Unsloth...", + "Loading checkpoint: /foo/bar", + "Starting export...", + ] + assert new_cursor == entries[-1]["seq"] + + +def test_sse_default_cursor_excludes_previous_run(orchestrator) -> None: + """After clear_logs(), lines from the PREVIOUS run must not leak + into the new run's SSE stream. Pre-fix this worked correctly + (clear_logs cleared the deque); the fix must preserve it. + """ + _append(orchestrator, "previous run line 1") + _append(orchestrator, "previous run line 2") + _append(orchestrator, "previous run line 3") + assert orchestrator.get_current_log_seq() == 3 + + orchestrator.clear_logs() + run_start = orchestrator.get_run_start_seq() + + _append(orchestrator, "new run line") + + entries, _ = orchestrator.get_logs_since(run_start) + assert [e["line"] for e in entries] == ["new run line"] + + +def test_clear_logs_twice_advances_run_start(orchestrator) -> None: + """Back-to-back clear_logs() calls (e.g. cleanup -> load -> + export in the same dialog session) must each re-anchor run_start + at the current seq, so successive runs each start with a fresh + low-water mark.""" + _append(orchestrator, "run 1 line a") + _append(orchestrator, "run 1 line b") + + orchestrator.clear_logs() + assert orchestrator.get_run_start_seq() == 2 + + _append(orchestrator, "run 2 line a") + _append(orchestrator, "run 2 line b") + _append(orchestrator, "run 2 line c") + + orchestrator.clear_logs() + assert orchestrator.get_run_start_seq() == 5 diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index aff56c3e6a..450691ddfd 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -42,7 +42,13 @@ export interface CheckpointListResponse { export interface ExportOperationResponse { success: boolean; message: string; - details?: Record | null; + /** + * Optional extras returned by the backend. The export endpoints set + * `details.output_path` to the resolved on-disk directory of the + * saved model when a local save was requested. Hub-only pushes leave + * `details` undefined. + */ + details?: { output_path?: string | null } & Record; } export async function fetchCheckpoints(): Promise { @@ -131,3 +137,172 @@ export async function cleanupExport(): Promise { const response = await authFetch("/api/export/cleanup", { method: "POST" }); return parseJson(response); } + +// ───────────────────────────────────────────────────────────────────── +// Live export log stream (Server-Sent Events) +// ───────────────────────────────────────────────────────────────────── + +export type ExportLogStream = "stdout" | "stderr" | "status"; + +export interface ExportLogEntry { + stream: ExportLogStream; + line: string; + ts: number | null; +} + +export type ExportLogEventName = "log" | "heartbeat" | "complete" | "error"; + +export interface ExportLogEvent { + event: ExportLogEventName; + id: number | null; + /** Present on `log` events. */ + entry?: ExportLogEntry; + /** Present on `error` events. */ + error?: string; +} + +interface ParsedSseMessage { + event: string; + id: number | null; + data: string; +} + +function parseSseMessage(raw: string): ParsedSseMessage | null { + const lines = raw.split(/\r?\n/); + let event = "message"; + let id: number | null = null; + const dataLines: string[] = []; + + for (const line of lines) { + if (!line) continue; + if (line.startsWith("event:")) { + event = line.slice(6).trim(); + continue; + } + if (line.startsWith("id:")) { + const value = Number(line.slice(3).trim()); + id = Number.isFinite(value) ? value : null; + continue; + } + if (line.startsWith("data:")) { + dataLines.push(line.slice(5).trimStart()); + continue; + } + // Comment lines (":heartbeat" etc.) are ignored per SSE spec. + } + + if (dataLines.length === 0) return null; + return { event, id, data: dataLines.join("\n") }; +} + +function isAbortError(error: unknown): boolean { + return error instanceof DOMException && error.name === "AbortError"; +} + +export async function streamExportLogs(options: { + signal: AbortSignal; + since?: number | null; + onOpen?: () => void; + onEvent: (event: ExportLogEvent) => void; +}): Promise { + const headers = new Headers(); + if (typeof options.since === "number") { + headers.set("Last-Event-ID", String(options.since)); + } + + const url = + typeof options.since === "number" + ? `/api/export/logs/stream?since=${options.since}` + : "/api/export/logs/stream"; + + const response = await authFetch(url, { + method: "GET", + headers, + signal: options.signal, + }); + + if (!response.ok) { + throw new Error(await readError(response)); + } + if (!response.body) { + throw new Error("Export log stream unavailable"); + } + + options.onOpen?.(); + + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + + try { + while (true) { + const { value, done } = await reader.read(); + if (done) return; + + buffer += decoder.decode(value, { stream: true }); + + let separatorIndex = buffer.search(/\r?\n\r?\n/); + while (separatorIndex >= 0) { + const rawEvent = buffer.slice(0, separatorIndex); + const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; + buffer = buffer.slice(separatorIndex + separatorLength); + + if (rawEvent.startsWith("retry:") || rawEvent.startsWith(":")) { + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + + const parsed = parseSseMessage(rawEvent); + if (!parsed) { + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + + try { + if (parsed.event === "log") { + const payload = JSON.parse(parsed.data) as { + stream?: ExportLogStream; + line?: string; + ts?: number | null; + }; + options.onEvent({ + event: "log", + id: parsed.id, + entry: { + stream: payload.stream ?? "stdout", + line: payload.line ?? "", + ts: payload.ts ?? null, + }, + }); + } else if (parsed.event === "heartbeat") { + options.onEvent({ event: "heartbeat", id: parsed.id }); + } else if (parsed.event === "complete") { + options.onEvent({ event: "complete", id: parsed.id }); + return; + } else if (parsed.event === "error") { + let errorMessage = "Export log stream error"; + try { + const payload = JSON.parse(parsed.data) as { error?: string }; + if (payload.error) errorMessage = payload.error; + } catch { + // fall through with default message + } + options.onEvent({ + event: "error", + id: parsed.id, + error: errorMessage, + }); + } + } catch (err) { + if (isAbortError(err)) return; + // Ignore malformed events, keep reading. + } + + separatorIndex = buffer.search(/\r?\n\r?\n/); + } + } + } catch (err) { + if (isAbortError(err)) return; + throw err; + } +} diff --git a/studio/frontend/src/features/export/components/export-dialog.tsx b/studio/frontend/src/features/export/components/export-dialog.tsx index bf7b680945..401be3ff4b 100644 --- a/studio/frontend/src/features/export/components/export-dialog.tsx +++ b/studio/frontend/src/features/export/components/export-dialog.tsx @@ -21,9 +21,199 @@ import { Switch } from "@/components/ui/switch"; import { AlertCircleIcon, ArrowRight01Icon, CheckmarkCircle02Icon, Key01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { AnimatePresence, motion } from "motion/react"; +import { useEffect, useRef, useState } from "react"; +import { streamExportLogs, type ExportLogEntry } from "../api/export-api"; import { collapseAnim } from "../anim"; import { EXPORT_METHODS, type ExportMethod } from "../constants"; +// Max log lines kept in the dialog's local state. Matches the backend +// ring buffer's maxlen so the UI shows the full scrollback captured +// server side. +const MAX_LOG_LINES = 4000; + +interface UseExportLogsResult { + lines: ExportLogEntry[]; + connected: boolean; + error: string | null; +} + +/** + * Subscribe to the live export log SSE stream while `exporting` is + * true, and accumulate lines in local state. Lines from a previous + * action are cleared: + * + * - when a new export starts (`exporting` flips to true), and + * - when the user switches export method, dialog opens fresh, or + * the dialog closes — so re-opening into a different action's + * screen doesn't show the prior screen's saved output. + */ +function useExportLogs( + exporting: boolean, + exportMethod: ExportMethod | null, + open: boolean, +): UseExportLogsResult { + const [lines, setLines] = useState([]); + const [connected, setConnected] = useState(false); + const [error, setError] = useState(null); + + // Reset log state whenever the user moves to a different screen -- + // either by switching export method or by reopening the dialog -- so + // each (open × method) tuple shows only its own run history. The + // streaming effect below additionally clears on new export start. + useEffect(() => { + setLines([]); + setError(null); + setConnected(false); + }, [exportMethod, open]); + + useEffect(() => { + if (!exporting) return; + + setLines([]); + setError(null); + + const abortCtrl = new AbortController(); + let cancelled = false; + // Track the highest seq we've observed on a `log` event so we can + // resume the stream via `since=` / `Last-Event-ID` after a drop. + // The backend's SSE `id:` field carries this as ExportLogEvent.id. + let lastSeq: number | null = null; + // Exponential backoff with jitter, capped. Reset on every + // successful connection so flaky networks don't accumulate delay. + let backoffMs = 500; + const MAX_BACKOFF_MS = 5000; + // Flipped by a terminal event (explicit `complete` from the + // backend or a non-transient error we choose not to retry). Stops + // the outer reconnect loop even if `exporting` is still true. + let terminated = false; + + const run = async () => { + while (!cancelled && !terminated) { + try { + await streamExportLogs({ + signal: abortCtrl.signal, + since: lastSeq, + onOpen: () => { + if (cancelled) return; + setConnected(true); + // Reset backoff on every successful connect so later + // drops don't inherit accumulated delay from earlier ones. + backoffMs = 500; + }, + onEvent: (event) => { + if (cancelled) return; + if (event.event === "log" && event.entry) { + if (typeof event.id === "number") { + lastSeq = event.id; + } + const entry = event.entry; + setLines((prev) => { + const next = prev.length >= MAX_LOG_LINES + ? prev.slice(prev.length - MAX_LOG_LINES + 1) + : prev.slice(); + next.push(entry); + return next; + }); + } else if (event.event === "complete") { + // Backend signalled the run is fully drained -- stop + // trying to reconnect even though `exporting` may not + // have flipped false yet on this tick. + terminated = true; + } else if (event.event === "error" && event.error) { + setError(event.error); + } + }, + }); + } catch (err: unknown) { + if (cancelled) return; + if (err instanceof DOMException && err.name === "AbortError") return; + setError(err instanceof Error ? err.message : String(err)); + // Fall through to the backoff path below; a fetch-level + // failure is retryable the same way a clean EOF is. + } + + setConnected(false); + if (cancelled || terminated) return; + + // Exponential backoff with jitter before reconnecting. The + // backend's ring buffer plus Last-Event-ID resume means we + // don't lose lines across the retry as long as the reconnect + // happens within the buffer's lifetime (~4000 lines). + const delay = backoffMs + Math.floor(Math.random() * 250); + backoffMs = Math.min(backoffMs * 2, MAX_BACKOFF_MS); + try { + await new Promise((resolve, reject) => { + if (abortCtrl.signal.aborted) { + reject(new DOMException("Aborted", "AbortError")); + return; + } + const timeoutId = window.setTimeout(resolve, delay); + abortCtrl.signal.addEventListener( + "abort", + () => { + window.clearTimeout(timeoutId); + reject(new DOMException("Aborted", "AbortError")); + }, + { once: true }, + ); + }); + } catch { + return; + } + } + }; + + // run()'s own try/catch handles every failure path we care about; + // swallow anything that somehow escapes so React's dev overlay + // doesn't flag an unhandled rejection on dialog close. + void run().catch(() => {}); + + return () => { + cancelled = true; + abortCtrl.abort(); + setConnected(false); + }; + }, [exporting]); + + return { lines, connected, error }; +} + +/** + * Tick every second while `exporting` is true and report elapsed + * seconds. Powers the "Working… 27s" badge in the log header so the + * panel doesn't look frozen during long single-step phases (cache + * file copy, GGUF conversion) when no new lines are arriving. + */ +function useElapsedSeconds(exporting: boolean): number { + const [elapsed, setElapsed] = useState(0); + useEffect(() => { + if (!exporting) { + setElapsed(0); + return; + } + const startedAt = Date.now(); + setElapsed(0); + const id = window.setInterval(() => { + setElapsed(Math.floor((Date.now() - startedAt) / 1000)); + }, 1000); + return () => window.clearInterval(id); + }, [exporting]); + return elapsed; +} + +function formatElapsed(seconds: number): string { + if (seconds < 60) return `${seconds}s`; + const m = Math.floor(seconds / 60); + const s = seconds % 60; + return `${m}m ${s.toString().padStart(2, "0")}s`; +} + +function formatLogLine(entry: ExportLogEntry): string { + // Strip trailing carriage returns that tqdm-style progress leaves + // in the stream so the scrollback doesn't render funky boxes. + return entry.line.replace(/\r+$/g, ""); +} + type Destination = "local" | "hub"; interface ExportDialogProps { @@ -49,6 +239,12 @@ interface ExportDialogProps { exporting: boolean; exportError: string | null; exportSuccess: boolean; + /** + * Resolved on-disk realpath of the most recent successful export. + * Surfaced on the Export Complete screen so users can find their + * model. Null when the export only pushed to the Hub. + */ + exportOutputPath: string | null; } export function ExportDialog({ @@ -74,7 +270,38 @@ export function ExportDialog({ exporting, exportError, exportSuccess, + exportOutputPath, }: ExportDialogProps) { + // Live log capture is only meaningful for export methods that run + // a slow subprocess operation with interesting stdout: merged and + // gguf. LoRA adapter export is a fast disk write and would just + // show a blank panel, so we hide it there. + const showLogPanel = + exportMethod === "merged" || exportMethod === "gguf"; + + const { lines: logLines, connected: logConnected, error: logError } = + useExportLogs(exporting && showLogPanel, exportMethod, open); + const elapsedSeconds = useElapsedSeconds(exporting && showLogPanel); + + const logScrollRef = useRef(null); + // Auto-scroll to bottom whenever a new line arrives, unless the + // user has scrolled up to read earlier output. + const [followTail, setFollowTail] = useState(true); + + useEffect(() => { + if (!followTail) return; + const el = logScrollRef.current; + if (el) el.scrollTop = el.scrollHeight; + }, [logLines, followTail]); + + const handleLogScroll = () => { + const el = logScrollRef.current; + if (!el) return; + const nearBottom = + el.scrollHeight - el.scrollTop - el.clientHeight < 24; + setFollowTail(nearBottom); + }; + return ( - { if (exporting) e.preventDefault(); }}> + { if (exporting) e.preventDefault(); }} + > {exportSuccess ? ( <>
-
+

Export Complete

-

+

{destination === "hub" ? "Model successfully pushed to Hugging Face Hub." : "Model saved locally."}

+ {exportOutputPath ? ( +
+ + Saved to + + + {exportOutputPath} + +
+ ) : null}
@@ -256,6 +499,78 @@ export function ExportDialog({
*/} + {/* Live export output panel */} + + {showLogPanel && (exporting || logLines.length > 0) && ( + +
+
+ +
+ + + {logConnected + ? "streaming" + : exporting + ? "connecting..." + : "idle"} + + {exporting && elapsedSeconds > 0 ? ( + + · {formatElapsed(elapsedSeconds)} + + ) : null} +
+
+
+ {logLines.length === 0 ? ( +
+ + + Waiting for worker output... + +
+ ) : ( + 
+                          {logLines.map((entry, idx) => (
+                            

+                              {formatLogLine(entry)}
+                            

+                          ))}
+
+ )} +
+ {logError && ( +

+ Log stream: {logError} +

+ )} +
+ + )} + +