Fix failing tests

Author: vishwakaria

src/sagemaker/pytorch/estimator.py

@@ -37,7 +37,7 @@ class PyTorch(Framework):
 
     _framework_name = "pytorch"
     LAUNCH_PYTORCH_DDP_ENV_NAME = "sagemaker_pytorch_ddp_enabled"
-    INSTANCE_TYPE = "sagemaker_instance_type"
+    INSTANCE_TYPE_ENV_NAME = "sagemaker_instance_type"
 
     def __init__(
         self,
@@ -207,13 +207,15 @@ def __init__(
             )
         self.framework_version = framework_version
         self.py_version = py_version
-        self.instance_type = instance_type
 
         if "enable_sagemaker_metrics" not in kwargs:
             # enable sagemaker metrics for PT v1.3 or greater:
             if self.framework_version and Version(self.framework_version) >= Version("1.3"):
                 kwargs["enable_sagemaker_metrics"] = True
 
+        if "instance_type" in kwargs:
+            self.instance_type = kwargs["instance_type"]
+
         super(PyTorch, self).__init__(
             entry_point, source_dir, hyperparameters, image_uri=image_uri, **kwargs
         )
@@ -231,17 +233,19 @@ def __init__(
         self.distribution = distribution or {}
 
     def _pytorch_distribution_configuration(self, distribution):
-        """Returns a dict of distribution config
+        """Returns a dict of distribution config for PyTorch training
+        
         Args:
-            None
+            distribution (dict): A dictionary with information on how to run distributed training.
         Returns:
-            dict containing torch ddp config
+            dict containing Pytorch DDP config
         """
         distribution_config = {}
         if "pytorchddp" in distribution:
             pytorch_ddp_enabled = distribution.get("pytorchddp").get("enabled", False)
             distribution_config[self.LAUNCH_PYTORCH_DDP_ENV_NAME] = pytorch_ddp_enabled
-            distribution_config[self.INSTANCE_TYPE] = self.instance_type
+            if self.instance_type is not None:
+                distribution_config[self.INSTANCE_TYPE_ENV_NAME] = self.instance_type
         else:
             distribution_config = self._distribution_configuration(distribution=distribution)
         return distribution_config

tests/data/pytorch_ddp/mnist_pt.py

@@ -24,7 +24,7 @@
 from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP
 import smdistributed.dataparallel.torch.distributed as dist
 
-dist.init_process_group()
+dist.init_process_group(backend="nccl")
 
 
 class Net(nn.Module):

tests/integ/test_pytorchddp.py

@@ -13,12 +13,9 @@
 from __future__ import absolute_import
 
 import os
-
 import pytest
-
 import sagemaker.utils
 import tests.integ as integ
-
 from sagemaker.pytorch import PyTorch
 from tests.integ import timeout
 from tests.integ.test_pytorch import _upload_training_data