Fix sa (#429)

* fix

* update connect

* add uniitest

paddleslim/common/controller_client.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import time
 import logging
 import socket
 from .log_helper import get_logger
@@ -31,6 +33,8 @@ class ControllerClient(object):
         client_name(str): Current client name, random generate for counting client number. Default: None.
     """
 
+    START = True
+
     def __init__(self,
                  server_ip=None,
                  server_port=None,
@@ -52,23 +56,58 @@ class ControllerClient(object):
             reward(float): The reward of tokens.
             iter(int): The iteration number of current client.
         """
+        ControllerClient.START = False
         socket_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        socket_client.connect((self.server_ip, self.server_port))
-        tokens = ",".join([str(token) for token in tokens])
-        socket_client.send("{}\t{}\t{}\t{}\t{}".format(
-            self._key, tokens, reward, iter, self._client_name).encode())
-        response = socket_client.recv(1024).decode()
-        if "ok" in response.strip('\n').split("\t"):
-            return True
+        errno = socket_client.connect_ex((self.server_ip, self.server_port))
+        if errno != 0:
+            _logger.info("Server is closed!!!")
+            os._exit(0)
         else:
-            return False
+            tokens = ",".join([str(token) for token in tokens])
+            socket_client.send("{}\t{}\t{}\t{}\t{}".format(
+                self._key, tokens, reward, iter, self._client_name).encode())
+            try:
+                response = socket_client.recv(1024).decode()
+                if "ok" in response.strip('\n').split("\t"):
+                    return True
+                else:
+                    return False
+            except Exception as err:
+                _logger.error(err)
+                os._exit(0)
 
     def next_tokens(self):
         """
         Get next tokens.
         """
-        socket_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        socket_client.connect((self.server_ip, self.server_port))
+        retry_cnt = 0
+
+        if ControllerClient.START:
+            while True:
+                socket_client = socket.socket(socket.AF_INET,
+                                              socket.SOCK_STREAM)
+                errno = socket_client.connect_ex(
+                    (self.server_ip, self.server_port))
+                if errno != 0:
+                    retry_cnt += 1
+                    _logger.info("Server is NOT ready, wait 10 second to retry")
+                    time.sleep(10)
+                else:
+                    break
+
+                if retry_cnt == 6:
+                    _logger.error(
+                        "Server is NOT ready in 1 minute, please check if it start"
+                    )
+                    os._exit(errno)
+
+        else:
+            socket_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            errno = socket_client.connect_ex((self.server_ip, self.server_port))
+            if errno != 0:
+                _logger.info("Server is closed")
+                os._exit(0)
+
         socket_client.send("next_tokens".encode())
         tokens = socket_client.recv(1024).decode()
         tokens = [int(token) for token in tokens.strip("\n").split(",")]
@@ -79,7 +118,11 @@ class ControllerClient(object):
         Request for current information.
         """
         socket_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        socket_client.connect((self.server_ip, self.server_port))
-        socket_client.send("current_info".encode())
-        current_info = socket_client.recv(1024).decode()
-        return eval(current_info)
+        errno = socket_client.connect_ex((self.server_ip, self.server_port))
+        if errno != 0:
+            _logger.info("Server is closed")
+            return None
+        else:
+            socket_client.send("current_info".encode())
+            current_info = socket_client.recv(1024).decode()
+            return eval(current_info)

tests/test_client_connect.py

+# Copyright (c) 2019  PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import sys
+sys.path.append("../")
+import os
+import time
+import signal
+import unittest
+import paddle.fluid as fluid
+from paddleslim.nas import SANAS
+from paddleslim.common.controller_client import ControllerClient
+import numpy as np
+from multiprocessing import Process
+import socket
+
+
+def start_client(configs, addr, port):
+    client_sanas = SANAS(
+        configs=configs,
+        server_addr=(addr, port),
+        save_checkpoint=None,
+        is_server=False)
+    for _ in range(2):
+        arch = client_sanas.next_archs()[0]
+        time.sleep(1)
+        client_sanas.reward(0.1)
+
+
+def start_server(configs, port):
+    server_sanas = SANAS(
+        configs=configs, server_addr=("", port), save_checkpoint=None)
+    server_sanas.next_archs()[0]
+    return server_sanas
+
+
+class TestClientConnect(unittest.TestCase):
+    def setUp(self):
+        self.configs = [('MobileNetV2BlockSpace', {'block_mask': [0]})]
+        self.port = np.random.randint(8337, 8773)
+        self.addr = socket.gethostbyname(socket.gethostname())
+
+    def test_client_start_first(self):
+        p = Process(
+            target=start_client, args=(self.configs, self.addr, self.port))
+        p.start()
+
+        start_server(self.configs, self.port)
+
+
+class TestClientConnectCase1(unittest.TestCase):
+    def setUp(self):
+        self.configs = [('MobileNetV2BlockSpace', {'block_mask': [0]})]
+        self.port = np.random.randint(8337, 8773)
+        self.addr = socket.gethostbyname(socket.gethostname())
+
+    def test_client_start_first(self):
+        p = Process(
+            target=start_client, args=(self.configs, self.addr, self.port))
+        p.start()
+
+        time.sleep(60)
+        server_sanas = start_server(self.configs, self.port)
+        os.kill(os.getpid(), 0)
+
+
+class TestClientConnectCase2(unittest.TestCase):
+    def setUp(self):
+        self.port = np.random.randint(8337, 8773)
+        self.addr = socket.gethostbyname(socket.gethostname())
+
+    def test_request_current_info(self):
+        client = ControllerClient(self.addr, self.port)
+        client.request_current_info()
+
+
+if __name__ == '__main__':
+    unittest.main()