HyukjinKwon commented on code in PR #59451:
URL: https://github.com/apache/spark/pull/59451#discussion_r4278429869
##########
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:
**Nit (P3):** This reservation only blocks binds on the same address (or the
wildcard) and outgoing-connection autobind. A socket bound to a different local
address can still take the same port. In production `address` is the executor
IP (`addrs[0]`), while `TCPStore` listens on the wildcard address. So a
loopback listener that gets port P can still make `TCPStore` fail with
`EADDRINUSE`.
I checked this on Linux with plain sockets. With `<host-ip>:P` reserved this
way, `bind` + `listen` on `127.0.0.1:P` succeeds, and an `SO_REUSEADDR` listen
on `[::]:P` / `0.0.0.0:P` then fails with errno 98. Loopback `bind(0)` calls
also still land on reserved ports at about the normal rate. Loopback
`connect()` never did.
Could this docstring, and the `# Other sockets cannot take the reserved
port.` comment in `test_reserve_free_port`, say what is actually blocked? Then
a future recurrence of this flake won't be ruled out because of the comment.
Reserving on the wildcard address would also close the gap on Linux, but it
changes bind semantics on BSD/macOS, so documenting the scope is the smaller
change.
**Verification:**
- **Inspection:** The helper docstring and the test comment state the
address-scoped guarantee accurately and make no claim broader than the
same-address assertion and TCPStore coexistence the test exercises.
##########
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:
**Nit (P3):** This test covers `_reserve_free_port` on its own, but nothing
enabled checks the part of the change that actually closes the race. That part
is partition 0 holding the socket across the torchrun call and closing it in
the `finally` of `wrapped_train_fn`. The distributed suites are skipped,
`test_log_streaming_server_start_failure_warns` patches
`_get_spark_task_function` out, and `test_connect_pipeline` only runs the path
without asserting on the port. So reverting `get_free_port` to
bind-and-release, or dropping the close loop, would keep every test green, and
the only symptom would be the intermittent `EADDRINUSE` coming back.
A small test in this mixin could do it. Build the task function with
`_get_spark_task_function`, then stub `BarrierTaskContext.get()` (one task on a
loopback address, `allGather` echoing its input), `LogStreamingClient`, the
partition-data setup, and the framework-wrapper call. Check that a plain bind
of `MASTER_ADDR:MASTER_PORT` fails while training runs, and succeeds once the
task function finishes, both when training returns and when it raises.
**Verification:**
- **Regression:** With the reservation released early (the merge-target
bind-and-release behavior), the new lifecycle test fails because a plain bind
of MASTER_ADDR:MASTER_PORT during the training call succeeds.
- **Behavior:** After the distributed task function completes, both on
success and when the training call raises, a plain bind of the same address and
port succeeds, showing the finally released the reservation.
--
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]
