def testWriterInitAndClose(self):
   logdir = self.get_temp_dir()
   get_total = lambda: len(summary_test_util.events_from_logdir(logdir))
   with summary_ops.always_record_summaries():
     writer = summary_ops.create_file_writer(
         logdir, max_queue=100, flush_millis=1000000)
     self.assertEqual(1, get_total())  # file_version Event
     # Calling init() again while writer is open has no effect
     writer.init()
     self.assertEqual(1, get_total())
     try:
       # Not using .as_default() to avoid implicit flush when exiting
       writer.set_as_default()
       summary_ops.scalar('one', 1.0, step=1)
       self.assertEqual(1, get_total())
       # Calling .close() should do an implicit flush
       writer.close()
       self.assertEqual(2, get_total())
       # Calling init() on a closed writer should start a new file
       time.sleep(1.1)  # Ensure filename has a different timestamp
       writer.init()
       files = sorted(gfile.Glob(os.path.join(logdir, '*tfevents*')))
       self.assertEqual(2, len(files))
       get_total = lambda: len(summary_test_util.events_from_file(files[1]))
       self.assertEqual(1, get_total())  # file_version Event
       summary_ops.scalar('two', 2.0, step=2)
       writer.close()
       self.assertEqual(2, get_total())
     finally:
       # Clean up by resetting default writer
       summary_ops.create_file_writer(None).set_as_default()

 def testWrite_usingDefaultStepVariable_fromLegacyGraph(self):
   logdir = self.get_temp_dir()
   try:
     with context.graph_mode():
       writer = summary_ops.create_file_writer(logdir)
       mystep = variables.Variable(0, dtype=dtypes.int64)
       summary_ops.set_step(mystep)
       with writer.as_default():
         write_op = summary_ops.write('tag', 1.0)
       first_assign_op = mystep.assign_add(1)
       second_assign_op = mystep.assign(10)
       with self.cached_session() as sess:
         sess.run(writer.init())
         sess.run(mystep.initializer)
         sess.run(write_op)
         sess.run(first_assign_op)
         sess.run(write_op)
         sess.run(second_assign_op)
         sess.run(write_op)
         sess.run(writer.flush())
     events = events_from_logdir(logdir)
     self.assertEqual(4, len(events))
     self.assertEqual(0, events[1].step)
     self.assertEqual(1, events[2].step)
     self.assertEqual(10, events[3].step)
   finally:
     # Reset to default state for other tests.
     summary_ops.set_step(None)

  def _test_summary_for_replica_zero_only(self, d):
    logdir = tempfile.mkdtemp()

    def run_fn():
      """Function executed for each replica."""
      with summary_writer.as_default():
        replica_id = ds_context.get_replica_context().replica_id_in_sync_group
        return summary_ops.write("a", replica_id)

    with self.cached_session() as sess, d.scope(), \
        summary_ops.always_record_summaries():
      # We need global_step because summary writing op *always* has global_step
      # as input, even when we always record summary or never record summary.
      global_step = training_util.get_or_create_global_step()
      if not context.executing_eagerly():
        # When executing eagerly, variables are initialized immediately after
        # creation, and its initializer will be None.
        global_step.initializer.run()
      summary_ops.set_step(0)
      summary_writer = summary_ops.create_file_writer(logdir)
      output = d.extended.call_for_each_replica(run_fn)
      unwrapped = d.unwrap(output)
      if not context.executing_eagerly():
        sess.run(summary_writer.init())
        sess.run(unwrapped)
        sess.run(summary_writer.close())

      events = _events_from_logdir(self, logdir)
      # There will be 2 entries: 1 summary file header entry, and 1 entry
      # written by replica 0.
      self.assertLen(events, 2)
      self.assertEqual(events[1].summary.value[0].tag, "a")
      self.assertEqual(events[1].summary.value[0].simple_value, 0.0)

  def testScalarSummaryNameScope(self):
    """Test record_summaries_every_n_global_steps and all_summaries()."""
    with ops.Graph().as_default(), self.cached_session() as sess:
      global_step = training_util.get_or_create_global_step()
      global_step.initializer.run()
      with ops.device('/cpu:0'):
        step_increment = state_ops.assign_add(global_step, 1)
      sess.run(step_increment)  # Increment global step from 0 to 1

      logdir = tempfile.mkdtemp()
      with summary_ops.create_file_writer(logdir, max_queue=0,
                                          name='t2').as_default():
        with summary_ops.record_summaries_every_n_global_steps(2):
          summary_ops.initialize()
          with ops.name_scope('scope'):
            summary_op = summary_ops.scalar('my_scalar', 2.0)

          # Neither of these should produce a summary because
          # global_step is 1 and "1 % 2 != 0"
          sess.run(summary_ops.all_summary_ops())
          sess.run(summary_op)
          events = summary_test_util.events_from_logdir(logdir)
          self.assertEqual(len(events), 1)

          # Increment global step from 1 to 2 and check that the summary
          # is now written
          sess.run(step_increment)
          sess.run(summary_ops.all_summary_ops())
          events = summary_test_util.events_from_logdir(logdir)
          self.assertEqual(len(events), 2)
          self.assertEqual(events[1].summary.value[0].tag, 'scope/my_scalar')

 def testWriterInitAndClose(self):
   logdir = self.get_temp_dir()
   with summary_ops.always_record_summaries():
     writer = summary_ops.create_file_writer(
         logdir, max_queue=100, flush_millis=1000000)
     with writer.as_default():
       summary_ops.scalar('one', 1.0, step=1)
   with self.cached_session() as sess:
     sess.run(summary_ops.summary_writer_initializer_op())
     get_total = lambda: len(summary_test_util.events_from_logdir(logdir))
     self.assertEqual(1, get_total())  # file_version Event
     # Running init() again while writer is open has no effect
     sess.run(writer.init())
     self.assertEqual(1, get_total())
     sess.run(summary_ops.all_summary_ops())
     self.assertEqual(1, get_total())
     # Running close() should do an implicit flush
     sess.run(writer.close())
     self.assertEqual(2, get_total())
     # Running init() on a closed writer should start a new file
     time.sleep(1.1)  # Ensure filename has a different timestamp
     sess.run(writer.init())
     sess.run(summary_ops.all_summary_ops())
     sess.run(writer.close())
     files = sorted(gfile.Glob(os.path.join(logdir, '*tfevents*')))
     self.assertEqual(2, len(files))
     self.assertEqual(2, len(summary_test_util.events_from_file(files[1])))

 def _init_writer(self):
   """Sets file writer."""
   if context.executing_eagerly():
     self.writer = summary_ops_v2.create_file_writer(self.log_dir)
   elif self.write_graph:
     self.writer = tf_summary.FileWriter(self.log_dir, K.get_session().graph)
   else:
     self.writer = tf_summary.FileWriter(self.log_dir)

 def run_metadata_graphs(self, *args, **kwargs):
   assert context.executing_eagerly()
   logdir = self.get_temp_dir()
   writer = summary_ops.create_file_writer(logdir)
   with writer.as_default():
     summary_ops.run_metadata_graphs(*args, **kwargs)
   writer.close()
   events = events_from_logdir(logdir)
   return events[1]

  def __init__(self, session, logdir, max_queue=10, flush_secs=120,
               filename_suffix=''):
    """Creates an `EventFileWriterV2` and an event file to write to.

    On construction, this calls `tf.contrib.summary.create_file_writer` within
    the graph from `session.graph` to look up a shared summary writer resource
    for `logdir` if one exists, and create one if not. Creating the summary
    writer resource in turn creates a new event file in `logdir` to be filled
    with `Event` protocol buffers passed to `add_event`. Graph ops to control
    this writer resource are added to `session.graph` during this init call;
    stateful methods on this class will call `session.run()` on these ops.

    Note that because the underlying resource is shared, it is possible that
    other parts of the code using the same session may interact independently
    with the resource, e.g. by flushing or even closing it. It is the caller's
    responsibility to avoid any undesirable sharing in this regard.

    The remaining arguments to the constructor (`flush_secs`, `max_queue`, and
    `filename_suffix`) control the construction of the shared writer resource
    if one is created. If an existing resource is reused, these arguments have
    no effect.  See `tf.contrib.summary.create_file_writer` for details.

    Args:
      session: A `tf.compat.v1.Session`. Session that will hold shared writer
        resource. The writer ops will be added to session.graph during this
        init call.
      logdir: A string. Directory where event file will be written.
      max_queue: 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.
      filename_suffix: A string. Every event file's name is suffixed with
        `filename_suffix`.
    """
    self._session = session
    self._logdir = logdir
    self._closed = False
    if not gfile.IsDirectory(self._logdir):
      gfile.MakeDirs(self._logdir)

    with self._session.graph.as_default():
      with ops.name_scope('filewriter'):
        file_writer = summary_ops_v2.create_file_writer(
            logdir=self._logdir,
            max_queue=max_queue,
            flush_millis=flush_secs * 1000,
            filename_suffix=filename_suffix)
        with summary_ops_v2.always_record_summaries(), file_writer.as_default():
          self._event_placeholder = array_ops.placeholder_with_default(
              constant_op.constant('unused', dtypes.string),
              shape=[])
          self._add_event_op = summary_ops_v2.import_event(
              self._event_placeholder)
        self._init_op = file_writer.init()
        self._flush_op = file_writer.flush()
        self._close_op = file_writer.close()
      self._session.run(self._init_op)

  def testSharedName(self):
    logdir = self.get_temp_dir()
    with summary_ops.always_record_summaries():
      # Create with default shared name (should match logdir)
      writer1 = summary_ops.create_file_writer(logdir)
      with writer1.as_default():
        summary_ops.scalar('one', 1.0, step=1)
      # Create with explicit logdir shared name (should be same resource/file)
      shared_name = 'logdir:' + logdir
      writer2 = summary_ops.create_file_writer(logdir, name=shared_name)
      with writer2.as_default():
        summary_ops.scalar('two', 2.0, step=2)
      # Create with different shared name (should be separate resource/file)
      writer3 = summary_ops.create_file_writer(logdir, name='other')
      with writer3.as_default():
        summary_ops.scalar('three', 3.0, step=3)

    with self.cached_session() as sess:
      # Run init ops across writers sequentially to avoid race condition.
      # TODO(nickfelt): fix race condition in resource manager lookup or create
      sess.run(writer1.init())
      sess.run(writer2.init())
      time.sleep(1.1)  # Ensure filename has a different timestamp
      sess.run(writer3.init())
      sess.run(summary_ops.all_summary_ops())
      sess.run([writer1.flush(), writer2.flush(), writer3.flush()])

    event_files = iter(sorted(gfile.Glob(os.path.join(logdir, '*tfevents*'))))

    # First file has tags "one" and "two"
    events = summary_test_util.events_from_file(next(event_files))
    self.assertEqual('brain.Event:2', events[0].file_version)
    tags = [e.summary.value[0].tag for e in events[1:]]
    self.assertItemsEqual(['one', 'two'], tags)

    # Second file has tag "three"
    events = summary_test_util.events_from_file(next(event_files))
    self.assertEqual('brain.Event:2', events[0].file_version)
    tags = [e.summary.value[0].tag for e in events[1:]]
    self.assertItemsEqual(['three'], tags)

    # No more files
    self.assertRaises(StopIteration, lambda: next(event_files))

 def keras_model(self, *args, **kwargs):
   logdir = self.get_temp_dir()
   writer = summary_ops.create_file_writer(logdir)
   with writer.as_default():
     summary_ops.keras_model(*args, **kwargs)
   writer.close()
   events = events_from_logdir(logdir)
   # The first event contains no summary values. The written content goes to
   # the second event.
   return events[1]

 def _init_writer(self, model):
   """Sets file writer."""
   if context.executing_eagerly():
     self.writer = summary_ops_v2.create_file_writer(self.log_dir)
     if not model.run_eagerly and self.write_graph:
       with self.writer.as_default():
         summary_ops_v2.graph(K.get_graph())
   elif self.write_graph:
     self.writer = tf_summary.FileWriter(self.log_dir, K.get_graph())
   else:
     self.writer = tf_summary.FileWriter(self.log_dir)

 def testSummaryName(self):
   logdir = self.get_temp_dir()
   writer = summary_ops.create_file_writer(logdir, max_queue=0)
   with writer.as_default(), summary_ops.always_record_summaries():
     summary_ops.scalar('scalar', 2.0, step=1)
   with self.cached_session() as sess:
     sess.run(summary_ops.summary_writer_initializer_op())
     sess.run(summary_ops.all_summary_ops())
   events = summary_test_util.events_from_logdir(logdir)
   self.assertEqual(2, len(events))
   self.assertEqual('scalar', events[1].summary.value[0].tag)

 def testEagerMemory(self):
   training_util.get_or_create_global_step()
   logdir = self.get_temp_dir()
   with summary_ops.create_file_writer(
       logdir, max_queue=0,
       name='t0').as_default(), summary_ops.always_record_summaries():
     summary_ops.generic('tensor', 1, '')
     summary_ops.scalar('scalar', 2.0)
     summary_ops.histogram('histogram', [1.0])
     summary_ops.image('image', [[[[1.0]]]])
     summary_ops.audio('audio', [[1.0]], 1.0, 1)

 def run_trace(self, f, step=1):
   assert context.executing_eagerly()
   logdir = self.get_temp_dir()
   writer = summary_ops.create_file_writer(logdir)
   summary_ops.trace_on(graph=True, profiler=False)
   with writer.as_default():
     f()
     summary_ops.trace_export(name='foo', step=step)
   writer.close()
   events = events_from_logdir(logdir)
   return events[1]

  def testSummaryGlobalStep(self):
    step = training_util.get_or_create_global_step()
    logdir = tempfile.mkdtemp()
    with summary_ops.create_file_writer(
        logdir, max_queue=0,
        name='t2').as_default(), summary_ops.always_record_summaries():

      summary_ops.scalar('scalar', 2.0, step=step)

      events = summary_test_util.events_from_logdir(logdir)
      self.assertEqual(len(events), 2)
      self.assertEqual(events[1].summary.value[0].tag, 'scalar')

 def testMaxQueue(self):
   logs = tempfile.mkdtemp()
   with summary_ops.create_file_writer(
       logs, max_queue=1, flush_millis=999999,
       name='lol').as_default(), summary_ops.always_record_summaries():
     get_total = lambda: len(summary_test_util.events_from_logdir(logs))
     # Note: First tf.Event is always file_version.
     self.assertEqual(1, get_total())
     summary_ops.scalar('scalar', 2.0, step=1)
     self.assertEqual(1, get_total())
     # Should flush after second summary since max_queue = 1
     summary_ops.scalar('scalar', 2.0, step=2)
     self.assertEqual(3, get_total())

  def testWriteSummaries(self):
    m = metrics.Mean()
    m([1, 10, 100])
    training_util.get_or_create_global_step()
    logdir = tempfile.mkdtemp()
    with summary_ops.create_file_writer(
        logdir, max_queue=0,
        name="t0").as_default(), summary_ops.always_record_summaries():
      m.result()  # As a side-effect will write summaries.

    events = summary_test_util.events_from_logdir(logdir)
    self.assertEqual(len(events), 2)
    self.assertEqual(events[1].summary.value[0].simple_value, 37.0)

  def testSharedName(self):
    logdir = self.get_temp_dir()
    with summary_ops.always_record_summaries():
      # Create with default shared name (should match logdir)
      writer1 = summary_ops.create_file_writer(logdir)
      with writer1.as_default():
        summary_ops.scalar('one', 1.0, step=1)
        summary_ops.flush()
      # Create with explicit logdir shared name (should be same resource/file)
      shared_name = 'logdir:' + logdir
      writer2 = summary_ops.create_file_writer(logdir, name=shared_name)
      with writer2.as_default():
        summary_ops.scalar('two', 2.0, step=2)
        summary_ops.flush()
      # Create with different shared name (should be separate resource/file)
      time.sleep(1.1)  # Ensure filename has a different timestamp
      writer3 = summary_ops.create_file_writer(logdir, name='other')
      with writer3.as_default():
        summary_ops.scalar('three', 3.0, step=3)
        summary_ops.flush()

    event_files = iter(sorted(gfile.Glob(os.path.join(logdir, '*tfevents*'))))

    # First file has tags "one" and "two"
    events = iter(summary_test_util.events_from_file(next(event_files)))
    self.assertEqual('brain.Event:2', next(events).file_version)
    self.assertEqual('one', next(events).summary.value[0].tag)
    self.assertEqual('two', next(events).summary.value[0].tag)
    self.assertRaises(StopIteration, lambda: next(events))

    # Second file has tag "three"
    events = iter(summary_test_util.events_from_file(next(event_files)))
    self.assertEqual('brain.Event:2', next(events).file_version)
    self.assertEqual('three', next(events).summary.value[0].tag)
    self.assertRaises(StopIteration, lambda: next(events))

    # No more files
    self.assertRaises(StopIteration, lambda: next(event_files))

 def testSummaryOps(self):
   training_util.get_or_create_global_step()
   logdir = tempfile.mkdtemp()
   with summary_ops.create_file_writer(
       logdir, max_queue=0,
       name='t0').as_default(), summary_ops.always_record_summaries():
     summary_ops.generic('tensor', 1, '')
     summary_ops.scalar('scalar', 2.0)
     summary_ops.histogram('histogram', [1.0])
     summary_ops.image('image', [[[[1.0]]]])
     summary_ops.audio('audio', [[1.0]], 1.0, 1)
     # The working condition of the ops is tested in the C++ test so we just
     # test here that we're calling them correctly.
     self.assertTrue(gfile.Exists(logdir))

 def testWrite_gpuDeviceContext(self):
   logdir = self.get_temp_dir()
   with context.eager_mode():
     with summary_ops.create_file_writer(logdir).as_default():
       with ops.device('/GPU:0'):
         value = constant_op.constant(42.0)
         step = constant_op.constant(12, dtype=dtypes.int64)
         summary_ops.write('tag', value, step=step).numpy()
   empty_metadata = summary_pb2.SummaryMetadata()
   events = events_from_logdir(logdir)
   self.assertEqual(2, len(events))
   self.assertEqual(12, events[1].step)
   self.assertEqual(42, to_numpy(events[1].summary.value[0]))
   self.assertEqual(empty_metadata, events[1].summary.value[0].metadata)