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 <danielhanchen@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
NuoFang
GitHub
parent
62e3de181f
commit
4cedeba8c2
4 changed files with 93 additions and 10 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue