Fix monitoring ctrlc hang

parsl/monitoring/db_manager.py

@@ -3,6 +3,7 @@
 import queue
 import os
 import time
+import datetime
 
 from parsl.dataflow.states import States
 from parsl.providers.error import OptionalModuleMissing
@@ -202,6 +203,8 @@ def __init__(self,
                  batching_threshold=99999,
                  ):
 
+        self.workflow_end = False
+        self.workflow_start_message = None
         self.logdir = logdir
         os.makedirs(self.logdir, exist_ok=True)
 
@@ -284,6 +287,7 @@ def start(self, priority_queue, node_queue, resource_queue):
                             self.logger.debug(
                                 "Inserting workflow start info to WORKFLOW table")
                             self._insert(table=WORKFLOW, messages=[msg])
+                            self.workflow_start_message = msg
                         else:                         # workflow end message
                             self.logger.debug(
                                 "Updating workflow end info to WORKFLOW table")
@@ -292,6 +296,8 @@ def start(self, priority_queue, node_queue, resource_queue):
                                                   'tasks_completed_count', 'time_completed',
                                                   'workflow_duration'],
                                          messages=[msg])
+                            self.workflow_end = True
+
                     else:                             # TASK_INFO message
                         all_messages.append(msg)
                         if msg['task_id'] in inserted_tasks:
@@ -389,22 +395,24 @@ def _migrate_logs_to_internal(self, logs_queue, queue_tag, kill_event):
     def _update(self, table, columns, messages):
         try:
             self.db.update(table=table, columns=columns, messages=messages)
-        except Exception:
+        except BaseException:
             self.logger.exception("Got exception when trying to update Table {}".format(table))
             try:
                 self.db.rollback()
             except Exception:
                 self.logger.exception("Rollback failed")
+            raise
 
     def _insert(self, table, messages):
         try:
             self.db.insert(table=table, messages=messages)
-        except Exception:
+        except BaseException:
             self.logger.exception("Got exception when trying to insert to Table {}".format(table))
             try:
                 self.db.rollback()
             except Exception:
                 self.logger.exception("Rollback failed")
+            raise
 
     def _get_messages_in_batch(self, msg_queue, interval=1, threshold=99999):
         messages = []
@@ -425,7 +433,19 @@ def _get_messages_in_batch(self, msg_queue, interval=1, threshold=99999):
     def close(self):
         if self.logger:
             self.logger.info(
-                "Finishing all the logging and terminating Database Manager.")
+                "Database Manager cleanup initiated.")
+        if not self.workflow_end and self.workflow_start_message:
+            if self.logger:
+                self.logger.info(
+                    "Logging workflow end info to database due to abnormal exit")
+            time_completed = datetime.datetime.now()
+            msg = {'time_completed': time_completed,
+                   'workflow_duration': (time_completed - self.workflow_start_message['time_began']).total_seconds()}
+            self.workflow_start_message.update(msg)
+            self._update(table=WORKFLOW,
+                         columns=['run_id', 'time_completed',
+                                  'workflow_duration'],
+                         messages=[self.workflow_start_message])
         self.batching_interval, self.batching_threshold = float(
             'inf'), float('inf')
         self._kill_event.set()
@@ -472,8 +492,13 @@ def dbm_starter(exception_q, priority_msgs, node_msgs, resource_msgs, *args, **k
     dbm.logger.info("Starting dbm in dbm starter")
     try:
         dbm.start(priority_msgs, node_msgs, resource_msgs)
+    except KeyboardInterrupt:
+        dbm.logger.exception("KeyboardInterrupt signal caught")
+        dbm.close()
+        raise
     except Exception as e:
         dbm.logger.exception("dbm.start exception")
         exception_q.put(("DBM", str(e)))
+        dbm.close()
 
     dbm.logger.info("End of dbm_starter")

parsl/monitoring/monitoring.py

@@ -207,8 +207,10 @@ def start(self, run_id):
 
         self.logger.debug("Initializing ZMQ Pipes to client")
         self.monitoring_hub_active = True
+        self.dfk_channel_timeout = 10000  # in milliseconds
         self._context = zmq.Context()
         self._dfk_channel = self._context.socket(zmq.DEALER)
+        self._dfk_channel.setsockopt(zmq.SNDTIMEO, self.dfk_channel_timeout)
         self._dfk_channel.set_hwm(0)
         self.dfk_port = self._dfk_channel.bind_to_random_port("tcp://{}".format(self.client_address),
                                                               min_port=self.client_port_range[0],
@@ -259,7 +261,11 @@ def start(self, run_id):
 
     def send(self, mtype, message):
         self.logger.debug("Sending message {}, {}".format(mtype, message))
-        return self._dfk_channel.send_pyobj((mtype, message))
+        try:
+            self._dfk_channel.send_pyobj((mtype, message))
+        except zmq.Again:
+            self.logger.exception(
+                "[MONITORING] The monitoring message sent from DFK to Hub timeouts after {}ms".format(self.dfk_channel_timeout))
 
     def close(self):
         if self.logger:
@@ -285,6 +291,7 @@ def close(self):
             if len(exception_msgs) == 0:
                 self.priority_msgs.put(("STOP", 0))
             self.dbm_proc.join()
+            self.logger.debug("Finished waiting for DBM termination")
 
     @staticmethod
     def monitor_wrapper(f,