Merge branch 'master' into dacorvo-fix-hf-vllm-image-uri

Author: mollyheamazon

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.

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):

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"

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

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 = {}