[PYTHON][MINOR] Use a unique per-run rendezvous id for distributed torch training

Distributed mode passed a fixed rendezvous id to torchrun, so concurrent
training runs sharing a rendezvous endpoint could cross-join each other's
rendezvous. Generate a unique id per run on the rank-0 barrier task, share
it with the other tasks via allGather, and require it when building the
torchrun command instead of falling back to a constant.

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-Authored-By: Holden Karau <holden@pigscanfly.ca>

Author: 3 people

python/pyspark/ml/deepspeed/tests/test_deepspeed_distributor.py

@@ -41,11 +41,12 @@ def _get_env_var(self, var_name: str, default_value: Any) -> Any:
         os.environ[var_name] = str(default_value)
         return default_value
 
-    def _get_env_variables_distributed(self) -> Tuple[Any, Any, Any]:
+    def _get_env_variables_distributed(self) -> Tuple[Any, Any, Any, Any]:
         master_addr = self._get_env_var("MASTER_ADDR", "127.0.0.1")
         master_port = self._get_env_var("MASTER_PORT", 2000)
         rank = self._get_env_var("RANK", 0)
-        return master_addr, master_port, rank
+        rdzv_id = self._get_env_var("PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID", "test-rdzv-id")
+        return master_addr, master_port, rank, rdzv_id
 
     def test_get_torchrun_args_local(self) -> None:
         number_of_processes = 5
@@ -60,12 +61,12 @@ def test_get_torchrun_args_local(self) -> None:
 
     def test_get_torchrun_args_distributed(self) -> None:
         number_of_processes = 5
-        master_addr, master_port, rank = self._get_env_variables_distributed()
+        master_addr, master_port, rank, rdzv_id = self._get_env_variables_distributed()
         expected_torchrun_args_distributed = [
             f"--nnodes={number_of_processes}",
             f"--node_rank={rank}",
             f"--rdzv_endpoint={master_addr}:{master_port}",
-            "--rdzv_id=0",
+            f"--rdzv_id={rdzv_id}",
         ]
         torchrun_args_distributed, process_per_node = DeepspeedTorchDistributor._get_torchrun_args(
             False, number_of_processes
@@ -131,12 +132,13 @@ def test_create_torchrun_command_distributed(self) -> None:
             distributed_master_address,
             distributed_master_port,
             distributed_rank,
+            distributed_rdzv_id,
         ) = self._get_env_variables_distributed()
         distributed_torchrun_args = [
             f"--nnodes={num_procs}",
             f"--node_rank={distributed_rank}",
             f"--rdzv_endpoint={distributed_master_address}:{distributed_master_port}",
-            "--rdzv_id=0",
+            f"--rdzv_id={distributed_rdzv_id}",
         ]
         with self.subTest(msg="Distributed training command verification with no extra args"):
             distributed_cmd_no_args_expected = [

python/pyspark/ml/torch/distributor.py

@@ -431,12 +431,20 @@ def _get_torchrun_args(local_mode: bool, num_processes: int) -> Tuple[List[Any],
         else:
             processes_per_node = 1
         node_rank = os.environ["RANK"]
+        # Set by set_torch_config; no constant fallback, since concurrent runs sharing a
+        # rendezvous endpoint would collide on a fixed id.
+        rdzv_id = os.environ.get("PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID")
+        if not rdzv_id:
+            raise RuntimeError(
+                "Missing PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID environment variable: a unique "
+                "per-run rendezvous id is required for distributed training."
+            )
 
         torchrun_args = [
             f"--nnodes={num_processes // processes_per_node}",
             f"--node_rank={node_rank}",
             f"--rdzv_endpoint={master_addr}:{master_port}",
-            "--rdzv_id=0",  # TODO: setup random ID that is gleaned from env variables
+            f"--rdzv_id={rdzv_id}",
         ]
         return torchrun_args, processes_per_node
 
@@ -660,6 +668,7 @@ def _get_spark_task_function(
         # Spark task program
         def wrapped_train_fn(iterator):  # type: ignore[no-untyped-def]
             import os
+            import secrets
             import pandas as pd
             import pyarrow
             from pyspark import BarrierTaskContext
@@ -687,6 +696,15 @@ def set_torch_config(context: "BarrierTaskContext") -> None:
 
                 os.environ["MASTER_ADDR"] = str(addrs[0])
                 os.environ["MASTER_PORT"] = str(get_free_port(addrs[0], context))
+                # Unique per run so that concurrent runs sharing a rendezvous endpoint
+                # do not collide.
+                rdzv_id = secrets.token_hex(16) if context.partitionId() == 0 else ""
+                rdzv_id = context.allGather(str(rdzv_id))[0]
+                if not rdzv_id:
+                    raise RuntimeError(
+                        "Failed to generate a shared rendezvous id for distributed training."
+                    )
+                os.environ["PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID"] = rdzv_id
                 os.environ["WORLD_SIZE"] = str(num_processes)
                 os.environ["NODE_RANK"] = str(context.partitionId())
                 os.environ["RANK"] = str(context.partitionId())

python/pyspark/ml/torch/tests/test_distributor.py

@@ -281,7 +281,21 @@ def test_create_torchrun_command(self) -> None:
         )
 
         distributed_mode_input_params = {"num_processes": 4, "local_mode": False}
-        input_env_vars = {"MASTER_ADDR": "localhost", "MASTER_PORT": "9350", "RANK": "3"}
+
+        # Without the per-run rendezvous id in the environment, command construction must
+        # fail instead of falling back to a fixed rendezvous id.
+        missing_rdzv_env_vars = {"MASTER_ADDR": "localhost", "MASTER_PORT": "9350", "RANK": "3"}
+        self.setup_env_vars(missing_rdzv_env_vars)
+        with self.assertRaisesRegex(RuntimeError, "PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID"):
+            TorchDistributor._create_torchrun_command(distributed_mode_input_params, train_path)
+        self.delete_env_vars(missing_rdzv_env_vars)
+
+        input_env_vars = {
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": "9350",
+            "RANK": "3",
+            "PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID": "0123456789abcdef",
+        }
 
         args_number = [1, 3]  # testing conversion to strings
         self.setup_env_vars(input_env_vars)
@@ -292,7 +306,7 @@ def test_create_torchrun_command(self) -> None:
             "--nnodes=4",
             "--node_rank=3",
             "--rdzv_endpoint=localhost:9350",
-            "--rdzv_id=0",
+            "--rdzv_id=0123456789abcdef",
             "--nproc_per_node=1",
             "train.py",
             "1",
@@ -313,13 +327,19 @@ def test_create_torchrun_command(self) -> None:
             "MASTER_ADDR": "11.22.33.44",
             "MASTER_PORT": "6677",
             "RANK": "1",
+            "PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID": "0123456789abcdef",
         },
     )
     def test_multi_gpu_node_get_torchrun_args(self):
         torchrun_args, processes_per_node = TorchDistributor._get_torchrun_args(False, 8)
         self.assertEqual(
             torchrun_args,
-            ["--nnodes=2", "--node_rank=1", "--rdzv_endpoint=11.22.33.44:6677", "--rdzv_id=0"],
+            [
+                "--nnodes=2",
+                "--node_rank=1",
+                "--rdzv_endpoint=11.22.33.44:6677",
+                "--rdzv_id=0123456789abcdef",
+            ],
         )
         self.assertEqual(processes_per_node, 4)