From a848091141898047bec4adf858336419d8f65a16 Mon Sep 17 00:00:00 2001
From: YH <100389977+yhna940@users.noreply.github.com>
Date: Tue, 28 Feb 2023 12:00:43 +0900
Subject: [PATCH] Fix port exception type (#2925)

---
 colossalai/utils/common.py | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

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