Merge branch 'master' into master

Author: rvasahu-amazon

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

src/sagemaker/serve/builder/model_builder.py

@@ -528,6 +528,8 @@ def _user_agent_decorator(self, func):
         def wrapper(*args, **kwargs):
             # Call the original function
             result = func(*args, **kwargs)
+            if "ModelBuilder" in result:
+                return result
             return result + " ModelBuilder"
 
         return wrapper

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(