From 304f0049e91c8a53ef755be89e24903ce293390f Mon Sep 17 00:00:00 2001 From: Timothy Wu Date: Wed, 21 Jan 2026 04:46:58 -0800 Subject: [PATCH] Fix: Don't apply default experiment config for pipelines in non-Eureka GA regions --- .../src/sagemaker/mlops/workflow/_utils.py | 30 +++++++++++++ .../src/sagemaker/mlops/workflow/pipeline.py | 11 ++++- .../unit/workflow/test_pipeline_class.py | 9 +++- .../test_pipeline_experiment_config.py | 42 +++++++++++++++++++ .../workflow/test_pipeline_mlflow_config.py | 1 + 5 files changed, 90 insertions(+), 3 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py index 3b98b49028..f8fd25e39e 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py @@ -58,6 +58,36 @@ --source_dir "${var_source_dir}" """ +# Static list of regions where Experiments (Eureka) is Generally Available. +# Note: Experiments is not expanding to new regions, so this list is static. +EUREKA_GA_REGIONS = frozenset([ + "us-east-1", # iad (N. Virginia) + "us-east-2", # cmh (Ohio) + "us-west-1", # sfo (N. California) + "us-west-2", # pdx (Oregon) + "ca-central-1", # yul (Montreal) + "eu-west-1", # dub (Dublin) + "eu-west-2", # lhr (London) + "eu-west-3", # cdg (Paris) + "eu-central-1", # fra (Frankfurt) + "eu-north-1", # arn (Stockholm) + "eu-south-1", # mxp (Milan) + "eu-south-2", # zaz (Spain) + "ap-northeast-1", # nrt (Tokyo) + "ap-northeast-2", # icn (Seoul) + "ap-northeast-3", # kix (Osaka) + "ap-southeast-1", # sin (Singapore) + "ap-southeast-2", # syd (Sydney) + "ap-southeast-3", # cgk (Jakarta) + "ap-south-1", # bom (Mumbai) + "ap-east-1", # hkg (Hong Kong) + "sa-east-1", # gru (São Paulo) + "af-south-1", # cpt (Cape Town) + "me-south-1", # bah (Bahrain) + "il-central-1", # tlv (Tel Aviv) + "cn-north-1", # bjs (Beijing) + "cn-northwest-1", # zhy (Ningxia) +]) class _RepackModelStep(TrainingStep): """Repacks model artifacts with custom inference entry points. diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py index 88d6b2f485..9ed8e93fb9 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py @@ -49,6 +49,7 @@ RESOURCE_NOT_FOUND_EXCEPTION, EXECUTION_TIME_PIPELINE_PARAMETER_FORMAT, ) +from sagemaker.mlops.workflow._utils import EUREKA_GA_REGIONS from sagemaker.mlops.workflow.lambda_step import LambdaOutput, LambdaStep from sagemaker.core.shapes.shapes import MlflowConfig from sagemaker.core.helper.pipeline_variable import ( @@ -116,6 +117,8 @@ def __init__( the same name already exists. By default, pipeline name is used as experiment name and execution id is used as the trial name. If set to None, no experiment or trial will be created automatically. + Note: The default experiment config is only applied in regions where + Experiments (Eureka) is Generally Available. mlflow_config (Optional[MlflowConfig]): If set, the pipeline will be configured with MLflow tracking for experiment tracking and model versioning. steps (Sequence[Union[Step, StepOutput]]): The list of the @@ -133,7 +136,6 @@ def __init__( """ self.name = name self.parameters = parameters if parameters else [] - self.pipeline_experiment_config = pipeline_experiment_config self.mlflow_config = mlflow_config self.steps = steps if steps else [] self.sagemaker_session = sagemaker_session if sagemaker_session else Session() @@ -146,6 +148,13 @@ def __init__( self._event_bridge_scheduler_helper = EventBridgeSchedulerHelper( self.sagemaker_session.boto_session.client("scheduler"), ) + self.pipeline_experiment_config = pipeline_experiment_config + + # Apply default experiment config only in Eureka GA regions + if pipeline_experiment_config is _DEFAULT_EXPERIMENT_CFG: + region = self.sagemaker_session.boto_region_name + if region not in EUREKA_GA_REGIONS: + self.pipeline_experiment_config = None @property def latest_pipeline_version_id(self): diff --git a/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py b/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py index 59502bed3f..434f7d7c13 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py +++ b/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py @@ -33,6 +33,7 @@ def mock_session(): session.sagemaker_client = Mock() session.boto_session = Mock() session.boto_session.client = Mock(return_value=Mock()) + session.boto_region_name = "us-east-1" session.local_mode = False session.sagemaker_config = {} session._append_sagemaker_config_tags = Mock(return_value=[]) @@ -59,8 +60,12 @@ class TestPipelineInit: def test_init_minimal(self): """Test Pipeline initialization with minimal parameters.""" with patch('sagemaker.mlops.workflow.pipeline.Session') as mock_session_class: - mock_session_class.return_value = Mock() - + mock_session = Mock() + mock_session.boto_region_name = "us-east-1" + mock_session.boto_session = Mock() + mock_session.boto_session.client = Mock(return_value=Mock()) + mock_session_class.return_value = mock_session + pipeline = Pipeline(name="test-pipeline") assert pipeline.name == "test-pipeline" diff --git a/sagemaker-mlops/tests/unit/workflow/test_pipeline_experiment_config.py b/sagemaker-mlops/tests/unit/workflow/test_pipeline_experiment_config.py index 32fe8c6261..7cb0e63d5e 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_pipeline_experiment_config.py +++ b/sagemaker-mlops/tests/unit/workflow/test_pipeline_experiment_config.py @@ -13,9 +13,12 @@ """Unit tests for workflow pipeline_experiment_config.""" from __future__ import absolute_import +from unittest.mock import Mock + from sagemaker.mlops.workflow.pipeline_experiment_config import ( PipelineExperimentConfig, PipelineExperimentConfigProperties ) +from sagemaker.mlops.workflow.pipeline import Pipeline, _DEFAULT_EXPERIMENT_CFG from sagemaker.core.workflow.execution_variables import ExecutionVariables @@ -41,3 +44,42 @@ def test_pipeline_experiment_config_with_execution_variables(): def test_pipeline_experiment_config_properties(): assert PipelineExperimentConfigProperties.EXPERIMENT_NAME.name == "ExperimentName" assert PipelineExperimentConfigProperties.TRIAL_NAME.name == "TrialName" + + +def _create_mock_session(region: str) -> Mock: + """Helper to create a mock SageMaker session with specified region.""" + mock_session = Mock() + mock_session.boto_region_name = region + mock_session.boto_session = Mock() + mock_session.boto_session.client = Mock(return_value=Mock()) + mock_session.local_mode = False + return mock_session + + +def test_default_config_applied_in_ga_region(): + """Default config applied when nothing provided in GA region.""" + mock_session = _create_mock_session("us-east-1") + pipeline = Pipeline(name="test-pipeline", sagemaker_session=mock_session) + assert pipeline.pipeline_experiment_config == _DEFAULT_EXPERIMENT_CFG + + +def test_no_default_config_in_non_ga_region(): + """No default config when nothing provided in non-GA region (THE FIX).""" + mock_session = _create_mock_session("us-gov-west-1") + pipeline = Pipeline(name="test-pipeline", sagemaker_session=mock_session) + assert pipeline.pipeline_experiment_config is None + + +def test_explicit_none_respected_in_ga_region(): + """None gets default config in GA region.""" + mock_session = _create_mock_session("us-east-1") + pipeline = Pipeline(name="test-pipeline", sagemaker_session=mock_session, pipeline_experiment_config=None) + assert pipeline.pipeline_experiment_config is None + + +def test_custom_config_respected(): + """Custom config respected regardless of region.""" + mock_session = _create_mock_session("us-east-1") + custom_config = PipelineExperimentConfig("my-experiment", "my-trial") + pipeline = Pipeline(name="test-pipeline", sagemaker_session=mock_session, pipeline_experiment_config=custom_config) + assert pipeline.pipeline_experiment_config == custom_config \ No newline at end of file diff --git a/sagemaker-mlops/tests/unit/workflow/test_pipeline_mlflow_config.py b/sagemaker-mlops/tests/unit/workflow/test_pipeline_mlflow_config.py index eb787bfb82..833d0b8f0e 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_pipeline_mlflow_config.py +++ b/sagemaker-mlops/tests/unit/workflow/test_pipeline_mlflow_config.py @@ -26,6 +26,7 @@ def mock_session(): """Create a mock SageMaker session for testing.""" session = Mock() session.boto_session.client.return_value = Mock() + session.boto_region_name = "us-east-1" session.sagemaker_client = Mock() session.local_mode = False session.sagemaker_config = {}