Set the nccl env to execute gpu test task. (intelligent-machine-learn…

…ing#1148)

* Set the nccl env to execute gpu test task.

* Fix test cases.

* Fix by comments.

dlrover/trainer/tests/torch/node_check_test.py

@@ -17,6 +17,7 @@
 
 from dlrover.trainer.torch.node_check.ascend_npu import main as npu_main
 from dlrover.trainer.torch.node_check.nvidia_gpu import main as gpu_main
+from dlrover.trainer.torch.node_check.nvidia_gpu import set_nccl_env
 from dlrover.trainer.torch.node_check.utils import mock_error
 
 
@@ -65,3 +66,10 @@ def test_mock_error(self):
  except ValueError:
  raised_error = True
  self.assertTrue(raised_error)
+
+ def test_set_nccl_env(self):
+ os.environ[
+ "NCCL_SETTINGS"
+ ] = "NCCL_DEBUG=INFO,NCCL_SOCKET_IFNAME=eth0,NCCL_IB_GID_INDEX=3"
+ set_nccl_env()
+ self.assertEqual(os.environ["NCCL_SOCKET_IFNAME"], "eth0")

dlrover/trainer/torch/node_check/nvidia_gpu.py

@@ -20,6 +20,13 @@
 from .utils import bm_allreduce, matmul, record_execution_time
 
 
+def set_nccl_env():
+ env_conf = os.getenv("NCCL_SETTINGS", "")
+ for item in env_conf.split(","):
+ k, v = item.split("=")
+ os.environ[k] = v
+
+
 @record_execution_time
 def main():
  use_cuda = torch.cuda.is_available()
@@ -39,4 +46,5 @@ def main():
 
 
 if __name__ == "__main__":
+ set_nccl_env()
  t = main()