From cd584fca81936142e069053b56433e66037362e2 Mon Sep 17 00:00:00 2001 From: yoshikisd Date: Sun, 9 Nov 2025 03:34:08 +0000 Subject: [PATCH] multigpu.setup now recognizes MASTER_ADDR and MASTER_PORT when setting up spawn-based jobs --- src/cdtools/tools/multigpu/multigpu.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/cdtools/tools/multigpu/multigpu.py b/src/cdtools/tools/multigpu/multigpu.py index 6051817..8156dbb 100644 --- a/src/cdtools/tools/multigpu/multigpu.py +++ b/src/cdtools/tools/multigpu/multigpu.py @@ -219,13 +219,19 @@ def setup(rank: int = None, """ # Make sure that the user explicitly defines parameters if spawn is used if get_launch_method() == 'spawn': + if init_method == 'env://' and None in (master_addr, master_port): + # We'll check if the master address/port is in the env variable + master_addr = os.environ.get('MASTER_ADDR') + master_port = os.environ.get('MASTER_PORT') + if None in (rank, world_size, master_addr, master_port): raise RuntimeError( 'torch.multiprocessing.spawn was detected as the launching \n' 'method, but either rank, world_size, master_addr, or \n' 'master_port has not been explicitly defined. Please ensure \n' - 'these parameters have been explicitly defined, or \n' - 'alternatively launch the multi-GPU job with torchrun.\n' + 'that either these parameters have been explicitly defined,\n' + 'MASTER_ADDR/MASTER_PORT have been defined as environment \n' + 'variables, or launch the multi-GPU job with torchrun.\n' ) elif init_method == 'env://': # Set up the environment variables