Prevent TCPServer get deleted too early (#67204)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67204
Fixes #66422
Fixes #66423
In the original test, all collectives are dummy local ones. As a
result, rank 0 could exit earlier than other ranks. However, the
`TCPStore` lives on rank 0, and other ranks might need to talk to
that store after rank 0 exits. This commit explicitly makes rank 0
wait for all other ranks to finish.
cc pietern mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse SciPioneer H-Huang
Test Plan: Imported from OSS
Reviewed By: albanD
Differential Revision: D31906802
Pulled By: mrshenli
fbshipit-source-id: 82745f5497d784ea3cea9df6bda537ec71380867
