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 and DataSet ARNs 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)

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?

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?

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.

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.

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

+            trainer.train(dry_run=True)
+
+
+class TestDryRunPassesWithValidInputs:


I would consider adding a test for serverful and hyperpod too. To make sure dry_runs in def train_hyperpod and train_serverful are also covered.

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.

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)

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(train): add dry_run=True to train()#6027

feat(train): add dry_run=True to train()#6027
amazeAmazing wants to merge 2 commits into
aws:master-nova-follow-upsfrom
amazeAmazing:feature/dry-run-validation

amazeAmazing commented Jul 15, 2026 •

edited

Loading

Uh oh!

Uh oh!

mujtaba1747 Jul 16, 2026

Uh oh!

mujtaba1747 Jul 16, 2026

Uh oh!

mujtaba1747 Jul 16, 2026

Uh oh!

mujtaba1747 Jul 16, 2026

Uh oh!

mujtaba1747 Jul 16, 2026

Uh oh!

mujtaba1747 Jul 16, 2026 •

edited

Loading

Uh oh!

mujtaba1747 Jul 16, 2026

Uh oh!

mujtaba1747 Jul 16, 2026

Uh oh!

mujtaba1747 Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		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):

		trainer.train(dry_run=True)


		class TestDryRunPassesWithValidInputs:



		@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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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