diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index 00b01d2156..db8d572a75 100644 --- a/src/sagemaker/estimator.py +++ b/src/sagemaker/estimator.py @@ -578,6 +578,17 @@ def __init__( self.dependencies = dependencies or [] self.uploaded_code: Optional[UploadedCode] = None + # Check that the user properly sets both subnet and secutiry_groupe_ids + if ( + subnets is not None + and security_group_ids is None + or security_group_ids is not None + and subnets is None + ): + raise RuntimeError( + "When setting up custom VPC, both subnets and security_group_ids must be set" + ) + if self.instance_type in ("local", "local_gpu"): if self.instance_type == "local_gpu" and self.instance_count > 1: raise RuntimeError("Distributed Training in Local GPU is not supported") diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py index ae75256794..437c150c8b 100644 --- a/tests/unit/test_estimator.py +++ b/tests/unit/test_estimator.py @@ -512,6 +512,24 @@ def test_framework_all_init_args(sagemaker_session): } +def test_subnets_without_security_groups(sagemaker_session): + with pytest.raises(RuntimeError): + DummyFramework( + entry_point=SCRIPT_PATH, + sagemaker_session=sagemaker_session, + subnets=["123"], + ) + + +def test_security_groups_without_subnets(sagemaker_session): + with pytest.raises(RuntimeError): + DummyFramework( + entry_point=SCRIPT_PATH, + sagemaker_session=sagemaker_session, + security_group_ids=["123"], + ) + + def test_framework_without_role_parameter(sagemaker_session): with pytest.raises(ValueError): DummyFramework(