def __init__(
self,
logdir,
max_queue_size=10,
flush_secs=120,
filename_base=None,
filename_suffix=""):
with warnings.catch_warnings():
warnings.simplefilter("ignore", Warning)
# pylint: disable=no-name-in-module
from tensorboard.summary.writer import event_file_writer as writelib
util.ensure_dir(logdir)
filename_base = filename_base or ("%010d.%s.%s.%s" % (
time.time(),
socket.gethostname(),
os.getpid(),
writelib._global_uid.get()))
filename = (
os.path.join(logdir, "events.out.tfevents.%s" % filename_base) +
filename_suffix)
self._writer = writelib._AsyncWriter(
writelib.RecordWriter(open(filename, "wb")),
max_queue_size, flush_secs)
event = writelib.event_pb2.Event(
wall_time=time.time(), file_version="brain.Event:2")
self.add_event(event)
self.flush()
def test_async_writer_write_once(self):
filename = os.path.join(self.get_temp_dir(), "async_writer_write_once")
w = _AsyncWriter(open(filename, 'wb'))
bytes_to_write = b"hello world"
w.write(bytes_to_write)
w.close()
with open(filename, 'rb') as f:
self.assertEqual(f.read(), bytes_to_write)
def test_async_writer_close_triggers_flush(self):
filename = os.path.join(self.get_temp_dir(), "async_writer_close_triggers_flush")
w = _AsyncWriter(open(filename, 'wb'))
bytes_to_write = b"x" * 64
w.write(bytes_to_write)
w.close()
with open(filename, 'rb') as f:
self.assertEqual(f.read(), bytes_to_write)
def test_async_writer_write_one_slot_queue(self):
filename = os.path.join(self.get_temp_dir(), "async_writer_write_one_slot_queue")
w = _AsyncWriter(open(filename, 'wb'), max_queue_size=1)
bytes_to_write = b"hello world"
repeat = 10 # faster
for i in range(repeat):
w.write(bytes_to_write)
w.close()
with open(filename, 'rb') as f:
self.assertEqual(f.read(), bytes_to_write * repeat)
def test_async_writer_write_queue_full(self):
filename = os.path.join(self.get_temp_dir(),
"async_writer_write_queue_full")
w = _AsyncWriter(open(filename, "wb"))
bytes_to_write = b"hello world"
repeat = 100
for i in range(repeat):
w.write(bytes_to_write)
w.close()
with open(filename, "rb") as f:
self.assertEqual(f.read(), bytes_to_write * repeat)
def test_write_after_async_writer_closed(self):
filename = os.path.join(self.get_temp_dir(), "write_after_async_writer_closed")
w = _AsyncWriter(open(filename, 'wb'))
bytes_to_write = b"x" * 64
w.write(bytes_to_write)
w.close()
with self.assertRaises(IOError):
w.write(bytes_to_write)
# nothing is written to the file after close
with open(filename, 'rb') as f:
self.assertEqual(f.read(), bytes_to_write)
def AsyncWriter(logdir,
max_queue_size=10,
flush_secs=120,
filename_base=None,
filename_suffix=""):
with warnings.catch_warnings():
warnings.simplefilter("ignore", Warning)
# pylint: disable=no-name-in-module
from tensorboard.summary.writer import event_file_writer
filename = event_filename(logdir,
filename_base=filename_base,
filename_suffix=filename_suffix)
return event_file_writer._AsyncWriter(
event_file_writer.RecordWriter(open(filename, "wb")), max_queue_size,
flush_secs)
