diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index ca126948bb..49d8b4c419 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -41,12 +41,6 @@ class CheckFormatResponse(BaseModel): warning: Optional[str] = None -class UploadDatasetRequest(BaseModel): - """Request for uploading a local training dataset file.""" - filename: str = Field(..., description="Original filename, e.g. my_data.jsonl") - content_base64: str = Field(..., description="Base64-encoded file bytes") - - class UploadDatasetResponse(BaseModel): """Response with stored dataset path for training.""" filename: str = Field(..., description="Original filename") diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 2faa7121ac..302ac0aa9d 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -2,13 +2,12 @@ Datasets API routes """ import base64 -import binascii import io import json import sys from pathlib import Path from uuid import uuid4 -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, UploadFile import logging # Add backend directory to path @@ -38,7 +37,6 @@ from models.datasets import ( CheckFormatResponse, LocalDatasetItem, LocalDatasetsResponse, - UploadDatasetRequest, UploadDatasetResponse, ) @@ -267,22 +265,12 @@ def _sanitize_filename(filename: str) -> str: return name -def _decode_base64_payload(content_base64: str) -> bytes: - raw = content_base64.strip() - if "," in raw and raw.lower().startswith("data:"): - raw = raw.split(",", 1)[1] - try: - return base64.b64decode(raw, validate=True) - except binascii.Error as exc: - raise HTTPException(status_code=400, detail="Invalid base64 payload") from exc - - @router.post("/upload", response_model=UploadDatasetResponse) -def upload_dataset( - payload: UploadDatasetRequest, +async def upload_dataset( + file: UploadFile, current_subject: str = Depends(get_current_subject), ) -> UploadDatasetResponse: - filename = _sanitize_filename(payload.filename) + 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)) @@ -291,16 +279,25 @@ def upload_dataset( detail=f"Unsupported file type: {ext}. Allowed: {allowed}", ) - file_bytes = _decode_base64_payload(payload.content_base64) - if not file_bytes: - raise HTTPException(status_code=400, detail="Empty upload payload") - + max_size_bytes = 512 * 1024 * 1024 DATASET_UPLOAD_DIR.mkdir(parents=True, exist_ok=True) - # Normalize extension to lowercase so downstream suffix checks work stem = Path(filename).stem stored_name = f"{uuid4().hex}_{stem}{ext}" stored_path = DATASET_UPLOAD_DIR / stored_name - stored_path.write_bytes(file_bytes) + + # 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)) diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index b81dd058f6..a72c21100a 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -334,18 +334,6 @@ export function DatasetSection() { hfResults.length, ); - const fileToBase64Payload = (file: File): Promise => - new Promise((resolve, reject) => { - const reader = new FileReader(); - reader.onload = () => { - const value = String(reader.result ?? ""); - const parts = value.split(","); - resolve(parts.length > 1 ? parts[1] : value); - }; - reader.onerror = () => reject(new Error("Failed to read file")); - reader.readAsDataURL(file); - }); - const [isUploading, setIsUploading] = useState(false); const handleUploadButtonClick = () => { @@ -367,11 +355,7 @@ export function DatasetSection() { setIsUploading(true); try { - const contentBase64 = await fileToBase64Payload(file); - const uploaded = await uploadTrainingDataset({ - filename: file.name, - contentBase64, - }); + const uploaded = await uploadTrainingDataset(file); selectLocalDataset(uploaded.stored_path); diff --git a/studio/frontend/src/features/training/api/datasets-api.ts b/studio/frontend/src/features/training/api/datasets-api.ts index 0f37ac53bf..f3098acc80 100644 --- a/studio/frontend/src/features/training/api/datasets-api.ts +++ b/studio/frontend/src/features/training/api/datasets-api.ts @@ -40,22 +40,15 @@ export async function checkDatasetFormat({ return res.json(); } -type UploadDatasetArgs = { - filename: string; - contentBase64: string; -}; +export async function uploadTrainingDataset( + file: File, +): Promise { + const form = new FormData(); + form.append("file", file); -export async function uploadTrainingDataset({ - filename, - contentBase64, -}: UploadDatasetArgs): Promise { const res = await authFetch("/api/datasets/upload", { method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - filename, - content_base64: contentBase64, - }), + body: form, }); if (!res.ok) {