feat(studio): implement S3 dataset loading (completes #5951) (#6222)

* feat(studio): add S3 dataset configuration foundation (#4539)

Add foundational types and configuration for S3 bucket dataset loading:

- Add S3Config type to frontend training types
- Add S3Config Pydantic model to backend training models
- Add "s3" as a DatasetSource option
- Add s3Config state and setS3Config action to training config store
- Add i18n translations for S3 configuration (English and Chinese)

This provides the type definitions and UI text for S3 integration.
Full implementation requires boto3 dependency and data loading logic.

Refs: #4539

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

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

* Wire S3 config into training pipeline and prevent secrets persistence

- Pass s3_config from request into training_kwargs so it flows to training subprocess
- Add s3Config to NON_PERSISTED_STATE_KEYS to prevent AWS secrets from being
  saved to localStorage

Addresses code review feedback on PR #5951.

* Exclude S3 config from database persistence to protect secrets

Filter out s3_config (which contains secret_access_key) from the
config_json stored in training_runs table, preventing AWS credentials
from being persisted to disk.

Addresses P1 security feedback on PR #5951.

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

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

* Re-raise HTTPException in start_training and defer s3 DatasetSource widening for PR #5951

* Redact s3_config from W&B run config and accept camelCase S3 credential aliases for PR #5951

* feat(studio): implement S3 dataset loading end-to-end

Builds the actual S3 loader on top of the hardened #5951 foundation,
turning the 501-gated scaffold into a working dataset source.

Backend:
- Add core/training/s3_dataset.py: lists and downloads supported dataset
  files (parquet/json/jsonl/csv) from an S3 bucket to a temp dir, using
  IAM-role or access-key credentials. boto3 is imported lazily (optional dep).
- Wire s3_config into UnslothTrainer.load_and_format_dataset (downloads then
  reuses the existing local-file path) and thread it through worker.py.
- Replace the 501 "not implemented" gate with a boto3-availability guard so
  S3 works when boto3 is present and fails clearly when it is not.
- Add boto3 to studio.txt requirements.
- Add tests/test_s3_dataset.py (8 tests) covering download/filtering,
  collisions, missing-boto3, and S3Config camelCase/IAM validation.

Frontend:
- Widen DatasetSource to include "s3"; add s3_config to the training payload
  type and mapper; add an S3 validation branch and selectS3Source store action.
- Add s3-config-form.tsx (bucket/region/prefix/keys/IAM toggle) reusing the
  existing studio.dataset.s3.* i18n strings.
- Add a Hugging Face / Local / Amazon S3 source toggle in dataset-section;
  the S3 config card replaces the dataset combobox when S3 is selected.
- Fix DatasetPreviewDialog to accept the widened DatasetSource type.

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

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

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

* Fix S3 dataset loader for PR #6222

* Fix S3 dataset edge cases for PR #6222

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

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

* Fix S3 IAM payload handling for PR #6222

* Block multimodal S3 datasets for PR #6222

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Ash <ash@MacBook-Pro.local>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>

studio/backend/core/training/resume.py

@@ -3,6 +3,7 @@
 
 """Helpers for validating resumable training outputs."""
 
+import json
 from pathlib import Path
 from typing import Optional
 
@@ -56,9 +57,29 @@ def normalize_resume_output_dir(path_value: str) -> str:
     return str(path)
 
 
+def _run_config(run: dict) -> dict:
+    raw_config = run.get("config_json")
+    if isinstance(raw_config, dict):
+        return raw_config
+    if not isinstance(raw_config, str) or not raw_config.strip():
+        return {}
+    try:
+        parsed = json.loads(raw_config)
+    except (json.JSONDecodeError, TypeError):
+        return {}
+    return parsed if isinstance(parsed, dict) else {}
+
+
+def _uses_s3_dataset(run: dict) -> bool:
+    config = _run_config(run)
+    return config.get("dataset_source") == "s3" or "s3_dataset" in config
+
+
 def can_resume_run(run: dict) -> bool:
     if run.get("resumed_later"):
         return False
+    if _uses_s3_dataset(run):
+        return False
 
     final_step = run.get("final_step")
     total_steps = run.get("total_steps")

studio/backend/core/training/s3_dataset.py

@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""
+S3 dataset loader.
+
+Downloads dataset files (parquet / json / jsonl / csv) from an AWS S3 bucket
+to a local temp directory so the existing local-file dataset path can consume
+them. boto3 is an optional dependency and is imported lazily — callers should
+gate on :func:`boto3_available` before invoking the loader.
+
+The S3 config dict mirrors ``models.training.S3Config.model_dump()`` (snake_case
+keys): bucket, region, prefix, access_key_id, secret_access_key, use_iam_role.
+Credentials are read once to build the client and never logged or persisted.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import shutil
+import tempfile
+from importlib.util import find_spec
+from typing import Callable, Optional
+
+logger = logging.getLogger(__name__)
+
+# Extensions the local-file loader (UnslothTrainer._loader_for_files) understands.
+SUPPORTED_EXTENSIONS = (".parquet", ".json", ".jsonl", ".csv")
+_JSON_EXTENSIONS = (".json", ".jsonl")
+_IGNORED_METADATA_FILENAMES = {
+    "dataset_info.json",
+    "metadata.json",
+    "schema.json",
+    "state.json",
+}
+
+
+class S3DownloadCancelled(RuntimeError):
+    """Raised when the caller cancels an S3 dataset download."""
+
+
+class S3DatasetDownload:
+    def __init__(
+        self,
+        files: list[str],
+        temp_dir: Optional[str] = None,
+    ):
+        self.files = files
+        self.temp_dir = temp_dir
+
+    def cleanup(self) -> None:
+        if not self.temp_dir:
+            return
+        shutil.rmtree(self.temp_dir, ignore_errors = True)
+        self.temp_dir = None
+
+
+def boto3_available() -> bool:
+    """True if boto3 can be imported (without importing it)."""
+    return find_spec("boto3") is not None
+
+
+def _build_s3_client(s3_config: dict):
+    """Create a boto3 S3 client from the config dict.
+
+    Uses explicit access keys when provided, otherwise falls back to the
+    default credential chain (IAM role / instance profile / env / shared creds).
+    """
+    import boto3  # lazy: optional dependency
+
+    region = s3_config.get("region") or "us-east-1"
+    use_iam_role = bool(s3_config.get("use_iam_role"))
+    access_key_id = s3_config.get("access_key_id")
+    secret_access_key = s3_config.get("secret_access_key")
+
+    if not use_iam_role and access_key_id and secret_access_key:
+        return boto3.client(
+            "s3",
+            region_name = region,
+            aws_access_key_id = access_key_id,
+            aws_secret_access_key = secret_access_key,
+        )
+    # IAM role / instance profile / ambient credentials
+    return boto3.client("s3", region_name = region)
+
+
+def _list_dataset_keys(client, bucket: str, prefix: Optional[str]) -> list[str]:
+    """List object keys under ``prefix`` that have a supported data extension."""
+    paginator = client.get_paginator("list_objects_v2")
+    list_kwargs = {"Bucket": bucket}
+    if prefix:
+        list_kwargs["Prefix"] = prefix
+
+    keys: list[str] = []
+    for page in paginator.paginate(**list_kwargs):
+        for obj in page.get("Contents", []):
+            key = obj["Key"]
+            if key.endswith("/"):
+                continue  # directory placeholder
+            if os.path.basename(key).lower() in _IGNORED_METADATA_FILENAMES:
+                continue
+            if key.lower().endswith(SUPPORTED_EXTENSIONS):
+                keys.append(key)
+    return keys
+
+
+def _extension_family(key: str) -> str:
+    ext = os.path.splitext(key)[1].lower()
+    if ext in _JSON_EXTENSIONS:
+        return "json"
+    return ext.lstrip(".")
+
+
+def _validate_single_extension_family(keys: list[str]) -> None:
+    families: list[str] = []
+    for key in keys:
+        family = _extension_family(key)
+        if family not in families:
+            families.append(family)
+
+    if len(families) <= 1:
+        return
+
+    raise ValueError(
+        "S3 prefix contains mixed dataset formats "
+        f"({', '.join(families)}). Keep one dataset format under the selected prefix."
+    )
+
+
+def _unique_local_path(target_dir: str, filename: str, used_paths: set[str]) -> str:
+    """Return an unused flattened path for an S3 object basename."""
+    stem, ext = os.path.splitext(filename)
+    candidate = os.path.join(target_dir, filename)
+    suffix = 1
+    while candidate in used_paths or os.path.exists(candidate):
+        candidate = os.path.join(target_dir, f"{stem}_{suffix}{ext}")
+        suffix += 1
+    used_paths.add(candidate)
+    return candidate
+
+
+def _raise_if_cancelled(cancel_callback: Optional[Callable[[], bool]]) -> None:
+    if cancel_callback is not None and cancel_callback():
+        raise S3DownloadCancelled("S3 dataset download cancelled")
+
+
+def prepare_s3_dataset_download(
+    s3_config: dict,
+    dest_dir: Optional[str] = None,
+    cancel_callback: Optional[Callable[[], bool]] = None,
+) -> S3DatasetDownload:
+    """Download supported dataset files from S3 to a local directory.
+
+    Returns the local files plus the owned temporary directory, when one was
+    created. Call ``cleanup()`` after the dataset loader has materialized data.
+
+    Raises ``RuntimeError`` if boto3 is missing, and ``ValueError`` if the
+    bucket/prefix contains no supported dataset files.
+    """
+    if not boto3_available():
+        raise RuntimeError("S3 dataset loading requires boto3. Install it with: pip install boto3")
+
+    bucket = s3_config.get("bucket")
+    if not bucket:
+        raise ValueError("s3_config.bucket is required")
+    prefix = s3_config.get("prefix")
+
+    _raise_if_cancelled(cancel_callback)
+    client = _build_s3_client(s3_config)
+
+    keys = _list_dataset_keys(client, bucket, prefix)
+    _raise_if_cancelled(cancel_callback)
+    if not keys:
+        where = f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}"
+        raise ValueError(
+            f"No supported dataset files ({', '.join(SUPPORTED_EXTENSIONS)}) "
+            f"found under {where}"
+        )
+
+    _validate_single_extension_family(keys)
+
+    owns_temp_dir = dest_dir is None
+    target_dir = dest_dir or tempfile.mkdtemp(prefix = "unsloth_s3_dataset_")
+    try:
+        os.makedirs(target_dir, exist_ok = True)
+
+        local_files: list[str] = []
+        used_paths: set[str] = set()
+        for key in keys:
+            _raise_if_cancelled(cancel_callback)
+            filename = os.path.basename(key)
+            local_path = _unique_local_path(target_dir, filename, used_paths)
+            download_kwargs = {}
+            if cancel_callback is not None:
+                download_kwargs["Callback"] = lambda _bytes: _raise_if_cancelled(cancel_callback)
+            client.download_file(bucket, key, local_path, **download_kwargs)
+            _raise_if_cancelled(cancel_callback)
+            local_files.append(local_path)
+    except Exception:
+        if owns_temp_dir:
+            shutil.rmtree(target_dir, ignore_errors = True)
+        raise
+
+    logger.info(
+        "Downloaded %d dataset file(s) from s3://%s/%s to %s",
+        len(local_files),
+        bucket,
+        prefix or "",
+        target_dir,
+    )
+    return S3DatasetDownload(
+        files = local_files,
+        temp_dir = target_dir if owns_temp_dir else None,
+    )
+
+
+def download_s3_dataset(
+    s3_config: dict,
+    dest_dir: Optional[str] = None,
+    cancel_callback: Optional[Callable[[], bool]] = None,
+) -> list[str]:
+    download = prepare_s3_dataset_download(
+        s3_config,
+        dest_dir = dest_dir,
+        cancel_callback = cancel_callback,
+    )
+    return download.files

studio/backend/core/training/trainer.py

@@ -2227,6 +2227,7 @@ class UnslothTrainer:
         dataset_slice_start: int = None,
         dataset_slice_end: int = None,
         is_cpt: bool = False,
+        s3_config: dict = None,
     ) -> Optional[tuple]:
         """
         Load and prepare a dataset for training.
@@ -2237,6 +2238,9 @@ class UnslothTrainer:
         Returns (dataset_info, eval_dataset) or None on error; eval_dataset
         may be None if no eval split is available.
         """
+        from core.training.s3_dataset import S3DownloadCancelled
+
+        s3_download = None
         try:
             dataset = None
             eval_dataset = None
@@ -2272,6 +2276,22 @@ class UnslothTrainer:
 
                 return result.dataset
 
+            # S3 datasets are downloaded to a local temp dir and then consumed
+            # through the same local-file path below.
+            if s3_config and not local_datasets:
+                from core.training.s3_dataset import prepare_s3_dataset_download
+
+                self._update_progress(status_message = "Downloading dataset from S3...")
+                s3_download = prepare_s3_dataset_download(
+                    s3_config,
+                    cancel_callback = lambda: self.should_stop,
+                )
+                local_datasets = s3_download.files
+                if self.should_stop:
+                    logger.info("Stopped during S3 download\n")
+                    return None
+                logger.info(f"Downloaded {len(local_datasets)} file(s) from S3\n")
+
             if local_datasets:
                 # Use load_dataset() for an Arrow-backed result; in-memory
                 # Dataset.from_list() has no cache and forces num_proc=1 during
@@ -2539,10 +2559,16 @@ class UnslothTrainer:
 
             return (dataset_info, eval_dataset)
 
+        except S3DownloadCancelled:
+            logger.info("Stopped during S3 download\n")
+            return None
         except Exception as e:
             logger.error(f"Error loading dataset: {e}")
             self._update_progress(error = str(e))
             return None
+        finally:
+            if s3_download is not None:
+                s3_download.cleanup()
 
     def _auto_detect_eval_split_from_hf(
         self, dataset_source: str, subset: str

studio/backend/core/training/training.py

@@ -40,6 +40,34 @@ logger = get_logger(__name__)
 _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$")
 
 
+def _sanitize_db_config(config: dict[str, Any]) -> dict[str, Any]:
+    db_config = {
+        k: v for k, v in config.items() if k not in {"hf_token", "wandb_token", "s3_config"}
+    }
+    s3_config = config.get("s3_config")
+    if hasattr(s3_config, "model_dump"):
+        s3_config = s3_config.model_dump()
+    if isinstance(s3_config, dict) and s3_config:
+        db_config["dataset_source"] = "s3"
+        db_config["s3_dataset"] = {
+            "bucket": s3_config.get("bucket"),
+            "region": s3_config.get("region"),
+            "prefix": s3_config.get("prefix"),
+            "use_iam_role": bool(s3_config.get("use_iam_role")),
+        }
+    return db_config
+
+
+def _s3_dataset_name(s3_dataset: Any) -> Optional[str]:
+    if not isinstance(s3_dataset, dict):
+        return None
+    bucket = s3_dataset.get("bucket")
+    if not bucket:
+        return None
+    prefix = s3_dataset.get("prefix")
+    return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}"
+
+
 def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None:
     """Remove only HF Trainer ``tmp-checkpoint-<step>/`` partials after a cancel.
 
@@ -236,6 +264,7 @@ class TrainingBackend:
             "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"),
             "trust_remote_code": kwargs.get("trust_remote_code", False),
             "gpu_ids": kwargs.get("gpu_ids"),
+            "s3_config": kwargs.get("s3_config"),
         }
 
         # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request.
@@ -309,7 +338,7 @@ class TrainingBackend:
         self._run_finalized = False
         self._db_run_created = False
         self._db_total_steps_set = False
-        self._db_config = {k: v for k, v in config.items() if k not in {"hf_token", "wandb_token"}}
+        self._db_config = _sanitize_db_config(config)
         self._db_started_at = datetime.now(timezone.utc).isoformat()
 
         # Assign subprocess handles after state reset.
@@ -732,8 +761,11 @@ class TrainingBackend:
         try:
             from storage.studio_db import create_run
 
-            dataset_name = self._db_config.get("hf_dataset") or next(
-                iter(self._db_config.get("local_datasets") or []), "unknown"
+            dataset_name = (
+                self._db_config.get("hf_dataset")
+                or next(iter(self._db_config.get("local_datasets") or []), None)
+                or _s3_dataset_name(self._db_config.get("s3_dataset"))
+                or "unknown"
             )
             create_run(
                 id = self.current_job_id,

studio/backend/core/training/worker.py

@@ -1336,6 +1336,32 @@ def _run_mlx_training(event_queue, stop_queue, config):
                 kwargs["message"] = sm
         event_queue.put({"type": event_type, "ts": time.time(), **kwargs})
 
+    _stop_save = [True]
+    _stop_requested = [False]
+    _trainer_ref = [None]
+
+    def _is_stop_requested():
+        return _stop_requested[0]
+
+    def _poll_stop():
+        while True:
+            try:
+                msg = stop_queue.get(timeout = 1.0)
+                if msg and msg.get("type") == "stop":
+                    _stop_save[0] = msg.get("save", True)
+                    _stop_requested[0] = True
+                    trainer = _trainer_ref[0]
+                    if trainer is not None:
+                        trainer.stop_requested = True
+                    return
+            except _queue.Empty:
+                continue
+            except (EOFError, OSError):
+                return
+
+    stop_thread = threading.Thread(target = _poll_stop, daemon = True)
+    stop_thread.start()
+
     _send("status", status_message = "Loading MLX libraries...")
 
     import mlx.core as mx
@@ -1515,6 +1541,26 @@ def _run_mlx_training(event_queue, stop_queue, config):
     elif config.get("local_datasets"):
         dataset = _load_local(config["local_datasets"])
         dataset = _slice(dataset)
+    elif config.get("s3_config"):
+        from core.training.s3_dataset import (
+            S3DownloadCancelled,
+            prepare_s3_dataset_download,
+        )
+
+        _send("status", status_message = "Downloading dataset from S3...")
+        try:
+            s3_download = prepare_s3_dataset_download(
+                config["s3_config"],
+                cancel_callback = _is_stop_requested,
+            )
+            try:
+                dataset = _load_local(s3_download.files)
+            finally:
+                s3_download.cleanup()
+        except S3DownloadCancelled:
+            _send("complete", output_dir = None, status_message = "Training cancelled")
+            return
+        dataset = _slice(dataset)
     else:
         raise ValueError("No dataset specified")
 
@@ -1693,6 +1739,9 @@ def _run_mlx_training(event_queue, stop_queue, config):
             eval_steps = eval_steps_val,
         ),
     )
+    _trainer_ref[0] = trainer
+    if _stop_requested[0]:
+        trainer.stop_requested = True
 
     # Tell the parent eval is configured so the frontend shows the eval chart
     if eval_dataset is not None and eval_steps_val > 0:
@@ -1733,7 +1782,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
             wandb_token = config.get("wandb_token")
             if wandb_token:
                 os.environ["WANDB_API_KEY"] = wandb_token
-            _wandb_sensitive = {"hf_token", "wandb_token"}
+            _wandb_sensitive = {"hf_token", "wandb_token", "s3_config"}
             wandb_run = _wandb.init(
                 project = config.get("wandb_project") or "unsloth-mlx",
                 config = {k: v for k, v in config.items() if k not in _wandb_sensitive},
@@ -1835,26 +1884,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
 
     trainer.add_eval_callback(_on_eval)
 
-    # ── 10. Stop signal polling ──
-    _stop_save = [True]  # mutable so thread can update; [save_flag]
-
-    def _poll_stop():
-        while True:
-            try:
-                msg = stop_queue.get(timeout = 1.0)
-                if msg and msg.get("type") == "stop":
-                    _stop_save[0] = msg.get("save", True)
-                    trainer.stop_requested = True
-                    return
-            except _queue.Empty:
-                continue
-            except (EOFError, OSError):
-                # Safe: pipe permanently broken, no more messages can arrive.
-                return
-
-    stop_thread = threading.Thread(target = _poll_stop, daemon = True)
-    stop_thread.start()
-
     # ── 11. Run training ──
     gc.collect()
     mx.synchronize()
@@ -2546,6 +2575,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
             dataset_slice_start = config.get("dataset_slice_start"),
             dataset_slice_end = config.get("dataset_slice_end"),
             is_cpt = _is_cpt_for_dataset,
+            s3_config = config.get("s3_config"),
         )
 
         if isinstance(dataset_result, tuple):
@@ -3010,20 +3040,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
         subset = config.get("subset") or None
         train_split = config.get("train_split", "train") or "train"
 
-        if hf_dataset and hf_dataset.strip():
-            hf_token = config.get("hf_token", "")
-            hf_token = hf_token if hf_token and hf_token.strip() else None
-            dataset = load_dataset(
-                hf_dataset.strip(),
-                subset,
-                split = train_split,
-                token = hf_token,
-            )
-        elif local_datasets:
-            # Load local file(s) — mirrors the non-embedding pipeline's directory
-            # handling so recipe outputs (parquet-files/) work.
+        def _load_local_embedding_dataset(dataset_paths: list[str]):
             all_files: list[str] = []
-            for dataset_file in local_datasets:
+            for dataset_file in dataset_paths:
                 file_path = (
                     dataset_file
                     if os.path.isabs(dataset_file)
@@ -3053,17 +3072,58 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
                 else:
                     all_files.append(file_path)
 
-            if all_files:
-                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]}")
-                dataset = load_dataset(loader, data_files = all_files, split = "train")
+            if not all_files:
+                raise ValueError("No local dataset files found")
+
+            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]}")
+            return load_dataset(loader, data_files = all_files, split = "train")
+
+        if hf_dataset and hf_dataset.strip():
+            hf_token = config.get("hf_token", "")
+            hf_token = hf_token if hf_token and hf_token.strip() else None
+            dataset = load_dataset(
+                hf_dataset.strip(),
+                subset,
+                split = train_split,
+                token = hf_token,
+            )
+        elif local_datasets:
+            dataset = _load_local_embedding_dataset(local_datasets)
+        elif config.get("s3_config"):
+            from core.training.s3_dataset import (
+                S3DownloadCancelled,
+                prepare_s3_dataset_download,
+            )
+
+            _send_status(event_queue, "Downloading dataset from S3...")
+            s3_download = None
+            try:
+                s3_download = prepare_s3_dataset_download(
+                    config["s3_config"],
+                    cancel_callback = lambda: _should_stop,
+                )
+                dataset = _load_local_embedding_dataset(s3_download.files)
+            except S3DownloadCancelled:
+                event_queue.put(
+                    {
+                        "type": "complete",
+                        "output_dir": None,
+                        "status_message": "Training cancelled",
+                        "ts": time.time(),
+                    }
+                )
+                return
+            finally:
+                if s3_download is not None:
+                    s3_download.cleanup()
         else:
             event_queue.put(
                 {

studio/backend/models/training.py

@@ -29,6 +29,43 @@ _MIN_VISION_IMAGE_SIZE = 256
 _MAX_VISION_IMAGE_SIZE = 2048
 
 
+class S3Config(BaseModel):
+    """S3 bucket configuration for loading datasets from AWS S3"""
+
+    # Accept both snake_case and the frontend's camelCase field names.
+    model_config = ConfigDict(populate_by_name = True)
+
+    bucket: str = Field(..., description = "S3 bucket name")
+    region: str = Field("us-east-1", description = "AWS region")
+    prefix: Optional[str] = Field(None, description = "Optional path prefix within bucket")
+    access_key_id: Optional[str] = Field(
+        None,
+        alias = "accessKeyId",
+        description = "AWS access key ID (optional if using IAM role)",
+    )
+    secret_access_key: Optional[str] = Field(
+        None,
+        alias = "secretAccessKey",
+        description = "AWS secret access key (optional if using IAM role)",
+    )
+    use_iam_role: bool = Field(
+        False,
+        alias = "useIamRole",
+        description = "Use IAM role credentials instead of access keys",
+    )
+
+    @model_validator(mode = "after")
+    def _check_credentials(self) -> "S3Config":
+        # Require either IAM role auth or a full key pair so credentials are
+        # never half-configured.
+        if not self.use_iam_role and not (self.access_key_id and self.secret_access_key):
+            raise ValueError(
+                "s3_config requires either use_iam_role=True or both "
+                "access_key_id and secret_access_key"
+            )
+        return self
+
+
 def _parse_lr(v: Any) -> float:
     """Parse learning_rate as a positive float strictly below _MAX_LR_VALUE."""
     if v is None:
@@ -338,6 +375,12 @@ class TrainingStartRequest(BaseModel):
         description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.",
     )
 
+    # S3 dataset source configuration
+    s3_config: Optional[S3Config] = Field(
+        None,
+        description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.",
+    )
+
     @model_validator(mode = "after")
     def _check_steps_or_epochs(self) -> "TrainingStartRequest":
         # Each accepts 0 as "use the other"; both 0 means nothing to train.

studio/backend/requirements/studio.txt

@@ -17,6 +17,7 @@ structlog>=24.1.0
 diceware
 ddgs
 cryptography>=42.0.0
+boto3>=1.34.0                    # optional: S3 dataset loading
 httpx>=0.27.0
 fastmcp>=3.0.2
 # RAG (knowledge bases, hybrid retrieval). sentence-transformers lives in

studio/backend/routes/training.py

@@ -125,6 +125,17 @@ async def start_training(
 
         backend = get_training_backend()
 
+        # S3 dataset loading needs the optional boto3 dependency. Reject early
+        # with a clear message so credentials are never accepted and then
+        # silently dropped on a host without boto3 installed.
+        if request.s3_config is not None:
+            from core.training.s3_dataset import boto3_available
+            if not boto3_available():
+                raise HTTPException(
+                    status_code = 501,
+                    detail = "S3 dataset loading requires boto3. Install it with: pip install boto3",
+                )
+
         # Check before mutating state.
         if backend.is_training_active():
             existing_job_id: Optional[str] = getattr(backend, "current_job_id", "")
@@ -235,6 +246,7 @@ async def start_training(
             "resume_from_checkpoint": request.resume_from_checkpoint,
             "trust_remote_code": request.trust_remote_code,
             "gpu_ids": request.gpu_ids,
+            "s3_config": request.s3_config.model_dump() if request.s3_config else None,
         }
 
         # Training page has no trust_remote_code toggle; as a safety net consult
@@ -293,6 +305,10 @@ async def start_training(
             error = None,
         )
 
+    except HTTPException:
+        # Deliberate rejections (S3 not implemented, resume validation) must
+        # reach the client with their original status, not a generic 500.
+        raise
     except ValueError as e:
         logger.warning("Rejected training GPU selection: %s", e)
         # Deliberate user-facing GPU-selection validation message.

studio/backend/storage/studio_db.py

@@ -652,7 +652,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
             SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
                    r.ended_at, r.total_steps, r.final_step, r.final_loss,
                    r.output_dir, r.duration_seconds, r.error_message,
-                   r.loss_sparkline, r.display_name,
+                   r.loss_sparkline, r.display_name, r.config_json,
                    CASE
                        WHEN r.status = 'stopped'
                             AND r.output_dir IS NOT NULL

studio/backend/tests/test_s3_dataset.py

@@ -0,0 +1,274 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Tests for the S3 dataset loader (core.training.s3_dataset).
+
+boto3 is optional and may be absent in CI, so the S3 client is mocked: a fake
+client provides a paginator over a synthetic bucket listing and writes files on
+download_file. No network or real AWS credentials are involved.
+"""
+
+import importlib.util
+import os
+from pathlib import Path
+
+import pytest
+
+# Load the modules under test directly by path. Importing them through their
+# packages (core.training / models) would execute heavy package __init__ chains
+# (structlog, torch, …) that aren't needed for these unit tests.
+_BACKEND = Path(__file__).resolve().parents[1]
+
+
+def _load(mod_name, rel_path):
+    spec = importlib.util.spec_from_file_location(mod_name, _BACKEND / rel_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+s3_dataset = _load("s3_dataset", "core/training/s3_dataset.py")
+S3Config = _load("models_training_s3", "models/training.py").S3Config
+
+
+class _FakePaginator:
+    def __init__(self, keys):
+        self._keys = keys
+
+    def paginate(self, **kwargs):
+        prefix = kwargs.get("Prefix")
+        contents = [{"Key": k} for k in self._keys if prefix is None or k.startswith(prefix)]
+        # Emit in two pages to exercise pagination handling.
+        mid = len(contents) // 2
+        yield {"Contents": contents[:mid]}
+        yield {"Contents": contents[mid:]}
+
+
+class _FakeS3Client:
+    def __init__(self, keys):
+        self._keys = keys
+        self.downloaded = []
+
+    def get_paginator(self, name):
+        assert name == "list_objects_v2"
+        return _FakePaginator(self._keys)
+
+    def download_file(self, bucket, key, local_path, **kwargs):
+        self.downloaded.append((bucket, key, local_path))
+        callback = kwargs.get("Callback")
+        if callback is not None:
+            callback(1)
+        with open(local_path, "w", encoding = "utf-8") as f:
+            f.write(f"content-of:{key}")
+
+
+@pytest.fixture
+def fake_client(monkeypatch):
+    """Force boto3_available True and stub the client builder."""
+    keys = [
+        "datasets/train.parquet",
+        "datasets/extra.parquet",
+        "datasets/notes.txt",  # filtered out (unsupported)
+        "datasets/subdir/",  # directory placeholder, skipped
+        "other/ignore.parquet",  # filtered out by prefix
+    ]
+    client = _FakeS3Client(keys)
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+    return client
+
+
+def _cfg(**overrides):
+    base = {
+        "bucket": "my-bucket",
+        "region": "us-east-1",
+        "prefix": "datasets/",
+        "access_key_id": "AKIA_TEST",
+        "secret_access_key": "secret",
+        "use_iam_role": False,
+    }
+    base.update(overrides)
+    return base
+
+
+def test_downloads_only_supported_files_under_prefix(fake_client, tmp_path):
+    files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
+    names = sorted(os.path.basename(f) for f in files)
+    # txt is unsupported, the directory placeholder is skipped, and the
+    # "other/" key is excluded by the prefix filter.
+    assert names == ["extra.parquet", "train.parquet"]
+    for f in files:
+        assert os.path.exists(f)
+
+
+def test_allows_json_and_jsonl_family(monkeypatch, tmp_path):
+    client = _FakeS3Client(["datasets/train.json", "datasets/extra.jsonl"])
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+
+    files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
+
+    assert sorted(os.path.basename(f) for f in files) == ["extra.jsonl", "train.json"]
+
+
+def test_ignores_common_json_metadata_files(monkeypatch, tmp_path):
+    client = _FakeS3Client(
+        [
+            "datasets/train.parquet",
+            "datasets/schema.json",
+            "datasets/metadata.json",
+            "datasets/dataset_info.json",
+        ]
+    )
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+
+    files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
+
+    assert [os.path.basename(f) for f in files] == ["train.parquet"]
+
+
+def test_raises_when_prefix_contains_mixed_formats(monkeypatch, tmp_path):
+    client = _FakeS3Client(["datasets/train.parquet", "datasets/stray.csv"])
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+
+    with pytest.raises(ValueError, match = "mixed dataset formats"):
+        s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
+
+    assert client.downloaded == []
+
+
+def test_raises_when_no_supported_files(monkeypatch, tmp_path):
+    client = _FakeS3Client(["datasets/readme.txt"])
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+    with pytest.raises(ValueError, match = "No supported dataset files"):
+        s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
+
+
+def test_raises_when_boto3_missing(monkeypatch, tmp_path):
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: False)
+    with pytest.raises(RuntimeError, match = "requires boto3"):
+        s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
+
+
+def test_basename_collisions_are_disambiguated(monkeypatch, tmp_path):
+    # Two keys share a basename under different sub-prefixes.
+    client = _FakeS3Client(["datasets/a/train.parquet", "datasets/b/train.parquet"])
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+    files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
+    assert len(files) == 2
+    assert len(set(files)) == 2  # no overwrite
+
+
+def test_basename_collision_skips_existing_generated_suffix(monkeypatch, tmp_path):
+    client = _FakeS3Client(
+        [
+            "datasets/a/train.parquet",
+            "datasets/b/train_1.parquet",
+            "datasets/c/train.parquet",
+        ]
+    )
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+
+    files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
+
+    assert [os.path.basename(f) for f in files] == [
+        "train.parquet",
+        "train_1.parquet",
+        "train_2.parquet",
+    ]
+    assert len(set(files)) == 3
+    assert (tmp_path / "train_1.parquet").read_text(encoding = "utf-8") == (
+        "content-of:datasets/b/train_1.parquet"
+    )
+    assert (tmp_path / "train_2.parquet").read_text(encoding = "utf-8") == (
+        "content-of:datasets/c/train.parquet"
+    )
+
+
+def test_download_handle_cleans_owned_temp_dir(monkeypatch, tmp_path):
+    target_dir = tmp_path / "owned-download"
+    client = _FakeS3Client(["datasets/train.parquet"])
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+    monkeypatch.setattr(s3_dataset.tempfile, "mkdtemp", lambda prefix: str(target_dir))
+
+    download = s3_dataset.prepare_s3_dataset_download(_cfg())
+
+    assert target_dir.exists()
+    assert download.files == [str(target_dir / "train.parquet")]
+    download.cleanup()
+    assert not target_dir.exists()
+
+
+def test_dest_dir_is_not_removed_by_cleanup(monkeypatch, tmp_path):
+    client = _FakeS3Client(["datasets/train.parquet"])
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+
+    download = s3_dataset.prepare_s3_dataset_download(_cfg(), dest_dir = str(tmp_path))
+
+    download.cleanup()
+    assert tmp_path.exists()
+    assert (tmp_path / "train.parquet").exists()
+
+
+def test_cancel_callback_aborts_and_removes_temp_dir(monkeypatch, tmp_path):
+    target_dir = tmp_path / "cancelled-download"
+    client = _FakeS3Client(["datasets/train.parquet"])
+    monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
+    monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
+    monkeypatch.setattr(s3_dataset.tempfile, "mkdtemp", lambda prefix: str(target_dir))
+    calls = 0
+
+    def cancel_after_download_starts():
+        nonlocal calls
+        calls += 1
+        return calls >= 4
+
+    with pytest.raises(s3_dataset.S3DownloadCancelled):
+        s3_dataset.prepare_s3_dataset_download(
+            _cfg(),
+            cancel_callback = cancel_after_download_starts,
+        )
+
+    assert not target_dir.exists()
+
+
+# ── S3Config model (camelCase aliases + credential validation) ──
+
+
+def test_s3config_accepts_camelcase_aliases():
+    cfg = S3Config.model_validate(
+        {
+            "bucket": "b",
+            "region": "eu-west-1",
+            "accessKeyId": "AKIA",
+            "secretAccessKey": "shh",
+        }
+    )
+    assert cfg.access_key_id == "AKIA"
+    assert cfg.secret_access_key == "shh"
+    # model_dump() yields snake_case for the loader.
+    assert cfg.model_dump()["access_key_id"] == "AKIA"
+
+
+def test_s3config_accepts_snake_case():
+    cfg = S3Config.model_validate(
+        {"bucket": "b", "access_key_id": "AKIA", "secret_access_key": "shh"}
+    )
+    assert cfg.access_key_id == "AKIA"
+
+
+def test_s3config_requires_credentials_or_iam():
+    with pytest.raises(ValueError):
+        S3Config.model_validate({"bucket": "b"})
+
+
+def test_s3config_iam_role_needs_no_keys():
+    cfg = S3Config.model_validate({"bucket": "b", "useIamRole": True})
+    assert cfg.use_iam_role is True

studio/backend/tests/test_training_resume.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Regression tests for resumable training run eligibility."""
+
+import importlib.util
+import json
+from pathlib import Path
+
+
+_BACKEND = Path(__file__).resolve().parents[1]
+
+
+def _load_resume_module():
+    spec = importlib.util.spec_from_file_location(
+        "training_resume_under_test",
+        _BACKEND / "core" / "training" / "resume.py",
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+resume = _load_resume_module()
+
+
+def _stopped_run(**overrides):
+    run = {
+        "status": "stopped",
+        "final_step": 5,
+        "total_steps": 10,
+        "output_dir": "/tmp/unsloth-output",
+        "resumed_later": False,
+        "config_json": json.dumps({"hf_dataset": "org/dataset"}),
+    }
+    run.update(overrides)
+    return run
+
+
+def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch):
+    monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
+
+    assert resume.can_resume_run(_stopped_run()) is True
+
+
+def test_can_resume_run_rejects_s3_dataset_source(monkeypatch):
+    monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
+
+    run = _stopped_run(
+        config_json = json.dumps(
+            {
+                "dataset_source": "s3",
+                "s3_dataset": {
+                    "bucket": "training-data",
+                    "prefix": "datasets/",
+                    "region": "us-east-1",
+                    "use_iam_role": True,
+                },
+            }
+        )
+    )
+
+    assert resume.can_resume_run(run) is False
+
+
+def test_can_resume_run_rejects_s3_metadata_marker(monkeypatch):
+    monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
+
+    run = _stopped_run(config_json = json.dumps({"s3_dataset": {"bucket": "training-data"}}))
+
+    assert resume.can_resume_run(run) is False
+
+
+def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path):
+    from storage import studio_db
+
+    monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+    monkeypatch.setattr(studio_db, "_schema_ready", False)
+    config_json = json.dumps({"dataset_source": "s3", "s3_dataset": {"bucket": "training-data"}})
+
+    studio_db.create_run(
+        id = "run-s3",
+        model_name = "unsloth/test-model",
+        dataset_name = "s3://training-data",
+        config_json = config_json,
+        started_at = "2026-01-01T00:00:00Z",
+        total_steps = 10,
+    )
+
+    result = studio_db.list_runs()
+
+    assert result["runs"][0]["config_json"] == config_json

studio/frontend/src/config/training.ts

@@ -141,6 +141,7 @@ export const DEFAULT_HYPERPARAMS = {
   finetuneAttentionModules: true,
   finetuneMLPModules: true,
   targetModules: TARGET_MODULES,
+  s3Config: null as import("@/types/training").S3Config | null,
 };
 
 export const MODEL_TYPE_TO_HF_TASK: Record<ModelType, PipelineType> = {

studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx

@@ -17,6 +17,7 @@ import { useTrainingActions, useTrainingConfigStore } from "@/features/training"
 import { checkDatasetFormat } from "@/features/training/api/datasets-api";
 import { isRawTextDatasetFormat } from "@/features/training/lib/training-methods";
 import type { CheckFormatResponse } from "@/features/training/types/datasets";
+import type { DatasetSource } from "@/types/training";
 import {
   AlertCircleIcon,
   CheckmarkCircle02Icon,
@@ -45,7 +46,7 @@ type DatasetPreviewDialogProps = {
   open: boolean;
   onOpenChange: (open: boolean) => void;
   datasetName: string | null;
-  datasetSource?: "huggingface" | "upload";
+  datasetSource?: DatasetSource;
   hfToken: string | null;
   datasetSubset?: string | null;
   datasetSplit?: string | null;

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

@@ -78,6 +78,7 @@ import {
 import { useShallow } from "zustand/react/shallow";
 import { DocumentUploadRedirectDialog } from "./document-upload-redirect-dialog";
 import { translate, useT } from "@/i18n";
+import { S3ConfigForm } from "./s3-config-form";
 
 const TRAINING_UPLOAD_EXTENSIONS = [
   ".csv",
@@ -154,6 +155,7 @@ export function DatasetSection() {
     datasetSource,
     selectHfDataset,
     selectLocalDataset,
+    selectS3Source,
     datasetFormat,
     setDatasetFormat,
     datasetSubset,
@@ -167,6 +169,8 @@ export function DatasetSection() {
     setUploadedEvalFile,
     hfToken,
     modelType,
+    isVisionModel,
+    isAudioModel,
     datasetSliceStart,
     setDatasetSliceStart,
     datasetSliceEnd,
@@ -177,6 +181,7 @@ export function DatasetSection() {
       datasetSource: s.datasetSource,
       selectHfDataset: s.selectHfDataset,
       selectLocalDataset: s.selectLocalDataset,
+      selectS3Source: s.selectS3Source,
       datasetFormat: s.datasetFormat,
       setDatasetFormat: s.setDatasetFormat,
       datasetSubset: s.datasetSubset,
@@ -190,6 +195,8 @@ export function DatasetSection() {
       setUploadedEvalFile: s.setUploadedEvalFile,
       hfToken: s.hfToken,
       modelType: s.modelType,
+      isVisionModel: s.isVisionModel,
+      isAudioModel: s.isAudioModel,
       datasetSliceStart: s.datasetSliceStart,
       setDatasetSliceStart: s.setDatasetSliceStart,
       datasetSliceEnd: s.datasetSliceEnd,
@@ -292,6 +299,11 @@ export function DatasetSection() {
   }
 
   const effectiveModelType = modelType ?? "text";
+  const isMultimodalModel =
+    effectiveModelType === "vision" ||
+    effectiveModelType === "audio" ||
+    isVisionModel ||
+    isAudioModel;
 
   const {
     results: hfResults,
@@ -373,6 +385,12 @@ export function DatasetSection() {
     selectLocalDataset,
   ]);
 
+  useEffect(() => {
+    if (datasetSource === "s3" && isMultimodalModel) {
+      selectHfDataset(dataset);
+    }
+  }, [dataset, datasetSource, isMultimodalModel, selectHfDataset]);
+
   const activeSourceTab = datasetSource === "upload" ? "local" : "huggingface";
   const comboboxItems =
     pickerTab === "huggingface" ? hfResultIds : localResultIds;
@@ -572,6 +590,35 @@ export function DatasetSection() {
         }`}
       >
         <div className="flex min-w-0 flex-col gap-4">
+          <Tabs
+            value={datasetSource}
+            onValueChange={(value) => {
+              if (value === datasetSource) return;
+              if (value === "huggingface") {
+                selectHfDataset(dataset);
+              } else if (value === "upload") {
+                selectLocalDataset(uploadedFile);
+              } else if (value === "s3") {
+                if (isMultimodalModel) return;
+                selectS3Source();
+              }
+            }}
+            className="w-full"
+          >
+            <TabsList className="w-full">
+              <TabsTrigger value="huggingface">Hugging Face</TabsTrigger>
+              <TabsTrigger value="upload">
+                {t("studio.dataset.localTab")}
+              </TabsTrigger>
+              {!isMultimodalModel && (
+                <TabsTrigger value="s3">Amazon S3</TabsTrigger>
+              )}
+            </TabsList>
+          </Tabs>
+
+          {datasetSource === "s3" && <S3ConfigForm />}
+
+          {datasetSource !== "s3" && (
           <div className="flex min-w-0 flex-col gap-2">
             <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
               {t("studio.dataset.chooseDataset")}
@@ -842,8 +889,10 @@ export function DatasetSection() {
               </p>
             )}
           </div>
+          )}
 
-          {isHfDatasetSelected ? (
+          {datasetSource !== "s3" &&
+            (isHfDatasetSelected ? (
             <HfDatasetSubsetSplitSelectors
               variant="studio"
               enabled={true}
@@ -918,7 +967,7 @@ export function DatasetSection() {
                 </div>
               </div>
             </div>
-          ) : null}
+          ) : null)}
 
           {datasetSource === "upload" && uploadedFile && (
             <div className="rounded-lg border bg-muted/20 px-3.5 py-3">
@@ -1102,111 +1151,115 @@ export function DatasetSection() {
             </CollapsibleContent>
           </Collapsible>
 
-          <div className="flex flex-col gap-3">
-            {selectedDatasetName ? (
-              <div className="flex items-center gap-3 rounded-lg border bg-muted/40 px-3.5 py-3">
-                <div className="rounded-md bg-indigo-500/10 p-1.5">
-                  <HugeiconsIcon
-                    icon={FileAttachmentIcon}
-                    className="size-4 text-indigo-500"
-                  />
-                </div>
-                <div className="flex-1 min-w-0">
-                  <p className="font-mono text-sm font-medium truncate">
-                    {datasetSource === "upload"
-                      ? (selectedLocalDataset?.label ??
-                        deriveLocalDatasetName(selectedDatasetName))
-                      : selectedDatasetName}
-                  </p>
-                  <p className="text-[10px] text-muted-foreground">
-                    {datasetSource === "upload" ? (
-                      uploadedFile ? (
-                        <>
-                          {t("studio.dataset.localDataset")}
-                          {selectedLocalRows != null
-                            ? t("studio.dataset.localDatasetRows", {
-                                count: selectedLocalRows.toLocaleString(),
-                              })
-                            : ""}
-                        </>
+          {datasetSource !== "s3" && (
+            <div className="flex flex-col gap-3">
+              {selectedDatasetName ? (
+                <div className="flex items-center gap-3 rounded-lg border bg-muted/40 px-3.5 py-3">
+                  <div className="rounded-md bg-indigo-500/10 p-1.5">
+                    <HugeiconsIcon
+                      icon={FileAttachmentIcon}
+                      className="size-4 text-indigo-500"
+                    />
+                  </div>
+                  <div className="flex-1 min-w-0">
+                    <p className="font-mono text-sm font-medium truncate">
+                      {datasetSource === "upload"
+                        ? (selectedLocalDataset?.label ??
+                          deriveLocalDatasetName(selectedDatasetName))
+                        : selectedDatasetName}
+                    </p>
+                    <p className="text-[10px] text-muted-foreground">
+                      {datasetSource === "upload" ? (
+                        uploadedFile ? (
+                          <>
+                            {t("studio.dataset.localDataset")}
+                            {selectedLocalRows != null
+                              ? t("studio.dataset.localDatasetRows", {
+                                  count: selectedLocalRows.toLocaleString(),
+                                })
+                              : ""}
+                          </>
+                        ) : (
+                          t("studio.dataset.localDataset")
+                        )
                       ) : (
-                        t("studio.dataset.localDataset")
-                      )
-                    ) : (
-                      <>
-                        {t("studio.dataset.huggingFaceDataset")}
-                        {datasetSubset && ` / ${datasetSubset}`}
-                        {datasetSplit && ` / ${datasetSplit}`}
-                      </>
-                    )}
-                  </p>
+                        <>
+                          {t("studio.dataset.huggingFaceDataset")}
+                          {datasetSubset && ` / ${datasetSubset}`}
+                          {datasetSplit && ` / ${datasetSplit}`}
+                        </>
+                      )}
+                    </p>
+                  </div>
+                  <Button
+                    variant="ghost"
+                    size="sm"
+                    className="shrink-0 text-xs"
+                    onClick={() => clearSelectionForTab(activeSourceTab)}
+                  >
+                    {t("studio.dataset.clear")}
+                  </Button>
                 </div>
-                <Button
-                  variant="ghost"
-                  size="sm"
-                  className="shrink-0 text-xs"
-                  onClick={() => clearSelectionForTab(activeSourceTab)}
+              ) : (
+                <button
+                  type="button"
+                  className={`flex w-full cursor-pointer items-center gap-3 rounded-lg border border-dashed px-3.5 py-3 text-left transition-colors ${
+                    isDatasetDragOver
+                      ? "border-indigo-500/70 bg-indigo-500/10"
+                      : "border-border bg-muted/20 hover:border-indigo-500/50 hover:bg-indigo-500/5"
+                  }`}
+                  disabled={isUploading}
+                  onClick={handleUploadButtonClick}
+                  onDrop={handleDatasetDrop}
+                  onDragOver={handleDatasetDragOver}
+                  onDragLeave={handleDatasetDragLeave}
                 >
-                  {t("studio.dataset.clear")}
+                  <HugeiconsIcon
+                    icon={CloudUploadIcon}
+                    className="pointer-events-none size-4 shrink-0 text-indigo-500"
+                  />
+                  <span className="pointer-events-none min-w-0">
+                    <span className="block text-xs font-medium text-foreground">
+                      {t("studio.dataset.dropFileOrClick")}
+                    </span>
+                    <span className="mt-0.5 block truncate text-[10px] text-muted-foreground">
+                      {TRAINING_DATASET_UPLOAD_LABEL} · up to{" "}
+                      {uploadLimitLabel}; {DOCUMENT_REDIRECT_LABEL}
+                    </span>
+                  </span>
+                </button>
+              )}
+
+              <div className="grid grid-cols-2 gap-2">
+                <Button
+                  variant="outline"
+                  size="sm"
+                  className="cursor-pointer gap-1.5"
+                  disabled={isUploading}
+                  onClick={handleUploadButtonClick}
+                >
+                  {isUploading ? (
+                    <Spinner className="size-3.5" />
+                  ) : (
+                    <HugeiconsIcon icon={CloudUploadIcon} className="size-3.5" />
+                  )}
+                  {isUploading
+                    ? t("studio.dataset.uploading")
+                    : t("studio.dataset.upload")}
+                </Button>
+                <Button
+                  variant="outline"
+                  size="sm"
+                  className="cursor-pointer gap-1.5"
+                  disabled={!selectedDatasetName}
+                  onClick={() => openPreview()}
+                >
+                  <HugeiconsIcon icon={ViewIcon} className="size-3.5" />
+                  {t("studio.dataset.viewDataset")}
                 </Button>
               </div>
-            ) : (
-              <button
-                type="button"
-                className={`flex w-full cursor-pointer items-center gap-3 rounded-lg border border-dashed px-3.5 py-3 text-left transition-colors ${
-                  isDatasetDragOver
-                    ? "border-indigo-500/70 bg-indigo-500/10"
-                    : "border-border bg-muted/20 hover:border-indigo-500/50 hover:bg-indigo-500/5"
-                }`}
-                disabled={isUploading}
-                onClick={handleUploadButtonClick}
-                onDrop={handleDatasetDrop}
-                onDragOver={handleDatasetDragOver}
-                onDragLeave={handleDatasetDragLeave}
-              >
-                <HugeiconsIcon
-                  icon={CloudUploadIcon}
-                  className="pointer-events-none size-4 shrink-0 text-indigo-500"
-                />
-                <span className="pointer-events-none min-w-0">
-                  <span className="block text-xs font-medium text-foreground">
-                    {t("studio.dataset.dropFileOrClick")}
-                  </span>
-                  <span className="mt-0.5 block truncate text-[10px] text-muted-foreground">
-                    {TRAINING_DATASET_UPLOAD_LABEL} · up to{" "}
-                    {uploadLimitLabel}; {DOCUMENT_REDIRECT_LABEL}
-                  </span>
-                </span>
-              </button>
-            )}
-
-            <div className="grid grid-cols-2 gap-2">
-              <Button
-                variant="outline"
-                size="sm"
-                className="cursor-pointer gap-1.5"
-                disabled={isUploading}
-                onClick={handleUploadButtonClick}
-              >
-                {isUploading ? (
-                  <Spinner className="size-3.5" />
-                ) : (
-                  <HugeiconsIcon icon={CloudUploadIcon} className="size-3.5" />
-                )}
-                {isUploading ? t("studio.dataset.uploading") : t("studio.dataset.upload")}
-              </Button>
-              <Button
-                variant="outline"
-                size="sm"
-                className="cursor-pointer gap-1.5"
-                disabled={!selectedDatasetName}
-                onClick={() => openPreview()}
-              >
-                <HugeiconsIcon icon={ViewIcon} className="size-3.5" />
-                {t("studio.dataset.viewDataset")}
-              </Button>
             </div>
-          </div>
+          )}
           <input
             ref={fileInputRef}
             type="file"

studio/frontend/src/features/studio/sections/s3-config-form.tsx

@@ -0,0 +1,142 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { Input } from "@/components/ui/input";
+import { Label } from "@/components/ui/label";
+import { Switch } from "@/components/ui/switch";
+import { useTrainingConfigStore } from "@/features/training";
+import { useT } from "@/i18n";
+import type { S3Config } from "@/types/training";
+import { useShallow } from "zustand/react/shallow";
+
+const DEFAULT_S3_CONFIG: S3Config = {
+  bucket: "",
+  region: "us-east-1",
+  prefix: "",
+  accessKeyId: "",
+  secretAccessKey: "",
+  useIamRole: false,
+};
+
+/**
+ * Inline S3 dataset configuration card. Shown in the dataset section when the
+ * selected source is "s3"; reads and writes the shared training-config store.
+ */
+export function S3ConfigForm() {
+  const t = useT();
+  const { s3Config, setS3Config } = useTrainingConfigStore(
+    useShallow((s) => ({
+      s3Config: s.s3Config,
+      setS3Config: s.setS3Config,
+    })),
+  );
+
+  const config = s3Config ?? DEFAULT_S3_CONFIG;
+
+  const update = (patch: Partial<S3Config>) => {
+    setS3Config({ ...config, ...patch });
+  };
+
+  const handleIamRoleChange = (useIamRole: boolean) => {
+    if (useIamRole) {
+      update({ useIamRole, accessKeyId: "", secretAccessKey: "" });
+      return;
+    }
+    update({ useIamRole });
+  };
+
+  return (
+    <div className="flex min-w-0 flex-col gap-3 rounded-lg border bg-muted/20 px-3.5 py-3">
+      <div>
+        <p className="text-xs font-medium text-foreground">
+          {t("studio.dataset.s3.title")}
+        </p>
+        <p className="text-[10px] text-muted-foreground/80">
+          {t("studio.dataset.s3.description")}
+        </p>
+      </div>
+
+      <div className="flex min-w-0 flex-col gap-1">
+        <Label htmlFor="s3-bucket" className="text-xs text-muted-foreground">
+          {t("studio.dataset.s3.bucket")}
+        </Label>
+        <Input
+          id="s3-bucket"
+          value={config.bucket}
+          onChange={(e) => update({ bucket: e.target.value })}
+          placeholder={t("studio.dataset.s3.bucketPlaceholder")}
+        />
+      </div>
+
+      <div className="flex min-w-0 gap-2">
+        <div className="flex min-w-0 flex-1 flex-col gap-1">
+          <Label htmlFor="s3-region" className="text-xs text-muted-foreground">
+            {t("studio.dataset.s3.region")}
+          </Label>
+          <Input
+            id="s3-region"
+            value={config.region}
+            onChange={(e) => update({ region: e.target.value })}
+            placeholder={t("studio.dataset.s3.regionPlaceholder")}
+          />
+        </div>
+        <div className="flex min-w-0 flex-1 flex-col gap-1">
+          <Label htmlFor="s3-prefix" className="text-xs text-muted-foreground">
+            {t("studio.dataset.s3.prefix")}
+          </Label>
+          <Input
+            id="s3-prefix"
+            value={config.prefix ?? ""}
+            onChange={(e) => update({ prefix: e.target.value })}
+            placeholder={t("studio.dataset.s3.prefixPlaceholder")}
+          />
+        </div>
+      </div>
+
+      <div className="flex items-center justify-between gap-2">
+        <Label htmlFor="s3-iam" className="text-xs text-muted-foreground">
+          {t("studio.dataset.s3.useIamRole")}
+        </Label>
+        <Switch
+          id="s3-iam"
+          checked={config.useIamRole ?? false}
+          onCheckedChange={handleIamRoleChange}
+        />
+      </div>
+
+      {!config.useIamRole && (
+        <>
+          <div className="flex min-w-0 flex-col gap-1">
+            <Label
+              htmlFor="s3-access-key"
+              className="text-xs text-muted-foreground"
+            >
+              {t("studio.dataset.s3.accessKeyId")}
+            </Label>
+            <Input
+              id="s3-access-key"
+              value={config.accessKeyId ?? ""}
+              onChange={(e) => update({ accessKeyId: e.target.value })}
+              placeholder={t("studio.dataset.s3.accessKeyIdPlaceholder")}
+            />
+          </div>
+          <div className="flex min-w-0 flex-col gap-1">
+            <Label
+              htmlFor="s3-secret-key"
+              className="text-xs text-muted-foreground"
+            >
+              {t("studio.dataset.s3.secretAccessKey")}
+            </Label>
+            <Input
+              id="s3-secret-key"
+              type="password"
+              value={config.secretAccessKey ?? ""}
+              onChange={(e) => update({ secretAccessKey: e.target.value })}
+              placeholder={t("studio.dataset.s3.secretAccessKeyPlaceholder")}
+            />
+          </div>
+        </>
+      )}
+    </div>
+  );
+}

studio/frontend/src/features/training/api/mappers.ts

@@ -17,6 +17,22 @@ function parseSliceValue(value: string | null): number | null {
   return num;
 }
 
+function buildS3PayloadConfig(config: TrainingConfigState) {
+  const s3 = config.datasetSource === "s3" ? config.s3Config : null;
+  if (!s3) {
+    return null;
+  }
+  if (s3.useIamRole) {
+    return {
+      bucket: s3.bucket,
+      region: s3.region,
+      prefix: s3.prefix,
+      useIamRole: s3.useIamRole,
+    };
+  }
+  return s3;
+}
+
 export function buildTrainingStartPayload(
   config: TrainingConfigState,
 ): TrainingStartRequest {
@@ -36,6 +52,7 @@ export function buildTrainingStartPayload(
     config.datasetSource === "upload" && config.uploadedFile
       ? [config.uploadedFile]
       : [];
+  const s3Config = buildS3PayloadConfig(config);
   let customFormatMapping: Record<string, unknown> | undefined =
     Object.keys(config.datasetManualMapping).length > 0
       ? { ...config.datasetManualMapping }
@@ -76,6 +93,7 @@ export function buildTrainingStartPayload(
       config.datasetSource === "upload" && config.uploadedEvalFile
         ? [config.uploadedEvalFile]
         : [],
+    s3_config: s3Config,
     format_type: config.datasetFormat,
     custom_format_mapping: customFormatMapping,
     num_epochs: config.epochs,

studio/frontend/src/features/training/lib/validation.ts

@@ -8,6 +8,31 @@ export interface StartValidationResult {
   message: string | null;
 }
 
+export function validateS3Source(
+  config: TrainingConfigState,
+): StartValidationResult {
+  if (
+    config.modelType === "vision" ||
+    config.modelType === "audio" ||
+    config.isVisionModel ||
+    config.isAudioModel
+  ) {
+    return {
+      ok: false,
+      message: "S3 datasets are not supported for vision or audio training yet.",
+    };
+  }
+  const s3 = config.s3Config;
+  if (!s3 || !s3.bucket.trim()) {
+    return { ok: false, message: "Enter an S3 bucket name first." };
+  }
+  const hasKeys = Boolean(s3.accessKeyId && s3.secretAccessKey);
+  if (!s3.useIamRole && !hasKeys) {
+    return { ok: false, message: "Provide S3 access keys or enable IAM role." };
+  }
+  return { ok: true, message: null };
+}
+
 export function validateTrainingConfig(
   config: TrainingConfigState,
 ): StartValidationResult {
@@ -23,10 +48,11 @@ export function validateTrainingConfig(
     if (!config.uploadedFile) {
       return { ok: false, message: "Select a local dataset first." };
     }
+  } else if (config.datasetSource === "s3") {
+    return validateS3Source(config);
   } else {
     return { ok: false, message: "Unsupported dataset source." };
   }
 
-
   return { ok: true, message: null };
 }

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

@@ -12,6 +12,7 @@ import { checkDatasetFormat } from "../api/datasets-api";
 import { checkVisionModel, getModelConfig } from "../api/models-api";
 import { mapBackendModelConfigToTrainingPatch } from "../lib/model-defaults";
 import { isRawTextDatasetFormat } from "../lib/training-methods";
+import { validateS3Source } from "../lib/validation";
 import type { BackendModelConfig } from "../api/models-api";
 import type { TrainingConfigState, TrainingConfigStore } from "../types/config";
 
@@ -124,6 +125,7 @@ const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set
   "isDatasetAudio",
   "trainOnCompletions",
   "maxPositionEmbeddings",
+  "s3Config",
 ]);
 
 function partializePersistedState(
@@ -148,9 +150,13 @@ function canProceedForStep(state: TrainingConfigState): boolean {
     case 2:
       return state.selectedModel !== null;
     case 3:
-      return state.datasetSource === "upload"
-        ? state.uploadedFile !== null
-        : state.dataset !== null;
+      if (state.datasetSource === "upload") {
+        return state.uploadedFile !== null;
+      }
+      if (state.datasetSource === "s3") {
+        return validateS3Source(state).ok;
+      }
+      return state.dataset !== null;
     case 4:
     case 5:
       return true;
@@ -569,6 +575,17 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             runDatasetCheck(uploadedFile, "train");
           }
         },
+        selectS3Source: () => {
+          _datasetCheckController?.abort();
+          _datasetCheckController = null;
+          _trainOnCompletionsManuallySet = false;
+          set({
+            datasetSource: "s3",
+            dataset: null,
+            uploadedFile: null,
+            ...resetDatasetState(),
+          });
+        },
         setDatasetFormat: (datasetFormat) =>
           set((state) => {
             if (state.trainingMethod === "cpt") {
@@ -746,6 +763,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         setFinetuneMLPModules: (finetuneMLPModules) =>
           set({ finetuneMLPModules }),
         setTargetModules: (targetModules) => set({ targetModules }),
+        setS3Config: (s3Config) => set({ s3Config }),
         canProceed: () => canProceedForStep(get()),
         reset: () => {
           _trainOnCompletionsManuallySet = false;

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

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
+import type { S3Config } from "@/types/training";
+
 export interface TrainingStartRequest {
   model_name: string;
   training_type: string;
@@ -18,6 +20,8 @@ export interface TrainingStartRequest {
   dataset_slice_end: number | null;
   local_datasets: string[];
   local_eval_datasets: string[];
+  /** S3 bucket configuration; only sent when the dataset source is "s3". */
+  s3_config?: S3Config | null;
   format_type: string;
   custom_format_mapping?: Record<string, unknown> | null;
   num_epochs: number;

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

@@ -6,6 +6,7 @@ import type {
   DatasetSource,
   GradientCheckpointing,
   ModelType,
+  S3Config,
   StepNumber,
   TrainingMethod,
 } from "@/types/training";
@@ -83,6 +84,7 @@ export interface TrainingConfigState {
   targetModules: string[];
   maxPositionEmbeddings: number | null;
   visionImageSize: number | null;
+  s3Config: S3Config | null;
 }
 
 export interface TrainingConfigActions {
@@ -98,6 +100,7 @@ export interface TrainingConfigActions {
   setDatasetSource: (source: DatasetSource) => void;
   selectHfDataset: (dataset: string | null) => void;
   selectLocalDataset: (file: string | null) => void;
+  selectS3Source: () => void;
   setDatasetFormat: (format: DatasetFormat) => void;
   setDataset: (dataset: string | null) => void;
   setDatasetSubset: (subset: string | null) => void;
@@ -147,6 +150,7 @@ export interface TrainingConfigActions {
   setFinetuneAttentionModules: (value: boolean) => void;
   setFinetuneMLPModules: (value: boolean) => void;
   setTargetModules: (value: string[]) => void;
+  setS3Config: (value: S3Config | null) => void;
   canProceed: () => boolean;
   reset: () => void;
   resetToModelDefaults: () => void;

studio/frontend/src/i18n/locales/en.ts

@@ -506,6 +506,28 @@ export const en = {
       preview: "Preview dataset",
       split: "Split",
       subset: "Subset",
+      s3: {
+        title: "S3 Configuration",
+        description: "Load .parquet, .json, .jsonl, or .csv datasets from Amazon S3",
+        bucket: "Bucket Name",
+        bucketPlaceholder: "my-training-data-bucket",
+        region: "AWS Region",
+        regionPlaceholder: "us-east-1",
+        prefix: "Path Prefix",
+        prefixPlaceholder: "datasets/whisper/",
+        prefixTooltip: "Optional path within the bucket to your dataset files",
+        accessKeyId: "Access Key ID",
+        accessKeyIdPlaceholder: "AKIAIOSFODNN7EXAMPLE",
+        secretAccessKey: "Secret Access Key",
+        secretAccessKeyPlaceholder: "Your AWS secret access key",
+        useIamRole: "Use IAM Role",
+        useIamRoleTooltip: "Use IAM role credentials instead of access keys (recommended for EC2/SageMaker)",
+        testConnection: "Test Connection",
+        connectionSuccess: "Successfully connected to S3 bucket",
+        connectionFailed: "Failed to connect to S3 bucket",
+        comingSoon: "S3 integration coming soon",
+        comingSoonDescription: "S3 dataset loading requires boto3. This feature is under development.",
+      },
     },
     params: {
       title: "Parameters",

studio/frontend/src/i18n/locales/zh-CN.ts

@@ -472,6 +472,28 @@ export const zhCN = {
       preview: "预览数据集",
       split: "切分",
       subset: "子集",
+      s3: {
+        title: "S3 配置",
+        description: "从 Amazon S3 加载 .parquet、.json、.jsonl 或 .csv 数据集",
+        bucket: "存储桶名称",
+        bucketPlaceholder: "my-training-data-bucket",
+        region: "AWS 区域",
+        regionPlaceholder: "us-east-1",
+        prefix: "路径前缀",
+        prefixPlaceholder: "datasets/whisper/",
+        prefixTooltip: "存储桶中数据集文件的可选路径",
+        accessKeyId: "访问密钥 ID",
+        accessKeyIdPlaceholder: "AKIAIOSFODNN7EXAMPLE",
+        secretAccessKey: "秘密访问密钥",
+        secretAccessKeyPlaceholder: "您的 AWS 秘密访问密钥",
+        useIamRole: "使用 IAM 角色",
+        useIamRoleTooltip: "使用 IAM 角色凭证而非访问密钥（推荐用于 EC2/SageMaker）",
+        testConnection: "测试连接",
+        connectionSuccess: "成功连接到 S3 存储桶",
+        connectionFailed: "无法连接到 S3 存储桶",
+        comingSoon: "S3 集成即将推出",
+        comingSoonDescription: "S3 数据集加载需要 boto3。此功能正在开发中。",
+      },
     },
     params: {
       title: "参数",

studio/frontend/src/types/training.ts

@@ -8,7 +8,17 @@ export function isAdapterMethod(method: TrainingMethod): boolean {
   return method === "lora" || method === "qlora" || method === "cpt";
 }
 export type StepNumber = 1 | 2 | 3 | 4 | 5;
-export type DatasetSource = "huggingface" | "upload";
+export type DatasetSource = "huggingface" | "upload" | "s3";
+
+/** S3 bucket configuration for loading datasets */
+export interface S3Config {
+  bucket: string;
+  region: string;
+  prefix?: string;
+  accessKeyId?: string;
+  secretAccessKey?: string;
+  useIamRole?: boolean;
+}
 export type DatasetFormat = "auto" | "alpaca" | "chatml" | "sharegpt" | "raw";
 export type GradientCheckpointing = "none" | "true" | "unsloth" | "mlx";