EVOLUTION-MANAGER

Edit File: event_file_writer.py

# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Writes events to disk in a logdir."""

import os
import queue
import socket
import threading
import time

from tensorboard.compat import tf
from tensorboard.compat.proto import event_pb2
from tensorboard.summary.writer.record_writer import RecordWriter

class AtomicCounter(object):
    def __init__(self, initial_value):
        self._value = initial_value
        self._lock = threading.Lock()

def get(self):
        with self._lock:
            try:
                return self._value
            finally:
                self._value += 1

_global_uid = AtomicCounter(0)

class EventFileWriter(object):
    """Writes `Event` protocol buffers to an event file.

The `EventFileWriter` class creates an event file in the specified
    directory, and asynchronously writes Event protocol buffers to the
    file. The Event file is encoded using the tfrecord format, which is
    similar to RecordIO.
    """

def __init__(
        self, logdir, max_queue_size=10, flush_secs=120, filename_suffix=""
    ):
        """Creates a `EventFileWriter` and an event file to write to.

On construction the summary writer creates a new event file in `logdir`.
        This event file will contain `Event` protocol buffers, which are written to
        disk via the add_event method.
        The other arguments to the constructor control the asynchronous writes to
        the event file:

Args:
          logdir: A string. Directory where event file will be written.
          max_queue_size: Integer. Size of the queue for pending events and summaries.
          flush_secs: Number. How often, in seconds, to flush the
            pending events and summaries to disk.
        """
        self._logdir = logdir
        tf.io.gfile.makedirs(logdir)
        self._file_name = (
            os.path.join(
                logdir,
                "events.out.tfevents.%010d.%s.%s.%s"
                % (
                    time.time(),
                    socket.gethostname(),
                    os.getpid(),
                    _global_uid.get(),
                ),
            )
            + filename_suffix
        )  # noqa E128
        self._general_file_writer = tf.io.gfile.GFile(self._file_name, "wb")
        self._async_writer = _AsyncWriter(
            RecordWriter(self._general_file_writer), max_queue_size, flush_secs
        )

# Initialize an event instance.
        _event = event_pb2.Event(
            wall_time=time.time(), file_version="brain.Event:2"
        )
        self.add_event(_event)
        self.flush()

def get_logdir(self):
        """Returns the directory where event file will be written."""
        return self._logdir

def add_event(self, event):
        """Adds an event to the event file.

Args:
          event: An `Event` protocol buffer.
        """
        if not isinstance(event, event_pb2.Event):
            raise TypeError(
                "Expected an event_pb2.Event proto, "
                " but got %s" % type(event)
            )
        self._async_writer.write(event.SerializeToString())

def flush(self):
        """Flushes the event file to disk.

Call this method to make sure that all pending events have been
        written to disk.
        """
        self._async_writer.flush()

def close(self):
        """Performs a final flush of the event file to disk, stops the
        write/flush worker and closes the file.

Call this method when you do not need the summary writer
        anymore.
        """
        self._async_writer.close()

class _AsyncWriter(object):
    """Writes bytes to a file."""

def __init__(self, record_writer, max_queue_size=20, flush_secs=120):
        """Writes bytes to a file asynchronously. An instance of this class
        holds a queue to keep the incoming data temporarily. Data passed to the
        `write` function will be put to the queue and the function returns
        immediately. This class also maintains a thread to write data in the
        queue to disk. The first initialization parameter is an instance of
        `tensorboard.summary.record_writer` which computes the CRC checksum and
        then write the combined result to the disk. So we use an async approach
        to improve performance.

Args:
            record_writer: A RecordWriter instance
            max_queue_size: Integer. Size of the queue for pending bytestrings.
            flush_secs: Number. How often, in seconds, to flush the
                pending bytestrings to disk.
        """
        self._writer = record_writer
        self._closed = False
        self._byte_queue = queue.Queue(max_queue_size)
        self._worker = _AsyncWriterThread(
            self._byte_queue, self._writer, flush_secs
        )
        self._lock = threading.Lock()
        self._worker.start()

def write(self, bytestring):
        """Enqueue the given bytes to be written asychronously."""
        with self._lock:
            if self._closed:
                raise IOError("Writer is closed")
            self._byte_queue.put(bytestring)

def flush(self):
        """Write all the enqueued bytestring before this flush call to disk.

Block until all the above bytestring are written.
        """
        with self._lock:
            if self._closed:
                raise IOError("Writer is closed")
            self._byte_queue.join()
            self._writer.flush()

def close(self):
        """Closes the underlying writer, flushing any pending writes first."""
        if not self._closed:
            with self._lock:
                if not self._closed:
                    self._closed = True
                    self._worker.stop()
                    self._writer.flush()
                    self._writer.close()

class _AsyncWriterThread(threading.Thread):
    """Thread that processes asynchronous writes for _AsyncWriter."""

def __init__(self, queue, record_writer, flush_secs):
        """Creates an _AsyncWriterThread.

Args:
          queue: A Queue from which to dequeue data.
          record_writer: An instance of record_writer writer.
          flush_secs: How often, in seconds, to flush the
            pending file to disk.
        """
        threading.Thread.__init__(self)
        self.daemon = True
        self._queue = queue
        self._record_writer = record_writer
        self._flush_secs = flush_secs
        # The first data will be flushed immediately.
        self._next_flush_time = 0
        self._has_pending_data = False
        self._shutdown_signal = object()

def stop(self):
        self._queue.put(self._shutdown_signal)
        self.join()

def run(self):
        # Here wait on the queue until an data appears, or till the next
        # time to flush the writer, whichever is earlier. If we have an
        # data, write it. If not, an empty queue exception will be raised
        # and we can proceed to flush the writer.
        while True:
            now = time.time()
            queue_wait_duration = self._next_flush_time - now
            data = None
            try:
                if queue_wait_duration > 0:
                    data = self._queue.get(True, queue_wait_duration)
                else:
                    data = self._queue.get(False)

if data is self._shutdown_signal:
                    return
                self._record_writer.write(data)
                self._has_pending_data = True
            except queue.Empty:
                pass
            finally:
                if data:
                    self._queue.task_done()

now = time.time()
            if now > self._next_flush_time:
                if self._has_pending_data:
                    # Small optimization - if there are no pending data,
                    # there's no need to flush, since each flush can be
                    # expensive (e.g. uploading a new file to a server).
                    self._record_writer.flush()
                    self._has_pending_data = False
                # Do it again in flush_secs.
                self._next_flush_time = now + self._flush_secs