Skip to content
Merged
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
6 changes: 3 additions & 3 deletions sagemaker-train/src/sagemaker/train/model_trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -331,7 +331,7 @@ def _populate_intelligent_defaults_from_training_job_space(self):
if not self.networking:
if default_enable_network_isolation is not None or default_vpc_config is not None:
self.networking = Networking(
default_enable_network_isolation=default_enable_network_isolation,
enable_network_isolation=default_enable_network_isolation,
subnets=self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_SUBNETS_PATH
),
Expand All @@ -347,8 +347,8 @@ def _populate_intelligent_defaults_from_training_job_space(self):
config_path=TRAINING_JOB_SUBNETS_PATH
)
if self.networking.security_group_ids is None:
self.networking.subnets = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_SUBNETS_PATH
self.networking.security_group_ids = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_SECURITY_GROUP_IDS_PATH
)

if not self.output_data_config:
Expand Down
115 changes: 115 additions & 0 deletions sagemaker-train/tests/unit/train/test_model_trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,10 +34,15 @@
MODEL_TRAINER,
_simple_path,
TRAINING_JOB_RESOURCE_CONFIG_PATH,
TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH,
TRAINING_JOB_VPC_CONFIG_PATH,
TRAINING_JOB_SUBNETS_PATH,
TRAINING_JOB_SECURITY_GROUP_IDS_PATH,
SAGEMAKER,
PYTHON_SDK,
MODULES,
)
from sagemaker.core.config.config_manager import SageMakerConfig
from sagemaker.core.helper.session_helper import Session
from sagemaker.train.model_trainer import ModelTrainer, Mode
from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE
Expand Down Expand Up @@ -1756,3 +1761,113 @@ def test_resolve_staging_bucket_returns_default_on_iam_error(model_trainer):

assert bucket == DEFAULT_BUCKET
assert prefix is None


# Networking intelligent defaults from the Training Job config space. These guard the two
# bugs fixed in _populate_intelligent_defaults_from_training_job_space: (1) the Networking
# constructor must be given `enable_network_isolation` (not the non-existent
# `default_enable_network_isolation`), and (2) an existing Networking object with an unset
# `security_group_ids` must be filled from the security-group-ids path, not the subnets path.

NETWORKING_DEFAULT_SUBNETS = ["subnet-000000000000"]
NETWORKING_DEFAULT_SECURITY_GROUP_IDS = ["sg-000000000000"]


def _make_config_mgr(values):
"""Return a real SageMakerConfig whose resolve_value_from_config is keyed by config_path.

A real instance (rather than a bare MagicMock) is required because ModelTrainer validates
that config_mgr is a SageMakerConfig on assignment.
"""
config_mgr = SageMakerConfig()
config_mgr.resolve_value_from_config = MagicMock(
side_effect=lambda **kwargs: values.get(kwargs["config_path"])
)
return config_mgr


def test_networking_intelligent_defaults_creates_networking(model_trainer):
"""A VpcConfig in the config space builds a valid Networking with the correct fields.

Regression for Bug 1: passing `default_enable_network_isolation` to Networking() raised
a pydantic ValidationError because that field does not exist.
"""
model_trainer.networking = None
model_trainer.config_mgr = _make_config_mgr(
{
TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH: True,
TRAINING_JOB_VPC_CONFIG_PATH: {
"subnets": NETWORKING_DEFAULT_SUBNETS,
"security_group_ids": NETWORKING_DEFAULT_SECURITY_GROUP_IDS,
},
TRAINING_JOB_SUBNETS_PATH: NETWORKING_DEFAULT_SUBNETS,
TRAINING_JOB_SECURITY_GROUP_IDS_PATH: NETWORKING_DEFAULT_SECURITY_GROUP_IDS,
}
)

model_trainer._populate_intelligent_defaults_from_training_job_space()

assert model_trainer.networking is not None
assert model_trainer.networking.enable_network_isolation is True
assert model_trainer.networking.subnets == NETWORKING_DEFAULT_SUBNETS
assert model_trainer.networking.security_group_ids == NETWORKING_DEFAULT_SECURITY_GROUP_IDS


def test_networking_intelligent_defaults_no_vpc_config_leaves_networking_unset(model_trainer):
"""No network isolation and no VpcConfig means Networking is not created."""
model_trainer.networking = None
model_trainer.config_mgr = _make_config_mgr({})

model_trainer._populate_intelligent_defaults_from_training_job_space()

assert model_trainer.networking is None


def test_networking_intelligent_defaults_fills_security_group_ids_on_existing(model_trainer):
"""An existing Networking with unset security_group_ids is filled from the SG path.

Regression for Bug 2: the missing security_group_ids branch mistakenly assigned to
`subnets` using the subnets path, so security_group_ids was never populated.
"""
model_trainer.networking = Networking(
enable_network_isolation=False,
subnets=["subnet-preexisting"],
security_group_ids=None,
)
model_trainer.config_mgr = _make_config_mgr(
{
TRAINING_JOB_SUBNETS_PATH: NETWORKING_DEFAULT_SUBNETS,
TRAINING_JOB_SECURITY_GROUP_IDS_PATH: NETWORKING_DEFAULT_SECURITY_GROUP_IDS,
}
)

model_trainer._populate_intelligent_defaults_from_training_job_space()

# security_group_ids is filled from the security-group-ids path...
assert model_trainer.networking.security_group_ids == NETWORKING_DEFAULT_SECURITY_GROUP_IDS
# ...and the pre-existing subnets are preserved (not overwritten).
assert model_trainer.networking.subnets == ["subnet-preexisting"]
model_trainer.config_mgr.resolve_value_from_config.assert_any_call(
config_path=TRAINING_JOB_SECURITY_GROUP_IDS_PATH
)


def test_networking_intelligent_defaults_fills_subnets_on_existing(model_trainer):
"""An existing Networking with unset subnets is filled from the subnets path."""
model_trainer.networking = Networking(
enable_network_isolation=False,
subnets=None,
security_group_ids=["sg-preexisting"],
)
model_trainer.config_mgr = _make_config_mgr(
{
TRAINING_JOB_SUBNETS_PATH: NETWORKING_DEFAULT_SUBNETS,
TRAINING_JOB_SECURITY_GROUP_IDS_PATH: NETWORKING_DEFAULT_SECURITY_GROUP_IDS,
}
)

model_trainer._populate_intelligent_defaults_from_training_job_space()

assert model_trainer.networking.subnets == NETWORKING_DEFAULT_SUBNETS
# pre-existing security_group_ids are preserved.
assert model_trainer.networking.security_group_ids == ["sg-preexisting"]
Loading