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)