From d3171d42d9304b93cb3cf843a8b8c58daffa0eb1 Mon Sep 17 00:00:00 2001 From: Mourya Baddam Date: Tue, 24 Feb 2026 01:50:06 +0000 Subject: [PATCH 1/2] fix: Apply default experiment config for pipelines only in regions with SageMaker Experiments --- src/sagemaker/workflow/_utils.py | 30 ++++++ src/sagemaker/workflow/pipeline.py | 9 ++ .../unit/sagemaker/workflow/test_pipeline.py | 93 +++++++++++++++++++ 3 files changed, 132 insertions(+) diff --git a/src/sagemaker/workflow/_utils.py b/src/sagemaker/workflow/_utils.py index b3796cac13..4be07361e3 100644 --- a/src/sagemaker/workflow/_utils.py +++ b/src/sagemaker/workflow/_utils.py @@ -65,6 +65,36 @@ --source_dir "${var_source_dir}" """ +# Static list of regions where SageMaker Experiments is Available. +# Note: Experiments is not expanding to new regions, so this list is static. +EXPERIMENTS_REGIONS = frozenset([ + "us-east-1", # US East (N. Virginia) + "us-east-2", # US East (Ohio) + "us-west-1", # US West (N. California) + "us-west-2", # US West (Oregon) + "ca-central-1", # Canada (Central) + "eu-west-1", # Europe (Ireland) + "eu-west-2", # Europe (London) + "eu-west-3", # Europe (Paris) + "eu-central-1", # Europe (Frankfurt) + "eu-north-1", # Europe (Stockholm) + "eu-south-1", # Europe (Milan) + "eu-south-2", # Europe (Spain) + "ap-northeast-1", # Asia Pacific (Tokyo) + "ap-northeast-2", # Asia Pacific (Seoul) + "ap-northeast-3", # Asia Pacific (Osaka) + "ap-southeast-1", # Asia Pacific (Singapore) + "ap-southeast-2", # Asia Pacific (Sydney) + "ap-southeast-3", # Asia Pacific (Jakarta) + "ap-south-1", # Asia Pacific (Mumbai) + "ap-east-1", # Asia Pacific (Hong Kong) + "sa-east-1", # South America (São Paulo) + "af-south-1", # Africa (Cape Town) + "me-south-1", # Middle East (Bahrain) + "il-central-1", # Israel (Tel Aviv) + "cn-north-1", # China (Beijing) + "cn-northwest-1", # China (Ningxia) +]) class _RepackModelStep(TrainingStep): """Repacks model artifacts with custom inference entry points. diff --git a/src/sagemaker/workflow/pipeline.py b/src/sagemaker/workflow/pipeline.py index 25d236f25e..539e0c25ad 100644 --- a/src/sagemaker/workflow/pipeline.py +++ b/src/sagemaker/workflow/pipeline.py @@ -34,6 +34,7 @@ from sagemaker.s3_utils import s3_path_join from sagemaker.session import Session from sagemaker.utils import resolve_value_from_config, retry_with_backoff, format_tags, Tags +from sagemaker.workflow._utils import EXPERIMENTS_REGIONS from sagemaker.workflow.callback_step import CallbackOutput, CallbackStep from sagemaker.workflow._event_bridge_client_helper import ( EventBridgeSchedulerHelper, @@ -99,6 +100,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 + SageMaker Experiments is Available. steps (Sequence[Union[Step, StepCollection, StepOutput]]): The list of the non-conditional steps associated with the pipeline. Any steps that are within the `if_steps` or `else_steps` of a `ConditionStep` cannot be listed in the steps of a @@ -127,6 +130,12 @@ def __init__( self.sagemaker_session.boto_session.client("scheduler"), ) + # Apply default experiment config only in regions where Experiments is available + if pipeline_experiment_config is _DEFAULT_EXPERIMENT_CFG: + region = self.sagemaker_session.boto_region_name + if region not in EXPERIMENTS_REGIONS: + self.pipeline_experiment_config = None + @property def latest_pipeline_version_id(self): """Retrieves the latest version id of this pipeline""" diff --git a/tests/unit/sagemaker/workflow/test_pipeline.py b/tests/unit/sagemaker/workflow/test_pipeline.py index 21ee5e75c7..acbb9e2c95 100644 --- a/tests/unit/sagemaker/workflow/test_pipeline.py +++ b/tests/unit/sagemaker/workflow/test_pipeline.py @@ -791,6 +791,99 @@ def test_pipeline_disable_experiment_config(sagemaker_session_mock): ) +@patch("sagemaker.workflow.pipeline.EXPERIMENTS_REGIONS", {"us-east-1", "us-west-2"}) +def test_pipeline_init_with_default_experiment_config_in_supported_region(sagemaker_session_mock): + """Test that default experiment config is preserved in regions where Experiments is available.""" + sagemaker_session_mock.boto_region_name = "us-east-1" + + pipeline = Pipeline( + name="MyPipeline", + steps=[CustomStep(name="MyStep", input_data="input")], + sagemaker_session=sagemaker_session_mock, + ) + + # Default experiment config should be preserved + assert pipeline.pipeline_experiment_config is not None + assert ( + pipeline.pipeline_experiment_config.experiment_name == ExecutionVariables.PIPELINE_NAME + ) + assert ( + pipeline.pipeline_experiment_config.trial_name + == ExecutionVariables.PIPELINE_EXECUTION_ID + ) + + +@patch("sagemaker.workflow.pipeline.EXPERIMENTS_REGIONS", {"us-east-1", "us-west-2"}) +def test_pipeline_init_with_default_experiment_config_in_unsupported_region( + sagemaker_session_mock, +): + """Test that default experiment config is disabled in regions where Experiments is not available.""" + sagemaker_session_mock.boto_region_name = "ap-southeast-4" # Not in EXPERIMENTS_REGIONS + + pipeline = Pipeline( + name="MyPipeline", + steps=[CustomStep(name="MyStep", input_data="input")], + sagemaker_session=sagemaker_session_mock, + ) + + # Default experiment config should be set to None + assert pipeline.pipeline_experiment_config is None + + +@patch("sagemaker.workflow.pipeline.EXPERIMENTS_REGIONS", {"us-east-1", "us-west-2"}) +def test_pipeline_init_with_explicit_experiment_config_in_unsupported_region( + sagemaker_session_mock, +): + """Test that explicitly set experiment config is preserved even in unsupported regions.""" + sagemaker_session_mock.boto_region_name = "ap-southeast-4" # Not in EXPERIMENTS_REGIONS + + explicit_config = PipelineExperimentConfig("MyExperiment", "MyTrial") + pipeline = Pipeline( + name="MyPipeline", + pipeline_experiment_config=explicit_config, + steps=[CustomStep(name="MyStep", input_data="input")], + sagemaker_session=sagemaker_session_mock, + ) + + # Explicitly set experiment config should be preserved + assert pipeline.pipeline_experiment_config is not None + assert pipeline.pipeline_experiment_config.experiment_name == "MyExperiment" + assert pipeline.pipeline_experiment_config.trial_name == "MyTrial" + + +@patch("sagemaker.workflow.pipeline.EXPERIMENTS_REGIONS", {"us-east-1", "us-west-2"}) +def test_pipeline_init_with_none_experiment_config_in_supported_region(sagemaker_session_mock): + """Test that explicitly setting experiment config to None is respected in supported regions.""" + sagemaker_session_mock.boto_region_name = "us-east-1" + + pipeline = Pipeline( + name="MyPipeline", + pipeline_experiment_config=None, + steps=[CustomStep(name="MyStep", input_data="input")], + sagemaker_session=sagemaker_session_mock, + ) + + # Explicitly set None should be preserved + assert pipeline.pipeline_experiment_config is None + + +@patch("sagemaker.workflow.pipeline.EXPERIMENTS_REGIONS", {"us-east-1"}) +def test_pipeline_definition_without_experiment_config_in_unsupported_region( + sagemaker_session_mock, +): + """Test that pipeline definition has no experiment config in unsupported regions.""" + sagemaker_session_mock.boto_region_name = "ap-southeast-4" # Not in EXPERIMENTS_REGIONS + + pipeline = Pipeline( + name="MyPipeline", + steps=[CustomStep(name="MyStep", input_data="input")], + sagemaker_session=sagemaker_session_mock, + ) + + definition = json.loads(pipeline.definition()) + assert definition["PipelineExperimentConfig"] is None + + def test_pipeline_list_executions(sagemaker_session_mock): sagemaker_session_mock.sagemaker_client.list_pipeline_executions.return_value = { "PipelineExecutionSummaries": [Mock()], From 3550fc3c1c2cbfb2965db6e30eafdddcacc3d9d7 Mon Sep 17 00:00:00 2001 From: Mourya Baddam Date: Tue, 24 Feb 2026 01:50:06 +0000 Subject: [PATCH 2/2] fix: Apply default experiment config for pipelines only in regions with SageMaker Experiments --- .../unit/sagemaker/workflow/test_pipeline.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tests/unit/sagemaker/workflow/test_pipeline.py b/tests/unit/sagemaker/workflow/test_pipeline.py index acbb9e2c95..baa8a6baa5 100644 --- a/tests/unit/sagemaker/workflow/test_pipeline.py +++ b/tests/unit/sagemaker/workflow/test_pipeline.py @@ -790,7 +790,6 @@ def test_pipeline_disable_experiment_config(sagemaker_session_mock): } ) - @patch("sagemaker.workflow.pipeline.EXPERIMENTS_REGIONS", {"us-east-1", "us-west-2"}) def test_pipeline_init_with_default_experiment_config_in_supported_region(sagemaker_session_mock): """Test that default experiment config is preserved in regions where Experiments is available.""" @@ -866,24 +865,6 @@ def test_pipeline_init_with_none_experiment_config_in_supported_region(sagemaker # Explicitly set None should be preserved assert pipeline.pipeline_experiment_config is None - -@patch("sagemaker.workflow.pipeline.EXPERIMENTS_REGIONS", {"us-east-1"}) -def test_pipeline_definition_without_experiment_config_in_unsupported_region( - sagemaker_session_mock, -): - """Test that pipeline definition has no experiment config in unsupported regions.""" - sagemaker_session_mock.boto_region_name = "ap-southeast-4" # Not in EXPERIMENTS_REGIONS - - pipeline = Pipeline( - name="MyPipeline", - steps=[CustomStep(name="MyStep", input_data="input")], - sagemaker_session=sagemaker_session_mock, - ) - - definition = json.loads(pipeline.definition()) - assert definition["PipelineExperimentConfig"] is None - - def test_pipeline_list_executions(sagemaker_session_mock): sagemaker_session_mock.sagemaker_client.list_pipeline_executions.return_value = { "PipelineExecutionSummaries": [Mock()],