Merge pull request #342 from unslothai/local-dataset

dataset upload

studio/backend/core/training/trainer.py

@@ -380,6 +380,9 @@ class UnslothTrainer:
                 self.is_audio = self._audio_type is not None
                 self.is_audio_vlm = False
 
+            if not self.is_audio and not self.is_audio_vlm:
+                self._cuda_audio_used = False
+
             # VLM: vision model with image dataset (mutually exclusive with audio paths)
             vision = is_vision_model(model_name) if not self.is_audio else False
             self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image
@@ -1799,18 +1802,20 @@ class UnslothTrainer:
             eval_enabled = eval_steps is not None and eval_steps > 0
 
             if local_datasets:
-                # Load local datasets
-                all_data = []
+                # Load local datasets using load_dataset() so the result is
+                # Arrow-backed (has cache files).  Dataset.from_list() creates
+                # an in-memory dataset with no cache, which forces num_proc=1
+                # during tokenization/map because sharding requires Arrow files.
+                all_files: list[str] = []
                 for dataset_file in local_datasets:
                     # dataset_file may already be an absolute path from routes/training.py
                     if os.path.isabs(dataset_file):
                         file_path = dataset_file
                     else:
                         # Fallback: try relative to assets/datasets
-                        file_path = _ASSETS_DATASETS_ROOT / dataset_file
+                        file_path = str(_ASSETS_DATASETS_ROOT / dataset_file)
 
                     file_path_obj = Path(file_path)
-                    file_path_str = str(file_path_obj)
 
                     if file_path_obj.is_dir():
                         parquet_dir = (
@@ -1820,36 +1825,41 @@ class UnslothTrainer:
                         )
                         parquet_files = sorted(parquet_dir.glob("*.parquet"))
                         if parquet_files:
-                            for parquet_file in parquet_files:
-                                df = pd.read_parquet(parquet_file)
-                                all_data.extend(df.to_dict("records"))
+                            all_files.extend(str(p) for p in parquet_files)
                             continue
+                        # Fall through to single-file detection for dirs with json/csv
+                        candidates: list[Path] = []
+                        for ext in ('.json', '.jsonl', '.csv', '.parquet'):
+                            candidates.extend(sorted(file_path_obj.glob(f"*{ext}")))
+                        if candidates:
+                            all_files.extend(str(c) for c in candidates)
+                            continue
+                        raise ValueError(f"No supported data files in directory: {file_path_obj}")
+                    else:
+                        all_files.append(str(file_path_obj))
 
-                    if file_path_str.endswith('.json'):
-                        with open(file_path_obj, 'r', encoding='utf-8') as f:
-                            data = json.load(f)
-                            if isinstance(data, list):
-                                all_data.extend(data)
-                            else:
-                                all_data.append(data)
-                    elif file_path_str.endswith('.csv'):
-                        df = pd.read_csv(file_path_obj)
-                        all_data.extend(df.to_dict('records'))
-                    elif file_path_str.endswith('.parquet'):
-                        df = pd.read_parquet(file_path_obj)
-                        all_data.extend(df.to_dict('records'))
-                        continue
+                if all_files:
+                    # Determine loader type from the first file extension
+                    first_ext = Path(all_files[0]).suffix.lower()
+                    if first_ext in ('.json', '.jsonl'):
+                        loader = 'json'
+                    elif first_ext == '.csv':
+                        loader = 'csv'
+                    elif first_ext == '.parquet':
+                        loader = 'parquet'
+                    else:
+                        raise ValueError(f"Unsupported local dataset format: {all_files[0]}")
 
-                if all_data:
-                    dataset = Dataset.from_list(all_data)
+                    dataset = load_dataset(loader, data_files=all_files, split='train')
 
                     # Check if stopped during dataset loading
                     if self.should_stop:
                         print("Stopped during dataset loading\n")
                         return None
 
-                    self._update_progress(status_message=f"Loaded {len(all_data)} samples from local files")
-                    print(f"Loaded {len(all_data)} samples from local files\n")
+                    self._update_progress(status_message=f"Loaded {len(dataset)} samples from local files")
+                    print(f"Loaded {len(dataset)} samples from local files\n")
+                    print(f"[DEBUG] Dataset cache_files: {dataset.cache_files}\n")
 
             elif dataset_source:
                 # Load from Hugging Face
@@ -2377,6 +2387,7 @@ class UnslothTrainer:
                 "dataset_num_proc": 1 if (self.is_audio or self.is_audio_vlm or self._cuda_audio_used) else safe_num_proc(max(1, os.cpu_count() // 4)),
                 "max_seq_length": training_args.get('max_seq_length', 2048),
             }
+            print(f"[DEBUG] dataset_num_proc={config_args['dataset_num_proc']} (is_audio={self.is_audio}, is_audio_vlm={self.is_audio_vlm}, _cuda_audio_used={self._cuda_audio_used})")
 
             # On Windows with transformers 5.x, disable DataLoader multiprocessing
             # to avoid issues with modified sys.path (.venv_t5) in spawned workers.

studio/backend/models/datasets.py

@@ -41,6 +41,12 @@ class CheckFormatResponse(BaseModel):
     warning: Optional[str] = None
 
 
+class UploadDatasetResponse(BaseModel):
+    """Response with stored dataset path for training."""
+    filename: str = Field(..., description="Original filename")
+    stored_path: str = Field(..., description="Absolute path stored on backend")
+
+
 class LocalDatasetItem(BaseModel):
     class Metadata(BaseModel):
         actual_num_records: Optional[int] = None

studio/backend/routes/datasets.py

@@ -6,7 +6,8 @@ import io
 import json
 import sys
 from pathlib import Path
-from fastapi import APIRouter, Depends, HTTPException
+from uuid import uuid4
+from fastapi import APIRouter, Depends, HTTPException, UploadFile
 import logging
 
 # Add backend directory to path
@@ -36,6 +37,7 @@ from models.datasets import (
     CheckFormatResponse,
     LocalDatasetItem,
     LocalDatasetsResponse,
+    UploadDatasetResponse,
 )
 
 
@@ -89,8 +91,10 @@ DATA_EXTS = (
     '.zip',
 )
 LOCAL_FILE_EXTS = ('.json', '.jsonl', '.csv', '.parquet')
+LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl", ".parquet"}
 BACKEND_ROOT = Path(__file__).resolve().parents[1]
 LOCAL_DATASETS_ROOT = BACKEND_ROOT / "assets" / "datasets"
+DATASET_UPLOAD_DIR = LOCAL_DATASETS_ROOT / "uploads"
 
 
 def _safe_read_metadata(path: Path) -> dict | None:
@@ -252,6 +256,50 @@ def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_s
     return preview_slice, total_rows
 
 
+def _sanitize_filename(filename: str) -> str:
+    name = Path(filename).name.strip().replace("\x00", "")
+    if not name:
+        return "dataset_upload"
+    return name
+
+
+@router.post("/upload", response_model=UploadDatasetResponse)
+async def upload_dataset(
+    file: UploadFile,
+    current_subject: str = Depends(get_current_subject),
+) -> UploadDatasetResponse:
+    filename = _sanitize_filename(file.filename or "dataset_upload")
+    ext = Path(filename).suffix.lower()
+    if ext not in LOCAL_UPLOAD_EXTS:
+        allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))
+        raise HTTPException(
+            status_code=400,
+            detail=f"Unsupported file type: {ext}. Allowed: {allowed}",
+        )
+
+    max_size_bytes = 512 * 1024 * 1024
+    DATASET_UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+    stem = Path(filename).stem
+    stored_name = f"{uuid4().hex}_{stem}{ext}"
+    stored_path = DATASET_UPLOAD_DIR / stored_name
+
+    # Stream file to disk in chunks to avoid holding entire file in memory
+    size = 0
+    with open(stored_path, "wb") as f:
+        while chunk := await file.read(1024 * 1024):
+            size += len(chunk)
+            if size > max_size_bytes:
+                stored_path.unlink(missing_ok=True)
+                raise HTTPException(status_code=413, detail="File too large (max 512MB)")
+            f.write(chunk)
+
+    if size == 0:
+        stored_path.unlink(missing_ok=True)
+        raise HTTPException(status_code=400, detail="Empty upload payload")
+
+    return UploadDatasetResponse(filename=filename, stored_path=str(stored_path))
+
+
 @router.get("/local", response_model=LocalDatasetsResponse)
 def list_local_datasets(
     current_subject: str = Depends(get_current_subject),

studio/frontend/src/features/studio/sections/dataset-section.tsx

@@ -37,6 +37,7 @@ import {
 } from "@/hooks";
 import {
   HfDatasetSubsetSplitSelectors,
+  uploadTrainingDataset,
   useDatasetPreviewDialogStore,
   useTrainingConfigStore,
 } from "@/features/training";
@@ -52,7 +53,8 @@ import {
   ViewIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { type ChangeEvent, useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { toast } from "sonner";
 import { useShallow } from "zustand/react/shallow";
 
 const SEARCH_INPUT_REASONS = new Set(["input-change", "input-paste", "input-clear"]);
@@ -72,7 +74,11 @@ function deriveLocalDatasetName(path: string): string {
   const parts = normalized.split("/").filter(Boolean);
   const parquetIndex = parts.lastIndexOf("parquet-files");
   if (parquetIndex > 0) return parts[parquetIndex - 1];
-  return parts[parts.length - 1] ?? path;
+  const basename = parts[parts.length - 1] ?? path;
+  // Strip UUID prefix from uploaded files (format: {32hex}_{original})
+  const uuidPrefixMatch = basename.match(/^[a-f0-9]{32}_(.+)$/);
+  if (uuidPrefixMatch) return uuidPrefixMatch[1];
+  return basename;
 }
 
 function formatUpdatedDate(timestamp: number | null): string {
@@ -286,6 +292,8 @@ export function DatasetSection() {
     if (datasetSource !== "upload") return;
     if (!uploadedFile) return;
     if (selectedLocalDataset) return;
+    // Don't clear if this is a direct file upload (not a recipe directory)
+    if (isLikelyLocalDatasetRef(uploadedFile)) return;
     selectLocalDataset(null);
   }, [
     datasetSource,
@@ -320,11 +328,49 @@ export function DatasetSection() {
   const selectedLocalUpdatedAt = selectedLocalDataset?.updated_at ?? null;
 
   const comboboxAnchorRef = useRef<HTMLDivElement>(null);
+  const fileInputRef = useRef<HTMLInputElement>(null);
   const { scrollRef, sentinelRef } = useInfiniteScroll(
     fetchMore,
     hfResults.length,
   );
 
+  const [isUploading, setIsUploading] = useState(false);
+
+  const handleUploadButtonClick = () => {
+    fileInputRef.current?.click();
+  };
+
+  const handleDatasetFileChange = async (event: ChangeEvent<HTMLInputElement>) => {
+    const file = event.target.files?.[0];
+    event.target.value = "";
+    if (!file) return;
+
+    const MAX_SIZE_BYTES = 512 * 1024 * 1024;
+    if (file.size > MAX_SIZE_BYTES) {
+      toast.error("File too large", {
+        description: "Maximum upload size is 512 MB.",
+      });
+      return;
+    }
+
+    setIsUploading(true);
+    try {
+      const uploaded = await uploadTrainingDataset(file);
+
+      selectLocalDataset(uploaded.stored_path);
+
+      toast.success("Dataset uploaded", {
+        description: uploaded.filename,
+      });
+    } catch (error) {
+      toast.error("Upload failed", {
+        description: error instanceof Error ? error.message : "Unknown error",
+      });
+    } finally {
+      setIsUploading(false);
+    }
+  };
+
   return (
     <div data-tour="studio-dataset" className="col-span-1 xl:col-span-4">
       <SectionCard
@@ -596,7 +642,7 @@ export function DatasetSection() {
               datasetEvalSplit={datasetEvalSplit}
               setDatasetEvalSplit={setDatasetEvalSplit}
             />
-          ) : datasetSource === "upload" ? (
+          ) : datasetSource === "upload" && selectedLocalDataset ? (
             <div className="rounded-lg border bg-muted/20 px-3.5 py-3">
               <div className="mb-3 flex items-center justify-between gap-3">
                 <div>
@@ -609,45 +655,39 @@ export function DatasetSection() {
                 </div>
               </div>
 
-              {uploadedFile ? (
-                <div className="flex flex-col gap-3">
-                  <div className="grid grid-cols-2 gap-x-4 gap-y-2 text-xs">
-                    <MetadataRow
-                      label="Rows"
-                      value={
-                        typeof selectedLocalRows === "number"
-                          ? selectedLocalRows.toLocaleString()
-                          : "--"
-                      }
-                    />
-                    <MetadataRow
-                      label="Columns"
-                      value={
-                        selectedLocalColumns.length > 0
-                          ? String(selectedLocalColumns.length)
-                          : "--"
-                      }
-                    />
-                    <MetadataRow
-                      label="Batches"
-                      value={
-                        typeof selectedLocalMetadata?.num_completed_batches === "number" &&
-                        typeof selectedLocalMetadata?.total_num_batches === "number"
-                          ? `${selectedLocalMetadata.num_completed_batches}/${selectedLocalMetadata.total_num_batches}`
-                          : "--"
-                      }
-                    />
-                    <MetadataRow
-                      label="Updated"
-                      value={formatUpdatedDate(selectedLocalUpdatedAt)}
-                    />
-                  </div>
+              <div className="flex flex-col gap-3">
+                <div className="grid grid-cols-2 gap-x-4 gap-y-2 text-xs">
+                  <MetadataRow
+                    label="Rows"
+                    value={
+                      typeof selectedLocalRows === "number"
+                        ? selectedLocalRows.toLocaleString()
+                        : "--"
+                    }
+                  />
+                  <MetadataRow
+                    label="Columns"
+                    value={
+                      selectedLocalColumns.length > 0
+                        ? String(selectedLocalColumns.length)
+                        : "--"
+                    }
+                  />
+                  <MetadataRow
+                    label="Batches"
+                    value={
+                      typeof selectedLocalMetadata?.num_completed_batches === "number" &&
+                      typeof selectedLocalMetadata?.total_num_batches === "number"
+                        ? `${selectedLocalMetadata.num_completed_batches}/${selectedLocalMetadata.total_num_batches}`
+                        : "--"
+                    }
+                  />
+                  <MetadataRow
+                    label="Updated"
+                    value={formatUpdatedDate(selectedLocalUpdatedAt)}
+                  />
                 </div>
-              ) : (
-                <p className="text-xs text-muted-foreground">
-                  Select a local dataset to view metadata.
-                </p>
-              )}
+              </div>
             </div>
           ) : null}
 
@@ -839,9 +879,15 @@ export function DatasetSection() {
                 variant="outline"
                 size="sm"
                 className="cursor-pointer gap-1.5"
+                disabled={isUploading}
+                onClick={handleUploadButtonClick}
               >
-                <HugeiconsIcon icon={CloudUploadIcon} className="size-3.5" />
-                Upload
+                {isUploading ? (
+                  <Spinner className="size-3.5" />
+                ) : (
+                  <HugeiconsIcon icon={CloudUploadIcon} className="size-3.5" />
+                )}
+                {isUploading ? "Uploading..." : "Upload"}
               </Button>
               <Button
                 variant="outline"
@@ -855,6 +901,15 @@ export function DatasetSection() {
               </Button>
             </div>
           </div>
+          <input
+            ref={fileInputRef}
+            type="file"
+            accept=".json,.jsonl,.csv,.parquet"
+            className="hidden"
+            onChange={(event) => {
+              void handleDatasetFileChange(event);
+            }}
+          />
       </div>
       </SectionCard>
     </div>

studio/frontend/src/features/training/api/datasets-api.ts

@@ -1,6 +1,7 @@
 import type {
   CheckFormatResponse,
   LocalDatasetsResponse,
+  UploadDatasetResponse,
 } from "../types/datasets";
 import { authFetch } from "@/features/auth";
 
@@ -39,6 +40,25 @@ export async function checkDatasetFormat({
   return res.json();
 }
 
+export async function uploadTrainingDataset(
+  file: File,
+): Promise<UploadDatasetResponse> {
+  const form = new FormData();
+  form.append("file", file);
+
+  const res = await authFetch("/api/datasets/upload", {
+    method: "POST",
+    body: form,
+  });
+
+  if (!res.ok) {
+    const body = await res.json().catch(() => null);
+    throw new Error(body?.detail || `Upload failed (${res.status})`);
+  }
+
+  return res.json();
+}
+
 export async function listLocalDatasets(): Promise<LocalDatasetsResponse> {
   const res = await authFetch("/api/datasets/local");
   if (!res.ok) {

studio/frontend/src/features/training/index.ts

@@ -7,6 +7,7 @@ export { useTrainingActions } from "./hooks/use-training-actions";
 export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle";
 export { HfDatasetSubsetSplitSelectors } from "./components/hf-dataset-subset-split-selectors";
 export { useDatasetPreviewDialogStore } from "./stores/dataset-preview-dialog-store";
+export { uploadTrainingDataset } from "./api/datasets-api";
 export { listLocalModels } from "./api/models-api";
 export type { LocalModelInfo } from "./api/models-api";
 export type { TrainingPhase } from "./types/runtime";

studio/frontend/src/features/training/stores/training-config-store.ts

@@ -285,6 +285,9 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             uploadedFile,
             ...resetDatasetState(),
           });
+          if (uploadedFile) {
+            runDatasetCheck(uploadedFile, "train");
+          }
         },
         setDatasetFormat: (datasetFormat) => set({ datasetFormat }),
         setDataset: (dataset) => {

studio/frontend/src/features/training/types/datasets.ts

@@ -15,6 +15,11 @@ export type CheckFormatResponse = {
   warning?: string | null;
 };
 
+export type UploadDatasetResponse = {
+  filename: string;
+  stored_path: string;
+};
+
 export type LocalDatasetInfo = {
   metadata?: {
     actual_num_records?: number | null;