Replaces ddp .spawn with subprocess (#2029)

* replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix
2026-09-12 12:40:20 +08:00 · 2020-06-01 11:00:32 -04:00
parent fd38f52e55
commit 82a20296e3
19 changed files with 283 additions and 174 deletions
@@ -117,6 +117,11 @@ import os
 import re
 from abc import ABC, abstractmethod
 from typing import Union
+import subprocess
+import sys
+from time import sleep
+import numpy as np
+from os.path import abspath

 import torch
 from pytorch_lightning import _logger as log
@@ -311,7 +316,7 @@ class TrainerDDPMixin(ABC):
        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

        # when slurm is managing the task it sets the visible devices
-        if not is_slurm_managing_tasks:
+        if not is_slurm_managing_tasks and 'CUDA_VISIBLE_DEVICES' not in os.environ:
            if isinstance(data_parallel_device_ids, int):
                id_str = ','.join(str(x) for x in list(range(data_parallel_device_ids)))
                os.environ["CUDA_VISIBLE_DEVICES"] = id_str
@@ -322,7 +327,74 @@ class TrainerDDPMixin(ABC):
        # don't make this debug... this is good UX
        log.info(f'CUDA_VISIBLE_DEVICES: [{os.environ["CUDA_VISIBLE_DEVICES"]}]')

-    def ddp_train(self, process_idx, model):
+    def __set_random_port(self):
+        """
+        When running DDP NOT managed by SLURM, the ports might collide
+        :return:
+        """
+        try:
+            default_port = os.environ['MASTER_PORT']
+        except Exception:
+            import random
+            default_port = random.randint(10000, 19000)
+            os.environ['MASTER_PORT'] = str(default_port)
+
+    def spawn_ddp_children(self, model):
+        self.__set_random_port()
+        port = os.environ['MASTER_PORT']
+
+        master_address = '127.0.0.1' if 'MASTER_ADDR' not in os.environ else os.environ['MASTER_ADDR']
+        os.environ['MASTER_PORT'] = f'{port}'
+        os.environ['MASTER_ADDR'] = f'{master_address}'
+
+        # allow the user to pass the node rank
+        node_rank = '0'
+        if 'NODE_RANK' in os.environ:
+            node_rank = os.environ['NODE_RANK']
+        if 'GROUP_RANK' in os.environ:
+            node_rank = os.environ['GROUP_RANK']
+
+        os.environ['NODE_RANK'] = node_rank
+        os.environ['LOCAL_RANK'] = '0'
+
+        # pull out the commands used to run the script and resolve the abs file path
+        command = sys.argv
+        full_path = abspath(command[0])
+        command[0] = full_path
+        command = ['python'] + command
+
+        # since this script sets the visible devices we replace the gpus flag with a number
+        num_gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',').__len__()
+
+        # if script called without a flag, pass in a flag anyhow
+        if '--gpus' not in command:
+            arg_gpus = len(self.gpus) if isinstance(self.gpus, list) else self.gpus
+            command += ['--gpus', arg_gpus]
+
+        gpu_flag_idx = command.index('--gpus')
+        command[gpu_flag_idx + 1] = f'{num_gpus}'
+
+        os.environ['WORLD_SIZE'] = f'{num_gpus * self.num_nodes}'
+
+        self.interactive_ddp_procs = []
+        for local_rank in range(1, self.num_processes):
+            env_copy = os.environ.copy()
+            env_copy['LOCAL_RANK'] = f'{local_rank}'
+
+            # import pdb; pdb.set_trace()
+            # start process
+            proc = subprocess.Popen(command, env=env_copy)
+            self.interactive_ddp_procs.append(proc)
+
+            # starting all processes at once can cause issues
+            # with dataloaders delay between 1-10 seconds
+            delay = np.random.uniform(1, 5, 1)[0]
+            sleep(delay)
+
+        local_rank = 0
+        self.ddp_train(local_rank, model, is_master=True)
+
+    def ddp_train(self, process_idx, model, is_master=False):
        """
        Entry point into a DP thread
        :param gpu_idx:
@@ -359,7 +431,14 @@ class TrainerDDPMixin(ABC):
        # MODEL
        # copy model to each gpu
        if self.on_gpu:
-            self.root_gpu = process_idx
+            gpu_idx = process_idx
+            if is_master:
+                # source of truth is cuda for gpu idx
+                gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
+                local_rank = int(os.environ['LOCAL_RANK'])
+                gpu_idx = int(gpus[local_rank])
+
+            self.root_gpu = gpu_idx
            torch.cuda.set_device(self.root_gpu)
            model.cuda(self.root_gpu)

@@ -388,9 +467,6 @@ class TrainerDDPMixin(ABC):
        # continue training routine
        self.run_pretrain_routine(model)

-        # when ddp ends, we save the model
-        self.save_spawn_weights(model)
-
    def save_spawn_weights(self, model):
        """
        Dump a temporary checkpoint after ddp ends to get weights out of the process