From 68c01ac2fc93c527108a62b5d2dfed1533f86e89 Mon Sep 17 00:00:00 2001
From: kuizhiqing <kuizhiqing@baidu.com>
Date: Fri, 19 Aug 2022 17:12:40 +0800
Subject: [PATCH] [LAUNCH] rewrite get free port strategy (#45247)

* rewrite get free port strategy

* hide the old one
---
 .../paddle/distributed/launch/context/node.py | 58 +++++++++++++++----
 1 file changed, 47 insertions(+), 11 deletions(-)

diff --git a/python/paddle/distributed/launch/context/node.py b/python/paddle/distributed/launch/context/node.py
index 6ee8fa6d10c..9c60ed43d47 100644
--- a/python/paddle/distributed/launch/context/node.py
+++ b/python/paddle/distributed/launch/context/node.py
@@ -14,6 +14,8 @@
 
 from .device import Device
 
+import os
+import random
 import socket
 import struct
 from contextlib import closing
@@ -28,6 +30,12 @@ class Node(object):
         self.free_ports = []
         self._allocated_ports = []
 
+        port_range = os.getenv('PORT_RANGE', '35100:64000')
+        port_range = port_range.split(':')
+        self._port_start = int(port_range[0])
+        self._port_end = int(port_range[1])
+        self._port_cur = random.randint(self._port_start, self._port_end)
+
     def get_host_ip(self):
         try:
             self.hostname = socket.gethostname()
@@ -44,21 +52,44 @@ class Node(object):
     def get_ports_occupied(self):
         return self.free_ports
 
+    def _get_free_port(self, port=0):
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER,
+                         struct.pack('ii', 1, 0))
+            try:
+                s.bind(('', port))
+                return s.getsockname()[1]
+            except:
+                return -1
+
+    def _update_port_cur(self):
+        self._port_cur += 1
+        if self._port_cur > self._port_end:
+            self._port_cur = self._port_start
+
     def get_free_port(self):
+        for _ in range(100):
+            ret = self._get_free_port(self._port_cur)
+            if ret > 0:
+                self._update_port_cur()
+                return ret
+            else:
+                self._update_port_cur()
+
+        return self._port_cur
+
+    '''
+    def get_free_port2(self):
         # for loop to avoid port conflict
         for _ in range(100):
-            with closing(socket.socket(socket.AF_INET,
-                                       socket.SOCK_STREAM)) as s:
-                s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER,
-                             struct.pack('ii', 1, 0))
-                s.bind(('', 0))
-                port = s.getsockname()[1]
-                if port in self._allocated_ports:
-                    continue
-                else:
-                    self._allocated_ports.append(port)
-                    return port
+            port = self._get_free_port()
+            if port in self._allocated_ports:
+                continue
+            else:
+                self._allocated_ports.append(port)
+                return port
         return port
+    '''
 
     @classmethod
     def is_server_ready(self, ip, port):
@@ -72,3 +103,8 @@ class Node(object):
                 return True
             else:
                 return False
+
+
+if __name__ == '__main__':
+    n = Node()
+    print(n.get_free_ports(10))
-- 
GitLab