fix: Apply default experiment config for pipelines only in regions with SageMaker Experiments

Author: Mourya Baddam

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.

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"""

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()],