Change: Raise error if custom VPC is halfway defined (#4269)

Author: martinRenou

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

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(