<commit_before># -*- coding: utf-8 -*-
# File: prof.py


import multiprocessing as mp
import numpy as np
import os
import time
import tensorflow as tf
from six.moves import map
from tensorflow.python.client import timeline

from ..tfutils.common import gpu_available_in_session
from ..utils import logger
from ..utils.concurrency import ensure_proc_terminate, start_proc_mask_signal
from ..utils.gpu import get_num_gpu
from ..utils.nvml import NVMLContext
from .base import Callback

__all__ = ['GPUUtilizationTracker', 'GraphProfiler', 'PeakMemoryTracker']


class GPUUtilizationTracker(Callback):
    """ Summarize the average GPU utilization within an epoch.

    It will start a process to run ``nvidia-smi`` every second
    within the epoch (the trigger_epoch time was not included),
    and write average utilization to monitors.

    This callback creates a process, therefore it's not safe to be used with MPI.
    """

    _chief_only = False

    def __init__(self, devices=None):
        """
        Args:
            devices (list[int]): physical GPU ids. If None, will use CUDA_VISIBLE_DEVICES
        """
        assert os.name != 'nt', "GPUUtilizationTracker does not support windows!"
        if devices is None:
            env = os.environ.get('CUDA_VISIBLE_DEVICES')
            if env is None:
                self._devices = list(range(get_num_gpu()))
                logger.warn("[GPUUtilizationTracker] Both devices and CUDA_VISIBLE_DEVICES are None! "
                            "Will monitor all {} visible GPUs!".format(len(self._devices)))
            else:
                if len(env):
                    self._devices = list(map(int, env.split(',')))
                else:
                    self._devices = []
        else:
            self._devices = devices
        assert len(self._devices), "[GPUUtilizationTracker] No GPU device given!"

    def _before_train(self):
        assert gpu_available_in_session(), "[GPUUtilizationTracker] needs GPU!"
        self._evt = mp.Event()
        self._stop_evt = mp.Event()
        self._queue = mp.Queue()
        self._proc = mp.Process(target=self.worker, args=(
            self._evt, self._queue, self._stop_evt))
        ensure_proc_terminate(self._proc)
        start_proc_mask_signal(self._proc)

    def _before_epoch(self):
        self._evt.set()

    def _after_epoch(self):
        while self._evt.is_set():   # unlikely
            pass
        self._evt.set()

    def _trigger_epoch(self):
        # Don't do this in after_epoch because
        # before,after_epoch are supposed to be extremely fast by design.
        stats = self._queue.get()
        for idx, dev in enumerate(self._devices):
            self.trainer.monitors.put_scalar('GPUUtil/{}'.format(dev), stats[idx])

    def _after_train(self):
        self._stop_evt.set()
        self._evt.set()
        self._proc.join()

    def worker(self, evt, rst_queue, stop_evt):
        while True:
            evt.wait()  # start epoch
            evt.clear()
            if stop_evt.is_set():   # or on exit
                return

            stats = np.zeros((len(self._devices),), dtype='f4')
            cnt = 0
            with NVMLContext() as ctx:
                while True:
                    time.sleep(1)

                    data = [ctx.device(i).utilization()['gpu'] for i in self._devices]
                    data = list(map(float, data))
                    stats += data
                    cnt += 1

                    if evt.is_set():    # stop epoch
                        if stop_evt.is_set():   # or on exit
                            return
                        evt.clear()
                        if cnt > 1:
                            # Ignore the last datapoint. Usually is zero, makes us underestimate the util.
                            stats -= data
                            cnt -= 1
                        rst_queue.put(stats / cnt)
                        break


# Can add more features from tfprof
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/profiler/README.md

class GraphProfiler(Callback):
    """
    Enable profiling by installing session hooks,
    and write tracing files / events / metadata to ``logger.get_logger_dir()``.

    The tracing files can be loaded from ``chrome://tracing``.
    The metadata files can be processed by
    `tfprof command line utils
    <https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/profiler/README.md>`_.
    The event is viewable from tensorboard.

    Tips:

    Note that the profiling is by default enabled for every step and is expensive.<commit_msg>Catch exception in GPUUtilization.worker (#1134)<commit_after># -*- coding: utf-8 -*-
# File: prof.py


import multiprocessing as mp
import numpy as np
import os
import time
import tensorflow as tf
from six.moves import map
from tensorflow.python.client import timeline

from ..tfutils.common import gpu_available_in_session
from ..utils import logger
from ..utils.concurrency import ensure_proc_terminate, start_proc_mask_signal
from ..utils.gpu import get_num_gpu
from ..utils.nvml import NVMLContext
from .base import Callback

__all__ = ['GPUUtilizationTracker', 'GraphProfiler', 'PeakMemoryTracker']


class GPUUtilizationTracker(Callback):
    """ Summarize the average GPU utilization within an epoch.

    It will start a process to run ``nvidia-smi`` every second
    within the epoch (the trigger_epoch time was not included),
    and write average utilization to monitors.

    This callback creates a process, therefore it's not safe to be used with MPI.
    """

    _chief_only = False

    def __init__(self, devices=None):
        """
        Args:
            devices (list[int]): physical GPU ids. If None, will use CUDA_VISIBLE_DEVICES
        """
        assert os.name != 'nt', "GPUUtilizationTracker does not support windows!"
        if devices is None:
            env = os.environ.get('CUDA_VISIBLE_DEVICES')
            if env is None:
                self._devices = list(range(get_num_gpu()))
                logger.warn("[GPUUtilizationTracker] Both devices and CUDA_VISIBLE_DEVICES are None! "
                            "Will monitor all {} visible GPUs!".format(len(self._devices)))
            else:
                if len(env):
                    self._devices = list(map(int, env.split(',')))
                else:
                    self._devices = []
        else:
            self._devices = devices
        assert len(self._devices), "[GPUUtilizationTracker] No GPU device given!"

    def _before_train(self):
        assert gpu_available_in_session(), "[GPUUtilizationTracker] needs GPU!"
        self._evt = mp.Event()
        self._stop_evt = mp.Event()
        self._queue = mp.Queue()
        self._proc = mp.Process(target=self.worker, args=(
            self._evt, self._queue, self._stop_evt))
        ensure_proc_terminate(self._proc)
        start_proc_mask_signal(self._proc)

    def _before_epoch(self):
        self._evt.set()

    def _after_epoch(self):
        while self._evt.is_set():   # unlikely
            pass
        self._evt.set()

    def _trigger_epoch(self):
        # 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])

    def _after_train(self):
        self._stop_evt.set()
        self._evt.set()
        self._proc.join()

    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
                    return

                stats = np.zeros((len(self._devices),), dtype='f4')
                cnt = 0
                with NVMLContext() as ctx:
                    while True:
                        time.sleep(1)

                        data = [ctx.device(i).utilization()['gpu'] for i in self._devices]
                        data = list(map(float, data))
                        stats += data
                        cnt += 1

                        if evt.is_set():    # stop epoch
                            if stop_evt.is_set():   # or on exit
                                return
                            evt.clear()
                            if cnt > 1:
                                # Ignore the last datapoint. Usually is zero, makes us underestimate the util.
                                stats -= data
                                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
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/profiler/README.md

class GraphProfiler(Callback):
    """
    Enable profiling by installing session hooks,
    and write tracing files / events / metadata to ``logger.get_logger_dir()``.

    The tracing files can be loaded from ``chrome://tracing``.
    The metadata files can be processed by
    `tfprof command line utils
    <https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/profiler/README.md>`_.
    The event is viewable from tensorboard.

    Tips:

    Note that the profiling is by default enabled for every step and is expensive.
{'fraction_non_alphanumeric': 0.09246715927750411, 'fraction_numerical': 0.002052545155993432, 'mean_word_length': 1.9037544696066746, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4200705', 'n_tokens_mistral': 2683, 'n_tokens_neox': 2603, 'n_words': 770}