Add unit tests

Author: Yongyan Rao

tests/unit/test_fw_utils.py

@@ -41,6 +41,13 @@ def cd(path):
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name="boto_session", region_name="us-west-2")
+    boto_mock.client.describe_instance_types = Mock(
+        return_value={
+            "InstanceTypes": [
+                {"GpuInfo": {}}
+            ]
+        }
+    )
     session_mock = Mock(
         name="sagemaker_session", boto_session=boto_mock, s3_client=None, s3_resource=None
     )
@@ -733,6 +740,24 @@ def test_validate_smdistributed_not_raises():
             )
 
 
+def test_validate_distribution_instance_no_smdistributed(sagemaker_session):
+    distribution = {}
+    instance_type = "mock_type"
+    fw_utils.validate_distribution_instance(sagemaker_session, distribution, instance_type)
+
+
+def test_validate_distribution_instance_no_modelparallel(sagemaker_session):
+    distribution = {"smdistributed": {}}
+    instance_type = "mock_type"
+    fw_utils.validate_distribution_instance(sagemaker_session, distribution, instance_type)
+
+
+def test_validate_distribution_instance_disabled_modelparallel(sagemaker_session):
+    distribution = {"smdistributed": {"modelparallel": {"enabled": False}}}
+    instance_type = "mock_type"
+    fw_utils.validate_distribution_instance(sagemaker_session, distribution, instance_type)
+
+
 def test_validate_smdistributed_raises():
     bad_args = [
         {"smdistributed": "dummy"},