fix: deriver master node from training environment

Author: satishpasumarthi

src/sagemaker_pytorch_container/training.py

@@ -47,7 +47,7 @@ def train(training_environment):
 
     _set_nccl_environment(training_environment.network_interface_name)
 
-    _set_distributed_environment(training_environment.hosts)
+    _set_distributed_environment(training_environment)
 
     mpi_enabled = training_environment.additional_framework_parameters.get('sagemaker_mpi_enabled')
 
@@ -88,15 +88,15 @@ def _dns_lookup(host):
     return socket.gethostbyname(host)
 
 
-def _set_distributed_environment(hosts):
+def _set_distributed_environment(training_env):
     """Set environment variable for distributed training.
 
     Args:
         hosts: list of hosts that are used for training.
     """
     # According to https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html
     # hosts are sorted lexicographically.
-    os.environ['MASTER_ADDR'] = hosts[0]
+    os.environ['MASTER_ADDR'] = training_env.master_hostname
     os.environ['MASTER_PORT'] = MASTER_PORT