aefe904d66* 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>
93 lines
2.6 KiB
Python
93 lines
2.6 KiB
Python
# 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
|