diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index cfc04e572e..96d593b6cf 100644 --- a/studio/backend/core/training/trainer.py +++ b/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. diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 19197d99ad..49d8b4c419 100644 --- a/studio/backend/models/datasets.py +++ b/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 diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 04345a944d..b06fa33942 100644 --- a/studio/backend/routes/datasets.py +++ b/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), diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index e522aaca9b..a72c21100a 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/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(null); + const fileInputRef = useRef(null); const { scrollRef, sentinelRef } = useInfiniteScroll( fetchMore, hfResults.length, ); + const [isUploading, setIsUploading] = useState(false); + + const handleUploadButtonClick = () => { + fileInputRef.current?.click(); + }; + + const handleDatasetFileChange = async (event: ChangeEvent) => { + 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 (
- ) : datasetSource === "upload" ? ( + ) : datasetSource === "upload" && selectedLocalDataset ? (
@@ -609,45 +655,39 @@ export function DatasetSection() {
- {uploadedFile ? ( -
-
- - 0 - ? String(selectedLocalColumns.length) - : "--" - } - /> - - -
+
+
+ + 0 + ? String(selectedLocalColumns.length) + : "--" + } + /> + +
- ) : ( -

- Select a local dataset to view metadata. -

- )} +
) : null} @@ -839,9 +879,15 @@ export function DatasetSection() { variant="outline" size="sm" className="cursor-pointer gap-1.5" + disabled={isUploading} + onClick={handleUploadButtonClick} > - - Upload + {isUploading ? ( + + ) : ( + + )} + {isUploading ? "Uploading..." : "Upload"}
+ { + void handleDatasetFileChange(event); + }} + /> diff --git a/studio/frontend/src/features/training/api/datasets-api.ts b/studio/frontend/src/features/training/api/datasets-api.ts index d056d92189..f3098acc80 100644 --- a/studio/frontend/src/features/training/api/datasets-api.ts +++ b/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 { + 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 { const res = await authFetch("/api/datasets/local"); if (!res.ok) { diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts index 9da25f88bb..abdba37388 100644 --- a/studio/frontend/src/features/training/index.ts +++ b/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"; diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 42f6b2a33b..3318a192d8 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -285,6 +285,9 @@ export const useTrainingConfigStore = create()( uploadedFile, ...resetDatasetState(), }); + if (uploadedFile) { + runDatasetCheck(uploadedFile, "train"); + } }, setDatasetFormat: (datasetFormat) => set({ datasetFormat }), setDataset: (dataset) => { diff --git a/studio/frontend/src/features/training/types/datasets.ts b/studio/frontend/src/features/training/types/datasets.ts index cee77f683b..41a71affaf 100644 --- a/studio/frontend/src/features/training/types/datasets.ts +++ b/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;