Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions src/sagemaker/workflow/_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
9 changes: 9 additions & 0 deletions src/sagemaker/workflow/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"""
Expand Down
74 changes: 74 additions & 0 deletions tests/unit/sagemaker/workflow/test_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -790,6 +790,80 @@ 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

def test_pipeline_list_executions(sagemaker_session_mock):
sagemaker_session_mock.sagemaker_client.list_pipeline_executions.return_value = {
Expand Down
Loading