Commit f20e411a authored by Yuxin Wu's avatar Yuxin Wu

Mask sigint in GPU util process (#516)

parent 5b2bf238
...@@ -13,7 +13,7 @@ from tensorflow.python.client import timeline ...@@ -13,7 +13,7 @@ from tensorflow.python.client import timeline
from .base import Callback from .base import Callback
from ..utils import logger from ..utils import logger
from ..utils.concurrency import ensure_proc_terminate, subproc_call from ..utils.concurrency import ensure_proc_terminate, subproc_call, start_proc_mask_signal
from ..utils.gpu import get_nr_gpu from ..utils.gpu import get_nr_gpu
__all__ = ['GPUUtilizationTracker', 'GraphProfiler', 'PeakMemoryTracker'] __all__ = ['GPUUtilizationTracker', 'GraphProfiler', 'PeakMemoryTracker']
...@@ -59,7 +59,7 @@ class GPUUtilizationTracker(Callback): ...@@ -59,7 +59,7 @@ class GPUUtilizationTracker(Callback):
self._proc = mp.Process(target=self.worker, args=( self._proc = mp.Process(target=self.worker, args=(
self._evt, self._queue, self._stop_evt)) self._evt, self._queue, self._stop_evt))
ensure_proc_terminate(self._proc) ensure_proc_terminate(self._proc)
self._proc.start() start_proc_mask_signal(self._proc)
def _before_epoch(self): def _before_epoch(self):
self._evt.set() self._evt.set()
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment