Preload and precompress static assets in CorePlugin

tensorboard/backend/http_util.py

@@ -119,18 +119,24 @@ def Respond(request,
   content = tf.compat.as_bytes(content, charset)
   if textual and not charset_match and mimetype not in _JSON_MIMETYPES:
     content_type += '; charset=' + charset
-  if (not content_encoding and textual and
-      _ALLOWS_GZIP_PATTERN.search(request.headers.get('Accept-Encoding', ''))):
+  gzip_accepted = _ALLOWS_GZIP_PATTERN.search(
+      request.headers.get('Accept-Encoding', ''))
+  # Automatically gzip uncompressed text data if accepted.
+  if textual and not content_encoding and gzip_accepted:
+    orig_len = len(content)
     out = six.BytesIO()
-    f = gzip.GzipFile(fileobj=out, mode='wb', compresslevel=3)
-    f.write(content)
-    f.close()
+    # Set mtime to zero to make payload for a given input deterministic.
+    with gzip.GzipFile(fileobj=out, mode='wb', compresslevel=3, mtime=0) as f:
+      f.write(content)
     content = out.getvalue()
     content_encoding = 'gzip'
-  if request.method == 'HEAD':
-    content = ''
-  headers = []
+  # Automatically unzip precompressed data if not accepted.
+  if content_encoding == 'gzip' and not gzip_accepted:
+    with gzip.GzipFile(fileobj=six.BytesIO(content), mode='rb') as f:
+      content = f.read()
+    content_encoding = None
 
+  headers = []
   headers.append(('Content-Length', str(len(content))))
   if content_encoding:
     headers.append(('Content-Encoding', content_encoding))
@@ -142,5 +148,8 @@ def Respond(request,
     headers.append(('Expires', '0'))
     headers.append(('Cache-Control', 'no-cache, must-revalidate'))
 
+  if request.method == 'HEAD':
+    content = None
+
   return wrappers.Response(
       response=content, status=code, headers=headers, content_type=content_type)

tensorboard/backend/http_util_test.py

@@ -36,14 +36,16 @@ def testHelloWorld(self):
     r = http_util.Respond(q, '<b>hello world</b>', 'text/html')
     self.assertEqual(r.status_code, 200)
     self.assertEqual(r.response, [six.b('<b>hello world</b>')])
+    self.assertEqual(r.headers.get('Content-Length'), '18')
 
   def testHeadRequest_doesNotWrite(self):
     builder = wtest.EnvironBuilder(method='HEAD')
     env = builder.get_environ()
     request = wrappers.Request(env)
     r = http_util.Respond(request, '<b>hello world</b>', 'text/html')
     self.assertEqual(r.status_code, 200)
-    self.assertEqual(r.response, [six.b('')])
+    self.assertEqual(r.response, [])
+    self.assertEqual(r.headers.get('Content-Length'), '18')
 
   def testPlainText_appendsUtf8ToContentType(self):
     q = wrappers.Request(wtest.EnvironBuilder().get_environ())
@@ -136,6 +138,20 @@ def testAcceptGzip_compressesResponse(self):
     self.assertEqual(
         r.response, [fall_of_hyperion_canto1_stanza1.encode('utf-8')])
 
+  def testAcceptGzip_alreadyCompressed_sendsPrecompressedResponse(self):
+    gzipped_text = _gzip('hello hello hello world')
+    e = wtest.EnvironBuilder(headers={'Accept-Encoding': 'gzip'}).get_environ()
+    q = wrappers.Request(e)
+    r = http_util.Respond(q, gzipped_text, 'text/plain', content_encoding='gzip')
+    self.assertEqual(r.response, [gzipped_text])  # Still singly zipped
+
+  def testPrecompressedResponse_noAcceptGzip_decompressesResponse(self):
+    orig_text = 'hello hello hello world'
+    gzipped_text = _gzip(orig_text)
+    q = wrappers.Request(wtest.EnvironBuilder().get_environ())
+    r = http_util.Respond(q, gzipped_text, 'text/plain', content_encoding='gzip')
+    self.assertEqual(r.response, [orig_text])
+
   def testJson_getsAutoSerialized(self):
     q = wrappers.Request(wtest.EnvironBuilder().get_environ())
     r = http_util.Respond(q, [1, 2, 3], 'application/json')
@@ -147,8 +163,16 @@ def testExpires_setsCruiseControl(self):
     self.assertEqual(r.headers.get('Cache-Control'), 'private, max-age=60')
 
 
+def _gzip(bs):
+  out = six.BytesIO()
+  with gzip.GzipFile(fileobj=out, mode='wb') as f:
+    f.write(bs)
+  return out.getvalue()
+
+
 def _gunzip(bs):
-  return gzip.GzipFile('', 'rb', 9, six.BytesIO(bs)).read()
+  with gzip.GzipFile(fileobj=six.BytesIO(bs), mode='rb') as f:
+    return f.read()
 
 
 if __name__ == '__main__':

tensorboard/plugins/core/core_plugin.py

@@ -19,9 +19,11 @@
 from __future__ import print_function
 
 import functools
+import gzip
 import mimetypes
 import zipfile
 
+import six
 import tensorflow as tf
 from werkzeug import utils
 from werkzeug import wrappers
@@ -65,12 +67,13 @@ def get_plugin_apps(self):
         '/images': self._redirect_to_index,
     }
     if self._assets_zip_provider:
-      apps['/'] = functools.partial(self._serve_asset, 'index.html')
       with self._assets_zip_provider() as fp:
         with zipfile.ZipFile(fp) as zip_:
-          for info in zip_.infolist():
-            path = info.filename
-            apps['/' + path] = functools.partial(self._serve_asset, path)
+          for path in zip_.namelist():
+            gzipped_asset_bytes = _gzip(zip_.read(path))
+            apps['/' + path] = functools.partial(
+                self._serve_asset, path, gzipped_asset_bytes)
+      apps['/'] = apps['/index.html']
     return apps
 
   @wrappers.Request.application
@@ -82,14 +85,11 @@ def _redirect_to_index(self, unused_request):
     return utils.redirect('/')
 
   @wrappers.Request.application
-  def _serve_asset(self, path, request):
-    """Serves a static asset from the zip file."""
+  def _serve_asset(self, path, gzipped_asset_bytes, request):
+    """Serves a pre-gzipped static asset from the zip file."""
     mimetype = mimetypes.guess_type(path)[0] or 'application/octet-stream'
-    with self._assets_zip_provider() as fp:
-      with zipfile.ZipFile(fp) as zip_:
-        with zip_.open(path) as file_:
-          html = file_.read()
-    return http_util.Respond(request, html, mimetype)
+    return http_util.Respond(
+        request, gzipped_asset_bytes, mimetype, content_encoding='gzip')
 
   @wrappers.Request.application
   def _serve_logdir(self, request):
@@ -129,3 +129,11 @@ def get_first_event_timestamp(run_name):
     }
     run_names.sort(key=first_event_timestamps.get)
     return http_util.Respond(request, run_names, 'application/json')
+
+
+def _gzip(bytestring):
+  out = six.BytesIO()
+  # Set mtime to zero for deterministic results across TensorBoard launches.
+  with gzip.GzipFile(fileobj=out, mode='wb', compresslevel=3, mtime=0) as f:
+    f.write(bytestring)
+  return out.getvalue()