Automated rollback of commit 568727eed199dba04e37f500265b50f96fed455e

PiperOrigin-RevId: 205875586

3 files changed, 30 insertions, 103 deletions

diff --git a/tensorflow/python/summary/writer/event_file_writer_v2.py b/tensorflow/python/summary/writer/event_file_writer_v2.py
index 262182d3b8..5c66c0f7a8 100644
--- a/tensorflow/python/summary/writer/event_file_writer_v2.py
+++ b/tensorflow/python/summary/writer/event_file_writer_v2.py
@@ -18,7 +18,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.python.client import session as tf_session
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
@@ -44,11 +43,11 @@ class EventFileWriterV2(object):
     """Creates an `EventFileWriterV2` and an event file to write to.
 
     On construction, this calls `tf.contrib.summary.create_file_writer` within
-    the default graph, which finds and returns a shared summary writer resource
-    for `logdir` if one exists, and creates one if not. Creating the summary
+    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 the default graph during this init call;
+    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
@@ -62,50 +61,38 @@ class EventFileWriterV2(object):
     no effect.  See `tf.contrib.summary.create_file_writer` for details.
 
     Args:
-      session: A `tf.Session`, or a callable that provides one which will be
-        called on-demand. The session will hold the shared writer resource.
+      session: A `tf.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`.
-
-    Raises:
-      ValueError: if `session` is not a `tf.Session` or a callable
     """
-    if isinstance(session, tf_session.SessionInterface):
-      self._session = lambda: session
-    elif callable(session):
-      self._session = session
-    else:
-      raise ValueError('session must be tf.Session or callable')
+    self._session = session
     self._logdir = logdir
-    self._initialized = False
     self._closed = False
     if not gfile.IsDirectory(self._logdir):
       gfile.MakeDirs(self._logdir)
 
-    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()
-
-  def _init_if_needed(self):
-    if not self._initialized:
-      self._session().run(self._init_op)
-      self._initialized = True
+    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 get_logdir(self):
     """Returns the directory where event file will be written."""
@@ -121,6 +108,7 @@ class EventFileWriterV2(object):
     """
     if self._closed:
       self._closed = False
+      self._session.run(self._init_op)
 
   def add_event(self, event):
     """Adds an event to the event file.
@@ -129,9 +117,8 @@ class EventFileWriterV2(object):
       event: An `Event` protocol buffer.
     """
     if not self._closed:
-      self._init_if_needed()
       event_pb = event.SerializeToString()
-      self._session().run(
+      self._session.run(
           self._add_event_op, feed_dict={self._event_placeholder: event_pb})
 
   def flush(self):
@@ -140,9 +127,7 @@ class EventFileWriterV2(object):
     Call this method to make sure that all pending events have been written to
     disk.
     """
-    if not self._closed:
-      self._init_if_needed()
-      self._session().run(self._flush_op)
+    self._session.run(self._flush_op)
 
   def close(self):
     """Flushes the event file to disk and close the file.
@@ -150,8 +135,6 @@ class EventFileWriterV2(object):
     Call this method when you do not need the summary writer anymore.
     """
     if not self._closed:
-      self._init_if_needed()
       self.flush()
-      self._session().run(self._close_op)
+      self._session.run(self._close_op)
       self._closed = True
-      self._initialized = False
diff --git a/tensorflow/python/summary/writer/writer.py b/tensorflow/python/summary/writer/writer.py
index 2a967ae3a5..aca084fc91 100644
--- a/tensorflow/python/summary/writer/writer.py
+++ b/tensorflow/python/summary/writer/writer.py
@@ -332,11 +332,8 @@ class FileWriter(SummaryToEventTransformer):
     the same shared resource name (which by default scoped to the logdir). If
     no such resource exists, one will be created using the remaining arguments
     to this constructor, but if one already exists those arguments are ignored.
-    In either case, ops will be added to the default graph to control the
+    In either case, ops will be added to `session.graph` to control the
     underlying file writer resource. See `tf.contrib.summary` for more details.
-    Instead of an actual `tf.Session`, this argument may also be a callable that
-    provides a `tf.Session` when invoked (e.g. `tf.get_default_session`), which
-    will be called on-demand when a session is needed.
 
     Args:
       logdir: A string. Directory where event file will be written.
@@ -347,8 +344,7 @@ class FileWriter(SummaryToEventTransformer):
       graph_def: DEPRECATED: Use the `graph` argument instead.
       filename_suffix: A string. Every event file's name is suffixed with
         `suffix`.
-      session: A `tf.Session` object or a callable that provides `tf.Session`
-        objects. See details above.
+      session: A `tf.Session` object. See details above.
 
     Raises:
       RuntimeError: If called with eager execution enabled.
diff --git a/tensorflow/python/summary/writer/writer_test.py b/tensorflow/python/summary/writer/writer_test.py
index 3380dea317..dc990c2602 100644
--- a/tensorflow/python/summary/writer/writer_test.py
+++ b/tensorflow/python/summary/writer/writer_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for writer.py."""
+"""Tests for training_coordinator.py."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -574,58 +574,6 @@ class SessionBasedFileWriterTestCase(FileWriterTestCase):
     # No more files
     self.assertRaises(StopIteration, lambda: next(event_paths))
 
-  def testSesssionArgument_callableProvider(self):
-    logdir = self.get_temp_dir()
-    setup_writer = summary_ops_v2.create_file_writer(logdir=logdir)
-    with summary_ops_v2.always_record_summaries(), setup_writer.as_default():
-      summary1 = summary_ops_v2.scalar("one", 0.0, step=0)
-      summary2 = summary_ops_v2.scalar("two", 0.0, step=0)
-    sess1 = session.Session()
-    sess1.run(setup_writer.init())
-    sess1.run(summary1)
-    sess1.run(setup_writer.flush())
-    time.sleep(1.1)  # Ensure filename has a different timestamp
-    sess2 = session.Session()
-    sess2.run(setup_writer.init())
-    sess2.run(summary2)
-    sess2.run(setup_writer.flush())
-
-    # Using get_default_session as session provider should make this FileWriter
-    # send its summaries to the current default session's shared summary writer
-    # resource (initializing it as needed).
-    test_writer = writer.FileWriter(
-        session=ops.get_default_session, logdir=logdir)
-    with sess1.as_default():
-      test_writer.add_summary(self._createTaggedSummary("won"), 1)
-      test_writer.flush()
-    with sess2.as_default():
-      test_writer.add_summary(self._createTaggedSummary("too"), 1)
-      test_writer.flush()
-
-    event_paths = iter(sorted(glob.glob(os.path.join(logdir, "event*"))))
-
-    # First file should have tags "one", "won"
-    events = summary_iterator.summary_iterator(next(event_paths))
-    self.assertEqual("brain.Event:2", next(events).file_version)
-    self.assertEqual("one", next(events).summary.value[0].tag)
-    self.assertEqual("won", next(events).summary.value[0].tag)
-    self.assertRaises(StopIteration, lambda: next(events))
-
-    # Second file should have tags "two", "too"
-    events = summary_iterator.summary_iterator(next(event_paths))
-    self.assertEqual("brain.Event:2", next(events).file_version)
-    self.assertEqual("two", next(events).summary.value[0].tag)
-    self.assertEqual("too", next(events).summary.value[0].tag)
-    self.assertRaises(StopIteration, lambda: next(events))
-
-    # No more files
-    self.assertRaises(StopIteration, lambda: next(event_paths))
-
-  def testSessionArgument_notSessionOrCallable(self):
-    logdir = self.get_temp_dir()
-    self.assertRaises(
-        ValueError, lambda: writer.FileWriter(session=[], logdir=logdir))
-
 
 class FileWriterCacheTest(test.TestCase):
   """FileWriterCache tests."""