# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Uploads a TensorBoard logdir to TensorBoard.dev.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import contextlib import functools import time import grpc import six from google.protobuf import message from tensorboard.compat.proto import graph_pb2 from tensorboard.compat.proto import summary_pb2 from tensorboard.compat.proto import types_pb2 from tensorboard.uploader.proto import server_info_pb2 from tensorboard.uploader.proto import experiment_pb2 from tensorboard.uploader.proto import write_service_pb2 from tensorboard.uploader import logdir_loader from tensorboard.uploader import util from tensorboard.backend import process_graph from tensorboard.backend.event_processing import directory_loader from tensorboard.backend.event_processing import event_file_loader from tensorboard.backend.event_processing import io_wrapper from tensorboard.plugins.graph import metadata as graphs_metadata from tensorboard.plugins.scalar import metadata as scalar_metadata from tensorboard.util import grpc_util from tensorboard.util import tb_logging from tensorboard.util import tensor_util # Minimum length of a logdir polling cycle in seconds. Shorter cycles will # sleep to avoid spinning over the logdir, which isn't great for disks and can # be expensive for network file systems. _MIN_LOGDIR_POLL_INTERVAL_SECS = 5 # Age in seconds of last write after which an event file is considered inactive. # TODO(@nfelt): consolidate with TensorBoard --reload_multifile default logic. _EVENT_FILE_INACTIVE_SECS = 4000 # Maximum length of a base-128 varint as used to encode a 64-bit value # (without the "msb of last byte is bit 63" optimization, to be # compatible with protobuf and golang varints). _MAX_VARINT64_LENGTH_BYTES = 10 logger = tb_logging.get_logger() class TensorBoardUploader(object): """Uploads a TensorBoard logdir to TensorBoard.dev.""" def __init__( self, writer_client, logdir, allowed_plugins, upload_limits, logdir_poll_rate_limiter=None, rpc_rate_limiter=None, tensor_rpc_rate_limiter=None, blob_rpc_rate_limiter=None, name=None, description=None, ): """Constructs a TensorBoardUploader. Args: writer_client: a TensorBoardWriterService stub instance logdir: path of the log directory to upload allowed_plugins: collection of string plugin names; events will only be uploaded if their time series's metadata specifies one of these plugin names upload_limits: instance of tensorboard.service.UploadLimits proto. logdir_poll_rate_limiter: a `RateLimiter` to use to limit logdir polling frequency, to avoid thrashing disks, especially on networked file systems rpc_rate_limiter: a `RateLimiter` to use to limit write RPC frequency. Note this limit applies at the level of single RPCs in the Scalar and Tensor case, but at the level of an entire blob upload in the Blob case-- which may require a few preparatory RPCs and a stream of chunks. Note the chunk stream is internally rate-limited by backpressure from the server, so it is not a concern that we do not explicitly rate-limit within the stream here. name: String name to assign to the experiment. description: String description to assign to the experiment. """ self._api = writer_client self._logdir = logdir self._allowed_plugins = frozenset(allowed_plugins) self._upload_limits = upload_limits self._name = name self._description = description self._request_sender = None if logdir_poll_rate_limiter is None: self._logdir_poll_rate_limiter = util.RateLimiter( _MIN_LOGDIR_POLL_INTERVAL_SECS ) else: self._logdir_poll_rate_limiter = logdir_poll_rate_limiter if rpc_rate_limiter is None: self._rpc_rate_limiter = util.RateLimiter( self._upload_limits.min_scalar_request_interval / 1000 ) else: self._rpc_rate_limiter = rpc_rate_limiter if tensor_rpc_rate_limiter is None: self._tensor_rpc_rate_limiter = util.RateLimiter( self._upload_limits.min_tensor_request_interval / 1000 ) else: self._tensor_rpc_rate_limiter = tensor_rpc_rate_limiter if blob_rpc_rate_limiter is None: self._blob_rpc_rate_limiter = util.RateLimiter( self._upload_limits.min_blob_request_interval / 1000 ) else: self._blob_rpc_rate_limiter = blob_rpc_rate_limiter active_filter = ( lambda secs: secs + _EVENT_FILE_INACTIVE_SECS >= time.time() ) directory_loader_factory = functools.partial( directory_loader.DirectoryLoader, loader_factory=event_file_loader.TimestampedEventFileLoader, path_filter=io_wrapper.IsTensorFlowEventsFile, active_filter=active_filter, ) self._logdir_loader = logdir_loader.LogdirLoader( self._logdir, directory_loader_factory ) def create_experiment(self): """Creates an Experiment for this upload session and returns the ID.""" logger.info("Creating experiment") request = write_service_pb2.CreateExperimentRequest( name=self._name, description=self._description ) response = grpc_util.call_with_retries( self._api.CreateExperiment, request ) self._request_sender = _BatchedRequestSender( response.experiment_id, self._api, allowed_plugins=self._allowed_plugins, upload_limits=self._upload_limits, rpc_rate_limiter=self._rpc_rate_limiter, tensor_rpc_rate_limiter=self._tensor_rpc_rate_limiter, blob_rpc_rate_limiter=self._blob_rpc_rate_limiter, ) return response.experiment_id def start_uploading(self): """Blocks forever to continuously upload data from the logdir. Raises: RuntimeError: If `create_experiment` has not yet been called. ExperimentNotFoundError: If the experiment is deleted during the course of the upload. """ if self._request_sender is None: raise RuntimeError( "Must call create_experiment() before start_uploading()" ) while True: self._logdir_poll_rate_limiter.tick() self._upload_once() def _upload_once(self): """Runs one upload cycle, sending zero or more RPCs.""" logger.info("Starting an upload cycle") sync_start_time = time.time() self._logdir_loader.synchronize_runs() sync_duration_secs = time.time() - sync_start_time logger.info("Logdir sync took %.3f seconds", sync_duration_secs) run_to_events = self._logdir_loader.get_run_events() self._request_sender.send_requests(run_to_events) def update_experiment_metadata( writer_client, experiment_id, name=None, description=None ): """Modifies user data associated with an experiment. Args: writer_client: a TensorBoardWriterService stub instance experiment_id: string ID of the experiment to modify name: If provided, modifies name of experiment to this value. description: If provided, modifies the description of the experiment to this value Raises: ExperimentNotFoundError: If no such experiment exists. PermissionDeniedError: If the user is not authorized to modify this experiment. InvalidArgumentError: If the server rejected the name or description, if, for instance, the size limits have changed on the server. """ logger.info("Modifying experiment %r", experiment_id) request = write_service_pb2.UpdateExperimentRequest() request.experiment.experiment_id = experiment_id if name is not None: logger.info("Setting exp %r name to %r", experiment_id, name) request.experiment.name = name request.experiment_mask.name = True if description is not None: logger.info( "Setting exp %r description to %r", experiment_id, description ) request.experiment.description = description request.experiment_mask.description = True try: grpc_util.call_with_retries(writer_client.UpdateExperiment, request) except grpc.RpcError as e: if e.code() == grpc.StatusCode.NOT_FOUND: raise ExperimentNotFoundError() if e.code() == grpc.StatusCode.PERMISSION_DENIED: raise PermissionDeniedError() if e.code() == grpc.StatusCode.INVALID_ARGUMENT: raise InvalidArgumentError(e.details()) raise def delete_experiment(writer_client, experiment_id): """Permanently deletes an experiment and all of its contents. Args: writer_client: a TensorBoardWriterService stub instance experiment_id: string ID of the experiment to delete Raises: ExperimentNotFoundError: If no such experiment exists. PermissionDeniedError: If the user is not authorized to delete this experiment. RuntimeError: On unexpected failure. """ logger.info("Deleting experiment %r", experiment_id) request = write_service_pb2.DeleteExperimentRequest() request.experiment_id = experiment_id try: grpc_util.call_with_retries(writer_client.DeleteExperiment, request) except grpc.RpcError as e: if e.code() == grpc.StatusCode.NOT_FOUND: raise ExperimentNotFoundError() if e.code() == grpc.StatusCode.PERMISSION_DENIED: raise PermissionDeniedError() raise class InvalidArgumentError(RuntimeError): pass class ExperimentNotFoundError(RuntimeError): pass class PermissionDeniedError(RuntimeError): pass class _OutOfSpaceError(Exception): """Action could not proceed without overflowing request budget. This is a signaling exception (like `StopIteration`) used internally by `_*RequestSender`; it does not mean that anything has gone wrong. """ pass class _BatchedRequestSender(object): """Helper class for building requests that fit under a size limit. This class maintains stateful request builders for each of the possible request types (scalars, tensors, and blobs). These accumulate batches independently, each maintaining its own byte budget and emitting a request when the batch becomes full. As a consequence, events of different types will likely be sent to the backend out of order. E.g., in the extreme case, a single tensor-flavored request may be sent only when the event stream is exhausted, even though many more recent scalar events were sent earlier. This class is not threadsafe. Use external synchronization if calling its methods concurrently. """ def __init__( self, experiment_id, api, allowed_plugins, upload_limits, rpc_rate_limiter, tensor_rpc_rate_limiter, blob_rpc_rate_limiter, ): # Map from `(run_name, tag_name)` to `SummaryMetadata` if the time # series is a scalar time series, else to `_NON_SCALAR_TIME_SERIES`. self._tag_metadata = {} self._allowed_plugins = frozenset(allowed_plugins) self._scalar_request_sender = _ScalarBatchedRequestSender( experiment_id, api, rpc_rate_limiter, upload_limits.max_scalar_request_size, ) self._tensor_request_sender = _TensorBatchedRequestSender( experiment_id, api, tensor_rpc_rate_limiter, upload_limits.max_tensor_request_size, upload_limits.max_tensor_point_size, ) self._blob_request_sender = _BlobRequestSender( experiment_id, api, blob_rpc_rate_limiter, upload_limits.max_blob_request_size, upload_limits.max_blob_size, ) def send_requests(self, run_to_events): """Accepts a stream of TF events and sends batched write RPCs. Each sent request will be batched, the size of each batch depending on the type of data (Scalar vs Tensor vs Blob) being sent. Args: run_to_events: Mapping from run name to generator of `tf.Event` values, as returned by `LogdirLoader.get_run_events`. Raises: RuntimeError: If no progress can be made because even a single point is too large (say, due to a gigabyte-long tag name). """ for (run_name, event, value) in self._run_values(run_to_events): time_series_key = (run_name, value.tag) # The metadata for a time series is memorized on the first event. # If later events arrive with a mismatching plugin_name, they are # ignored with a warning. metadata = self._tag_metadata.get(time_series_key) first_in_time_series = False if metadata is None: first_in_time_series = True metadata = value.metadata self._tag_metadata[time_series_key] = metadata plugin_name = metadata.plugin_data.plugin_name if value.HasField("metadata") and ( plugin_name != value.metadata.plugin_data.plugin_name ): logger.warning( "Mismatching plugin names for %s. Expected %s, found %s.", time_series_key, metadata.plugin_data.plugin_name, value.metadata.plugin_data.plugin_name, ) continue if plugin_name not in self._allowed_plugins: if first_in_time_series: logger.info( "Skipping time series %r with unsupported plugin name %r", time_series_key, plugin_name, ) continue if metadata.data_class == summary_pb2.DATA_CLASS_SCALAR: self._scalar_request_sender.add_event( run_name, event, value, metadata ) elif metadata.data_class == summary_pb2.DATA_CLASS_TENSOR: self._tensor_request_sender.add_event( run_name, event, value, metadata ) elif metadata.data_class == summary_pb2.DATA_CLASS_BLOB_SEQUENCE: self._blob_request_sender.add_event( run_name, event, value, metadata ) self._scalar_request_sender.flush() self._tensor_request_sender.flush() self._blob_request_sender.flush() def _run_values(self, run_to_events): """Helper generator to create a single stream of work items. Note that `dataclass_compat` may emit multiple variants of the same event, for backwards compatibility. Thus this stream should be filtered to obtain the desired version of each event. Here, we ignore any event that does not have a `summary` field. Furthermore, the events emitted here could contain values that do not have `metadata.data_class` set; these too should be ignored. In `_send_summary_value(...)` above, we switch on `metadata.data_class` and drop any values with an unknown (i.e., absent or unrecognized) `data_class`. """ # Note that this join in principle has deletion anomalies: if the input # stream contains runs with no events, or events with no values, we'll # lose that information. This is not a problem: we would need to prune # such data from the request anyway. for (run_name, events) in six.iteritems(run_to_events): for event in events: _filter_graph_defs(event) for value in event.summary.value: yield (run_name, event, value) class _ScalarBatchedRequestSender(object): """Helper class for building requests that fit under a size limit. This class accumulates a current request. `add_event(...)` may or may not send the request (and start a new one). After all `add_event(...)` calls are complete, a final call to `flush()` is needed to send the final request. This class is not threadsafe. Use external synchronization if calling its methods concurrently. """ def __init__(self, experiment_id, api, rpc_rate_limiter, max_request_size): if experiment_id is None: raise ValueError("experiment_id cannot be None") self._experiment_id = experiment_id self._api = api self._rpc_rate_limiter = rpc_rate_limiter self._byte_budget_manager = _ByteBudgetManager(max_request_size) self._runs = {} # cache: map from run name to `Run` proto in request self._tags = ( {} ) # cache: map from `(run, tag)` to `Tag` proto in run in request self._new_request() def _new_request(self): """Allocates a new request and refreshes the budget.""" self._request = write_service_pb2.WriteScalarRequest() self._runs.clear() self._tags.clear() self._request.experiment_id = self._experiment_id self._byte_budget_manager.reset(self._request) def add_event(self, run_name, event, value, metadata): """Attempts to add the given event to the current request. If the event cannot be added to the current request because the byte budget is exhausted, the request is flushed, and the event is added to the next request. """ try: self._add_event_internal(run_name, event, value, metadata) except _OutOfSpaceError: self.flush() # Try again. This attempt should never produce OutOfSpaceError # because we just flushed. try: self._add_event_internal(run_name, event, value, metadata) except _OutOfSpaceError: raise RuntimeError("add_event failed despite flush") def _add_event_internal(self, run_name, event, value, metadata): run_proto = self._runs.get(run_name) if run_proto is None: run_proto = self._create_run(run_name) self._runs[run_name] = run_proto tag_proto = self._tags.get((run_name, value.tag)) if tag_proto is None: tag_proto = self._create_tag(run_proto, value.tag, metadata) self._tags[(run_name, value.tag)] = tag_proto self._create_point(tag_proto, event, value) def flush(self): """Sends the active request after removing empty runs and tags. Starts a new, empty active request. """ request = self._request _prune_empty_tags_and_runs(request) if not request.runs: return self._rpc_rate_limiter.tick() with _request_logger(request, request.runs): try: # TODO(@nfelt): execute this RPC asynchronously. grpc_util.call_with_retries(self._api.WriteScalar, request) except grpc.RpcError as e: if e.code() == grpc.StatusCode.NOT_FOUND: raise ExperimentNotFoundError() logger.error("Upload call failed with error %s", e) self._new_request() def _create_run(self, run_name): """Adds a run to the live request, if there's space. Args: run_name: String name of the run to add. Returns: The `WriteScalarRequest.Run` that was added to `request.runs`. Raises: _OutOfSpaceError: If adding the run would exceed the remaining request budget. """ run_proto = self._request.runs.add(name=run_name) self._byte_budget_manager.add_run(run_proto) return run_proto def _create_tag(self, run_proto, tag_name, metadata): """Adds a tag for the given value, if there's space. Args: run_proto: `WriteScalarRequest.Run` proto to which to add a tag. tag_name: String name of the tag to add (as `value.tag`). metadata: TensorBoard `SummaryMetadata` proto from the first occurrence of this time series. Returns: The `WriteScalarRequest.Tag` that was added to `run_proto.tags`. Raises: _OutOfSpaceError: If adding the tag would exceed the remaining request budget. """ tag_proto = run_proto.tags.add(name=tag_name) tag_proto.metadata.CopyFrom(metadata) self._byte_budget_manager.add_tag(tag_proto) return tag_proto def _create_point(self, tag_proto, event, value): """Adds a scalar point to the given tag, if there's space. Args: tag_proto: `WriteScalarRequest.Tag` proto to which to add a point. event: Enclosing `Event` proto with the step and wall time data. value: Scalar `Summary.Value` proto with the actual scalar data. Raises: _OutOfSpaceError: If adding the point would exceed the remaining request budget. """ point = tag_proto.points.add() point.step = event.step # TODO(@nfelt): skip tensor roundtrip for Value with simple_value set point.value = tensor_util.make_ndarray(value.tensor).item() util.set_timestamp(point.wall_time, event.wall_time) try: self._byte_budget_manager.add_point(point) except _OutOfSpaceError: tag_proto.points.pop() raise class _TensorBatchedRequestSender(object): """Helper class for building WriteTensor() requests that fit under a size limit. This class accumulates a current request. `add_event(...)` may or may not send the request (and start a new one). After all `add_event(...)` calls are complete, a final call to `flush()` is needed to send the final request. This class is not threadsafe. Use external synchronization if calling its methods concurrently. """ def __init__( self, experiment_id, api, rpc_rate_limiter, max_request_size, max_tensor_point_size, ): if experiment_id is None: raise ValueError("experiment_id cannot be None") self._experiment_id = experiment_id self._api = api self._rpc_rate_limiter = rpc_rate_limiter self._byte_budget_manager = _ByteBudgetManager(max_request_size) self._max_tensor_point_size = max_tensor_point_size self._runs = {} # cache: map from run name to `Run` proto in request self._tags = ( {} ) # cache: map from `(run, tag)` to `Tag` proto in run in request self._new_request() def _new_request(self): """Allocates a new request and refreshes the budget.""" self._request = write_service_pb2.WriteTensorRequest() self._runs.clear() self._tags.clear() self._request.experiment_id = self._experiment_id self._byte_budget_manager.reset(self._request) def add_event(self, run_name, event, value, metadata): """Attempts to add the given event to the current request. If the event cannot be added to the current request because the byte budget is exhausted, the request is flushed, and the event is added to the next request. """ try: self._add_event_internal(run_name, event, value, metadata) except _OutOfSpaceError: self.flush() # Try again. This attempt should never produce OutOfSpaceError # because we just flushed. try: self._add_event_internal(run_name, event, value, metadata) except _OutOfSpaceError: raise RuntimeError("add_event failed despite flush") def _add_event_internal(self, run_name, event, value, metadata): run_proto = self._runs.get(run_name) if run_proto is None: run_proto = self._create_run(run_name) self._runs[run_name] = run_proto tag_proto = self._tags.get((run_name, value.tag)) if tag_proto is None: tag_proto = self._create_tag(run_proto, value.tag, metadata) self._tags[(run_name, value.tag)] = tag_proto self._create_point(tag_proto, event, value) def flush(self): """Sends the active request after removing empty runs and tags. Starts a new, empty active request. """ request = self._request _prune_empty_tags_and_runs(request) if not request.runs: return self._rpc_rate_limiter.tick() with _request_logger(request, request.runs): try: grpc_util.call_with_retries(self._api.WriteTensor, request) except grpc.RpcError as e: if e.code() == grpc.StatusCode.NOT_FOUND: raise ExperimentNotFoundError() logger.error("Upload call failed with error %s", e) self._new_request() def _create_run(self, run_name): """Adds a run to the live request, if there's space. Args: run_name: String name of the run to add. Returns: The `WriteTensorRequest.Run` that was added to `request.runs`. Raises: _OutOfSpaceError: If adding the run would exceed the remaining request budget. """ run_proto = self._request.runs.add(name=run_name) self._byte_budget_manager.add_run(run_proto) return run_proto def _create_tag(self, run_proto, tag_name, metadata): """Adds a tag for the given value, if there's space. Args: run_proto: `WriteTensorRequest.Run` proto to which to add a tag. tag_name: String name of the tag to add (as `value.tag`). metadata: TensorBoard `SummaryMetadata` proto from the first occurrence of this time series. Returns: The `WriteTensorRequest.Tag` that was added to `run_proto.tags`. Raises: _OutOfSpaceError: If adding the tag would exceed the remaining request budget. """ tag_proto = run_proto.tags.add(name=tag_name) tag_proto.metadata.CopyFrom(metadata) self._byte_budget_manager.add_tag(tag_proto) return tag_proto def _create_point(self, tag_proto, event, value): """Adds a tensor point to the given tag, if there's space. Args: tag_proto: `WriteTensorRequest.Tag` proto to which to add a point. event: Enclosing `Event` proto with the step and wall time data. value: Tensor `Summary.Value` proto with the actual tensor data. Raises: _OutOfSpaceError: If adding the point would exceed the remaining request budget. """ point = tag_proto.points.add() point.step = event.step point.value.CopyFrom(value.tensor) util.set_timestamp(point.wall_time, event.wall_time) if point.value.ByteSize() > self._max_tensor_point_size: logger.warning( "Tensor too large; skipping. " "Size %d exceeds limit of %d bytes.", point.value.ByteSize(), self._max_tensor_point_size, ) tag_proto.points.pop() return self._validate_tensor_value( value.tensor, value.tag, event.step, event.wall_time ) try: self._byte_budget_manager.add_point(point) except _OutOfSpaceError: tag_proto.points.pop() raise def _validate_tensor_value(self, tensor_proto, tag, step, wall_time): """Validate a TensorProto by attempting to parse it.""" try: tensor_util.make_ndarray(tensor_proto) except ValueError as error: raise ValueError( "The uploader failed to upload a tensor. This seems to be " "due to a malformation in the tensor, which may be caused by " "a bug in the process that wrote the tensor.\n\n" "The tensor has tag '%s' and is at step %d and wall_time %.6f.\n\n" "Original error:\n%s" % (tag, step, wall_time, error) ) class _ByteBudgetManager(object): """Helper class for managing the request byte budget for certain RPCs. This should be used for RPCs that organize data by Runs, Tags, and Points, specifically WriteScalar and WriteTensor. Any call to add_run(), add_tag(), or add_point() may raise an _OutOfSpaceError, which is non-fatal. It signals to the caller that they should flush the current request and begin a new one. For more information on the protocol buffer encoding and how byte cost can be calculated, visit: https://developers.google.com/protocol-buffers/docs/encoding """ def __init__(self, max_bytes): # The remaining number of bytes that we may yet add to the request. self._byte_budget = None # type: int self._max_bytes = max_bytes def reset(self, base_request): """Resets the byte budget and calculates the cost of the base request. Args: base_request: Base request. Raises: _OutOfSpaceError: If the size of the request exceeds the entire request byte budget. """ self._byte_budget = self._max_bytes self._byte_budget -= base_request.ByteSize() if self._byte_budget < 0: raise RuntimeError("Byte budget too small for base request") def add_run(self, run_proto): """Integrates the cost of a run proto into the byte budget. Args: run_proto: The proto representing a run. Raises: _OutOfSpaceError: If adding the run would exceed the remaining request budget. """ cost = ( # The size of the run proto without any tag fields set. run_proto.ByteSize() # The size of the varint that describes the length of the run # proto. We can't yet know the final size of the run proto -- we # haven't yet set any tag or point values -- so we can't know the # final size of this length varint. We conservatively assume it is # maximum size. + _MAX_VARINT64_LENGTH_BYTES # The size of the proto key. + 1 ) if cost > self._byte_budget: raise _OutOfSpaceError() self._byte_budget -= cost def add_tag(self, tag_proto): """Integrates the cost of a tag proto into the byte budget. Args: tag_proto: The proto representing a tag. Raises: _OutOfSpaceError: If adding the tag would exceed the remaining request budget. """ cost = ( # The size of the tag proto without any tag fields set. tag_proto.ByteSize() # The size of the varint that describes the length of the tag # proto. We can't yet know the final size of the tag proto -- we # haven't yet set any point values -- so we can't know the final # size of this length varint. We conservatively assume it is maximum # size. + _MAX_VARINT64_LENGTH_BYTES # The size of the proto key. + 1 ) if cost > self._byte_budget: raise _OutOfSpaceError() self._byte_budget -= cost def add_point(self, point_proto): """Integrates the cost of a point proto into the byte budget. Args: point_proto: The proto representing a point. Raises: _OutOfSpaceError: If adding the point would exceed the remaining request budget. """ submessage_cost = point_proto.ByteSize() cost = ( # The size of the point proto. submessage_cost # The size of the varint that describes the length of the point # proto. + _varint_cost(submessage_cost) # The size of the proto key. + 1 ) if cost > self._byte_budget: raise _OutOfSpaceError() self._byte_budget -= cost class _BlobRequestSender(object): """Uploader for blob-type event data. Unlike the other types, this class does not accumulate events in batches; every blob is sent individually and immediately. Nonetheless we retain the `add_event()`/`flush()` structure for symmetry. This class is not threadsafe. Use external synchronization if calling its methods concurrently. """ def __init__( self, experiment_id, api, rpc_rate_limiter, max_blob_request_size, max_blob_size, ): if experiment_id is None: raise ValueError("experiment_id cannot be None") self._experiment_id = experiment_id self._api = api self._rpc_rate_limiter = rpc_rate_limiter self._max_blob_request_size = max_blob_request_size self._max_blob_size = max_blob_size # Start in the empty state, just like self._new_request(). self._run_name = None self._event = None self._value = None self._metadata = None def _new_request(self): """Declares the previous event complete.""" self._run_name = None self._event = None self._value = None self._metadata = None def add_event( self, run_name, event, value, metadata, ): """Attempts to add the given event to the current request. If the event cannot be added to the current request because the byte budget is exhausted, the request is flushed, and the event is added to the next request. """ if self._value: raise RuntimeError("Tried to send blob while another is pending") self._run_name = run_name self._event = event # provides step and possibly plugin_name self._value = value # TODO(soergel): should we really unpack the tensor here, or ship # it wholesale and unpack server side, or something else? # TODO(soergel): can we extract the proto fields directly instead? self._blobs = tensor_util.make_ndarray(self._value.tensor) if self._blobs.ndim == 1: self._metadata = metadata self.flush() else: logger.warning( "A blob sequence must be represented as a rank-1 Tensor. " "Provided data has rank %d, for run %s, tag %s, step %s ('%s' plugin) .", self._blobs.ndim, run_name, self._value.tag, self._event.step, metadata.plugin_data.plugin_name, ) # Skip this upload. self._new_request() def flush(self): """Sends the current blob sequence fully, and clears it to make way for the next. """ if self._value: blob_sequence_id = self._get_or_create_blob_sequence() logger.info( "Sending %d blobs for sequence id: %s", len(self._blobs), blob_sequence_id, ) sent_blobs = 0 for seq_index, blob in enumerate(self._blobs): # Note the _send_blob() stream is internally flow-controlled. # This rate limit applies to *starting* the stream. self._rpc_rate_limiter.tick() sent_blobs += self._send_blob(blob_sequence_id, seq_index, blob) logger.info( "Sent %d of %d blobs for sequence id: %s", sent_blobs, len(self._blobs), blob_sequence_id, ) self._new_request() def _get_or_create_blob_sequence(self): request = write_service_pb2.GetOrCreateBlobSequenceRequest( experiment_id=self._experiment_id, run=self._run_name, tag=self._value.tag, step=self._event.step, final_sequence_length=len(self._blobs), metadata=self._metadata, ) util.set_timestamp(request.wall_time, self._event.wall_time) with _request_logger(request): try: # TODO(@nfelt): execute this RPC asynchronously. response = grpc_util.call_with_retries( self._api.GetOrCreateBlobSequence, request ) blob_sequence_id = response.blob_sequence_id except grpc.RpcError as e: if e.code() == grpc.StatusCode.NOT_FOUND: raise ExperimentNotFoundError() logger.error("Upload call failed with error %s", e) # TODO(soergel): clean up raise return blob_sequence_id def _send_blob(self, blob_sequence_id, seq_index, blob): """Tries to send a single blob for a given index within a blob sequence. The blob will not be sent if it was sent already, or if it is too large. Returns: The number of blobs successfully sent (i.e., 1 or 0). """ # TODO(soergel): retry and resume logic if len(blob) > self._max_blob_size: logger.warning( "Blob too large; skipping. Size %d exceeds limit of %d bytes.", len(blob), self._max_blob_size, ) return 0 request_iterator = self._write_blob_request_iterator( blob_sequence_id, seq_index, blob ) upload_start_time = time.time() count = 0 # TODO(soergel): don't wait for responses for greater throughput # See https://stackoverflow.com/questions/55029342/handling-async-streaming-request-in-grpc-python try: for response in self._api.WriteBlob(request_iterator): count += 1 # TODO(soergel): validate responses? probably not. pass upload_duration_secs = time.time() - upload_start_time logger.info( "Upload for %d chunks totaling %d bytes took %.3f seconds (%.3f MB/sec)", count, len(blob), upload_duration_secs, len(blob) / upload_duration_secs / (1024 * 1024), ) return 1 except grpc.RpcError as e: if e.code() == grpc.StatusCode.ALREADY_EXISTS: logger.error("Attempted to re-upload existing blob. Skipping.") return 0 else: logger.info("WriteBlob RPC call got error %s", e) raise def _write_blob_request_iterator(self, blob_sequence_id, seq_index, blob): # For now all use cases have the blob in memory already. # In the future we may want to stream from disk; that will require # refactoring here. # TODO(soergel): compute crc32c's to allow server-side data validation. for offset in range(0, len(blob), self._max_blob_request_size): chunk = blob[offset : offset + self._max_blob_request_size] finalize_object = offset + self._max_blob_request_size >= len(blob) request = write_service_pb2.WriteBlobRequest( blob_sequence_id=blob_sequence_id, index=seq_index, data=chunk, offset=offset, crc32c=None, finalize_object=finalize_object, final_crc32c=None, blob_bytes=len(blob), ) yield request @contextlib.contextmanager def _request_logger(request, runs=None): upload_start_time = time.time() request_bytes = request.ByteSize() logger.info("Trying request of %d bytes", request_bytes) yield upload_duration_secs = time.time() - upload_start_time if runs: logger.info( "Upload for %d runs (%d bytes) took %.3f seconds", len(runs), request_bytes, upload_duration_secs, ) else: logger.info( "Upload of (%d bytes) took %.3f seconds", request_bytes, upload_duration_secs, ) def _varint_cost(n): """Computes the size of `n` encoded as an unsigned base-128 varint. This should be consistent with the proto wire format: Args: n: A non-negative integer. Returns: An integer number of bytes. """ result = 1 while n >= 128: result += 1 n >>= 7 return result def _prune_empty_tags_and_runs(request): for (run_idx, run) in reversed(list(enumerate(request.runs))): for (tag_idx, tag) in reversed(list(enumerate(run.tags))): if not tag.points: del run.tags[tag_idx] if not run.tags: del request.runs[run_idx] def _filter_graph_defs(event): for v in event.summary.value: if v.metadata.plugin_data.plugin_name != graphs_metadata.PLUGIN_NAME: continue if v.tag == graphs_metadata.RUN_GRAPH_NAME: data = list(v.tensor.string_val) filtered_data = [_filtered_graph_bytes(x) for x in data] filtered_data = [x for x in filtered_data if x is not None] if filtered_data != data: new_tensor = tensor_util.make_tensor_proto( filtered_data, dtype=types_pb2.DT_STRING ) v.tensor.CopyFrom(new_tensor) def _filtered_graph_bytes(graph_bytes): try: graph_def = graph_pb2.GraphDef().FromString(graph_bytes) # The reason for the RuntimeWarning catch here is b/27494216, whereby # some proto parsers incorrectly raise that instead of DecodeError # on certain kinds of malformed input. Triggering this seems to require # a combination of mysterious circumstances. except (message.DecodeError, RuntimeWarning): logger.warning( "Could not parse GraphDef of size %d. Skipping.", len(graph_bytes), ) return None # Use the default filter parameters: # limit_attr_size=1024, large_attrs_key="_too_large_attrs" process_graph.prepare_graph_for_ui(graph_def) return graph_def.SerializeToString()