HyukjinKwon commented on code in PR #59451:
URL: https://github.com/apache/spark/pull/59451#discussion_r4280203284
##########
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:
**Nit (P3):** Thanks for adding this test. The in-training check does fail
if the port is released early. This release check, though, runs only after
`list(task_fn(...))` has run the generator to completion, or after
`assertRaisesRegex` has cleared the traceback frames. By then the frame that
holds `reserved_port_sockets` is gone and CPython has already closed the
socket, so deleting the `for sock in reserved_port_sockets: sock.close()` loop
in the `finally` keeps both subtests green. I reproduced this with the same
closure/generator/`finally` shape.
Checking while the generator or exception is still alive would cover the
close loop, for example:
```python
gen = task_fn(iter([]))
next(gen) # partition 0's output batch, yielded after the finally
bind_master_port()
list(gen)
```
and for the raising case, call `bind_master_port()` inside an `except
RuntimeError:` block instead of after `assertRaisesRegex` exits.
**Verification:**
- **Regression:** With the explicit close of the reservation socket removed,
the lifecycle test fails on both the success and raising training paths.
- **Behavior:** With the change as written, the port is bindable as soon as
training returns or raises, while the task generator or exception is still
alive, and the test passes.
##########
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:
**Nit (P3):** This sentence only holds for binds that don't set
`SO_REUSEADDR`. The reserving socket sets it and never listens, so on Linux
another socket that also sets `SO_REUSEADDR` can bind **and listen** on this
exact address and port while the reservation is held. I checked with plain
sockets: with the reservation on `127.0.0.1:P`, a plain bind on `127.0.0.1:P`
or `0.0.0.0:P` fails with `EADDRINUSE`, but an `SO_REUSEADDR` bind+listen on
`127.0.0.1:P` succeeds. After that, an `SO_REUSEADDR` wildcard listen on `P`
(what `TCPStore` does) fails with errno 98. That is the same failure this PR
fixes, and it can happen to any server that explicitly binds the port with
`SO_REUSEADDR`, which many stacks set by default.
My earlier comment's "same-address or wildcard binds" had the same gap,
sorry about that. Something like this would be accurate:
> The reservation only blocks binds that do not set SO_REUSEADDR on the same
address or the wildcard address, and outgoing connections picking the port. A
socket that sets SO_REUSEADDR, or one bound to another local address such as
loopback, can still take the port.
The `# Binds on the same address cannot take the reserved port.` comment in
`test_reserve_free_port` could say "Plain binds" for the same reason, since
that is what the assertion checks.
**Verification:**
- **Inspection:** The docstring and test comment describe only the bind
cases the reservation actually blocks, and they agree with the stated TCPStore
SO_REUSEADDR coexistence.
--
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]
