diff --git a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py b/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py index 756982375f..419ef2115d 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py +++ b/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,6 +714,7 @@ 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: if options_dict: return FineTuningOptions(options_dict), model_arn, is_gated_model else: @@ -969,7 +1005,8 @@ 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: @@ -977,6 +1014,8 @@ def _create_mlflow_config(sagemaker_session, mlflow_resource_arn=None, 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 @@ -984,7 +1023,9 @@ def _create_mlflow_config(sagemaker_session, mlflow_resource_arn=None, # 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 diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py index 47ce435705..20879e202f 100644 --- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py +++ b/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() diff --git a/sagemaker-train/src/sagemaker/train/evaluate/benchmark_evaluator.py b/sagemaker-train/src/sagemaker/train/evaluate/benchmark_evaluator.py index 4a8a2ce793..96ababa87a 100644 --- a/sagemaker-train/src/sagemaker/train/evaluate/benchmark_evaluator.py +++ b/sagemaker-train/src/sagemaker/train/evaluate/benchmark_evaluator.py @@ -971,5 +971,5 @@ def _evaluate_hyperpod(self, subtask=None): return self._submit_hyperpod_eval_job( override_parameters=override_parameters, - base_job_name=f"eval-{self.benchmark.value}", + base_job_name=self.base_eval_name or f"eval-{self.benchmark.value}", ) diff --git a/sagemaker-train/src/sagemaker/train/evaluate/custom_scorer_evaluator.py b/sagemaker-train/src/sagemaker/train/evaluate/custom_scorer_evaluator.py index 15bde55f20..46766ba45d 100644 --- a/sagemaker-train/src/sagemaker/train/evaluate/custom_scorer_evaluator.py +++ b/sagemaker-train/src/sagemaker/train/evaluate/custom_scorer_evaluator.py @@ -787,5 +787,5 @@ def _evaluate_hyperpod(self): return self._submit_hyperpod_eval_job( override_parameters=override_parameters, - base_job_name="custom-eval", + base_job_name=self.base_eval_name or "custom-eval", ) diff --git a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py index 507d38202a..8794a2a625 100644 --- a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py +++ b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py @@ -272,15 +272,18 @@ def train( self, training_dataset: Optional[Union[str, DataSet]] = None, wait: bool = True, + dry_run: bool = False, ) -> AgentRFTJob: """Launch an Agentic RFT job. Args: training_dataset: Training dataset override. wait: If True (default), block until job reaches terminal status. + dry_run: If True, runs validation without submitting a job. + Returns None on success. Returns: - AgentRFTJob instance for tracking the job. + AgentRFTJob instance for tracking the job, or None if dry_run=True. """ sagemaker_session = TrainDefaults.get_sagemaker_session( sagemaker_session=self.sagemaker_session @@ -301,7 +304,11 @@ def train( if training_dataset is not None: self.training_dataset = training_dataset - job_config_doc = self._build_job_config_document() + job_config_doc = self._build_job_config_document(dry_run=dry_run) + + if dry_run: + logger.info("Dry-run validation passed. No job submitted.") + return None tags = _get_jumpstart_tags(self._model_name, get_sagemaker_hub_name()) @@ -361,14 +368,14 @@ def attach(cls, job_name: str, session=None) -> AgentRFTJob: # ---- Private: JobConfigDocument construction ---- - def _build_job_config_document(self) -> str: + def _build_job_config_document(self, dry_run: bool = False) -> str: """Build the JobConfigDocument JSON string conforming to v1_0_0 schema.""" config = { "AgentConfig": self._build_agent_config(), "InputDataConfig": self._build_input_data_config(), "OutputDataConfig": self._build_output_data_config(), "ModelPackageConfig": self._build_model_package_config(), - "TrainingConfig": self._build_training_config(), + "TrainingConfig": self._build_training_config(dry_run=dry_run), } if self.networking: config["VpcConfig"] = { @@ -444,12 +451,12 @@ def _build_model_package_config(self) -> dict: ) return config - def _build_training_config(self) -> dict: + def _build_training_config(self, dry_run: bool = False) -> dict: hyperparameters = getattr(self, "_final_hyperparameters", {}) config = { "BaseModelArn": self._model_arn, } - mlflow_config = self._build_mlflow_config() + mlflow_config = self._build_mlflow_config(dry_run=dry_run) if mlflow_config: config["MlflowConfig"] = mlflow_config if self.accept_eula is not None: @@ -462,7 +469,7 @@ def _build_training_config(self) -> dict: config["HyperParameters"] = user_set return config - def _build_mlflow_config(self) -> Optional[dict]: + def _build_mlflow_config(self, dry_run: bool = False) -> Optional[dict]: arn = ( self.mlflow_app_arn.arn if isinstance(self.mlflow_app_arn, MlflowApp) @@ -472,7 +479,9 @@ def _build_mlflow_config(self) -> Optional[dict]: session = self.sagemaker_session or TrainDefaults.get_sagemaker_session( sagemaker_session=self.sagemaker_session ) - arn = _resolve_mlflow_resource_arn(session, None, min_mlflow_version=MIN_MLFLOW_VERSION) + arn = _resolve_mlflow_resource_arn( + session, None, min_mlflow_version=MIN_MLFLOW_VERSION, dry_run=dry_run + ) if not arn: return None logger.info("MLflow resource ARN: %s", arn) diff --git a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py index 8054282a17..15e580ebe6 100644 --- a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py +++ b/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() diff --git a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py index 19cce18226..4831d7bf4a 100644 --- a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py +++ b/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() diff --git a/sagemaker-train/src/sagemaker/train/sft_trainer.py b/sagemaker-train/src/sagemaker/train/sft_trainer.py index b8c147c9f2..005cdec054 100644 --- a/sagemaker-train/src/sagemaker/train/sft_trainer.py +++ b/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() diff --git a/sagemaker-train/tests/unit/train/common_utils/test_mlflow_dry_run.py b/sagemaker-train/tests/unit/train/common_utils/test_mlflow_dry_run.py new file mode 100644 index 0000000000..81b63c6ead --- /dev/null +++ b/sagemaker-train/tests/unit/train/common_utils/test_mlflow_dry_run.py @@ -0,0 +1,206 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for MLflow dry_run behavior in _resolve_mlflow_resource_arn.""" +import logging +from unittest.mock import Mock, patch, MagicMock + +import pytest + +from sagemaker.train.common_utils.finetune_utils import ( + _resolve_mlflow_resource_arn, + _create_mlflow_config, +) + + +class TestResolveMlflowDryRunSkipsCreation: + """When dry_run=True, _resolve_mlflow_resource_arn must never create or wait.""" + + @patch("sagemaker.train.common_utils.finetune_utils._create_mlflow_app") + @patch("sagemaker.train.common_utils.finetune_utils._get_current_domain_id") + @patch("sagemaker.train.common_utils.finetune_utils._get_prod_sm_client") + def test_no_apps_dry_run_skips_creation( + self, mock_client, mock_domain, mock_create_app + ): + """dry_run=True with zero apps returns None without calling _create_mlflow_app.""" + mock_paginator = MagicMock() + mock_paginator.paginate.return_value = [{"Summaries": []}] + mock_client.return_value.get_paginator.return_value = mock_paginator + mock_domain.return_value = None + + mock_session = Mock() + + result = _resolve_mlflow_resource_arn(mock_session, dry_run=True) + + assert result is None + mock_create_app.assert_not_called() + + @patch("sagemaker.train.common_utils.finetune_utils._create_mlflow_app") + @patch("sagemaker.train.common_utils.finetune_utils._get_current_domain_id") + @patch("sagemaker.train.common_utils.finetune_utils._get_prod_sm_client") + def test_no_apps_non_dry_run_creates_app( + self, mock_client, mock_domain, mock_create_app + ): + """Without dry_run, zero apps triggers _create_mlflow_app.""" + mock_paginator = MagicMock() + mock_paginator.paginate.return_value = [{"Summaries": []}] + mock_client.return_value.get_paginator.return_value = mock_paginator + mock_domain.return_value = None + mock_create_app.return_value = "arn:aws:sagemaker:us-east-1:123:mlflow-app/new" + + mock_session = Mock() + + result = _resolve_mlflow_resource_arn(mock_session, dry_run=False) + + assert result == "arn:aws:sagemaker:us-east-1:123:mlflow-app/new" + mock_create_app.assert_called_once() + + @patch("sagemaker.train.common_utils.finetune_utils._create_mlflow_app") + @patch("sagemaker.train.common_utils.finetune_utils._get_current_domain_id") + @patch("sagemaker.train.common_utils.finetune_utils._get_prod_sm_client") + def test_creating_app_dry_run_skips_wait( + self, mock_client, mock_domain, mock_create_app + ): + """dry_run=True with an app in 'Creating' state returns ARN without waiting.""" + creating_app = { + "Arn": "arn:aws:sagemaker:us-east-1:123:mlflow-app/creating", + "Status": "Creating", + "AccountDefaultStatus": "ENABLED", + "MlflowVersion": "3.4.0", + } + mock_paginator = MagicMock() + mock_paginator.paginate.return_value = [{"Summaries": [creating_app]}] + mock_client.return_value.get_paginator.return_value = mock_paginator + mock_domain.return_value = None + + mock_session = Mock() + + result = _resolve_mlflow_resource_arn(mock_session, dry_run=True) + + assert result == "arn:aws:sagemaker:us-east-1:123:mlflow-app/creating" + mock_create_app.assert_not_called() + + @patch("sagemaker.train.common_utils.finetune_utils._create_mlflow_app_as_upgrade") + @patch("sagemaker.train.common_utils.finetune_utils._get_current_domain_id") + @patch("sagemaker.train.common_utils.finetune_utils._get_prod_sm_client") + def test_version_below_minimum_dry_run_skips_upgrade( + self, mock_client, mock_domain, mock_upgrade + ): + """dry_run=True with app below min version returns ARN without upgrading.""" + old_app = { + "Arn": "arn:aws:sagemaker:us-east-1:123:mlflow-app/old", + "Status": "Created", + "AccountDefaultStatus": "ENABLED", + "MlflowVersion": "2.0.0", + } + mock_paginator = MagicMock() + mock_paginator.paginate.return_value = [{"Summaries": [old_app]}] + mock_client.return_value.get_paginator.return_value = mock_paginator + mock_domain.return_value = None + + mock_session = Mock() + + result = _resolve_mlflow_resource_arn( + mock_session, min_mlflow_version="3.10", dry_run=True + ) + + assert result == "arn:aws:sagemaker:us-east-1:123:mlflow-app/old" + mock_upgrade.assert_not_called() + + @patch("sagemaker.train.common_utils.finetune_utils._get_current_domain_id") + @patch("sagemaker.train.common_utils.finetune_utils._get_prod_sm_client") + def test_ready_app_dry_run_returns_normally(self, mock_client, mock_domain): + """dry_run=True with an app in 'Updated' state returns ARN normally.""" + ready_app = { + "Arn": "arn:aws:sagemaker:us-east-1:123:mlflow-app/ready", + "Status": "Updated", + "AccountDefaultStatus": "ENABLED", + "MlflowVersion": "3.4.0", + } + mock_paginator = MagicMock() + mock_paginator.paginate.return_value = [{"Summaries": [ready_app]}] + mock_client.return_value.get_paginator.return_value = mock_paginator + mock_domain.return_value = None + + mock_session = Mock() + + result = _resolve_mlflow_resource_arn(mock_session, dry_run=True) + + assert result == "arn:aws:sagemaker:us-east-1:123:mlflow-app/ready" + + +class TestResolveMlflowDryRunWarnings: + """Verify appropriate warnings are logged during dry_run.""" + + @patch("sagemaker.train.common_utils.finetune_utils._get_current_domain_id") + @patch("sagemaker.train.common_utils.finetune_utils._get_prod_sm_client") + def test_no_apps_logs_warning(self, mock_client, mock_domain, caplog): + """Warns that job submission would create an app.""" + mock_paginator = MagicMock() + mock_paginator.paginate.return_value = [{"Summaries": []}] + mock_client.return_value.get_paginator.return_value = mock_paginator + mock_domain.return_value = None + + mock_session = Mock() + + with caplog.at_level(logging.WARNING): + _resolve_mlflow_resource_arn(mock_session, dry_run=True) + + assert "No MLflow app exists" in caplog.text + assert "would create a new app" in caplog.text + + @patch("sagemaker.train.common_utils.finetune_utils._get_current_domain_id") + @patch("sagemaker.train.common_utils.finetune_utils._get_prod_sm_client") + def test_creating_app_logs_warning(self, mock_client, mock_domain, caplog): + """Warns that job submission would block on a Creating app.""" + creating_app = { + "Arn": "arn:aws:sagemaker:us-east-1:123:mlflow-app/creating", + "Status": "Creating", + "AccountDefaultStatus": "ENABLED", + "MlflowVersion": "3.4.0", + } + mock_paginator = MagicMock() + mock_paginator.paginate.return_value = [{"Summaries": [creating_app]}] + mock_client.return_value.get_paginator.return_value = mock_paginator + mock_domain.return_value = None + + mock_session = Mock() + + with caplog.at_level(logging.WARNING): + _resolve_mlflow_resource_arn(mock_session, dry_run=True) + + assert "Creating" in caplog.text + assert "would block" in caplog.text + + +class TestCreateMlflowConfigDryRun: + """Verify _create_mlflow_config passes dry_run through.""" + + @patch("sagemaker.train.common_utils.finetune_utils._resolve_mlflow_resource_arn") + def test_passes_dry_run_to_resolve(self, mock_resolve): + """dry_run flag is forwarded to _resolve_mlflow_resource_arn.""" + mock_resolve.return_value = None + mock_session = Mock() + + _create_mlflow_config(mock_session, dry_run=True) + + mock_resolve.assert_called_once_with(mock_session, None, dry_run=True) + + @patch("sagemaker.train.common_utils.finetune_utils._resolve_mlflow_resource_arn") + def test_dry_run_false_by_default(self, mock_resolve): + """dry_run defaults to False.""" + mock_resolve.return_value = None + mock_session = Mock() + + _create_mlflow_config(mock_session) + + mock_resolve.assert_called_once_with(mock_session, None, dry_run=False) diff --git a/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py b/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py index a69f201624..291c7cc79e 100644 --- a/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py +++ b/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py @@ -719,3 +719,71 @@ def test_paginates(self, mock_session_cls): result = MultiTurnRLTrainer.list_bedrock_agentcore_runtimes() assert len(result) == 2 assert mock_client.list_agent_runtimes.call_count == 2 + + +class TestDryRun: + """Test dry_run=True skips job submission and MLflow creation.""" + + def _make_trainer(self): + """Create a trainer with mocked internals for dry_run testing.""" + trainer = object.__new__(MultiTurnRLTrainer) + trainer.agent_env = BEDROCK_AGENT_ARN + trainer.bedrock_agentcore_qualifier = "DEFAULT" + trainer.s3_output_path = S3_OUTPUT + trainer.output_model_package_group = MPG_ARN + trainer.intermediate_checkpoint_model_package_group = "arn:aws:sagemaker:us-west-2:123456789012:model-package-group/ckpt-mpg" + trainer.mlflow_app_arn = None # Force MLflow resolution + trainer.mlflow_experiment_name = None + trainer.mlflow_run_name = None + trainer.accept_eula = True + trainer.kms_key_arn = None + trainer.networking = None + trainer.model = "test-model-id" + trainer.validation_dataset = None + trainer._model_arn = MODEL_ARN + trainer._model_name = "test-model" + trainer.training_dataset = S3_DATA + trainer.hyperparameters = MagicMock() + trainer.hyperparameters.to_dict.return_value = {} + trainer.hyperparameters._specs = {} + trainer._hp_defaults = {} + trainer._final_hyperparameters = {} + mock_session = MagicMock() + mock_session.sagemaker_config = {"SchemaVersion": "1.0"} + mock_session.boto_session.region_name = "us-west-2" + trainer.sagemaker_session = mock_session + trainer.role = "arn:aws:iam::123456789012:role/TestRole" + trainer.base_job_name = "test-mtrl" + trainer._recipe_path = None + trainer._overrides = None + trainer._resolved_recipe_cache = None + return trainer + + @patch("sagemaker.train.multi_turn_rl_trainer._resolve_mlflow_resource_arn") + @patch("sagemaker.train.multi_turn_rl_trainer.Job") + @patch("sagemaker.train.multi_turn_rl_trainer.TrainDefaults.get_role") + def test_dry_run_skips_job_creation(self, mock_get_role, mock_job_cls, mock_resolve_mlflow): + """dry_run=True returns None without calling Job.create.""" + mock_resolve_mlflow.return_value = None + mock_get_role.return_value = "arn:aws:iam::123456789012:role/TestRole" + trainer = self._make_trainer() + + result = trainer.train(dry_run=True) + + assert result is None + mock_job_cls.create.assert_not_called() + + @patch("sagemaker.train.multi_turn_rl_trainer._resolve_mlflow_resource_arn") + @patch("sagemaker.train.multi_turn_rl_trainer.Job") + @patch("sagemaker.train.multi_turn_rl_trainer.TrainDefaults.get_role") + def test_dry_run_passes_flag_to_mlflow_resolver(self, mock_get_role, mock_job_cls, mock_resolve_mlflow): + """dry_run=True is forwarded to _resolve_mlflow_resource_arn.""" + mock_resolve_mlflow.return_value = None + mock_get_role.return_value = "arn:aws:iam::123456789012:role/TestRole" + trainer = self._make_trainer() + + trainer.train(dry_run=True) + + # Verify dry_run=True was passed + call_kwargs = mock_resolve_mlflow.call_args[1] + assert call_kwargs["dry_run"] is True