From aefe904d66ec026a559812048436c9db7dd4e557 Mon Sep 17 00:00:00 2001 From: ashzak Date: Fri, 12 Jun 2026 07:52:04 -0500 Subject: [PATCH] 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) * [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 Co-authored-by: Ash Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: wasimysaid --- studio/backend/core/training/resume.py | 21 ++ studio/backend/core/training/s3_dataset.py | 228 +++++++++++++++ studio/backend/core/training/trainer.py | 26 ++ studio/backend/core/training/training.py | 38 ++- studio/backend/core/training/worker.py | 150 +++++++--- studio/backend/models/training.py | 43 +++ studio/backend/requirements/studio.txt | 1 + studio/backend/routes/training.py | 16 + studio/backend/storage/studio_db.py | 2 +- studio/backend/tests/test_s3_dataset.py | 274 ++++++++++++++++++ studio/backend/tests/test_training_resume.py | 93 ++++++ studio/frontend/src/config/training.ts | 1 + .../sections/dataset-preview-dialog.tsx | 3 +- .../studio/sections/dataset-section.tsx | 255 +++++++++------- .../studio/sections/s3-config-form.tsx | 142 +++++++++ .../src/features/training/api/mappers.ts | 18 ++ .../src/features/training/lib/validation.ts | 28 +- .../training/stores/training-config-store.ts | 24 +- .../src/features/training/types/api.ts | 4 + .../src/features/training/types/config.ts | 4 + studio/frontend/src/i18n/locales/en.ts | 22 ++ studio/frontend/src/i18n/locales/zh-CN.ts | 22 ++ studio/frontend/src/types/training.ts | 12 +- 23 files changed, 1271 insertions(+), 156 deletions(-) create mode 100644 studio/backend/core/training/s3_dataset.py create mode 100644 studio/backend/tests/test_s3_dataset.py create mode 100644 studio/backend/tests/test_training_resume.py create mode 100644 studio/frontend/src/features/studio/sections/s3-config-form.tsx diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py index 165c1c2cf..2a4a19861 100644 --- a/studio/backend/core/training/resume.py +++ b/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") diff --git a/studio/backend/core/training/s3_dataset.py b/studio/backend/core/training/s3_dataset.py new file mode 100644 index 000000000..3d05d19c7 --- /dev/null +++ b/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 diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 085f999dd..018b66ea9 100644 --- a/studio/backend/core/training/trainer.py +++ b/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 diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 2764101d9..6dd42976c 100644 --- a/studio/backend/core/training/training.py +++ b/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-/`` 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, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index c7c9003a0..1120744a2 100644 --- a/studio/backend/core/training/worker.py +++ b/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( { diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 4a3303c73..ca0459117 100644 --- a/studio/backend/models/training.py +++ b/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. diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 9bffb8154..96fef6047 100644 --- a/studio/backend/requirements/studio.txt +++ b/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 diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index d7687ffde..281f03bca 100644 --- a/studio/backend/routes/training.py +++ b/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. diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index da8d9b5e6..85cfacbc2 100644 --- a/studio/backend/storage/studio_db.py +++ b/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 diff --git a/studio/backend/tests/test_s3_dataset.py b/studio/backend/tests/test_s3_dataset.py new file mode 100644 index 000000000..f47db565f --- /dev/null +++ b/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 diff --git a/studio/backend/tests/test_training_resume.py b/studio/backend/tests/test_training_resume.py new file mode 100644 index 000000000..91fdac996 --- /dev/null +++ b/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 diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index c0a60d0de..873e9aa20 100644 --- a/studio/frontend/src/config/training.ts +++ b/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 = { diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx index f337d5100..c467880b6 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/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; diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 3677b8e1b..970e5246a 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/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() { }`} >
+ { + 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" + > + + Hugging Face + + {t("studio.dataset.localTab")} + + {!isMultimodalModel && ( + Amazon S3 + )} + + + + {datasetSource === "s3" && } + + {datasetSource !== "s3" && (
{t("studio.dataset.chooseDataset")} @@ -842,8 +889,10 @@ export function DatasetSection() {

)}
+ )} - {isHfDatasetSelected ? ( + {datasetSource !== "s3" && + (isHfDatasetSelected ? (
- ) : null} + ) : null)} {datasetSource === "upload" && uploadedFile && (
@@ -1102,111 +1151,115 @@ export function DatasetSection() { -
- {selectedDatasetName ? ( -
-
- -
-
-

- {datasetSource === "upload" - ? (selectedLocalDataset?.label ?? - deriveLocalDatasetName(selectedDatasetName)) - : selectedDatasetName} -

-

- {datasetSource === "upload" ? ( - uploadedFile ? ( - <> - {t("studio.dataset.localDataset")} - {selectedLocalRows != null - ? t("studio.dataset.localDatasetRows", { - count: selectedLocalRows.toLocaleString(), - }) - : ""} - + {datasetSource !== "s3" && ( +

+ {selectedDatasetName ? ( +
+
+ +
+
+

+ {datasetSource === "upload" + ? (selectedLocalDataset?.label ?? + deriveLocalDatasetName(selectedDatasetName)) + : selectedDatasetName} +

+

+ {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}`} - - )} -

+ <> + {t("studio.dataset.huggingFaceDataset")} + {datasetSubset && ` / ${datasetSubset}`} + {datasetSplit && ` / ${datasetSplit}`} + + )} +

+
+
- + )} + +
+ +
- ) : ( - - )} - -
- -
-
+ )} ({ + s3Config: s.s3Config, + setS3Config: s.setS3Config, + })), + ); + + const config = s3Config ?? DEFAULT_S3_CONFIG; + + const update = (patch: Partial) => { + setS3Config({ ...config, ...patch }); + }; + + const handleIamRoleChange = (useIamRole: boolean) => { + if (useIamRole) { + update({ useIamRole, accessKeyId: "", secretAccessKey: "" }); + return; + } + update({ useIamRole }); + }; + + return ( +
+
+

+ {t("studio.dataset.s3.title")} +

+

+ {t("studio.dataset.s3.description")} +

+
+ +
+ + update({ bucket: e.target.value })} + placeholder={t("studio.dataset.s3.bucketPlaceholder")} + /> +
+ +
+
+ + update({ region: e.target.value })} + placeholder={t("studio.dataset.s3.regionPlaceholder")} + /> +
+
+ + update({ prefix: e.target.value })} + placeholder={t("studio.dataset.s3.prefixPlaceholder")} + /> +
+
+ +
+ + +
+ + {!config.useIamRole && ( + <> +
+ + update({ accessKeyId: e.target.value })} + placeholder={t("studio.dataset.s3.accessKeyIdPlaceholder")} + /> +
+
+ + update({ secretAccessKey: e.target.value })} + placeholder={t("studio.dataset.s3.secretAccessKeyPlaceholder")} + /> +
+ + )} +
+ ); +} diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 1223888b1..ead719f82 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/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 | 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, diff --git a/studio/frontend/src/features/training/lib/validation.ts b/studio/frontend/src/features/training/lib/validation.ts index 700de1c1a..fcba1dcab 100644 --- a/studio/frontend/src/features/training/lib/validation.ts +++ b/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 }; } 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 7de78096f..75f23cf40 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/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 = 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()( 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()( setFinetuneMLPModules: (finetuneMLPModules) => set({ finetuneMLPModules }), setTargetModules: (targetModules) => set({ targetModules }), + setS3Config: (s3Config) => set({ s3Config }), canProceed: () => canProceedForStep(get()), reset: () => { _trainOnCompletionsManuallySet = false; diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index 5c12e7118..8490d5ee6 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/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 | null; num_epochs: number; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index f40f053e8..5ef1e2d13 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/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; diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 19863d42b..9e0ec15e5 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/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", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index 5173f9fa7..72dad973c 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/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: "参数", diff --git a/studio/frontend/src/types/training.ts b/studio/frontend/src/types/training.ts index feca65ebd..a09c92878 100644 --- a/studio/frontend/src/types/training.ts +++ b/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";