From 368b19b23775f9822e2fcff2c256d65b81f9fcfc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 13 Jun 2026 04:06:17 -0700 Subject: [PATCH] Studio: fix training output dir escaping outputs root for models on another drive (#6293) * Studio: derive training output dir from model basename for local-drive models A LoRA/QLoRA run started from a model loaded by absolute path (common when models live on a non-system drive, e.g. G:\modelsAI\...\gemma-4-12B-it) seeded the default output dir with that full path. resolve_output_dir then raised "path escapes root ... is not under the studio outputs folder", so training could not start from a model stored off the system drive. Add default_run_dir_name(): Hugging Face repo ids keep their namespace (org/model becomes org_model), while local paths collapse to their final component so an absolute source path can no longer leak into the output dir. Use it at the three worker derivation sites and add a regression test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: cap run dir name length and drop redundant resolve Apply PR review feedback: length-cap the auto-generated output dir component so an unusually long model name stays under the filesystem name limit, and drop the redundant double resolve_output_dir at the embedding site so all three derivation sites match. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/training/worker.py | 15 +++-- .../tests/test_default_output_dir_name.py | 65 +++++++++++++++++++ studio/backend/utils/paths/__init__.py | 2 + studio/backend/utils/paths/storage_roots.py | 18 +++++ 4 files changed, 93 insertions(+), 7 deletions(-) create mode 100644 studio/backend/tests/test_default_output_dir_name.py diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 1120744a2d..4a6d182da8 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1687,12 +1687,12 @@ def _run_mlx_training(event_queue, stop_queue, config): warmup_steps = 5 # ── 5. Build output dir ── + # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it + from utils.paths import resolve_output_dir, ensure_dir, default_run_dir_name + output_dir = config.get("output_dir", "") if not output_dir: - output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}" - # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it - from utils.paths import resolve_output_dir, ensure_dir - + output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) @@ -2450,6 +2450,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> resolve_output_dir, resolve_tensorboard_dir, datasets_root, + default_run_dir_name, ) import transformers @@ -2773,7 +2774,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> resume_from_checkpoint ) if not output_dir: - output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}" + output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) @@ -2924,7 +2925,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> from datasets import Dataset from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset from transformers import TrainerCallback - from utils.paths import datasets_root, resolve_output_dir + from utils.paths import datasets_root, resolve_output_dir, default_run_dir_name except ImportError as e: event_queue.put( { @@ -3182,7 +3183,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> resume_from_checkpoint ) if not output_dir: - output_dir = str(resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}")) + output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" output_dir = str(resolve_output_dir(output_dir)) num_epochs = config.get("num_epochs", 2) diff --git a/studio/backend/tests/test_default_output_dir_name.py b/studio/backend/tests/test_default_output_dir_name.py new file mode 100644 index 0000000000..d8a7f5ae21 --- /dev/null +++ b/studio/backend/tests/test_default_output_dir_name.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Auto-generated training output dir names stay inside outputs_root. + +Regression for local-model training: a model loaded by absolute path (e.g. +``G:\\modelsAI\\...\\gemma-4-12B-it`` on a non-system drive) used to seed the +default run dir with that full path, so ``resolve_output_dir`` raised +``path escapes root`` because the result was not under ``/outputs``. +""" + +import importlib.util +from pathlib import Path + +import pytest + + +_BACKEND_DIR = Path(__file__).resolve().parent.parent + + +def _load_storage_roots(): + path = _BACKEND_DIR / "utils/paths/storage_roots.py" + spec = importlib.util.spec_from_file_location("storage_roots_under_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_repo_id_keeps_namespace(): + sr = _load_storage_roots() + assert sr.default_run_dir_name("unsloth/gemma-3-4b") == "unsloth_gemma-3-4b" + assert sr.default_run_dir_name("gemma-3-4b") == "gemma-3-4b" + + +def test_local_paths_collapse_to_basename(): + sr = _load_storage_roots() + assert sr.default_run_dir_name(r"G:\modelsAI\gguf\test\gemma-4-12B-it") == "gemma-4-12B-it" + assert sr.default_run_dir_name("/data/models/gemma-3-4b") == "gemma-3-4b" + assert sr.default_run_dir_name("~/models/gemma-3-4b") == "gemma-3-4b" + assert sr.default_run_dir_name("C:/Users/me/models/gemma-3-4b") == "gemma-3-4b" + + +def test_empty_falls_back_to_model(): + sr = _load_storage_roots() + assert sr.default_run_dir_name("") == "model" + assert sr.default_run_dir_name(" ") == "model" + + +def test_very_long_name_is_capped(): + sr = _load_storage_roots() + name = sr.default_run_dir_name("a" * 500) + assert 0 < len(name) <= 200 + + +def test_derived_name_resolves_under_outputs_root(tmp_path, monkeypatch): + sr = _load_storage_roots() + outputs = tmp_path / "outputs" + outputs.mkdir() + monkeypatch.setattr(sr, "outputs_root", lambda: outputs) + + name = sr.default_run_dir_name(r"G:\modelsAI\gguf\test\gemma-4-12B-it") + resolved = sr.resolve_output_dir(f"{name}_1781327234") + assert resolved == outputs / "gemma-4-12B-it_1781327234" + # No escape: the absolute G: source no longer leaks into the output path. + assert "modelsAI" not in str(resolved) diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index eff9b64678..e812fa9e61 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -41,6 +41,7 @@ from .storage_roots import ( ensure_dir, ensure_studio_directories, resolve_under_root, + default_run_dir_name, resolve_output_dir, resolve_export_dir, resolve_export_write_dir, @@ -88,6 +89,7 @@ __all__ = [ "ensure_dir", "ensure_studio_directories", "resolve_under_root", + "default_run_dir_name", "resolve_output_dir", "resolve_export_dir", "resolve_export_write_dir", diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index d336bc2e71..c718f38ffb 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -5,6 +5,7 @@ from __future__ import annotations import json import os +import re import sys from pathlib import Path, PurePosixPath, PureWindowsPath import tempfile @@ -384,6 +385,23 @@ def resolve_under_root( return candidate +def default_run_dir_name(model_name: str) -> str: + # Folder-safe run name for an auto-created output dir. Repo ids keep their + # namespace (org/model -> org_model); local paths (incl. G:\dir\model) + # collapse to their final component so an absolute source can't escape + # outputs_root. Length-capped to stay under the filesystem name limit. + raw = str(model_name or "").strip() + is_path = ( + "\\" in raw + or raw.startswith(("/", "~", ".")) + or os.path.isabs(raw) + or (len(raw) >= 2 and raw[1] == ":") + ) + base = PureWindowsPath(raw).name if is_path else raw.replace("/", "_") + base = re.sub(r"[^A-Za-z0-9._-]+", "_", base)[:200].strip("._-") + return base or "model" + + def resolve_output_dir(path_value: str | None = None) -> Path: return resolve_under_root( path_value,