assert on number of threads left running after each test in config local test

Author: benclifford

parsl/executors/high_throughput/executor.py

@@ -329,6 +329,9 @@ def __init__(self,
             launch_cmd = DEFAULT_LAUNCH_CMD
         self.launch_cmd = launch_cmd
 
+        self._queue_management_thread_exit = threading.Event()
+        self._queue_management_thread: Optional[threading.Thread] = None
+
     radio_mode = "htex"
 
     def _warn_deprecated(self, old: str, new: str):
@@ -450,9 +453,9 @@ def _queue_management_worker(self):
         """
         logger.debug("Queue management worker starting")
 
-        while not self.bad_state_is_set:
+        while not self.bad_state_is_set and not self._queue_management_thread_exit.is_set():
             try:
-                msgs = self.incoming_q.get()
+                msgs = self.incoming_q.get(timeout_ms=1000)
 
             except IOError as e:
                 logger.exception("Caught broken queue with exception code {}: {}".format(e.errno, e))
@@ -465,57 +468,55 @@ def _queue_management_worker(self):
             else:
 
                 if msgs is None:
-                    logger.debug("Got None, exiting")
-                    return
+                    continue
+
+                for serialized_msg in msgs:
+                    try:
+                        msg = pickle.loads(serialized_msg)
+                    except pickle.UnpicklingError:
+                        raise BadMessage("Message received could not be unpickled")
 
-                else:
-                    for serialized_msg in msgs:
+                    if msg['type'] == 'heartbeat':
+                        continue
+                    elif msg['type'] == 'result':
                         try:
-                            msg = pickle.loads(serialized_msg)
-                        except pickle.UnpicklingError:
-                            raise BadMessage("Message received could not be unpickled")
+                            tid = msg['task_id']
+                        except Exception:
+                            raise BadMessage("Message received does not contain 'task_id' field")
+
+                        if tid == -1 and 'exception' in msg:
+                            logger.warning("Executor shutting down due to exception from interchange")
+                            exception = deserialize(msg['exception'])
+                            self.set_bad_state_and_fail_all(exception)
+                            break
 
-                        if msg['type'] == 'heartbeat':
-                            continue
-                        elif msg['type'] == 'result':
+                        task_fut = self.tasks.pop(tid)
+
+                        if 'result' in msg:
+                            result = deserialize(msg['result'])
+                            task_fut.set_result(result)
+
+                        elif 'exception' in msg:
                             try:
-                                tid = msg['task_id']
-                            except Exception:
-                                raise BadMessage("Message received does not contain 'task_id' field")
-
-                            if tid == -1 and 'exception' in msg:
-                                logger.warning("Executor shutting down due to exception from interchange")
-                                exception = deserialize(msg['exception'])
-                                self.set_bad_state_and_fail_all(exception)
-                                break
-
-                            task_fut = self.tasks.pop(tid)
-
-                            if 'result' in msg:
-                                result = deserialize(msg['result'])
-                                task_fut.set_result(result)
-
-                            elif 'exception' in msg:
-                                try:
-                                    s = deserialize(msg['exception'])
-                                    # s should be a RemoteExceptionWrapper... so we can reraise it
-                                    if isinstance(s, RemoteExceptionWrapper):
-                                        try:
-                                            s.reraise()
-                                        except Exception as e:
-                                            task_fut.set_exception(e)
-                                    elif isinstance(s, Exception):
-                                        task_fut.set_exception(s)
-                                    else:
-                                        raise ValueError("Unknown exception-like type received: {}".format(type(s)))
-                                except Exception as e:
-                                    # TODO could be a proper wrapped exception?
-                                    task_fut.set_exception(
-                                        DeserializationError("Received exception, but handling also threw an exception: {}".format(e)))
-                            else:
-                                raise BadMessage("Message received is neither result or exception")
+                                s = deserialize(msg['exception'])
+                                # s should be a RemoteExceptionWrapper... so we can reraise it
+                                if isinstance(s, RemoteExceptionWrapper):
+                                    try:
+                                        s.reraise()
+                                    except Exception as e:
+                                        task_fut.set_exception(e)
+                                elif isinstance(s, Exception):
+                                    task_fut.set_exception(s)
+                                else:
+                                    raise ValueError("Unknown exception-like type received: {}".format(type(s)))
+                            except Exception as e:
+                                # TODO could be a proper wrapped exception?
+                                task_fut.set_exception(
+                                    DeserializationError("Received exception, but handling also threw an exception: {}".format(e)))
                         else:
-                            raise BadMessage("Message received with unknown type {}".format(msg['type']))
+                            raise BadMessage("Message received is neither result or exception")
+                    else:
+                        raise BadMessage("Message received with unknown type {}".format(msg['type']))
 
         logger.info("Queue management worker finished")
 
@@ -815,13 +816,19 @@ def shutdown(self, timeout: float = 10.0):
 
         logger.info("Attempting HighThroughputExecutor shutdown")
 
+        logger.info("Terminating interchange and queue management thread")
+        self._queue_management_thread_exit.set()
         self.interchange_proc.terminate()
         try:
             self.interchange_proc.wait(timeout=timeout)
         except subprocess.TimeoutExpired:
-            logger.info("Unable to terminate Interchange process; sending SIGKILL")
+            logger.warning("Unable to terminate Interchange process; sending SIGKILL")
             self.interchange_proc.kill()
 
+        logger.info("Waiting for queue management thread exit")
+        if self._queue_management_thread:
+            self._queue_management_thread.join()
+
         logger.info("closing context sockets")
         # this might block if there are outstanding messages (eg if the interchange
         # has gone away... probably something to do with zmq.LINGER sockopt to remove

parsl/executors/high_throughput/zmq_pipes.py

@@ -205,12 +205,21 @@ def __init__(self, ip_address, port_range, cert_dir: Optional[str] = None):
         self.port = self.results_receiver.bind_to_random_port("tcp://{}".format(ip_address),
                                                               min_port=port_range[0],
                                                               max_port=port_range[1])
+        self.poller = zmq.Poller()
+        self.poller.register(self.results_receiver, zmq.POLLIN)
 
-    def get(self):
+    def get(self, timeout_ms=None):
+        """Get a message from the queue, returning None if timeout expires
+        without a message. timeout is measured in milliseconds.
+        """
         logger.debug("Waiting for ResultsIncoming message")
-        m = self.results_receiver.recv_multipart()
-        logger.debug("Received ResultsIncoming message")
-        return m
+        socks = dict(self.poller.poll(timeout=timeout_ms))
+        if self.results_receiver in socks and socks[self.results_receiver] == zmq.POLLIN:
+            m = self.results_receiver.recv_multipart()
+            logger.debug("Received ResultsIncoming message")
+            return m
+        else:
+            return None
 
     def close(self):
         self.results_receiver.close()

parsl/tests/conftest.py

@@ -180,6 +180,8 @@ def load_dfk_session(request, pytestconfig, tmpd_cwd_session):
     config = pytestconfig.getoption('config')[0]
 
     if config != 'local':
+        assert threading.active_count() == 1, "precondition: only one thread can be running before this test: " + repr(threading.enumerate())
+
         spec = importlib.util.spec_from_file_location('', config)
         module = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(module)
@@ -207,6 +209,9 @@ def load_dfk_session(request, pytestconfig, tmpd_cwd_session):
             raise RuntimeError("DFK changed unexpectedly during test")
         dfk.cleanup()
         assert DataFlowKernelLoader._dfk is None
+
+        assert threading.active_count() == 1, "test left threads running: " + repr(threading.enumerate())
+
     else:
         yield
 
@@ -232,6 +237,7 @@ def load_dfk_local_module(request, pytestconfig, tmpd_cwd_session):
         start_fds = this_process.num_fds()
         logger.error(f"BENC: start open fds: {start_fds}")
 
+        assert threading.active_count() == 1, "precondition: only one thread can be running before this test"
         local_setup = getattr(request.module, "local_setup", None)
         local_teardown = getattr(request.module, "local_teardown", None)
         local_config = getattr(request.module, "local_config", None)
@@ -266,6 +272,8 @@ def load_dfk_local_module(request, pytestconfig, tmpd_cwd_session):
         end_fds = this_process.num_fds()
         logger.error(f"BENC: end open fds: {end_fds} (vs start {start_fds}")
 
+        assert threading.active_count() == 1, "test left threads running: " + repr(threading.enumerate())
+
     else:
         yield
 

parsl/tests/test_htex/test_htex.py

@@ -116,7 +116,7 @@ def kill_interchange(*args, **kwargs):
         assert mock_ix_proc.wait.called
         assert {"timeout": 10} == mock_ix_proc.wait.call_args[1]
         if timeout_expires:
-            assert "Unable to terminate Interchange" in mock_logs[1][0][0]
+            assert "Unable to terminate Interchange" in mock_logs[3][0][0]
             assert mock_ix_proc.kill.called
         assert "Attempting" in mock_logs[0][0][0]
         assert "Finished" in mock_logs[-1][0][0]