From 59daeca31528ad20932921e687d2a3ae042991a1 Mon Sep 17 00:00:00 2001 From: Dan Smilkov Date: Fri, 4 Nov 2016 06:34:23 -0800 Subject: [PATCH] Fix performance issues with the Embedding Projector plugin. - Setup data handlers in the class c-tor instead of on every GET request. - Use only 1 instance of the plugin, instead of creating a new instance on every GET Request. - Memorize readers, and other data, and share it across the multiple threads of TB server. Previously # of checkpoint readers = # of GET requests, ups. - Checkpoints can be big (order of 1G). Reading it takes a while. Only start reading checkpoints if someone explicitly opens the embeddings tab. Previously the checkpoint was read on TB startup, which slows down TensorBoard. Change: 138187941 --- tensorflow/tensorboard/backend/handler.py | 75 +++++----- tensorflow/tensorboard/plugins/__init__.py | 2 +- tensorflow/tensorboard/plugins/base_plugin.py | 9 -- .../tensorboard/plugins/projector/plugin.py | 133 ++++++++++-------- 4 files changed, 111 insertions(+), 108 deletions(-) diff --git a/tensorflow/tensorboard/backend/handler.py b/tensorflow/tensorboard/backend/handler.py index 4201d1a26a8..e016684665c 100644 --- a/tensorflow/tensorboard/backend/handler.py +++ b/tensorflow/tensorboard/backend/handler.py @@ -24,13 +24,11 @@ from __future__ import division from __future__ import print_function import csv +import functools import imghdr -import json import mimetypes import os -import six -import six as _six from six import StringIO from six.moves import BaseHTTPServer from six.moves import urllib @@ -101,8 +99,41 @@ class TensorboardHandler(BaseHTTPServer.BaseHTTPRequestHandler): def __init__(self, multiplexer, logdir, *args): self._multiplexer = multiplexer self._logdir = logdir + self._setup_data_handlers() BaseHTTPServer.BaseHTTPRequestHandler.__init__(self, *args) + def _setup_data_handlers(self): + self.data_handlers = { + DATA_PREFIX + LOGDIR_ROUTE: self._serve_logdir, + DATA_PREFIX + SCALARS_ROUTE: self._serve_scalars, + DATA_PREFIX + GRAPH_ROUTE: self._serve_graph, + DATA_PREFIX + RUN_METADATA_ROUTE: self._serve_run_metadata, + DATA_PREFIX + HISTOGRAMS_ROUTE: self._serve_histograms, + DATA_PREFIX + COMPRESSED_HISTOGRAMS_ROUTE: + self._serve_compressed_histograms, + DATA_PREFIX + IMAGES_ROUTE: self._serve_images, + DATA_PREFIX + INDIVIDUAL_IMAGE_ROUTE: self._serve_image, + DATA_PREFIX + AUDIO_ROUTE: self._serve_audio, + DATA_PREFIX + INDIVIDUAL_AUDIO_ROUTE: self._serve_individual_audio, + DATA_PREFIX + RUNS_ROUTE: self._serve_runs, + '/app.js': self._serve_js + } + + # Serve the routes from the registered plugins using their name as the route + # prefix. For example if plugin z has two routes /a and /b, they will be + # served as /data/plugin/z/a and /data/plugin/z/b. + for name in REGISTERED_PLUGINS: + try: + plugin = REGISTERED_PLUGINS[name] + plugin_handlers = plugin.get_plugin_handlers( + self._multiplexer.RunPaths(), self._logdir) + except Exception as e: # pylint: disable=broad-except + logging.warning('Plugin %s failed. Exception: %s', name, str(e)) + continue + for route, handler in plugin_handlers.items(): + path = DATA_PREFIX + PLUGIN_PREFIX + '/' + name + route + self.data_handlers[path] = functools.partial(handler, self) + def respond(self, *args, **kwargs): """Delegates to http.Respond.""" http.Respond(self, *args, **kwargs) @@ -436,6 +467,7 @@ class TensorboardHandler(BaseHTTPServer.BaseHTTPRequestHandler): def _serve_static_file(self, path): """Serves the static file located at the given path. + Args: path: The path of the static file, relative to the tensorboard/ directory. """ @@ -483,39 +515,6 @@ class TensorboardHandler(BaseHTTPServer.BaseHTTPRequestHandler): if clean_path.endswith('/'): clean_path = clean_path[:-1] - data_handlers = { - DATA_PREFIX + LOGDIR_ROUTE: self._serve_logdir, - DATA_PREFIX + SCALARS_ROUTE: self._serve_scalars, - DATA_PREFIX + GRAPH_ROUTE: self._serve_graph, - DATA_PREFIX + RUN_METADATA_ROUTE: self._serve_run_metadata, - DATA_PREFIX + HISTOGRAMS_ROUTE: self._serve_histograms, - DATA_PREFIX + COMPRESSED_HISTOGRAMS_ROUTE: - self._serve_compressed_histograms, - DATA_PREFIX + IMAGES_ROUTE: self._serve_images, - DATA_PREFIX + INDIVIDUAL_IMAGE_ROUTE: self._serve_image, - DATA_PREFIX + AUDIO_ROUTE: self._serve_audio, - DATA_PREFIX + INDIVIDUAL_AUDIO_ROUTE: self._serve_individual_audio, - DATA_PREFIX + RUNS_ROUTE: self._serve_runs, - '/app.js': self._serve_js - } - - # Serve the routes from the registered plugins using their name as the route - # prefix. For example if plugin z has two routes /a and /b, they will be - # served as /data/plugin/z/a and /data/plugin/z/b. - for name in REGISTERED_PLUGINS: - try: - plugin = REGISTERED_PLUGINS[name]() - # Initialize the plugin by passing the main http handler. - plugin.initialize(self) - plugin_handlers = plugin.get_plugin_handlers( - self._multiplexer.RunPaths(), self._logdir) - except Exception as e: # pylint: disable=broad-except - logging.warning('Plugin %s failed. Exception: %s', name, str(e)) - continue - for route, handler in six.iteritems(plugin_handlers): - path = DATA_PREFIX + PLUGIN_PREFIX + '/' + name + route - data_handlers[path] = handler - query_params = urlparse.parse_qs(parsed_url.query) # parse_qs returns a list of values for each key; we're only interested in # the first. @@ -528,8 +527,8 @@ class TensorboardHandler(BaseHTTPServer.BaseHTTPRequestHandler): return query_params[key] = query_params[key][0] - if clean_path in data_handlers: - data_handlers[clean_path](query_params) + if clean_path in self.data_handlers: + self.data_handlers[clean_path](query_params) elif clean_path in TAB_ROUTES: self._serve_index(query_params) else: diff --git a/tensorflow/tensorboard/plugins/__init__.py b/tensorflow/tensorboard/plugins/__init__.py index 819baa25f08..defcd304fa4 100644 --- a/tensorflow/tensorboard/plugins/__init__.py +++ b/tensorflow/tensorboard/plugins/__init__.py @@ -19,4 +19,4 @@ from __future__ import print_function from tensorflow.tensorboard.plugins.projector.plugin import ProjectorPlugin # Map of registered plugins in TensorBoard. -REGISTERED_PLUGINS = {'projector': ProjectorPlugin} +REGISTERED_PLUGINS = {'projector': ProjectorPlugin()} diff --git a/tensorflow/tensorboard/plugins/base_plugin.py b/tensorflow/tensorboard/plugins/base_plugin.py index 0b9199b4daf..efb445a6ab3 100644 --- a/tensorflow/tensorboard/plugins/base_plugin.py +++ b/tensorflow/tensorboard/plugins/base_plugin.py @@ -30,15 +30,6 @@ class TBPlugin(object): """TensorBoard plugin interface. Every plugin must extend from this class.""" __metaclass__ = ABCMeta - def initialize(self, handler): - """Initializes the plugin. - - Args: - handler: The tensorboard http handler that has methods that are used - by plugins such as serving json or gzip response. - """ - self.handler = handler - @abstractmethod def get_plugin_handlers(self, run_paths, logdir): """Returns a set of http handlers that the plugin implements. diff --git a/tensorflow/tensorboard/plugins/projector/plugin.py b/tensorflow/tensorboard/plugins/projector/plugin.py index 4479edbd5c2..c293d813c7a 100644 --- a/tensorflow/tensorboard/plugins/projector/plugin.py +++ b/tensorflow/tensorboard/plugins/projector/plugin.py @@ -64,13 +64,18 @@ def _read_tensor_file(fpath): class ProjectorPlugin(TBPlugin): """Embedding projector.""" - def get_plugin_handlers(self, run_paths, logdir): - self.configs, self.config_fpaths = self._read_config_files(run_paths, - logdir) + def __init__(self): + self._handlers = None self.readers = {} - self._augment_configs_with_checkpoint_info() + self.run_paths = None + self.logdir = None + self._configs = None + self.old_num_run_paths = None - return { + def get_plugin_handlers(self, run_paths, logdir): + self.run_paths = run_paths + self.logdir = logdir + self._handlers = { RUNS_ROUTE: self._serve_runs, CONFIG_ROUTE: self._serve_config, TENSOR_ROUTE: self._serve_tensor, @@ -78,9 +83,26 @@ class ProjectorPlugin(TBPlugin): BOOKMARKS_ROUTE: self._serve_bookmarks, SPRITE_IMAGE_ROUTE: self._serve_sprite_image } + return self._handlers + + @property + def configs(self): + """Returns a map of run paths to `ProjectorConfig` protos.""" + if self._run_paths_changed(): + self._configs, self.config_fpaths = self._read_config_files( + self.run_paths, self.logdir) + self._augment_configs_with_checkpoint_info() + return self._configs + + def _run_paths_changed(self): + num_run_paths = len(list(self.run_paths.keys())) + if num_run_paths != self.old_num_run_paths: + self.old_num_run_paths = num_run_paths + return True + return False def _augment_configs_with_checkpoint_info(self): - for run, config in self.configs.items(): + for run, config in self._configs.items(): # Find the size of the embeddings that are associated with a tensor file. for embedding in config.embeddings: if embedding.tensor_path and not embedding.tensor_shape: @@ -111,18 +133,18 @@ class ProjectorPlugin(TBPlugin): # Remove configs that do not have any valid (2D) tensors. runs_to_remove = [] - for run, config in self.configs.items(): + for run, config in self._configs.items(): if not config.embeddings: runs_to_remove.append(run) for run in runs_to_remove: - del self.configs[run] + del self._configs[run] del self.config_fpaths[run] - def _read_config_files(self, run_paths, logdir): + def _read_config_files(self, run_paths, summary_logdir): # If there are no summary event files, the projector can still work, # thus treating the `logdir` as the model checkpoint directory. if not run_paths: - run_paths['.'] = logdir + run_paths['.'] = summary_logdir configs = {} config_fpaths = {} @@ -164,7 +186,7 @@ class ProjectorPlugin(TBPlugin): if run in self.readers: return self.readers[run] - config = self.configs[run] + config = self._configs[run] reader = None if config.model_checkpoint_path: try: @@ -201,48 +223,45 @@ class ProjectorPlugin(TBPlugin): return info return None - def _serve_runs(self, query_params): + def _serve_runs(self, request, query_params): """Returns a list of runs that have embeddings.""" - self.handler.respond(list(self.configs.keys()), 'application/json') + request.respond(list(self.configs.keys()), 'application/json') - def _serve_config(self, query_params): + def _serve_config(self, request, query_params): run = query_params.get('run') if run is None: - self.handler.respond('query parameter "run" is required', - 'text/plain', 400) + request.respond('query parameter "run" is required', 'text/plain', 400) return if run not in self.configs: - self.handler.respond('Unknown run: %s' % run, 'text/plain', 400) + request.respond('Unknown run: %s' % run, 'text/plain', 400) return config = self.configs[run] - self.handler.respond(json_format.MessageToJson(config), 'application/json') + request.respond(json_format.MessageToJson(config), 'application/json') - def _serve_metadata(self, query_params): + def _serve_metadata(self, request, query_params): run = query_params.get('run') if run is None: - self.handler.respond('query parameter "run" is required', - 'text/plain', 400) + request.respond('query parameter "run" is required', 'text/plain', 400) return name = query_params.get('name') if name is None: - self.handler.respond('query parameter "name" is required', - 'text/plain', 400) + request.respond('query parameter "name" is required', 'text/plain', 400) return if run not in self.configs: - self.handler.respond('Unknown run: %s' % run, 'text/plain', 400) + request.respond('Unknown run: %s' % run, 'text/plain', 400) return config = self.configs[run] fpath = self._get_metadata_file_for_tensor(name, config) if not fpath: - self.handler.respond( + request.respond( 'No metadata file found for tensor %s in the config file %s' % (name, self.config_fpaths[run]), 'text/plain', 400) return if not file_io.file_exists(fpath) or file_io.is_directory(fpath): - self.handler.respond('%s is not a file' % fpath, 'text/plain', 400) + request.respond('%s is not a file' % fpath, 'text/plain', 400) return num_header_rows = 0 @@ -256,23 +275,21 @@ class ProjectorPlugin(TBPlugin): num_header_rows = 1 if len(lines) >= LIMIT_NUM_POINTS + num_header_rows: break - self.handler.respond(''.join(lines), 'text/plain') + request.respond(''.join(lines), 'text/plain') - def _serve_tensor(self, query_params): + def _serve_tensor(self, request, query_params): run = query_params.get('run') if run is None: - self.handler.respond('query parameter "run" is required', - 'text/plain', 400) + request.respond('query parameter "run" is required', 'text/plain', 400) return name = query_params.get('name') if name is None: - self.handler.respond('query parameter "name" is required', - 'text/plain', 400) + request.respond('query parameter "name" is required', 'text/plain', 400) return if run not in self.configs: - self.handler.respond('Unknown run: %s' % run, 'text/plain', 400) + request.respond('Unknown run: %s' % run, 'text/plain', 400) return reader = self._get_reader_for_run(run) @@ -282,19 +299,19 @@ class ProjectorPlugin(TBPlugin): # See if there is a tensor file in the config. embedding = self._get_embedding(name, config) if not embedding or not embedding.tensor_path: - self.handler.respond('Tensor %s has no tensor_path in the config' % - name, 'text/plain', 400) + request.respond('Tensor %s has no tensor_path in the config' % + name, 'text/plain', 400) return if not file_io.file_exists(embedding.tensor_path): - self.handler.respond('Tensor file %s does not exist' % - embedding.tensor_path, 'text/plain', 400) + request.respond('Tensor file %s does not exist' % + embedding.tensor_path, 'text/plain', 400) return tensor = _read_tensor_file(embedding.tensor_path) else: if not reader.has_tensor(name): - self.handler.respond('Tensor %s not found in checkpoint dir %s' % - (name, config.model_checkpoint_path), - 'text/plain', 400) + request.respond('Tensor %s not found in checkpoint dir %s' % + (name, config.model_checkpoint_path), + 'text/plain', 400) return tensor = reader.get_tensor(name) @@ -302,75 +319,71 @@ class ProjectorPlugin(TBPlugin): tensor = tensor[:LIMIT_NUM_POINTS] # Stream it as TSV. tsv = '\n'.join(['\t'.join([str(val) for val in row]) for row in tensor]) - self.handler.respond(tsv, 'text/tab-separated-values') + request.respond(tsv, 'text/tab-separated-values') - def _serve_bookmarks(self, query_params): + def _serve_bookmarks(self, request, query_params): run = query_params.get('run') if not run: - self.handler.respond('query parameter "run" is required', 'text/plain', - 400) + request.respond('query parameter "run" is required', 'text/plain', 400) return name = query_params.get('name') if name is None: - self.handler.respond('query parameter "name" is required', 'text/plain', - 400) + request.respond('query parameter "name" is required', 'text/plain', 400) return if run not in self.configs: - self.handler.respond('Unknown run: %s' % run, 'text/plain', 400) + request.respond('Unknown run: %s' % run, 'text/plain', 400) return config = self.configs[run] fpath = self._get_bookmarks_file_for_tensor(name, config) if not fpath: - self.handler.respond( + request.respond( 'No bookmarks file found for tensor %s in the config file %s' % (name, self.config_fpaths[run]), 'text/plain', 400) return if not file_io.file_exists(fpath) or file_io.is_directory(fpath): - self.handler.respond('%s is not a file' % fpath, 'text/plain', 400) + request.respond('%s is not a file' % fpath, 'text/plain', 400) return bookmarks_json = None with file_io.FileIO(fpath, 'r') as f: bookmarks_json = f.read() - self.handler.respond(bookmarks_json, 'application/json') + request.respond(bookmarks_json, 'application/json') - def _serve_sprite_image(self, query_params): + def _serve_sprite_image(self, request, query_params): run = query_params.get('run') if not run: - self.handler.respond('query parameter "run" is required', 'text/plain', - 400) + request.respond('query parameter "run" is required', 'text/plain', 400) return name = query_params.get('name') if name is None: - self.handler.respond('query parameter "name" is required', 'text/plain', - 400) + request.respond('query parameter "name" is required', 'text/plain', 400) return if run not in self.configs: - self.handler.respond('Unknown run: %s' % run, 'text/plain', 400) + request.respond('Unknown run: %s' % run, 'text/plain', 400) return config = self.configs[run] embedding_info = self._get_embedding(name, config) if not embedding_info or not embedding_info.sprite.image_path: - self.handler.respond( + request.respond( 'No sprite image file found for tensor %s in the config file %s' % (name, self.config_fpaths[run]), 'text/plain', 400) return fpath = embedding_info.sprite.image_path if not file_io.file_exists(fpath) or file_io.is_directory(fpath): - self.handler.respond('%s does not exist or is directory' % fpath, - 'text/plain', 400) + request.respond( + '%s does not exist or is directory' % fpath, 'text/plain', 400) return f = file_io.FileIO(fpath, 'r') encoded_image_string = f.read() f.close() image_type = imghdr.what(None, encoded_image_string) mime_type = _IMGHDR_TO_MIMETYPE.get(image_type, _DEFAULT_IMAGE_MIMETYPE) - self.handler.respond(encoded_image_string, mime_type) + request.respond(encoded_image_string, mime_type)