fix timeout error on TCPStore creation (#147)

Summary:
Pull Request resolved: #147

Async torchsnapshotting was causing a timeout error when creating the TCPStore within PendingSnapshot on an ondemand

Using `torch.distributed.elastic.utils.distributed's` `get_socket_with_port()` resolves the issue.

Reviewed By: daniellepintz

Differential Revision: D48072665

fbshipit-source-id: a573a146f33ecec5f91ed800984e9b0f95cb2974

Author: JKSenthil

torchsnapshot/dist_store.py

@@ -5,11 +5,11 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import socket
 from datetime import timedelta
 from typing import Dict, Optional
 
 import torch.distributed as dist
+from torch.distributed.elastic.utils.distributed import get_socket_with_port
 
 from .pg_wrapper import PGWrapper
 
@@ -63,10 +63,8 @@ def create_store(pg_wrapper: PGWrapper) -> dist.Store:
     """
     if pg_wrapper.get_rank() == 0:
         # Find a free port
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        sock.bind((socket.gethostname(), 0))
-        master_addr, master_port = sock.getsockname()
+        sock = get_socket_with_port()
+        master_addr, master_port, _, _ = sock.getsockname()
         sock.close()
         # Broadcast master address/port to peers
         obj_list = [master_addr, master_port]