From 4cedeba8c241dbfa42362d398746cd4e1ae98e2a Mon Sep 17 00:00:00 2001 From: NuoFang <137974286+NuoFang6@users.noreply.github.com> Date: Sun, 22 Mar 2026 21:11:24 +0800 Subject: [PATCH] fix(studio): prevent ModuleNotFoundError in dataset.map() on Windows (#4473) * fix(studio): prevent ModuleNotFoundError in dataset.map() on Windows On Windows, dataset.map() uses "spawn", which requires workers to import compiled modules from disk. Previously, clear_unsloth_compiled_cache() deleted the entire directory, causing workers to crash when looking for UnslothSFTTrainer.py. Changes: 1. Added `preserve_patterns` to cache cleanup to keep `Unsloth*Trainer.py` on Windows while clearing model-specific files. 2. Added the cache directory to PYTHONPATH for spawn workers. Linux/macOS behavior is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix spawn-platform coverage, CWD path mismatch, and race condition for PR #4473 - Extend platform guard from win32-only to include macOS (also uses spawn since Python 3.8, same ModuleNotFoundError would occur) - Replace fragile CWD-based PYTHONPATH registration with centralized register_compiled_cache_on_path() that uses the same __file__-relative _CACHE_DIRS already used by cache_cleanup -- fixes path mismatch when studio is launched from a directory other than the repo root - Move PYTHONPATH registration to the top of _train_worker(), before any dataset.map() call (previously it ran late in config assembly, after dataset formatting which also calls dataset.map()) - Update inference.py model-unload to preserve trainer files on spawn platforms, preventing a race where unloading a model via inference tab would delete UnslothSFTTrainer.py while training workers are importing it * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix cache-dir precedence reversal in register_compiled_cache_on_path() Iterating _CACHE_DIRS in forward order while calling insert(0) each time reverses the declared priority: later entries shadow earlier ones. When multiple compiled-cache directories exist, spawned workers could import a stale trainer from the wrong cache. Fix: iterate in reverse so that the highest-priority entry (first in _CACHE_DIRS) is inserted last and ends up at position 0 in sys.path and PYTHONPATH. * fix: harden worker-count helpers against cpu_count=None and desired<=0 - safe_num_proc: guard os.cpu_count() with `or 1`, clamp multi-GPU path with max(1, min(4, desired)), clamp return with max(1, desired) - safe_thread_num_proc: same os.cpu_count() guard and return clamp - Add regression tests (31 L1 unit + 10 sandbox edge-case tests) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove regression tests from PR --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Daniel Han --- studio/backend/core/inference/inference.py | 12 +++- studio/backend/core/training/trainer.py | 14 ++++- studio/backend/utils/cache_cleanup.py | 71 ++++++++++++++++++++-- studio/backend/utils/hardware/hardware.py | 6 +- 4 files changed, 93 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index a95fa93daa..6cb077f4a9 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -549,10 +549,18 @@ class InferenceBackend: # Clear GPU memory cache clear_gpu_cache() - # Remove stale compiled cache so the next model gets a fresh one + # Remove stale compiled cache so the next model gets a fresh one. + # On spawn-based platforms, preserve trainer files so that any + # concurrent training dataset.map() workers can still import them. + import sys as _sys from utils.cache_cleanup import clear_unsloth_compiled_cache - clear_unsloth_compiled_cache() + _preserve = ( + ["Unsloth*Trainer.py"] + if _sys.platform in ("win32", "darwin") + else None + ) + clear_unsloth_compiled_cache(preserve_patterns = _preserve) logger.info(f"Model '{model_name}' successfully unloaded.") return True diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 9b320c93a8..5f504bbdf4 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -529,7 +529,10 @@ class UnslothTrainer: # Remove stale compiled cache so the new model gets a fresh one from utils.cache_cleanup import clear_unsloth_compiled_cache - clear_unsloth_compiled_cache() + _preserve = ( + ["Unsloth*Trainer.py"] if sys.platform in ("win32", "darwin") else None + ) + clear_unsloth_compiled_cache(preserve_patterns = _preserve) # Detect audio model type dynamically (config.json + tokenizer) self._audio_type = detect_audio_type(model_name, hf_token) # audio_vlm is detected as an audio_type now, handle it separately @@ -2718,6 +2721,15 @@ class UnslothTrainer: def _train_worker(self, dataset: Dataset, **training_args): """Worker function for training (runs in separate thread)""" try: + # On spawn-based platforms (Windows, macOS), register all known + # compiled-cache directories on sys.path and PYTHONPATH before any + # dataset.map() call so spawned workers can import dynamically + # compiled modules such as UnslothSFTTrainer. + if sys.platform in ("win32", "darwin"): + from utils.cache_cleanup import register_compiled_cache_on_path + + register_compiled_cache_on_path() + # Store training parameters for metrics calculation self.batch_size = training_args.get("batch_size", 2) self.max_seq_length = training_args.get("max_seq_length", 2048) diff --git a/studio/backend/utils/cache_cleanup.py b/studio/backend/utils/cache_cleanup.py index b3ffbcfc05..4c8e6239a0 100644 --- a/studio/backend/utils/cache_cleanup.py +++ b/studio/backend/utils/cache_cleanup.py @@ -6,13 +6,16 @@ Utility for cleaning up the Unsloth compiled cache directory. The unsloth_compiled_cache is created by unsloth_zoo/compiler.py during FastModel.from_pretrained() and contains model-type-specific compiled Python -files. It should be cleared between model loads to avoid stale artefacts. +files. It should be selectively cleared between model loads to avoid stale +artefacts, while preserving model-agnostic components (like Trainers) needed +by spawned subprocesses. """ import shutil import structlog from loggers import get_logger from pathlib import Path +from typing import List, Optional logger = get_logger(__name__) @@ -27,9 +30,69 @@ _CACHE_DIRS = [ ] -def clear_unsloth_compiled_cache() -> None: - """Remove every known unsloth_compiled_cache directory (idempotent).""" +def get_existing_cache_dirs() -> List[Path]: + """Return known compiled-cache directories that currently exist on disk.""" + return [d for d in _CACHE_DIRS if d.exists()] + + +def register_compiled_cache_on_path() -> None: + """Add all existing compiled-cache directories to sys.path and PYTHONPATH. + + This ensures spawned workers (on platforms using the 'spawn' start method, + i.e. Windows and macOS) can import dynamically compiled modules such as + UnslothSFTTrainer. + """ + import os + import sys + + pypath = os.environ.get("PYTHONPATH", "") + pypath_entries = [p for p in pypath.split(os.pathsep) if p] + + # Iterate in reverse so that earlier _CACHE_DIRS entries (higher priority) + # are inserted last and therefore end up first in sys.path / PYTHONPATH. + for cache_dir in reversed(get_existing_cache_dirs()): + resolved = str(cache_dir.resolve()) + if resolved not in sys.path: + sys.path.insert(0, resolved) + if resolved not in pypath_entries: + pypath_entries.insert(0, resolved) + + os.environ["PYTHONPATH"] = os.pathsep.join(pypath_entries) + + +def clear_unsloth_compiled_cache(preserve_patterns: Optional[List[str]] = None) -> None: + """ + Remove compiled files from the cache directory (idempotent). + + Args: + preserve_patterns: A list of glob patterns for files to keep + (e.g., ["Unsloth*Trainer.py"]). If None or empty, + the entire cache directory is deleted (legacy behavior). + """ for cache_dir in _CACHE_DIRS: - if cache_dir.exists(): + if not cache_dir.exists(): + continue + + if preserve_patterns: + logger.info( + f"Cleaning unsloth compiled cache (preserving {preserve_patterns}): " + f"{cache_dir}" + ) + + for item in cache_dir.iterdir(): + if item.is_file(): + # Check if the file matches any of the patterns we want to keep + preserve = any(item.match(pattern) for pattern in preserve_patterns) + if not preserve: + try: + item.unlink() + except OSError as e: + logger.debug(f"Could not delete {item}: {e}") + + elif item.is_dir(): + # Always clear __pycache__ and other subdirectories + shutil.rmtree(item, ignore_errors = True) + else: + # Legacy behavior: nuke the entire directory logger.info(f"Removing unsloth compiled cache: {cache_dir}") shutil.rmtree(cache_dir, ignore_errors = True) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index b055c4e52d..61ee8a0967 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -521,14 +521,14 @@ def safe_num_proc(desired: Optional[int] = None) -> int: visible = get_visible_gpu_count() if visible > 1: - capped = min(4, desired) + capped = max(1, min(4, desired)) logger.info( f"Multi-GPU detected ({visible} visible GPUs) -- " f"capping num_proc {desired} -> {capped} to avoid fork deadlocks" ) return capped - return desired + return max(1, desired) def safe_thread_num_proc(desired: Optional[int] = None) -> int: @@ -551,7 +551,7 @@ def safe_thread_num_proc(desired: Optional[int] = None) -> int: if desired is None or not isinstance(desired, int): desired = max(1, (os.cpu_count() or 1) // 3) - return desired + return max(1, desired) def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]: