Studio: training survives a non-writable HF datasets cache (#6148)

* Studio: training survives a non-writable HF datasets cache

A shared HF datasets cache can contain subtrees owned by another user
(for example populated by an earlier root-run job). datasets then dies
with "[Errno 13] Permission denied: ..._builder.lock" while locking
the cached builder and the training run fails. load_dataset in the
training worker and trainer now goes through a wrapper that catches the
EACCES and rebuilds the dataset in a Studio-owned cache under
cache_root()/hf-datasets, logging the fallback.

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

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

* Scope the HF_DATASETS_CACHE override to the fallback load

* Route non-streaming dataset preview loads through the cache-safe wrapper

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

studio/backend/core/training/trainer.py

@@ -58,7 +58,8 @@ from pathlib import Path
 from typing import Optional, Callable
 from dataclasses import dataclass
 import pandas as pd
-from datasets import Dataset, load_dataset
+from datasets import Dataset
+from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
 
 from core.inference.llama_cpp import _hf_offline_if_dns_dead
 from utils.models import is_vision_model, detect_audio_type

studio/backend/core/training/worker.py

@@ -1353,7 +1353,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
             "(unsloth_zoo.mlx.loader / unsloth_zoo.mlx.trainer). Reinstall via "
             "install.sh on Apple Silicon."
         ) from e
-    from datasets import load_dataset
+    from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
 
     if mx.metal.is_available():
         info = mx.device_info()
@@ -2879,7 +2879,8 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
         )
         from sentence_transformers.losses import MultipleNegativesRankingLoss
         from sentence_transformers.training_args import BatchSamplers
-        from datasets import load_dataset, Dataset
+        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
     except ImportError as e:

studio/backend/hub/services/datasets/formatting.py

@@ -223,7 +223,10 @@ def _load_processed_hf_preview_slice(
     if not _is_valid_repo_id(request.dataset_name):
         return None
     try:
-        from datasets import DownloadConfig, load_dataset
+        from datasets import DownloadConfig
+
+        # Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper.
+        from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
     except Exception:
         return None
 

studio/backend/hub/services/datasets/local.py

@@ -224,7 +224,8 @@ def _stream_file_preview_slice(path: Path, preview_size: int):
 
 
 def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
-    from datasets import load_dataset
+    # Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper.
+    from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
 
     if dataset_path.is_dir():
         parquet_dir = (

studio/backend/routes/datasets.py

@@ -412,7 +412,8 @@ def _build_local_dataset_items() -> list[LocalDatasetItem]:
 
 
 def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
-    from datasets import load_dataset
+    # Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper.
+    from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
 
     if dataset_path.is_dir():
         parquet_dir = (

studio/backend/utils/datasets/cache_safe.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Permission-safe wrapper around datasets.load_dataset.
+
+A shared HF datasets cache can contain subtrees owned by another user (for
+example populated by an earlier root-run job). datasets then raises
+"[Errno 13] Permission denied: ..._builder.lock" while locking the cached
+builder, killing the training run even though the dataset itself is fine.
+Retry such loads in a Studio-owned cache so the run proceeds; the worst case
+is one rebuild of the dataset in the fallback location.
+"""
+
+import logging
+import os
+
+from utils.paths.storage_roots import cache_root
+
+logger = logging.getLogger(__name__)
+
+
+def studio_datasets_cache() -> str:
+    path = cache_root() / "hf-datasets"
+    path.mkdir(parents = True, exist_ok = True)
+    return str(path)
+
+
+def load_dataset_cache_safe(*args, **kwargs):
+    """datasets.load_dataset, retried in a Studio-owned cache on EACCES."""
+    from datasets import load_dataset
+    try:
+        return load_dataset(*args, **kwargs)
+    except PermissionError as error:
+        fallback = studio_datasets_cache()
+        logger.warning(
+            "HF datasets cache is not writable (%s); rebuilding in %s",
+            error,
+            fallback,
+        )
+        kwargs["cache_dir"] = fallback
+        # Nested builders consult the env var while the load runs; restore it
+        # after so other datasets keep trying the shared cache first.
+        old_env = os.environ.get("HF_DATASETS_CACHE")
+        os.environ["HF_DATASETS_CACHE"] = fallback
+        try:
+            return load_dataset(*args, **kwargs)
+        finally:
+            if old_env is None:
+                os.environ.pop("HF_DATASETS_CACHE", None)
+            else:
+                os.environ["HF_DATASETS_CACHE"] = old_env