Yicong-Huang commented on code in PR #59451:
URL: https://github.com/apache/spark/pull/59451#discussion_r4288560961
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -448,6 +449,25 @@ def _get_torchrun_args(local_mode: bool, num_processes:
int) -> Tuple[List[Any],
]
return torchrun_args, processes_per_node
+ @staticmethod
+ def _reserve_free_port(address: str) -> socket.socket:
+ """
+ Bind a socket to a free port on the given address and return it.
+
+ Closing the socket right away would let another socket take the port
before
+ torchrun's TCPStore binds it. The returned socket is bound but not
listening and
+ sets SO_REUSEADDR, so TCPStore, which also sets SO_REUSEADDR, can
still listen on
+ the port while the socket keeps it reserved.
+
+ The reservation only blocks binds on the same address or the wildcard
address, and
Review Comment:
Fixed in 591746b08b9: the docstring now uses your wording (only binds
without SO_REUSEADDR on the same/wildcard address are blocked), and the test
comment says "Plain binds".
##########
python/pyspark/ml/torch/tests/test_distributor.py:
##########
@@ -300,6 +300,78 @@ class StopTraining(Exception):
self.assertIn("Start torch distributor log streaming server failed",
message)
self.assertIn("Address already in use", message)
+ def test_reserve_free_port(self) -> None:
+ import socket
+ from datetime import timedelta
+
+ from torch.distributed import TCPStore
+
+ sock = TorchDistributor._reserve_free_port("127.0.0.1")
+ try:
+ port = sock.getsockname()[1]
+ # Binds on the same address cannot take the reserved port.
+ other = socket.socket()
+ try:
+ with self.assertRaises(OSError):
+ other.bind(("127.0.0.1", port))
+ finally:
+ other.close()
+ # torchrun's TCPStore can still listen on it.
+ store = TCPStore("127.0.0.1", port, 1, True,
timeout=timedelta(seconds=30))
+ store.set("key", "value")
+ self.assertEqual(store.get("key"), b"value")
+ finally:
+ sock.close()
+
+ def test_master_port_reserved_until_training_finishes(self) -> None:
+ import socket
+
+ from pyspark.taskcontext import BarrierTaskContext
+
+ def bind_master_port() -> None:
+ sock = socket.socket()
+ try:
+ sock.bind((os.environ["MASTER_ADDR"],
int(os.environ["MASTER_PORT"])))
+ finally:
+ sock.close()
+
+ def train_and_check_reserved(*args: Any, **kwargs: Any) -> str:
+ with self.assertRaises(OSError):
+ bind_master_port()
+ return "done"
+
+ def train_check_reserved_and_raise(*args: Any, **kwargs: Any) -> str:
+ train_and_check_reserved()
+ raise RuntimeError("training failed")
+
+ dist = TorchDistributor(num_processes=1, local_mode=False,
use_gpu=False)
+ dist.driver_address = "127.0.0.1"
+ dist.log_streaming_server_port = -1
+ dist.log_streaming_auth_secret = None
+ context = MagicMock()
+ context.partitionId.return_value = 0
+ context.getTaskInfos.return_value = [MagicMock(address="127.0.0.1:0")]
+ context.allGather.side_effect = lambda message: [message]
+
+ for train_fn in [train_and_check_reserved,
train_check_reserved_and_raise]:
+ with (
+ self.subTest(train_fn=train_fn.__name__),
+ patch.dict(os.environ),
+ patch.object(BarrierTaskContext, "get", return_value=context),
+ patch("pyspark.ml.torch.distributor.LogStreamingClient"),
+ patch.object(
+ TorchDistributor, "_get_output_from_framework_wrapper",
side_effect=train_fn
+ ),
+ ):
+ task_fn = dist._get_spark_task_function(None, MagicMock(),
None, None)
+ if train_fn is train_check_reserved_and_raise:
+ with self.assertRaisesRegex(RuntimeError, "training
failed"):
+ list(task_fn(iter([])))
+ else:
+ self.assertEqual(len(list(task_fn(iter([])))), 1)
+ # The reservation is released once the task function finishes.
+ bind_master_port()
Review Comment:
Fixed in 591746b08b9: the release check now runs right after `next(gen)` on
the success path and inside `except RuntimeError:` on the raising path, while
the reservation socket is still alive. Verified both subtests fail with the
close loop removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
