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,