Wait for UCXX listener addresses before distribution - #1103
Wait for UCXX listener addresses before distribution#1103fallintoplace wants to merge 4 commits into
Conversation
| if (listener_addresses_distributed_) { | ||
| return; | ||
| } | ||
| listener_addresses_distributed_ = true; |
There was a problem hiding this comment.
I think we need to wait for the listener-address map to be complete before setting this flag and snapshotting the maps. The root barrier currently only waits for rank_to_endpoint_ to reach nranks(), but endpoint registration and listener-address registration are separate for the HostPortPair path: listener_callback() registers the endpoint when the rank connects, and the RegisterListenerAddress control message is handled later. That means this can snapshot a partial rank_to_listener_address_, mark distribution as done, and never send the missing addresses.
Could we wait for both rank_to_endpoint_.size() and rank_to_listener_address_.size() to reach nranks() before distributing, and also assert both invariants here while holding both mutexes before setting listener_addresses_distributed_ = true? I’d also update the doc comment from "after all ranks have connected" to "after all ranks have connected and registered listener addresses.".
There was a problem hiding this comment.
Thank you for your attention, could you please check again?
| // The root needs endpoints and listener addresses for all ranks to continue. | ||
| while (rank_ == 0 && !all_barrier_maps_registered()) { | ||
| progress_worker(); | ||
| } |
There was a problem hiding this comment.
This change is necessary I think, and fixes a race bug.
| } | ||
|
|
||
| if (rank_ == 0) { | ||
| auto const rank_to_endpoint = rank_to_endpoint_snapshot(); |
There was a problem hiding this comment.
I think this one is not however. For the following reason: we can only reach this line once all_barrier_maps_registered() returns true. But, by definition that call can only return true once all other threads have finished modifying the rank_to_endpoint_ map having registered the relevant endpoints. IOW, as soon as all_barrier_maps_registered() returns true, we are guaranteed that there are going to be no more modifications of rank_to_endpoint_ and hence it is safe to read without the lock because there will be no data races.
Do you agree?
| * Must only be called on the root rank, after all ranks have connected and | ||
| * registered listener addresses. | ||
| */ | ||
| void distribute_listener_addresses() { | ||
| if (listener_addresses_distributed_) { | ||
| return; | ||
| RankToEndpointMap rank_to_endpoint; | ||
| RankToListenerAddressMap rank_to_listener_address; | ||
| { | ||
| std::scoped_lock lock(endpoints_mutex_, listener_mutex_); | ||
| if (listener_addresses_distributed_) { | ||
| return; |
There was a problem hiding this comment.
Here, again, the precondition for this function, I think requires that all modifications of rank_to_endpoint_ and rank_to_listener_address_ have been committed ,and so there is no possibility of a data race. What am I missing?
There was a problem hiding this comment.
Thank you for your review, I’m happy to simplify this and remove the snapshots later today.
There was a problem hiding this comment.
No I think there is some misunderstanding here. I contend that these changes weren't necessary (for the reasons in #1103 (comment)) but I am not sure. So I asked for a second opinion given that you'd observed this potential race.
It might be that there is a race and my understanding of how the code fits together is wrong, so I would like to understand one way (or the other) to improve overall understanding
There was a problem hiding this comment.
Yes, I agree that after all_barrier_maps_registered() returns true, the later direct reads are safe by the protocol, so I removed the snapshots.
Looking at the history, this seems to come from #934 / 33a8511: that change moved listener-address lookup from on-demand QueryListenerAddress to pre-distribution during the first barrier, but the root barrier readiness check stayed as the older endpoint-only check.
The ordering hole is in the HostPortPair path: root records the endpoint in listener_callback(), but the listener address is sent later via RegisterListenerAddress after the non-root receives its assigned rank. So the old barrier could see all endpoints, distribute with a partial listener-address map, and mark distribution done.
I think the necessary part is just waiting for both maps, plus checking those invariants before setting listener_addresses_distributed_.
fallintoplace
changed the title
3 participants
Summary
Fix UCXX root barrier setup so listener-address distribution only starts after the root has both the endpoint and listener address for every rank.
In the
HostPortPairpath, endpoint registration and listener-address registration are separate:listener_callback()records the endpoint when a rank connects, andRegisterListenerAddresscan be handled later. Previously the root barrier waited only forrank_to_endpoint_to reachnranks(), sodistribute_listener_addresses()could run with a partialrank_to_listener_address_, mark distribution complete, and skip late addresses.This change makes the root wait for both maps to reach
nranks(), documents the stronger precondition, and checks the invariants under the existing map mutexes before settinglistener_addresses_distributed_.Validation
git diff --checkpre-commit run --files cpp/src/communicator/ucxx.cppFailed to find nvcc.