Fix distill exit bug (#109)

Author: wangxicoding

python/paddle_edl/distill/distill_reader.py

@@ -100,6 +100,7 @@ def __init__(self, ins, predicts):
         self._service_name = None
 
         # reader worker args
+        self._reader_worker = None
         self._reader = None
         self._reader_type = None
         self._reader_out_queue = None
@@ -167,6 +168,7 @@ def _start_reader_worker(self):
                     self._reader_cond, ))
             reader_worker.daemon = True
             reader_worker.start()
+            self._reader_worker = reader_worker
             self._is_reader_start = True
         else:
             with self._reader_cond:
@@ -209,7 +211,8 @@ def _start_predict_worker_pool(self):
                     self._require_num,
                     self._predict_stop_events,
                     self._get_servers,
-                    self._predict_manage_stop_event, ))
+                    self._predict_manage_stop_event,
+                    self._predict_cond, ))
             self._predict_manage_thread.daemon = True
             self._predict_manage_thread.start()
 
@@ -388,3 +391,19 @@ def __call__(self):
                 self._reader_type, self._predict_out_queue,
                 self._fetch_stop_event, self._task_semaphore):
             yield data
+
+    def __del__(self):
+        if not self._is_args_init:
+            return
+
+        # stop reader worker
+        with self._reader_cond:
+            self._reader_stop_event.set()
+            self._reader_cond.notify()
+
+        self._predict_manage_stop_event.set()
+
+        for i in range(20):
+            if self._reader_worker.is_alive() or \
+               self._predict_manage_thread.is_alive():
+                time.sleep(1)

python/paddle_edl/distill/distill_worker.py

@@ -15,6 +15,9 @@
 import logging
 import numpy as np
 import os
+import signal
+import six
+import sys
 import time
 
 from paddle_serving_client import Client
@@ -53,7 +56,7 @@ def __init__(self, server_id, server, stop_event_id, state=PENDING):
 
 def predict_manage_worker(process, server_queue, server_result_queue,
                           require_num, predict_stop_events, get_servers_fun,
-                          stop_event):
+                          stop_event, predict_cond):
     """ thread that manage predict worker """
     num_shutdown_process = [0]
 
@@ -129,6 +132,34 @@ def shutdown_one_process():
         except queue.Empty:
             pass
 
+    def clean_queue(data_queue):
+        while True:
+            try:
+                data_queue.get_nowait()
+            except Exception:
+                break
+
+    clean_queue(server_queue)
+    clean_queue(server_result_queue)
+
+    with predict_cond:
+        for predict_stop_event in predict_stop_events:
+            predict_stop_event.set()
+        predict_cond.notify_all()
+
+    for i in range(require_num):
+        shutdown_one_process()
+        clean_queue(server_result_queue)
+
+    for i in range(20):
+        shutdown_process = 0
+        for p in process:
+            if not p.is_alive():
+                shutdown_process += 1
+        if shutdown_process == len(process):
+            break
+        time.sleep(1)
+
 
 class _PoisonPill:
     def __init__(self, feed_count, predict_count=0):
@@ -284,22 +315,38 @@ def __del__(self):
 def predict_worker(server_queue, server_result_queue, working_predict_count,
                    in_queue, out_queue, feeds, fetchs, conf_file, stop_events,
                    predict_lock, global_finished_task, predict_cond):
-    while True:
-        # get server
-        server_item = server_queue.get()
-        if server_item is None:
-            server_queue.put(None)  # poison_pill
-            return
-
-        # predict
-        success = predict_loop(server_item, working_predict_count, in_queue,
-                               out_queue, feeds, fetchs, conf_file,
-                               stop_events, predict_lock, global_finished_task,
-                               predict_cond)
+    signal_exit = [False, ]
+
+    # Define signal handler function
+    def predict_signal_handle(signum, frame):
+        signal_exit[0] = True
+        exit(0)
+
+    # register signal.SIGTERM's handler
+    signal.signal(signal.SIGTERM, predict_signal_handle)
+
+    try:
+        while True:
+            # get server
+            server_item = server_queue.get()
+            if server_item is None:
+                server_result_queue.put(None)
+                return
 
-        server_item.state = ServerItem.FINISHED if success else ServerItem.ERROR
-        server_result_queue.put(server_item)
-        logger.info('Stopped server={}'.format(server_item.server))
+            # predict
+            success = predict_loop(server_item, working_predict_count,
+                                   in_queue, out_queue, feeds, fetchs,
+                                   conf_file, stop_events, predict_lock,
+                                   global_finished_task, predict_cond)
+
+            server_item.state = ServerItem.FINISHED if success else ServerItem.ERROR
+            server_result_queue.put(server_item)
+            logger.info('Stopped server={}'.format(server_item.server))
+    except Exception as e:
+        if signal_exit[0] is True:
+            pass
+        else:
+            six.reraise(*sys.exc_info())
 
 
 def predict_loop(server_item, working_predict_count, in_queue, out_queue,
@@ -365,7 +412,8 @@ def predict_loop(server_item, working_predict_count, in_queue, out_queue,
                     out_queue.put(poison_pill)  # poison consumer
                 else:
                     in_queue.put(poison_pill)  # poison other predict worker
-
+                if stop_event.is_set():
+                    break
                 # wait next reader iter or last failed predict job
                 predict_cond.wait()
 
@@ -435,22 +483,40 @@ def reader_worker(reader, reader_type, teacher_batch_size, out_queue,
     # consumer may recv out-of-order task(3, 1, 2) before task(0), consumer will store then,
     # when task(0) is completed and consumer recv it, it will release semaphore,
     # reader go on working.
+
+    signal_exit = [False, ]
+
+    def reader_signal_handle(signum, frame):
+        signal_exit[0] = True
+        exit(0)
+
+    # register signal.SIGTERM's handler
+    signal.signal(signal.SIGTERM, reader_signal_handle)
+
     read_func_map = {
         ReaderType.SAMPLE: read_sample,
         ReaderType.SAMPLE_LIST: read_sample_list,
         ReaderType.BATCH: read_batch
     }
     read_func = read_func_map[reader_type]
 
-    while not stop_event.is_set():
-        task_size = read_func(reader, teacher_batch_size, out_queue,
-                              task_semaphore)
-
-        poison_pill = _PoisonPill(task_size)
-        with reader_cond:
-            out_queue.put(poison_pill)
-            # wait next reader iter
-            reader_cond.wait()
+    try:
+        while not stop_event.is_set():
+            task_size = read_func(reader, teacher_batch_size, out_queue,
+                                  task_semaphore)
+
+            poison_pill = _PoisonPill(task_size)
+            with reader_cond:
+                out_queue.put(poison_pill)
+                if stop_event.is_set():
+                    break
+                # wait next reader iter
+                reader_cond.wait()
+    except Exception as e:
+        if signal_exit[0] is True:
+            pass
+        else:
+            six.reraise(*sys.exc_info())
 
 
 def read_sample(reader, teacher_batch_size, out_queue, task_semaphore):