Improve threads + taking in account review

Author: Hartorn

giskard/ml_worker/testing/registry/registry.py

@@ -1,3 +1,5 @@
+from typing import Dict, Optional
+
 import hashlib
 import importlib.util
 import inspect
@@ -6,16 +8,15 @@
 import sys
 import uuid
 from pathlib import Path
-from typing import Optional, Dict
 
 import cloudpickle
 
 from giskard.core.core import SavableMeta
-from giskard.settings import expand_env_var, settings
+from giskard.settings import settings
 
 
 def find_plugin_location():
-    return Path(expand_env_var(settings.home)) / "plugins"
+    return settings.home_dir / "plugins"
 
 
 logger = logging.getLogger(__name__)

giskard/utils/__init__.py

@@ -19,7 +19,10 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-class WorkerPool:
+NOT_STARTED = "Pool is not started"
+
+
+class SingletonWorkerPool:
     "Utility class to wrap a Process pool"
 
     def __init__(self):
@@ -30,9 +33,7 @@ def start(self, max_workers: int = None):
         if self.pool is not None:
             return
         self.max_workers = max(max_workers, settings.min_workers) if max_workers is not None else os.cpu_count()
-        LOGGER.info("Starting worker pool with %s workers...", self.max_workers)
         self.pool = WorkerPoolExecutor(nb_workers=self.max_workers)
-        LOGGER.info("Pool is started")
 
     def shutdown(self, wait=True):
         if self.pool is None:
@@ -41,21 +42,21 @@ def shutdown(self, wait=True):
 
     def schedule(self, fn, args=None, kwargs=None, timeout=None) -> Future:
         if self.pool is None:
-            raise ValueError("Pool is not started")
+            raise RuntimeError(NOT_STARTED)
         return self.pool.schedule(fn, args=args, kwargs=kwargs, timeout=timeout)
 
     def submit(self, *args, **kwargs) -> Future:
         if self.pool is None:
-            raise ValueError("Pool is not started")
+            raise RuntimeError(NOT_STARTED)
         return self.pool.submit(*args, **kwargs)
 
     def map(self, *args, **kwargs):
         if self.pool is None:
-            raise ValueError("Pool is not started")
+            raise RuntimeError(NOT_STARTED)
         return self.pool.map(*args, **kwargs)
 
 
-POOL = WorkerPool()
+POOL = SingletonWorkerPool()
 
 
 def start_pool(max_workers: int = None):

giskard/utils/worker_pool.py

@@ -12,6 +12,7 @@
 from multiprocessing import Process, Queue, SimpleQueue, cpu_count, get_context
 from multiprocessing.context import SpawnContext, SpawnProcess
 from multiprocessing.managers import SyncManager
+from queue import Empty
 from threading import Thread
 from time import sleep
 from uuid import uuid4
@@ -72,13 +73,12 @@ class GiskardTask:
 @dataclass(frozen=True)
 class GiskardResult:
     id: str
+    logs: str
     result: Any = None
     exception: Any = None
-    logs: str = None
 
 
 def _process_worker(tasks_queue: SimpleQueue, tasks_results: SimpleQueue, running_process: Dict[str, str]):
-    # TODO(Bazire): For now, worker stops because the queues will be closed. Not the cleanest...
     pid = os.getpid()
 
     while True:
@@ -93,7 +93,7 @@ def _process_worker(tasks_queue: SimpleQueue, tasks_results: SimpleQueue, runnin
                 handler = logging.StreamHandler(f)
                 logging.getLogger().addHandler(handler)
                 try:
-                    LOGGER.debug("Doing task %", task.id)
+                    LOGGER.debug("Doing task %s", task.id)
                     running_process[task.id] = pid
                     result = task.callable(*task.args, **task.kwargs)
                     to_return = GiskardResult(id=task.id, result=result, logs=f.getvalue())
@@ -119,7 +119,8 @@ class PoolState(Enum):
 
 
 class WorkerPoolExecutor(Executor):
-    def __init__(self, nb_workers: Optional[int] = None):
+    def __init__(self, nb_workers: Optional[int] = None, name: Optional[str] = None):
+        self._prefix = f"{name}_" if name is not None else "giskard_pool_"
         if nb_workers is None:
             nb_workers = cpu_count()
         if nb_workers <= 0:
@@ -135,33 +136,42 @@ def __init__(self, nb_workers: Optional[int] = None):
         # Mapping of the running tasks and worker pids
         self._running_process: Dict[str, str] = self._manager.dict()
         # Mapping of the running tasks and worker pids
-        self._with_timeout_tasks: List[Tuple[str, float]] = []
-        # Queue with tasks to run
-        self._pending_tasks_queue: SimpleQueue[GiskardTask] = self._mp_context.SimpleQueue()
+        self._with_timeout_tasks: List[TimeoutData] = []
         # Queue with tasks to run
+        self._pending_tasks_queue: Queue[GiskardTask] = self._mp_context.Queue()
+        # Queue with tasks to be consumed asap
         # As in ProcessPool, add one more to avoid idling process
-        self._running_tasks_queue: Queue[GiskardTask] = self._mp_context.Queue(maxsize=self._nb_workers + 1)
+        self._running_tasks_queue: Queue[Optional[GiskardTask]] = self._mp_context.Queue(maxsize=self._nb_workers + 1)
         # Queue with results to notify
-        self._tasks_results: SimpleQueue[GiskardResult] = self._mp_context.SimpleQueue()
+        self._tasks_results: Queue[GiskardResult] = self._mp_context.Queue()
         # Mapping task_id with future
         self._futures_mapping: Dict[str, Future] = dict()
+        LOGGER.debug("Starting threads for the WorkerPoolExecutor")
 
         self._threads = [
-            Thread(name="giskard_pool_" + target.__name__, target=target, daemon=True, args=[self], kwargs=None)
+            Thread(name=f"{self._prefix}{target.__name__}", target=target, daemon=True, args=[self], kwargs=None)
             for target in [_killer_thread, _feeder_thread, _results_thread]
         ]
         for t in self._threads:
             t.start()
+        LOGGER.debug("Threads started, spawning workers...")
+        LOGGER.debug("Starting the pool with %s", {self._nb_workers})
 
         # Startup workers
         for _ in range(self._nb_workers):
             self._spawn_worker()
+        LOGGER.info("WorkerPoolExecutor is started")
+
+    def health_check(self):
+        if any([not p.is_alive() for p in self._processes.values()]):
+            LOGGER.warning("At least one process died for an unknown reason, marking pool as broken")
+            self._state = PoolState.BROKEN
 
     def _spawn_worker(self):
         # Daemon means process are linked to main one, and will be stopped if current process is stopped
         p = self._mp_context.Process(
             target=_process_worker,
-            name="giskard_worker_process",
+            name=f"{self._prefix}_worker_process",
             args=(self._running_tasks_queue, self._tasks_results, self._running_process),
             daemon=True,
         )
@@ -190,7 +200,7 @@ def schedule(
         self._futures_mapping[task.id] = res
         self._pending_tasks_queue.put(task)
         if timeout is not None:
-            self._with_timeout_tasks.append((task.id, time.monotonic() + timeout))
+            self._with_timeout_tasks.append(TimeoutData(task.id, time.monotonic() + timeout))
         return res
 
     def shutdown(self, wait=True, timeout: float = 5):
@@ -207,8 +217,14 @@ def shutdown(self, wait=True, timeout: float = 5):
         while not self._running_tasks_queue.empty():
             try:
                 self._running_tasks_queue.get_nowait()
-            except BaseException:
-                pass
+            except ValueError as e:
+                # This happens if queues is closed
+                LOGGER.warning("Running task queue is already closed")
+                LOGGER.exception(e)
+            except Empty as e:
+                # May happen if a process consume an element
+                LOGGER.warning("Queue was empty, skipping")
+                LOGGER.exception(e)
         # Try to nicely stop the worker, by adding None into the running tasks
         for _ in range(self._nb_workers):
             self._running_tasks_queue.put(None, timeout=1)
@@ -235,92 +251,97 @@ def shutdown(self, wait=True, timeout: float = 5):
         return exit_codes
 
 
+def _safe_get(queue: Queue, executor: WorkerPoolExecutor, timeout: float = 1) -> Tuple[Any, bool]:
+    try:
+        result = queue.get(timeout=1)
+    except Empty:
+        result = None
+    except ValueError as e:
+        # If queue is closed
+        if executor._state not in FINAL_STATES:
+            LOGGER.error("Queue is closed, and executor not in final state")
+            executor._state = PoolState.BROKEN
+            raise e
+        return None, True
+    if executor._state in FINAL_STATES:
+        return None, True
+    return result, False
+
+
 def _results_thread(
     executor: WorkerPoolExecutor,
 ):
     # Goal of this thread is to feed the running tasks from pending one as soon as possible
-    while True:
-        while executor._state not in FINAL_STATES and (executor._tasks_results.empty()):
-            # TODO(Bazire): find a way to improve this ?
-            # Cannot use select, since we want to be windows compatible
-            sleep(0.01)
-        if executor._state in FINAL_STATES:
+    # while True:
+    while executor._state not in FINAL_STATES:
+        result, should_stop = _safe_get(executor._tasks_results, executor)
+        if should_stop:
             return
-        result = executor._tasks_results.get()
-        future = executor._futures_mapping.get(result.id)
-        if future.cancelled():
-            try:
-                del executor._futures_mapping[result.id]
-            except BaseException:
-                pass
+        if result is None:
+            continue
+
+        future = executor._futures_mapping.pop(result.id, None)
+        if future is None or future.cancelled():
             continue
         future.logs = result.logs
         if result.exception is None:
             future.set_result(result.result)
         else:
-            # TODO(Bazire): improve to get Traceback
             future.set_exception(RuntimeError(result.exception))
-        try:
-            del executor._futures_mapping[result.id]
-        except BaseException:
-            pass
 
 
 def _feeder_thread(
     executor: WorkerPoolExecutor,
 ):
     # Goal of this thread is to feed the running tasks from pending one as soon as possible
-    while True:
-        while executor._state not in FINAL_STATES and (
-            executor._running_tasks_queue.full() or executor._pending_tasks_queue.empty()
-        ):
-            # TODO(Bazire): find a way to improve this ?
-            # Cannot use select, since we want to be windows compatible
-            sleep(0.01)
-        if executor._state in FINAL_STATES:
+    while executor._state not in FINAL_STATES:
+        task, should_stop = _safe_get(executor._pending_tasks_queue, executor)
+        if should_stop:
             return
-        task = executor._pending_tasks_queue.get()
+        if task is None:
+            continue
+
         future = executor._futures_mapping.get(task.id)
-        if future is not None and future.set_running_or_notify_cancel():
+        if future is None:
+            continue
+        if future.set_running_or_notify_cancel():
             executor._running_tasks_queue.put(task)
-        elif future is not None:
+        else:
             # Future has been cancelled already, nothing to do
-            del executor._futures_mapping[task.id]
+            executor._futures_mapping.pop(task.id, False)
 
 
 def _killer_thread(
     executor: WorkerPoolExecutor,
 ):
-    while True:
+    while executor._state not in FINAL_STATES:
         while len(executor._with_timeout_tasks) == 0 and executor._state not in FINAL_STATES:
             # No need to be too active
             sleep(1)
+            executor.health_check()
         if executor._state in FINAL_STATES:
             return
 
-        clean_up = []
+        clean_up: List[TimeoutData] = []
         exception = None
         for timeout_data in executor._with_timeout_tasks:
-            task_id, end_time = timeout_data
-            if task_id not in executor._futures_mapping:
+            if timeout_data.id not in executor._futures_mapping:
                 # Task is already completed, do not wait for it
                 clean_up.append(timeout_data)
-            elif time.monotonic() > end_time:
+            elif time.monotonic() > timeout_data.end_time:
                 # Task has timed out, we should kill it
                 try:
-                    pid = executor._running_process.get(task_id)
-                    if pid is None:
-                        # Task must have finished
-                        continue
-                    future = executor._futures_mapping.get(task_id)
-                    if not future.cancel():
+                    future = executor._futures_mapping.pop(timeout_data.id, None)
+                    if future is not None and not future.cancel():
                         LOGGER.warning("Killing a timed out process")
                         future.set_exception(TimeoutError("Task took too long"))
-                        p = executor._processes[pid]
-                        del executor._processes[pid]
-                        _stop_processes([p])
-                        executor._spawn_worker()
+                        pid = executor._running_process.pop(timeout_data.id, None)
+                        if pid is not None:
+                            p = executor._processes.pop(pid)
+                            _stop_processes([p])
+                            executor._spawn_worker()
                 except BaseException as e:
+                    # This is probably an OSError, but we want to be extra safe
                     LOGGER.warning("Unexpected error when killing a timed out process, pool is broken")
                     LOGGER.exception(e)
                     exception = e