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

sagemaker-train/src/sagemaker/train/base_trainer.py

@@ -32,6 +32,7 @@
     _validate_hyperparameter_values,
     _get_smhp_replicas_enum,
 )
+from sagemaker.train.common_utils.data_utils import validate_data_path_exists
 from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics
 from sagemaker.train.common_utils.mlflow_config_utils import resolve_mlflow_tracking_fields
 from sagemaker.train.common_utils.validator import validate_hyperpod_compute
@@ -598,7 +599,7 @@ def _validate_instance_count(self, instance_count, sagemaker_session):
         return smhp_replicas_enum
 
     @abstractmethod
-    def train(self, input_data_config: List[InputData], wait: bool = True, logs: bool = True, wait_timeout: Optional[int] = None):
+    def train(self, input_data_config: List[InputData], wait: bool = True, logs: bool = True, wait_timeout: Optional[int] = None, dry_run: bool = False):
         """Common training method that calls the specific implementation."""
         pass
 
@@ -614,7 +615,7 @@ def _get_extra_smtj_hyperparameters(self) -> Dict[str, Any]:
         return {}
 
     def _train_serverful_smtj(self, training_dataset=None, validation_dataset=None,
-                    wait=True, wait_timeout=None, poll=5):
+                    wait=True, wait_timeout=None, poll=5, dry_run=False):
         """Execute training on serverful SageMaker Training Job (SMTJ) compute.
 
         Uses ModelTrainer.from_recipe() with the model's recipe template from
@@ -967,6 +968,20 @@ def _yaml_safe_default(value):
             base_job_name=base_job_name,
         )
 
+        # Validate data paths exist before submission
+        if resolved_training_dataset:
+            validate_data_path_exists(
+                resolved_training_dataset, sagemaker_session, label="training dataset"
+            )
+        if resolved_validation_dataset:
+            validate_data_path_exists(
+                resolved_validation_dataset, sagemaker_session, label="validation dataset"
+            )
+
+        if dry_run:
+            logger.info("Dry-run validation passed. No job submitted.")
+            return None
+
         # Execute training
         model_trainer.train(
             wait=wait,
@@ -1085,7 +1100,7 @@ def _resolve_checkpoint_from_manifest(
         return checkpoint_path
 
     def _train_hyperpod(self, training_dataset=None, validation_dataset=None,
-                        wait=True, wait_timeout=None, poll=5):
+                        wait=True, wait_timeout=None, poll=5, dry_run=False):
         """Execute training on a SageMaker HyperPod cluster.
 
         Uses the HyperPod CLI to connect to the cluster and submit a training job
@@ -1228,6 +1243,20 @@ def _train_hyperpod(self, training_dataset=None, validation_dataset=None,
         if getattr(self, 'model_source', None):
             override_parameters["recipes.run.model_name_or_path"] = self.model_source
 
+        # Validate data paths exist before submission
+        if resolved_training_dataset:
+            validate_data_path_exists(
+                resolved_training_dataset, sagemaker_session, label="training dataset"
+            )
+        if resolved_validation_dataset:
+            validate_data_path_exists(
+                resolved_validation_dataset, sagemaker_session, label="validation dataset"
+            )
+
+        if dry_run:
+            logger.info("Dry-run validation passed. No job submitted.")
+            return None
+
         # Submit job
         start_job_cmd = [
             "hyperpod", "start-job",

sagemaker-train/src/sagemaker/train/common_utils/data_utils.py

@@ -94,6 +94,117 @@ def load_file_content(
             raise FileLoadError(f"Failed to read file {file_path}: {e}")
 
 
+def validate_data_path_exists(
+    data_path: Union[str, "DataSet"],
+    sagemaker_session,
+    label: str = "data",
+) -> None:
+    """Validate that a data path (S3 URI, dataset ARN, or DataSet object) exists and is accessible.
+
+    Called inline during dry_run to catch bad paths before job submission.
+
+    Args:
+        data_path: S3 URI, SageMaker hub-content DataSet ARN, or DataSet object to validate.
+        sagemaker_session: SageMaker session (provides boto_session).
+        label: Human-readable label for error messages.
+
+    Raises:
+        ValueError: If the path does not exist or is inaccessible.
+    """
+    # Handle DataSet objects — extract the ARN for validation
+    if isinstance(data_path, DataSet):
+        data_path = data_path.arn
+
+    # Handle SageMaker hub-content DataSet ARNs
+    if re.match(r"^arn:aws(?:-[a-z]+)*:sagemaker:.+/DataSet/", data_path):
+        _validate_dataset_arn_exists(data_path, sagemaker_session, label=label)
+        return
+
+    # Handle S3 URIs
+    parts = _parse_s3_uri(data_path)
+    if parts is None:
+        raise ValueError(
+            f"Invalid {label} path format: {data_path}. "
+            f"Expected an S3 URI (s3://bucket/key) or a DataSet ARN."
+        )
+
+    bucket, key = parts
+    s3 = sagemaker_session.boto_session.client("s3")
+
+    try:
+        resp = s3.list_objects_v2(Bucket=bucket, Prefix=key, MaxKeys=1)
+        if resp.get("KeyCount", 0) == 0:
+            raise ValueError(
+                f"S3 {label} path does not exist: {data_path}"
+            )
+    except ClientError as e:
+        code = e.response["Error"]["Code"]
+        if code == "403" or "AccessDenied" in str(e):
+            # Caller may not have access but the execution role might —
+            # log a warning and allow the job to proceed.
+            logger.warning(
+                "Cannot verify S3 %s path %s from caller identity "
+                "(AccessDenied). The execution role may still have access.",
+                label, data_path,
+            )
+        else:
+            raise ValueError(f"Error accessing S3 {label} path {data_path}: {e}")
+
+
+def _validate_dataset_arn_exists(
+    dataset_arn: str,
+    sagemaker_session,
+    label: str = "data",
+) -> None:
+    """Validate that a SageMaker hub-content DataSet ARN exists.
+
+    Args:
+        dataset_arn: ARN like arn:aws:sagemaker:<region>:<account>:hub-content/<hub>/DataSet/<name>/<version>
+        sagemaker_session: SageMaker session (provides boto_session).
+        label: Human-readable label for error messages.
+
+    Raises:
+        ValueError: If the dataset ARN cannot be described.
+    """
+
+    pattern = (
+        r"^arn:aws(?:-[a-z]+)*:sagemaker:([^:]+):(\d+):hub-content/"
+        r"([^/]+)/DataSet/([^/]+)/([\d\.]+)$"
+    )
+    match = re.match(pattern, dataset_arn)
+    if not match:
+        raise ValueError(
+            f"Invalid {label} DataSet ARN format: {dataset_arn}"
+        )
+
+    region, _, hub_name, content_name, content_version = match.groups()
+    sm_client = sagemaker_session.sagemaker_client
+
+    try:
+        sm_client.describe_hub_content(
+            HubName=hub_name,
+            HubContentType="DataSet",
+            HubContentName=content_name,
+            HubContentVersion=content_version,
+        )
+    except ClientError as e:
+        code = e.response["Error"]["Code"]
+        if code == "ResourceNotFound" or "does not exist" in str(e).lower():
+            raise ValueError(
+                f"{label.capitalize()} DataSet does not exist: {dataset_arn}"
+            )
+        elif code == "AccessDeniedException" or "AccessDenied" in str(e):
+            logger.warning(
+                "Cannot verify %s DataSet %s from caller identity "
+                "(AccessDenied). The execution role may still have access.",
+                label, dataset_arn,
+            )
+        else:
+            raise ValueError(
+                f"Error validating {label} DataSet {dataset_arn}: {e}"
+            )
+
+
 def _has_multimodal_content(record: dict) -> bool:
     """Check if a single record contains multimodal content."""
     if "messages" not in record:

sagemaker-train/src/sagemaker/train/cpt_trainer.py

@@ -186,6 +186,7 @@ def train(
         wait: bool = True,
         wait_timeout: Optional[int] = None,
         poll: int = 5,
+        dry_run: bool = False,
     ):
         """Execute the CPT training job on HyperPod.
 
@@ -200,9 +201,13 @@ def train(
                 Maximum time in seconds to wait.
             poll (int):
                 Polling interval in seconds. Defaults to 5.
+            dry_run (bool):
+                If True, runs validation without submitting a job. 
+                Returns None on success, raises on validation failure.
+                Defaults to False.
 
         Returns:
-            str: The HyperPod job name.
+            str: The HyperPod job name, or None if dry_run=True.
 
         Raises:
             ValueError: If compute is not configured or recipe is missing.
@@ -252,4 +257,5 @@ def train(
             wait=wait,
             wait_timeout=wait_timeout,
             poll=poll,
+            dry_run=dry_run,
         )

sagemaker-train/src/sagemaker/train/dpo_trainer.py

@@ -24,7 +24,7 @@
     _validate_eula_for_gated_model,
     _validate_hyperparameter_values
 )
-from sagemaker.train.common_utils.data_utils import is_multimodal_data
+from sagemaker.train.common_utils.data_utils import is_multimodal_data, validate_data_path_exists
 from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter
 from sagemaker.core.telemetry.constants import Feature
 from sagemaker.train.constants import get_sagemaker_hub_name
@@ -217,7 +217,8 @@ def train(self,
               validation_dataset: Optional[Union[str, DataSet]] = None,
               wait: bool = True,
               wait_timeout: Optional[int] = None,
-              poll: int = 5):
+              poll: int = 5,
+              dry_run: bool = False):
         """Execute the DPO training job.
 
         Parameters:
@@ -234,9 +235,13 @@ def train(self,
                 If None, uses the default timeout from the wait utility.
             poll (int):
                 Polling interval in seconds for checking training job status. Defaults to 5.
+            dry_run (bool):
+                If True, runs all validation (IAM, hyperparameters, infrastructure, data paths)
+                without submitting a job. Returns None on success, raises on validation failure.
+                Defaults to False.
 
         Returns:
-            TrainingJob: The SageMaker training job object.
+            TrainingJob: The SageMaker training job object, or None if dry_run=True.
         """
         # Dispatch based on compute type
         if isinstance(self.compute, HyperPodCompute):
@@ -246,6 +251,7 @@ def train(self,
                 wait=wait,
                 wait_timeout=wait_timeout,
                 poll=poll,
+                dry_run=dry_run,
             )
         elif isinstance(self.compute, TrainingJobCompute):
             return self._train_serverful_smtj(
@@ -254,6 +260,7 @@ def train(self,
                 wait=wait,
                 wait_timeout=wait_timeout,
                 poll=poll,
+                dry_run=dry_run,
             )
 
         # Default: serverless compute (None)
@@ -340,6 +347,22 @@ def train(self,
         if self.stopping_condition is not None:
             create_args["stopping_condition"] = self.stopping_condition
 
+        # Validate data paths exist before submission
+        effective_training = training_dataset or self.training_dataset
+        effective_validation = validation_dataset or self.validation_dataset
+        if effective_training:
+            validate_data_path_exists(
+                effective_training, sagemaker_session, label="training dataset"
+            )
+        if effective_validation:
+            validate_data_path_exists(
+                effective_validation, sagemaker_session, label="validation dataset"
+            )
+
+        if dry_run:
+            logger.info("Dry-run validation passed. No job submitted.")
+            return None
+
         try:
             training_job = TrainingJob.create(**create_args)
         except Exception as e:

sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py

@@ -1034,15 +1034,22 @@ def get_resolved_recipe(self) -> Dict[str, Any]:
         object.__setattr__(self, '_resolved_recipe_cache', resolved)
         return copy.deepcopy(resolved)
 
-    def evaluate(self) -> Any:
+    def evaluate(self, dry_run: bool = False) -> Any:
         """Create and start an evaluation execution.
 
         This method must be implemented by subclasses to define the specific
         evaluation logic for different evaluation types (benchmark, custom scorer,
         LLM-as-judge, etc.).
 
+        Args:
+            dry_run (bool):
+                If True, runs all validation (IAM, model resolution, data paths)
+                without submitting the evaluation. Returns None on success, raises
+                on validation failure. Defaults to False.
+
         Returns:
-            EvaluationPipelineExecution: The created evaluation execution object.
+            EvaluationPipelineExecution: The created evaluation execution object,
+            or None if dry_run=True.
 
         Raises:
             NotImplementedError: This is an abstract method that must be implemented by subclasses.

sagemaker-train/src/sagemaker/train/evaluate/benchmark_evaluator.py

@@ -591,7 +591,7 @@ def _get_benchmark_template_additions(self, eval_subtask: Optional[Union[str, Li
         return benchmark_context
 
     @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="BenchMarkEvaluator.evaluate")
-    def evaluate(self, subtask: Optional[Union[str, List[str]]] = None) -> EvaluationPipelineExecution:
+    def evaluate(self, subtask: Optional[Union[str, List[str]]] = None, dry_run: bool = False) -> EvaluationPipelineExecution:
         """Create and start a benchmark evaluation job.
 
         Supports multiple compute backends via the ``compute`` parameter set at
@@ -601,12 +601,17 @@ def evaluate(self, subtask: Optional[Union[str, List[str]]] = None) -> Evaluatio
         - **HyperPod**: Submits to a HyperPod cluster via the HyperPod CLI.
 
         Args:
-            subtask (Optional[Union[str, list[str]]]): Optional subtask(s) to evaluate. If not provided, 
-                uses the subtasks from constructor. Can be a single subtask string, a list of 
-                subtasks, or 'ALL' to run all subtasks.
+            subtask (Optional[Union[str, list[str]]]): Optional subtask(s) to evaluate.
+                If not provided, uses the subtasks from constructor. Can be a single
+                subtask string, a list of subtasks, or 'ALL' to run all subtasks.
+            dry_run (bool):
+                If True, runs all validation (IAM, model resolution, data paths)
+                without submitting the evaluation. Returns None on success, raises
+                on validation failure. Defaults to False.
 
         Returns:
-            EvaluationPipelineExecution: The created benchmark evaluation execution.
+            EvaluationPipelineExecution: The created benchmark evaluation execution,
+            or None if dry_run=True.
 
         Example:
 
@@ -702,7 +707,11 @@ def evaluate(self, subtask: Optional[Union[str, List[str]]] = None) -> Evaluatio
 
         # Generate execution name
         name = self.base_eval_name or f"benchmark-eval-{self.benchmark.value}"
-
+
+        if dry_run:
+            _logger.info("Dry-run validation passed. No evaluation submitted.")
+            return None
+
         # Start execution
         return self._start_execution(
             eval_type=EvalType.BENCHMARK,