Catch exception in GPUUtilization.worker (#1134)

tensorpack/callbacks/prof.py

@@ -75,6 +75,9 @@ class GPUUtilizationTracker(Callback):
         # Don't do this in after_epoch because
         # before,after_epoch are supposed to be extremely fast by design.
         stats = self._queue.get()
+        if stats == -1:
+            from ..train.base import StopTraining
+            raise StopTraining("GPUUtilizationTracker.worker has failed.")
         for idx, dev in enumerate(self._devices):
             self.trainer.monitors.put_scalar('GPUUtil/{}'.format(dev), stats[idx])
 
@@ -85,6 +88,7 @@ class GPUUtilizationTracker(Callback):
 
     def worker(self, evt, rst_queue, stop_evt):
         while True:
+            try:
                 evt.wait()  # start epoch
                 evt.clear()
                 if stop_evt.is_set():   # or on exit
@@ -111,6 +115,10 @@ class GPUUtilizationTracker(Callback):
                                 cnt -= 1
                             rst_queue.put(stats / cnt)
                             break
+            except Exception:
+                logger.exception("Exception in GPUUtilizationTracker.worker")
+                rst_queue.put(-1)
+                return
 
 
 # Can add more features from tfprof