fix(dry_run): skip MLflow app creation and warn on unsupported instance types

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

@@ -31,6 +31,7 @@
     get_recipe_s3_uri,
     _validate_hyperparameter_values,
     _get_smhp_replicas_enum,
+    _warn_unsupported_instance_type,
 )
 from sagemaker.train.common_utils.data_utils import validate_data_path_exists
 from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics
@@ -1157,6 +1158,12 @@ def _yaml_safe_default(value):
                 resolved_validation_dataset, sagemaker_session, label="validation dataset"
             )
 
+        # Warn if instance type is not in the model's supported list
+        _warn_unsupported_instance_type(
+            compute.instance_type,
+            recipe_supported_types=getattr(self.hyperparameters, '_supported_instance_types', None),
+        )
+
         if dry_run:
             logger.info("Dry-run validation passed. No job submitted.")
             return None
@@ -1436,6 +1443,13 @@ def _train_hyperpod(self, training_dataset=None, validation_dataset=None,
                 resolved_validation_dataset, sagemaker_session, label="validation dataset"
             )
 
+        # Warn if instance type is not in the model's supported list
+        if compute.instance_type:
+            _warn_unsupported_instance_type(
+                compute.instance_type,
+                recipe_supported_types=getattr(self.hyperparameters, '_supported_instance_types', None),
+            )
+
         if dry_run:
             logger.info("Dry-run validation passed. No job submitted.")
             return None

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

@@ -6,7 +6,6 @@
 import logging
 import json
 from typing import Any, Dict, Optional, Union
-import time
 import boto3
 from sagemaker.core.resources import ModelPackage, ModelPackageGroup
 from sagemaker.core.helper.session_helper import Session
@@ -188,7 +187,12 @@ def _get_prod_sm_client(sagemaker_session) -> "boto3.client":
     return boto3.client("sagemaker", region_name=region)
 
 
-def _resolve_mlflow_resource_arn(sagemaker_session, mlflow_resource_arn: Optional[str] = None, min_mlflow_version: Optional[str] = None) -> Optional[str]:
+def _resolve_mlflow_resource_arn(
+    sagemaker_session,
+    mlflow_resource_arn: Optional[str] = None,
+    min_mlflow_version: Optional[str] = None,
+    dry_run: bool = False,
+) -> Optional[str]:
     """Resolve MLflow resource ARN using default experience logic.
 
     All MLflow API calls use a raw boto3 client against prod (no custom endpoint),
@@ -199,6 +203,8 @@ def _resolve_mlflow_resource_arn(sagemaker_session, mlflow_resource_arn: Optiona
         mlflow_resource_arn: Explicit ARN to use (returned as-is if provided).
         min_mlflow_version: Minimum required MLflow version (e.g. "3.10").
             If the resolved app's version is below this, a new app is created.
+        dry_run: If True, only performs read-only checks (list/describe) without
+            creating new apps or waiting for apps in Creating status.
     """
     if mlflow_resource_arn:
         return mlflow_resource_arn
@@ -240,10 +246,24 @@ def _resolve_mlflow_resource_arn(sagemaker_session, mlflow_resource_arn: Optiona
                 logger.warning("Resolved MLflow app %s is in failed state: %s. Skipping.",
                                resolved_arn, resolved_app.get("Status"))
                 resolved_app = None
+            elif dry_run and resolved_app.get("Status") in ["Creating", "Updating"]:
+                logger.warning(
+                    "dry_run: MLflow app %s is in '%s' state. "
+                    "Job submission would block until the app is ready.",
+                    resolved_arn, resolved_app.get("Status"),
+                )
+                return resolved_arn
 
         # Version check: if resolved app is below min version, create a new one as default
         if resolved_app and min_mlflow_version and not _mlflow_version_meets_minimum_dict(resolved_app, min_mlflow_version):
             resolved_arn = resolved_app["Arn"]
+            if dry_run:
+                logger.warning(
+                    "dry_run: MLflow app %s has version below %s. "
+                    "Job submission would create a new app (may take several minutes).",
+                    resolved_arn, min_mlflow_version,
+                )
+                return resolved_arn
             logger.info(
                 "Existing MLflow app %s has version below %s. Creating new app as default.",
                 resolved_arn, min_mlflow_version
@@ -258,6 +278,21 @@ def _resolve_mlflow_resource_arn(sagemaker_session, mlflow_resource_arn: Optiona
         if resolved_app:
             return resolved_app["Arn"]
 
+        # In dry_run mode, don't create a new app — just warn and return None
+        if dry_run:
+            if mlflow_apps_list:
+                # Apps exist but none are in a ready/usable state
+                logger.warning(
+                    "dry_run: No MLflow app in ready state found. "
+                    "Job submission would create a new app (may take several minutes)."
+                )
+            else:
+                logger.warning(
+                    "dry_run: No MLflow app exists. "
+                    "Job submission would create a new app (may take several minutes)."
+                )
+            return None
+
         # Create new app
         new_arn = _create_mlflow_app(sagemaker_session)
         if new_arn:
@@ -679,10 +714,21 @@ def _get_fine_tuning_options_and_model_arn(model_name: str, customization_techni
             except Exception as e:
                 logger.debug(f"Could not fetch subscription recipe override_params: {type(e).__name__}: {e}")
 
+        # Build union of supported instance types from both sources:
+        # 1. SupportedInstanceTypes on the recipe entry itself (OSS models)
+        # 2. instance_type.enum in the override params JSON (Nova models)
+        supported_instance_types = set(recipe.get("SupportedInstanceTypes") or [])
+        override_instance_enum = options_dict.get("instance_type", {})
+        if isinstance(override_instance_enum, dict):
+            supported_instance_types.update(override_instance_enum.get("enum") or [])
+
         if options_dict:
-            return FineTuningOptions(options_dict), model_arn, is_gated_model
+            ft_options = FineTuningOptions(options_dict)
         else:
-            return FineTuningOptions({}), model_arn, is_gated_model
+            ft_options = FineTuningOptions({})
+
+        ft_options._supported_instance_types = sorted(supported_instance_types) if supported_instance_types else None
+        return ft_options, model_arn, is_gated_model
 
     except Exception as e:
         logger.debug("Exception getting fine-tuning options: %s", e)
@@ -969,22 +1015,27 @@ def _create_model_package_config(model_package_group_name, model, sagemaker_sess
 
 
 def _create_mlflow_config(sagemaker_session, mlflow_resource_arn=None, 
-                       mlflow_experiment_name=None, mlflow_run_name=None):
+                       mlflow_experiment_name=None, mlflow_run_name=None,
+                       dry_run=False):
     """Create MLflow configuration with resolved resource ARN.
     
     Args:
         sagemaker_session: SageMaker session for resolving MLflow ARN
         mlflow_resource_arn: MLflow resource ARN (if None, uses default experience)
         mlflow_experiment_name: MLflow experiment name
         mlflow_run_name: MLflow run name
+        dry_run: If True, only performs read-only checks without creating
+            new MLflow apps or waiting for apps in Creating status.
     
     Returns:
         MlflowConfig object or None if no MLflow resource ARN is resolved
     """
 
 
     # Derive mlflow_resource_arn with default experience
-    resolved_mlflow_arn = _resolve_mlflow_resource_arn(sagemaker_session, mlflow_resource_arn)
+    resolved_mlflow_arn = _resolve_mlflow_resource_arn(
+        sagemaker_session, mlflow_resource_arn, dry_run=dry_run
+    )
     logger.info(f"MLflow resource ARN: {resolved_mlflow_arn}")
 
     # Create MlflowConfig using shapes
@@ -1157,6 +1208,63 @@ def _validate_s3_path_exists(s3_path: str, sagemaker_session):
         raise ValueError(f"Failed to validate/create S3 path '{s3_path}': {str(e)}")
 
 
+def _warn_unsupported_instance_type(
+    instance_type: str,
+    recipe_supported_types=None,
+    override_instance_enum=None,
+):
+    """Log a warning if instance_type is not in the union of supported types.
+
+    Args:
+        instance_type: The user's selected instance type.
+        recipe_supported_types: SupportedInstanceTypes from the Hub recipe entry.
+        override_instance_enum: instance_type.enum from the override params JSON.
+    """
+    supported = set(recipe_supported_types or [])
+    supported.update(override_instance_enum or [])
+
+    if not supported:
+        return  # No metadata available — can't validate
+
+    if instance_type not in supported:
+        logger.warning(
+            "Instance type '%s' is not in the model's supported instance types: %s. "
+            "The training job may fail at submission time.",
+            instance_type,
+            sorted(supported),
+        )
+
+
+def _get_supported_instance_types(
+    model_name: str,
+    customization_technique: str,
+    training_type,
+    sagemaker_session,
+    platform: str = "smtj",
+):
+    """Fetch supported instance types for a model/technique/platform from Hub.
+
+    Returns the union of SupportedInstanceTypes (recipe entry) and
+    instance_type.enum (override params JSON), or None if unavailable.
+    """
+    try:
+        recipe_entry, override_spec = _get_recipe_entry_and_override_spec(
+            model_name=model_name,
+            customization_technique=customization_technique,
+            training_type=training_type,
+            sagemaker_session=sagemaker_session,
+            platform=platform,
+        )
+        supported = set(recipe_entry.get("SupportedInstanceTypes") or [])
+        instance_type_spec = override_spec.get("instance_type")
+        if isinstance(instance_type_spec, dict):
+            supported.update(instance_type_spec.get("enum") or [])
+        return sorted(supported) if supported else None
+    except Exception as e:
+        logger.debug("Could not fetch supported instance types: %s", e)
+        return None
+
+
 def _validate_hyperparameter_values(hyperparameters: dict):
     """Validate hyperparameter values for allowed characters."""
     import re

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

@@ -286,7 +286,6 @@ def train(self,
         )
 
         logger.info(f"Training Job Name: {current_training_job_name}")
-        print(f"Training Job Name: {current_training_job_name}")
 
         #data
         input_data_config = _create_input_data_config(training_dataset or self.training_dataset,
@@ -316,6 +315,7 @@ def train(self,
             mlflow_resource_arn=self.mlflow_resource_arn,
             mlflow_experiment_name=self.mlflow_experiment_name,
             mlflow_run_name=self.mlflow_run_name,
+            dry_run=dry_run,
         )
 
         final_hyperparameters = self.hyperparameters.to_dict()

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

@@ -31,6 +31,8 @@
 from sagemaker.train.common_utils.finetune_utils import (
     _resolve_mlflow_resource_arn,
     _is_nova_model,
+    _get_supported_instance_types,
+    _warn_unsupported_instance_type,
 )
 from sagemaker.train.common_utils.recipe_utils import resolve_recipe, get_resolved_recipe_from_context
 from sagemaker.train.common_utils.validator import validate_hyperpod_compute
@@ -1779,6 +1781,20 @@ def _submit_hyperpod_eval_job(self, override_parameters=None, base_job_name=None
             )
             _logger.info(f"Auto-resolved HyperPod eval recipe from Hub: {recipe_name}")
 
+            # Warn if instance type is not in the model's supported list
+            if compute.instance_type:
+                supported = _get_supported_instance_types(
+                    model_name=model_name,
+                    customization_technique="Evaluation",
+                    training_type="FULL",
+                    sagemaker_session=self.sagemaker_session,
+                    platform="hyperpod",
+                )
+                _warn_unsupported_instance_type(
+                    compute.instance_type,
+                    recipe_supported_types=supported,
+                )
+
         # Build base override parameters
         base_overrides = {}
         if compute.instance_type:

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

@@ -22,6 +22,7 @@
 from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter, TelemetryParamType
 from sagemaker.train.common_utils.telemetry_params import BASE_EVALUATOR_TELEMETRY_PARAMS
 from sagemaker.core.telemetry.constants import Feature
+from sagemaker.train.common_utils.finetune_utils import _warn_unsupported_instance_type
 from sagemaker.train.constants import get_sagemaker_hub_name
 
 _logger = logging.getLogger(__name__)
@@ -819,6 +820,13 @@ def _evaluate_serverful_smtj(self, subtask=None):
         # Fall back to first available SMTJ eval recipe if no benchmark match
         recipe_metadata = benchmark_recipes[0] if benchmark_recipes else smtj_eval_recipes[0]
 
+        # Warn if instance type is not in the model's supported list
+        if self.compute and self.compute.instance_type:
+            _warn_unsupported_instance_type(
+                self.compute.instance_type,
+                recipe_supported_types=recipe_metadata.get("SupportedInstanceTypes"),
+            )
+
         # Get recipe template S3 URI and image URI
         training_image = self.training_image or recipe_metadata.get("SmtjImageUri")
 

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

@@ -18,6 +18,7 @@
 from sagemaker.train.common_utils.telemetry_params import BASE_EVALUATOR_TELEMETRY_PARAMS
 from sagemaker.core.telemetry.constants import Feature
 from sagemaker.train.common_utils.data_utils import validate_data_path_exists
+from sagemaker.train.common_utils.finetune_utils import _warn_unsupported_instance_type
 from sagemaker.train.constants import get_sagemaker_hub_name
 from sagemaker.train.defaults import TrainDefaults
 
@@ -606,6 +607,13 @@ def _evaluate_serverful_smtj(self):
         ]
         recipe_metadata = custom_scorer_recipes[0] if custom_scorer_recipes else smtj_eval_recipes[0]
 
+        # Warn if instance type is not in the model's supported list
+        if self.compute and self.compute.instance_type:
+            _warn_unsupported_instance_type(
+                self.compute.instance_type,
+                recipe_supported_types=recipe_metadata.get("SupportedInstanceTypes"),
+            )
+
         # Resolve training image
         training_image = self.training_image
         if not training_image:

sagemaker-train/src/sagemaker/train/rlaif_trainer.py

@@ -276,6 +276,7 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati
             mlflow_resource_arn=self.mlflow_resource_arn,
             mlflow_experiment_name=self.mlflow_experiment_name,
             mlflow_run_name=self.mlflow_run_name,
+            dry_run=dry_run,
         )
 
         final_hyperparameters = self.hyperparameters.to_dict()

sagemaker-train/src/sagemaker/train/rlvr_trainer.py

@@ -477,6 +477,7 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None,
             mlflow_resource_arn=self.mlflow_resource_arn,
             mlflow_experiment_name=self.mlflow_experiment_name,
             mlflow_run_name=self.mlflow_run_name,
+            dry_run=dry_run,
         )
 
         final_hyperparameters = self.hyperparameters.to_dict()

sagemaker-train/src/sagemaker/train/sft_trainer.py

@@ -368,6 +368,7 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati
             mlflow_resource_arn=self.mlflow_resource_arn,
             mlflow_experiment_name=self.mlflow_experiment_name,
             mlflow_run_name=self.mlflow_run_name,
+            dry_run=dry_run,
         )
 
         final_hyperparameters = self.hyperparameters.to_dict()