Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix segfault from passing invalid input to tf.summary.flush()
This removes an obsolete legacy codepath from tf.summary.flush(). Because that codepath was still used by tpu_estimator.py, for now, we preserve a copy of the legacy codepath and update tpu_estimator.py to use it.

PiperOrigin-RevId: 441606843
  • Loading branch information
nfelt authored and tensorflower-gardener committed Apr 13, 2022
1 parent 59419fd commit dbdd98c
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 4 deletions.
10 changes: 6 additions & 4 deletions tensorflow/python/kernel_tests/summary_ops/summary_ops_test.py
Expand Up @@ -985,10 +985,12 @@ def testFlushFunction(self):
self.assertEqual(3, get_total())
summary_ops.flush(writer=writer)
self.assertEqual(4, get_total())
summary_ops.write('tag', 1, step=0)
self.assertEqual(4, get_total())
summary_ops.flush(writer=writer._resource) # pylint:disable=protected-access
self.assertEqual(5, get_total())

# Regression test for b/228097117.
def testFlushFunction_disallowsInvalidWriterInput(self):
with context.eager_mode():
with self.assertRaisesRegex(ValueError, 'Invalid argument to flush'):
summary_ops.flush(writer=())

@test_util.assert_no_new_tensors
def testNoMemoryLeak_graphMode(self):
Expand Down
23 changes: 23 additions & 0 deletions tensorflow/python/ops/summary_ops_v2.py
Expand Up @@ -1111,12 +1111,35 @@ def flush(writer=None, name=None):
Returns:
The created `tf.Operation`.
"""
del name # unused
if writer is None:
writer = _summary_state.writer
if writer is None:
return control_flow_ops.no_op()
if isinstance(writer, SummaryWriter):
return writer.flush()
raise ValueError("Invalid argument to flush(): %r" % (writer,))


def legacy_raw_flush(writer=None, name=None):
"""Legacy version of flush() that accepts a raw resource tensor for `writer`.
Do not use this function in any new code. Not supported and not part of the
public TF APIs.
Args:
writer: The `tf.summary.SummaryWriter` to flush. If None, the current
default writer will be used instead; if there is no current writer, this
returns `tf.no_op`. For this legacy version only, also accepts a raw
resource tensor pointing to the underlying C++ writer resource.
name: Ignored legacy argument for a name for the operation.
Returns:
The created `tf.Operation`.
"""
if writer is None or isinstance(writer, SummaryWriter):
# Forward to the TF2 implementation of flush() when possible.
return flush(writer, name)
else:
# Legacy fallback in case we were passed a raw resource tensor.
with ops.device("cpu:0"):
Expand Down

0 comments on commit dbdd98c

Please sign in to comment.