Yicong-Huang commented on code in PR #59451:
URL: https://github.com/apache/spark/pull/59451#discussion_r4279540509
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -448,6 +449,21 @@ 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.
Review Comment:
Fixed in 11a07cdca62: the docstring now states the reservation only covers
same-address/wildcard binds and outgoing connects, and the test comment is
narrowed to same-address binds. Kept the bind address as is to avoid the
BSD/macOS semantics change.
##########
python/pyspark/ml/torch/tests/test_distributor.py:
##########
@@ -300,6 +300,29 @@ 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:
Review Comment:
Added `test_master_port_reserved_until_training_finishes` in 11a07cdca62,
which builds the task function with stubbed
`BarrierTaskContext`/`LogStreamingClient`/training and checks a plain bind of
`MASTER_ADDR:MASTER_PORT` fails during training and succeeds afterwards, on
both return and raise. Verified it fails when the reservation is released early.
--
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]
