ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/routes/export.py
Anmol Mishra 554c289538
 fix: respect absolute export paths to prevent cross-drive copy failures (WinError 112) (#6088) 
* fix: allow absolute save_directory in export paths to prevent cross-drive copy failures

The GGUF export pipeline (and all other export flows) forced every
save_directory through resolve_export_dir(), which always resolved
the path under exports_root() — typically ~/.unsloth/studio/exports/
on the system drive (C: on Windows).

When a user selected an output directory on a different drive (E:):
1. The absolute path was rejected at the Pydantic validator level.
2. Even if it got through, resolve_export_dir would re-resolve it
   under C:\Users\.unsloth\studio\exports\.
3. After GGUF conversion completed on E:, the relocation step would
   try to move/copy the finished files to C:, causing:
   - WinError 17 (cross-drive move failure when shutil.move falls
     through to a cross-filesystem copy)
   - WinError 112 (disk full on C:)

Fix both layers:
- _validate_save_directory: accept absolute paths (they represent an
  explicit user choice of output location).
- resolve_export_dir, resolve_output_dir, resolve_tensorboard_dir:
  return absolute paths as-is instead of forcing them under the
  default root. Keep the existing safety checks (null bytes, '..'
  segments) and fall through to resolve_under_root for relative paths.

Fixes: https://github.com/unslothai/unsloth/issues/6082

* refactor: centralize user path validation into _resolve_user_path helper

Addresses code review feedback: the null-byte, '..', and absolute-path
checks were duplicated across resolve_output_dir, resolve_export_dir,
and resolve_tensorboard_dir. Extract a single _resolve_user_path helper
that all three delegate to.

No behavioral change — pure consolidation.

* fix: address code review — contain destructive cleanup and scope absolute paths

Address all review feedback from gemini-code-assist:

1. P1: destructive subdirectory cleanup (export_gguf)
   The flattening loop in export_gguf previously rmtree'd every
   subdirectory under abs_save_dir. When targeting an existing user
   directory on a different drive (#6082), this could nuke unrelated
   subdirectories. Now snapshot existing subdirectories before the
   export and only clean up dirs created during this run.

2. P2: keep scan/read endpoints contained
   Only resolve_export_dir accepts absolute paths (export is a write
   path where user picks location). Reverted resolve_output_dir and
   resolve_tensorboard_dir to use resolve_under_root directly — these
   are used by scan/read/training endpoints that must stay contained
   under their respective roots.

3. Centralization feedback
   Removed the _resolve_user_path helper since it's no longer needed
   with the narrowed scope. resolve_export_dir has the absolute path
   logic inline with a clear docstring.

* fix: skip pre-existing subdirs in GGUF flatten loop and clean stale export intermediates

Two issues caught in code review (chatgpt-codex-connector):

1. The flattening loop moved ALL .gguf files from ALL subdirectories
   into abs_save_dir, including pre-existing unrelated user subdirs.
   Now skip pre-existing subdirs entirely unless they are known
   export-owned intermediates (model/, model_gguf/).

2. After a failed export, known export-owned subdirectories (model/,
   model_gguf/) were snapshotted as pre-existing on retry and never
   cleaned up. These are now always cleaned up regardless, since they
   are known intermediates created by the export pipeline.

* fix: separate write vs read export paths, guard same-dir rmtree

Three issues caught in code review (chatgpt-codex-connector):

1. P1: scan endpoint containment
   resolve_export_dir was changed to accept absolute paths, but it's
   also used by scan/read endpoints (routes/models.py) that must stay
   contained under exports_root(). Split into:
   - resolve_export_dir: contained, used by scans
   - resolve_export_write_dir: accepts absolute paths, used by export
     backend only

2. P1: same-directory rmtree
   When a non-PEFT checkpoint's gguf_dir resolves to the same path as
   abs_save_dir (user selected the checkpoint's gguf output as their
   export directory), shutil.rmtree(gguf_dir) would delete the user's
   chosen output directory. Now skip relocation when both paths resolve
   to the same location.

3. P1: pre-existing subdir flatten loop
   Reverted _EXPORT_OWNED_SUBDIRS logic — 'model/' and 'model_gguf/'
   are common directory names in shared model folders and don't prove
   export ownership. Now only clean up subdirs that didn't exist before
   the export started.

* fix: remove dead _EXPORT_OWNED_SUBDIRS and fix _export_details for absolute paths

Two fixes from review comments:

1. Remove unused _EXPORT_OWNED_SUBDIRS declaration (leftover from
   previous iteration that was intentionally removed).

2. _export_details now returns the full absolute path when the export
   target is outside exports_root(), instead of truncating to basename.
   Users who export to E:\ can now see the full destination path in
   the success dialog.

* fix: use unique tmp dir for GGUF intermediates to avoid overwriting user dirs

When exporting to an absolute destination that already contains a
model/ subdirectory (e.g. a shared models folder), the hard-coded
model_save_path would overwrite files in that unrelated directory.

Use _tmp_model_<uuid> as the intermediate path instead, so user
directories are never touched. The tmp dir is created as a new subdir
of abs_save_dir and cleaned up by the flatten loop after GGUF files
are relocated.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix GGUF local export paths for PR #6088

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address GGUF export follow-ups for PR #6088

* Clean GGUF temp dirs on export failure for PR #6088

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix/adjust export path tests for PR #6088

* Fix/adjust export path review findings for PR #6088

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix/adjust home export path handling for PR #6088

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
2026-06-12 12:52:57 +02:00

483 lines
17 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Export API routes: checkpoint discovery and model export operations."""
import asyncio
import json
import os
import sys
import time
from pathlib import Path
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
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
from auth.authentication import get_current_subject
from utils.utils import safe_error_detail
try:
from core.export import get_export_backend
except ImportError:
parent_backend = backend_path.parent / "backend"
if str(parent_backend) not in sys.path:
sys.path.insert(0, str(parent_backend))
from core.export import get_export_backend
from models import (
LoadCheckpointRequest,
ExportStatusResponse,
ExportOperationResponse,
ExportMergedModelRequest,
ExportBaseModelRequest,
ExportGGUFRequest,
ExportLoRAAdapterRequest,
)
router = APIRouter()
logger = get_logger(__name__)
@router.post("/load-checkpoint", response_model = ExportOperationResponse)
async def load_checkpoint(
request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject)
):
"""Load a checkpoint into the export backend (ExportBackend.load_checkpoint)."""
try:
# Free GPU memory: shut down running inference/training subprocesses
# before loading the export checkpoint (they'd compete for VRAM).
try:
from core.inference import get_inference_backend
inf = get_inference_backend()
if inf.active_model_name:
logger.info(
"Unloading inference model '%s' to free GPU memory for export",
inf.active_model_name,
)
inf._shutdown_subprocess()
inf.active_model_name = None
inf.models.clear()
except Exception as e:
logger.warning("Could not unload inference model: %s", e)
try:
from core.training import get_training_backend
trn = get_training_backend()
if trn.is_training_active():
logger.info("Stopping active training to free GPU memory for export")
trn.stop_training()
# Wait for the training subprocess to exit, else it may still hold GPU memory.
for _ in range(60): # up to 30s
if not trn.is_training_active():
break
await asyncio.sleep(0.5)
else:
logger.warning("Training subprocess did not exit within 30s, proceeding anyway")
except Exception as e:
logger.warning("Could not stop training: %s", e)
backend = get_export_backend()
# Run in a worker thread (spawns and waits on a subprocess, can take
# minutes) so the event loop stays free to serve the live log SSE stream.
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,
trust_remote_code = request.trust_remote_code,
)
if not success:
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success = True, message = message)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error loading checkpoint: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to load checkpoint",
)
@router.post("/cleanup", response_model = ExportOperationResponse)
async def cleanup_export_memory(current_subject: str = Depends(get_current_subject)):
"""Cleanup export-related models from memory (ExportBackend.cleanup_memory)."""
try:
backend = get_export_backend()
success = await asyncio.to_thread(backend.cleanup_memory)
if not success:
raise HTTPException(
status_code = 500,
detail = "Memory cleanup failed. See server logs for details.",
)
return ExportOperationResponse(
success = True,
message = "Memory cleanup completed successfully",
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error during export memory cleanup: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to cleanup export memory",
)
@router.get("/status", response_model = ExportStatusResponse)
async def get_export_status(current_subject: str = Depends(get_current_subject)):
"""Get export backend status (loaded checkpoint, model type, PEFT flag)."""
try:
backend = get_export_backend()
return ExportStatusResponse(
current_checkpoint = backend.current_checkpoint,
is_vision = bool(getattr(backend, "is_vision", False)),
is_peft = bool(getattr(backend, "is_peft", False)),
)
except Exception as e:
logger.error(f"Error getting export status: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to get export status",
)
def _try_register_external_export(path: Path) -> tuple[bool, Optional[str]]:
"""Best-effort registration so absolute exports show up in local scans."""
try:
from storage.studio_db import add_scan_folder
folder = add_scan_folder(str(path))
return True, str(folder.get("path") or path)
except Exception as exc:
logger.warning("Could not register export scan folder %s: %s", path, exc)
return False, None
def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]:
"""Return relative export paths, keeping external absolute paths visible."""
if not output_path:
return None
try:
from utils.paths.storage_roots import exports_root
path = Path(output_path)
# If it's outside exports_root, return the full absolute path
# so users can find their files on a different drive.
if path.is_absolute():
try:
path.resolve().relative_to(exports_root().resolve())
except ValueError:
registered, registered_path = _try_register_external_export(path)
return {
"output_path": str(path),
"scan_folder_registered": registered,
"scan_folder_path": registered_path,
}
rel = os.path.relpath(output_path, exports_root())
return {"output_path": rel}
except Exception:
return {"output_path": output_path}
@router.post("/export/merged", response_model = ExportOperationResponse)
async def export_merged_model(
request: ExportMergedModelRequest, current_subject: str = Depends(get_current_subject)
):
"""Export a merged PEFT model (16-bit or 4-bit), optionally pushing to Hub.
Wraps ExportBackend.export_merged_model.
"""
try:
backend = get_export_backend()
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,
repo_id = request.repo_id,
hf_token = request.hf_token,
private = request.private,
)
if not success:
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(
success = True,
message = message,
details = _export_details(output_path),
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error exporting merged model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to export merged model",
)
@router.post("/export/base", response_model = ExportOperationResponse)
async def export_base_model(
request: ExportBaseModelRequest, current_subject: str = Depends(get_current_subject)
):
"""Export a non-PEFT base model, optionally pushing to Hub.
Wraps ExportBackend.export_base_model.
"""
try:
backend = get_export_backend()
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,
hf_token = request.hf_token,
private = request.private,
base_model_id = request.base_model_id,
)
if not success:
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(
success = True,
message = message,
details = _export_details(output_path),
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error exporting base model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to export base model",
)
@router.post("/export/gguf", response_model = ExportOperationResponse)
async def export_gguf(
request: ExportGGUFRequest, current_subject: str = Depends(get_current_subject)
):
"""Export the current model to GGUF format, optionally pushing to Hub.
Wraps ExportBackend.export_gguf.
"""
try:
backend = get_export_backend()
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,
repo_id = request.repo_id,
hf_token = request.hf_token,
)
if not success:
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(
success = True,
message = message,
details = _export_details(output_path),
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error exporting GGUF model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to export GGUF model",
)
@router.post("/export/lora", response_model = ExportOperationResponse)
async def export_lora_adapter(
request: ExportLoRAAdapterRequest, current_subject: str = Depends(get_current_subject)
):
"""Export only the LoRA adapter (if the loaded model is PEFT).
Wraps ExportBackend.export_lora_adapter.
"""
try:
backend = get_export_backend()
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,
hf_token = request.hf_token,
private = request.private,
)
if not success:
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(
success = True,
message = message,
details = _export_details(output_path),
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error exporting LoRA adapter: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to export LoRA adapter",
)
# Live export log stream (Server-Sent Events).
#
# The export worker's stdout/stderr is piped to the orchestrator as log
# entries (core/export/worker.py, orchestrator.py); this endpoint streams
# them to the browser for a live terminal panel during export operations.
#
# Shape follows routes/training.py::stream_training_progress: each event
# carries id/event/data, 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 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` : once the worker is idle and no new lines arrived for
~1 second. Clients should close.
- `error` : unrecoverable server-side error
Each event's `id:` field is the log entry's monotonic seq number so the
browser can resume via `Last-Event-ID` on reconnect.
"""
backend = get_export_backend()
# Starting cursor: explicit `since` wins, then Last-Event-ID on reconnect,
# else the run-start snapshot so the client sees every line since the run
# began even if the SSE connection opened after the export-kickoff POST.
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
# 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():
# Let the reader thread drain trailing lines printed just
# before the worker signalled 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: end 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": safe_error_detail(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",
},
)
Reference in a new issue View git blame Copy permalink