From 5c8e28b864ba210d19b9bc75e65029dcc0e720ec Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Mon, 20 Jul 2026 12:19:36 -0700 Subject: [PATCH] fix(train): correct Networking field names in ModelTrainer intelligent defaults Fixes two bugs in _populate_intelligent_defaults_from_training_job_space() in model_trainer.py that prevent ModelTrainer from working when sagemaker_config has a VpcConfig set. Bug 1: Networking() was passed default_enable_network_isolation, which is not a field in sagemaker-core 2.x. The correct field is enable_network_isolation. This raised a pydantic ValidationError on every train() call when VPC config was present. Bug 2: When updating an existing Networking object, security_group_ids was mistakenly assigned the subnets value from TRAINING_JOB_SUBNETS_PATH instead of security_group_ids from TRAINING_JOB_SECURITY_GROUP_IDS_PATH. Adds unit tests covering both bug paths (fresh Networking construction and updating an existing Networking object). Fix originally from #5866; unit tests added here. --- .../src/sagemaker/train/model_trainer.py | 6 +- .../tests/unit/train/test_model_trainer.py | 115 ++++++++++++++++++ 2 files changed, 118 insertions(+), 3 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 420a5543a8..ea3d62d145 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -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 ), @@ -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: diff --git a/sagemaker-train/tests/unit/train/test_model_trainer.py b/sagemaker-train/tests/unit/train/test_model_trainer.py index bc3a9885a6..b2da7cd2a4 100644 --- a/sagemaker-train/tests/unit/train/test_model_trainer.py +++ b/sagemaker-train/tests/unit/train/test_model_trainer.py @@ -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 @@ -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"]