unsloth/studio/backend/utils/datasets
Andrew Barnes 2c5d3c48ec
fix: subprocess crash during map operation on Windows (#4507)
* fix: handle Windows subprocess crash during dataset.map()

Windows uses spawn (not fork) for multiprocessing. Spawned workers
cannot resolve Unsloth's dynamically compiled cache modules from
unsloth_compiled_cache/, causing ModuleNotFoundError and RuntimeError
during dataset.map() tokenization.

Add two platform-guarded patches for sys.platform == "win32":
1. Force HF_DATASETS_MULTITHREADING_MAX_WORKERS=1 and set spawn method
2. Monkey-patch Dataset.map() to force num_proc=None

Fixes #4490

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review: extend spawn fix to macOS, add multiprocess fallback

- Change platform checks from sys.platform == "win32" to
  sys.platform != "linux" so macOS (also spawn-based) is covered
- Wrap multiprocess import in try/except falling back to stdlib
  multiprocessing when the multiprocess package isn't installed
- Rename _win32_safe_map to _spawn_safe_map to reflect broader scope

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: replace global Dataset.map monkey-patch with targeted num_proc routing

The previous approach had issues: Patch 1 set HF_DATASETS_MULTITHREADING_MAX_WORKERS
and forced set_start_method (dead code on platforms already using spawn), and Patch 2
globally monkey-patched Dataset.map() (too broad, missed Dataset.filter()).

Replace with a two-layer fix:

1. Studio layer: Add dataset_map_num_proc() that returns None on spawn platforms
   (Windows, macOS). Unlike num_proc=1 which still creates Pool(1) and spawns a
   worker, num_proc=None runs Dataset.map()/filter() truly in-process.
   Update all dataset.map() callsites to use it. ThreadPoolExecutor callers
   (format_conversion.py) keep using safe_num_proc() since threads are unaffected.

2. Root-cause layer: Propagate UNSLOTH_COMPILE_LOCATION via PYTHONPATH on spawn
   platforms so spawned workers can import compiled modules. Mirrors the .venv_t5
   pattern in worker.py. Does not import unsloth_zoo.compiler (heavy torch/triton
   imports). Completely skipped on Linux.

Also extend safe_num_proc() to return 1 on macOS (was only guarding Windows),
and narrow the transformers 5.x dataloader guard from != "linux" to explicit
("win32", "darwin").

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: add safe_thread_num_proc() for ThreadPoolExecutor callsites

safe_num_proc() correctly caps to 1 on macOS/Windows for process-based
multiprocessing, but format_conversion.py reuses it for ThreadPoolExecutor
workers. Threads share address space and are unaffected by spawn, so
capping to 1 makes image URL downloads sequential -- a real regression.

Add safe_thread_num_proc() that skips the platform guard but keeps the
cpu_count heuristic, and switch both ThreadPoolExecutor callsites in
format_conversion.py to use it.

* fix: remove double-wrap in dataset_num_proc + fix num_proc=1 in datasets route

- trainer.py:3009: Replace safe_num_proc(max(1, os.cpu_count() // 4))
  with max(1, (os.cpu_count() or 1) // 4) to avoid double-wrapping
  inside dataset_map_num_proc which already calls safe_num_proc
- trainer.py:15-20: Clarify comment on PYTHONPATH propagation
- datasets.py:445: Change num_proc=1 to num_proc=None for 10-row
  preview slice (avoids unnecessary multiprocessing overhead)

* fix: guard os.cpu_count() against None in worker-count helpers

os.cpu_count() can return None on some platforms. Use (os.cpu_count() or 1)
to prevent TypeError in safe_num_proc() and safe_thread_num_proc().

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-03-22 05:21:09 -07:00
..
__init__.py Update license headers 2026-03-12 17:23:10 +00:00
chat_templates.py fix: subprocess crash during map operation on Windows (#4507) 2026-03-22 05:21:09 -07:00
data_collators.py Final cleanup 2026-03-12 18:28:04 +00:00
dataset_utils.py studio: web search, KV cache dtype, training progress, inference fixes 2026-03-17 00:30:01 -07:00
format_conversion.py fix: subprocess crash during map operation on Windows (#4507) 2026-03-22 05:21:09 -07:00
format_detection.py Final cleanup 2026-03-12 18:28:04 +00:00
llm_assist.py feat(chat): server-side timings, context display & source hover cards (#4467) 2026-03-20 23:42:01 -07:00
model_mappings.py studio: improve onboarding UX, tooltips, and training defaults (#4355) 2026-03-17 07:46:07 -07:00
vlm_processing.py Final cleanup 2026-03-12 18:28:04 +00:00