Fix unit test failures

Author: satishpasumarthi

src/sagemaker/fw_utils.py

@@ -859,27 +859,28 @@ def validate_distribution_for_instance_type(instance_type, distribution):
         instance_type (str): A string representing the type of training instance selected.
         distribution (dict): A dictionary with information to enable distributed training.
     """
-    match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
     err_msg = ""
-    if match and match[1].startswith("trn"):
-        keys = list(distribution.keys())
-        if len(keys) == 0:
-            return
-        if len(keys) == 1:
-            distribution_strategy = keys[0]
-            if distribution_strategy != "torch_distributed":
+    if isinstance(instance_type, str):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match and match[1].startswith("trn"):
+            keys = list(distribution.keys())
+            if len(keys) == 0:
+                return
+            if len(keys) == 1:
+                distribution_strategy = keys[0]
+                if distribution_strategy != "torch_distributed":
+                    err_msg += (
+                        f"Provided distribution strategy {distribution_strategy} is not supported"
+                        " for Trainium instances.\n"
+                        "Please specify one of the following supported distribution strategies:"
+                        f" {TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES} \n"
+                    )
+            elif len(keys) > 1:
                 err_msg += (
-                    f"Provided distribution strategy {distribution_strategy} is not supported for"
-                    " Trainium instances.\n"
+                    "Multiple distribution strategies are not supported for Trainium instances.\n"
                     "Please specify one of the following supported distribution strategies:"
-                    f" {TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES} \n"
+                    f" {TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES} "
                 )
-        elif len(keys) > 1:
-            err_msg += (
-                "Multiple distribution strategies are not supported for Trainium instances.\n"
-                "Please specify one of the following supported distribution strategies:"
-                f" {TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES} "
-            )
 
     if err_msg:
         raise ValueError(err_msg)

src/sagemaker/pytorch/estimator.py

@@ -230,12 +230,13 @@ def __init__(
             if self.framework_version and Version(self.framework_version) >= Version("1.3"):
                 kwargs["enable_sagemaker_metrics"] = True
 
-        if "entry_point" not in kwargs:
-            kwargs["entry_point"] = entry_point
-
         super(PyTorch, self).__init__(
             entry_point, source_dir, hyperparameters, image_uri=image_uri, **kwargs
         )
+
+        if "entry_point" not in kwargs:
+            kwargs["entry_point"] = entry_point
+
         if distribution is not None:
             distribution = validate_distribution(
                 distribution,

tests/integ/test_torch_distributed.py

@@ -24,13 +24,9 @@
 
 
 @pytest.mark.skip(
-    reason="This test is skipped for now due ML capacity error."
+    reason="Disabling until the launch of SM Trainium containers"
     "This test should be re-enabled later."
 )
-@pytest.mark.skipif(
-    integ.test_region() not in integ.DATA_PARALLEL_TESTING_REGIONS,
-    reason="Only allow this test to run in IAD and CMH to limit usage of ml.trn1.2xlarge",
-)
 def test_torch_distributed_trn1_pt_mnist(
     sagemaker_session,
     torch_distributed_framework_version,

tests/unit/test_fw_utils.py

@@ -655,7 +655,7 @@ def test_validate_distribution_not_raises():
                 None,  # framework_version
                 None,  # py_version
                 "custom-container",
-                {"instance_type": instance_type},  # kwargs
+                {"instance_type": instance_type, "entry_point": "train.py"},  # kwargs
             )
 
     for framework in frameworks:
@@ -683,7 +683,7 @@ def test_validate_distribution_not_raises():
                 None,  # framework_version
                 None,  # py_version
                 "custom-container",
-                {},  # kwargs
+                {"entry_point": "train.py"},  # kwargs
             )
 
 
@@ -723,7 +723,7 @@ def test_validate_distribution_raises():
                     None,  # framework_version
                     None,  # py_version
                     "custom-container",
-                    {"instance_type": instance_type},  # kwargs
+                    {"instance_type": instance_type, "entry_point": "train.py"},  # kwargs
                 )
 
     for framework in frameworks:
@@ -952,7 +952,7 @@ def test_validate_torch_distributed_not_raises():
 
     # Case 1: Framework is PyTorch, but distribution is not torch_distributed
     torch_distributed_disabled = {"torch_distributed": {"enabled": False}}
-    fw_utils.validate_pytorch_distribution(
+    fw_utils.validate_torch_distributed_distribution(
         instance_type="ml.trn1.2xlarge",
         distribution=torch_distributed_disabled,
         framework_version="1.11.0",