feat(train): add dry_run=True to train() by amazeAmazing · Pull Request #6027 · aws/sagemaker-python-sdk

amazeAmazing · 2026-07-15T17:50:04Z

Summary

Add dry_run=True parameter to all trainers (SFT, DPO, RLVR, RLAIF) and evaluators (BenchMark, CustomScorer, LLMAsJudge). When dry_run=True, all validation runs but no job is submitted and no compute is consumed.

What it does

All existing inline validation runs as normal (IAM role, hyperparameters, recipe constraints, model resolution)
validate_data_path_exists() validates S3 URIs, DataSet ARNs, and DataSet objects unconditionally before submission
dry_run=True short-circuits before the API call, returning None
Raises with a clear error message on validation failure

Changes

Commit 1: feat(train): add dry_run=True to train()

data_utils.py — new validate_data_path_exists() (S3 via list_objects_v2, DataSet ARN via describe_hub_content)
base_trainer.py — add dry_run to train(), _train_serverful_smtj(), _train_hyperpod()
sft/dpo/rlvr/rlaif_trainer.py — add dry_run, validate data paths, short-circuit
Trainer notebook examples (SFT, DPO, RLVR, RLAIF)
Unit tests in existing test files
Integration test (test_dry_run_integration.py)

Commit 2: feat(evaluate): add dry_run=True to evaluate()

base_evaluator.py — add dry_run to evaluate() signature
benchmark/custom_scorer/llm_as_judge_evaluator.py — add dry_run, validate dataset paths (custom scorer + LLM-as-judge), short-circuit before _start_execution()
Evaluator notebook examples (benchmark, custom_scorer, llm_as_judge)

Commit 3: fix(dry_run): support DataSet objects, deduplicate ARN validation, expand coverage

Addresses PR review feedback:

DataSet object support: validate_data_path_exists() now accepts Union[str, DataSet]. When a DataSet object is passed, it extracts .arn and validates the ARN exists via describe_hub_content.
Deduplicated ARN validation: Removed inline ARN validation from validate_data_path_exists() and delegated to the existing _validate_dataset_arn_exists() helper.
Unified AccessDenied behavior: _validate_dataset_arn_exists() now logs a warning on AccessDenied (same as the S3 path) instead of raising, since the execution role may still have access even if the caller identity does not.
Removed isinstance(..., str) guards: All trainers (SFT, DPO, RLVR, RLAIF) and evaluators (CustomScorer, LLMAsJudge) now pass DataSet objects through to validation rather than skipping them.
CPTTrainer dry_run: Added dry_run: bool = False parameter to CPTTrainer.train(), threaded through to _train_hyperpod().
ModelTrainer dry_run: Added dry_run: bool = False parameter to ModelTrainer.train(). When True, runs _create_training_job_args() (all validation/resolution) then returns None without submitting.
Integration test improvements:
- valid_dataset fixture checks if data exists before uploading (no teardown delete)
- Added nonexistent_dataset_s3, nonexistent_dataset_arn, nonexistent_dataset_obj fixtures
- Added tests validating failure with ARN strings and DataSet objects
- Added TestDryRunServerful class covering SFT, DPO, RLVR with TrainingJobCompute
Unit test additions:
- test_dataset_object_extracts_arn — validates DataSet object ARN is extracted and described
- test_dataset_object_not_found_raises — validates ValueError on nonexistent DataSet ARN

Testing

Integration test uploads sample data to default bucket, validates S3 path failures, confirms result is None, verifies no job created via SageMaker API

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Add dry_run parameter to all trainers (SFT, DPO, RLVR, RLAIF). When dry_run=True: - All existing validation runs inline (IAM role, hyperparameters, recipe constraints, infrastructure availability) - Returns None without submitting a job or consuming compute - Raises with clear error message on validation failure Additionally, validate_data_path_exists() is called unconditionally (regardless of dry_run) before job submission to catch non-existent S3 paths or dataset ARNs early. Design follows nova-forge-sdk pattern: validation always runs as part of the normal code path, dry_run short-circuits before the actual TrainingJob.create API call. Changes: - data_utils.py: add validate_data_path_exists() utility (S3 + DataSet ARN) - base_trainer.py: add dry_run to abstract train(), _train_serverful_smtj(), and _train_hyperpod() - sft/dpo/rlvr/rlaif_trainer.py: add dry_run param, pass through to shared methods, short-circuit serverless path - Notebook examples added to SFT, DPO, RLVR, RLAIF notebooks - Unit tests added to existing test files - Integration test added

Add dry_run parameter to BaseEvaluator.evaluate() and all subclasses (BenchMarkEvaluator, CustomScorerEvaluator, LLMAsJudgeEvaluator). When dry_run=True: - All existing validation runs (IAM role, model resolution, recipe, pipeline rendering) - Dataset S3 path / DataSet ARN validated via validate_data_path_exists() - Returns None without submitting a pipeline execution - Raises on validation failure Dataset validation runs unconditionally (not just during dry_run) for CustomScorerEvaluator and LLMAsJudgeEvaluator which accept user datasets. Changes: - base_evaluator.py: add dry_run to evaluate() signature - benchmark_evaluator.py: add dry_run, short-circuit before _start_execution() - custom_scorer_evaluator.py: add dry_run, validate dataset, short-circuit - llm_as_judge_evaluator.py: add dry_run, validate dataset, short-circuit - Notebook examples added to benchmark, custom_scorer, llm_as_judge notebooks

mujtaba1747 · 2026-07-16T22:08:10Z

            raise FileLoadError(f"Failed to read file {file_path}: {e}")


+def validate_data_path_exists(


Will have to check this but some trainers may also support a Dataset object. In that case the arn is stored as dataset.arn.

mujtaba1747 · 2026-07-16T22:15:41Z

+        ValueError: If the path does not exist or is inaccessible.
+    """
+    # Handle SageMaker hub-content DataSet ARNs
+    if data_path.startswith("arn:aws:sagemaker:") and "/DataSet/" in data_path:


Can we just call _validate_dataset_arn_exists that's defined below to avoid duplicate code?

Thanks for the callout - will make the change

mujtaba1747 · 2026-07-16T22:25:10Z

-
+
+        # Validate dataset path exists
+        if hasattr(self, 'dataset') and self.dataset and isinstance(self.dataset, str):


Will have to check if Dataset object is also supported by Evaluator.

mujtaba1747 · 2026-07-16T22:40:54Z

    @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="RLVRTrainer.train")
    def train(self, training_dataset: Optional[Union[str, DataSet]] = None,
-              validation_dataset: Optional[Union[str, DataSet]] = None, wait: bool = True, wait_timeout: Optional[int] = None, poll: int = 5):
+              validation_dataset: Optional[Union[str, DataSet]] = None, wait: bool = True, wait_timeout: Optional[int] = None, poll: int = 5, dry_run: bool = False):


Does CPTTrainer also require additional dry_run: bool parameter?

Thanks for the callout - will make the change

mujtaba1747 · 2026-07-16T22:54:08Z

+        ]},
+    ]
+    body = "\n".join(json.dumps(s) for s in samples)
+    s3.put_object(Bucket=bucket, Key=DATASET_KEY, Body=body.encode("utf-8"))


nit: Why create and delete on every test run. Just check if the object exists at that path. If it doesn't, create it.

Thanks for the callout - will make the change

mujtaba1747 · 2026-07-16T23:00:19Z

            The MLflow experiment name for organizing runs.
        mlflow_run_name (Optional[str]):
            The MLflow run name for this training job.
        training_dataset (Optional[Union[str, DataSet]]):


Yes, we support Dataset object as well. So dry_run needs to support Dataset object too. Should be an easy change IMO.

Thanks for the callout - will make the change

mujtaba1747 · 2026-07-16T23:04:50Z

+            return None
+
        # Execute training
        model_trainer.train(


Why didn't we add dry_run feature to Model trainer class. Customers can submit training jobs using model_trainer directly as well? I'm ok with not adding it if there's a good reason.

Thanks for the callout - will make the change! That was a miss on my part

mujtaba1747 · 2026-07-16T23:46:42Z

+
+
+@pytest.fixture(scope="module")
+def nonexistent_dataset(sagemaker_session):


Should run the test with nonexistent_dataset_s3 and nonexistent_dataset_obj (ie Dataset object or Dataset arn)

Thanks for the callout - will make the change

…pand coverage - validate_data_path_exists() now accepts Union[str, DataSet]; extracts .arn from DataSet objects for validation - Removed duplicate ARN validation logic; delegates to _validate_dataset_arn_exists() - _validate_dataset_arn_exists() warns on AccessDenied instead of raising (execution role may still have access) - Removed isinstance(..., str) guards in all trainers and evaluators so DataSet objects flow through validation - Added dry_run=True parameter to CPTTrainer.train() - Added dry_run=True parameter to ModelTrainer.train() - Integration test: valid_dataset fixture no longer re-creates on every run - Integration test: added nonexistent_dataset_arn and nonexistent_dataset_obj fixtures - Integration test: added TestDryRunServerful class (serverful compute path) - Unit test: added test_dataset_object_extracts_arn, test_dataset_object_not_found_raises

mujtaba1747 · 2026-07-17T17:10:59Z

+    """
+
+    pattern = (
+        r"^arn:aws:sagemaker:([^:]+):(\d+):hub-content/"


ARN needs to support all aws partitions like China, ISO etc. Should be a quick change in regex pattern.

Thank you for the callout! Will update

- Update ARN regex to aws(?:-[a-z]+)* to match aws-cn, aws-us-gov, aws-iso, aws-iso-b partitions - Use regex guard in validate_data_path_exists() for consistent matching - Add unit tests for each partition (standard, China, GovCloud, ISO, ISO-B) and invalid partition rejection

amazeAmazing force-pushed the feature/dry-run-validation branch 2 times, most recently from 164a14c to 7b06162 Compare July 15, 2026 17:55

amazeAmazing closed this Jul 15, 2026

amazeAmazing reopened this Jul 15, 2026

amazeAmazing force-pushed the feature/dry-run-validation branch 9 times, most recently from 008794d to 16c783a Compare July 15, 2026 22:18

amazeAmazing changed the title ~~feat(train): add dry_run=True to train() and deploy()~~ feat(train): add dry_run=True to train() Jul 15, 2026

amazeAmazing marked this pull request as ready for review July 15, 2026 22:30

amazeAmazing force-pushed the feature/dry-run-validation branch 6 times, most recently from f6bd43d to 8a3936b Compare July 15, 2026 23:39

ehsu3 reviewed Jul 16, 2026

View reviewed changes

Comment thread sagemaker-train/src/sagemaker/train/common_utils/data_utils.py Outdated

ehsu3 approved these changes Jul 16, 2026

View reviewed changes

amazeAmazing force-pushed the feature/dry-run-validation branch from 8a3936b to 6375702 Compare July 16, 2026 18:43

mujtaba1747 reviewed Jul 16, 2026

View reviewed changes

Comment thread sagemaker-train/tests/integ/train/test_dry_run_integration.py

mujtaba1747 reviewed Jul 16, 2026

View reviewed changes

mujtaba1747 reviewed Jul 17, 2026

View reviewed changes

amazeAmazing force-pushed the feature/dry-run-validation branch from a935adb to bb7028d Compare July 17, 2026 19:17

mujtaba1747 approved these changes Jul 17, 2026

View reviewed changes

mujtaba1747 merged commit f055206 into aws:master-nova-follow-ups Jul 20, 2026
1 check passed

		raise FileLoadError(f"Failed to read file {file_path}: {e}")


		def validate_data_path_exists(



		# Validate dataset path exists
		if hasattr(self, 'dataset') and self.dataset and isinstance(self.dataset, str):



		@pytest.fixture(scope="module")
		def nonexistent_dataset(sagemaker_session):

Uh oh!

Conversation

amazeAmazing commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What it does

Changes

Testing

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mujtaba1747 Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

amazeAmazing commented Jul 15, 2026 •

edited

Loading

mujtaba1747 Jul 16, 2026 •

edited

Loading