diff --git a/colossalai/utils/common.py b/colossalai/utils/common.py index e35b29c2a..e15981140 100644 --- a/colossalai/utils/common.py +++ b/colossalai/utils/common.py @@ -50,16 +50,20 @@ def ensure_path_exists(filename: str): Path(dirpath).mkdir(parents=True, exist_ok=True) -def free_port(): +def free_port() -> int: + """Get a free port on localhost. + + Returns: + int: A free port on localhost. + """ while True: + port = random.randint(20000, 65000) try: - sock = socket.socket() - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - port = random.randint(20000, 65000) - sock.bind(('localhost', port)) - sock.close() - return port - except Exception: + with socket.socket() as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("localhost", port)) + return port + except OSError: continue