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]: