Update driver_pytorch.py to fix NCCL_SOCKET_IFNAME issue.

Author: qiuosier

ads/jobs/templates/driver_pytorch.py

@@ -270,6 +270,7 @@ def find_self_ip(self):
                 ip = snics[0].address
                 if ipaddress.ip_address(ip) in ipaddress.ip_network(cidr):
                     logger.info("Node IP address: %s", ip)
+                    # Specify the network interface for NCCL/GLOO
                     os.environ["GLOO_SOCKET_IFNAME"] = interface
                     os.environ["NCCL_SOCKET_IFNAME"] = interface
                     return ip
@@ -579,13 +580,22 @@ def save_deepspeed_env(self):
                 # as the .deepspeed_env file is parsed line by line.
                 if not v or "\n" in v:
                     continue
+                # Ignore variables that are node specific
+                # The network interface name for each job run is a unique string, e.g. ens300f0v1604
+                if k in ["NCCL_SOCKET_IFNAME", "GLOO_SOCKET_IFNAME", "JOB_RUN_OCID"]:
+                    continue
                 # Quote the value if it contains space
                 # Environment variable containing space may not be exported correctly when using pdsh
                 # https://github.com/microsoft/DeepSpeed/blob/v0.9.2/deepspeed/launcher/multinode_runner.py#L79
                 if " " in v:
                     v = shlex.quote(v)
 
                 f.write(f"{k}={v}\n")
+            # The following are required for specifying the network interface to be used by NCCL/GLOO
+            # The value should be the prefix of the expected network interface name
+            # https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-socket-ifname
+            f.write("NCCL_SOCKET_IFNAME=ens\n")
+            f.write("GLOO_SOCKET_IFNAME=ens\n")
         logger.debug("Environment variables saved to %s", self.ENV_FILE)
         self.run_command(f"cat {self.ENV_FILE}")