data: add dispatching data provider (#3847)

wchargin · web-flow · commit 03b7edf797bc · 2020-07-20T13:11:08.000-07:00
Summary: A new `DispatchingDataProvider` composes multiple sub-providers under different experiment ID prefixes. Thus, if `FooDataProvider` expects experiment IDs like `123` and `BarDataProvider` expects IDs like `abc`, you could use a `DispatchingDataProvider` to simultaneously expose experiment IDs like `foo:123` and `bar:abc`. Most method implementations of `DispatchingDataProvider` simply involve determining the appropriate sub-provider and forwarding all arguments. We implement those with a tiny bit of metaprogramming (since there are nine such methods), stopping short of [more drastic measures][ddd]. [ddd]: https://wchargin.github.io/posts/dubious-dynamic-delegation-in-python/ Test Plan: Unit tests included. For Google-internal use, <http://cl/321472166> shows an application of this data provider. wchargin-branch: data-dispatching-provider
diff --git a/tensorboard/data/BUILD b/tensorboard/data/BUILD
@@ -40,3 +40,28 @@ py_test(
         "@org_pythonhosted_six",
     ],
 )
+
+py_library(
+    name = "dispatching_provider",
+    srcs = ["dispatching_provider.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":provider",
+        "//tensorboard:errors",
+    ],
+)
+
+py_test(
+    name = "dispatching_provider_test",
+    size = "small",
+    srcs = ["dispatching_provider_test.py"],
+    srcs_version = "PY3",
+    tags = ["support_notf"],
+    deps = [
+        ":dispatching_provider",
+        ":provider",
+        "//tensorboard:context",
+        "//tensorboard:errors",
+        "//tensorboard:test",
+    ],
+)
diff --git a/tensorboard/data/dispatching_provider.py b/tensorboard/data/dispatching_provider.py
@@ -0,0 +1,216 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""A data provider that dispatches to sub-providers based on prefix."""
+
+import base64
+import json
+
+from tensorboard import errors
+from tensorboard.data import provider
+
+
+# Separator between prefix and sub-ID.
+_SEPARATOR = ":"
+
+
+class DispatchingDataProvider(provider.DataProvider):
+    """Data provider that dispatches to sub-providers based on prefix.
+
+    If you have one data provider that talks to the Foo service with IDs
+    like `123` and another that talks to the Bar service with IDs like
+    `3a28213a`, then this data provider lets you talk to both of them
+    with IDs like `foo:123` and `bar:3a28213a`, respectively. The part
+    before the colon is the *prefix*, and identifies the sub-provider to
+    which to dispatch. The part after the colon is the *sub-ID*, and is
+    passed verbatim to the sub-provider. The sub-ID may contain any
+    characters, including colons, so the sub-providers may themselves be
+    hierarchical.
+
+    Optionally, an `unprefixed_provider` can be specified as default in
+    the case where an experiment ID does not contain a colon. Note that
+    this is not used as a fallback when the prefix is simply not one of
+    the registered prefixes; that will always be an error.
+    """
+
+    # Implementation note: this data provider provides a simple
+    # pass-through for most methods, but has extra logic for methods
+    # related to blob keys, where we need to annotate or extract the
+    # associated sub-provider.
+
+    def __init__(self, providers, unprefixed_provider=None):
+        """Initialize a `DispatchingDataProvider`.
+
+        Args:
+          providers: Dict mapping prefix (`str`) to sub-provider
+            instance (`provider.DataProvider`). Keys will appear in
+            experiment IDs and so must be URL-safe.
+          unprefixed_provider: Optional `provider.DataProvider` instance
+            to use with experiment IDs that do not have a prefix.
+
+        Raises:
+          ValueError: If any of the provider keys contains a colon,
+            which would make it impossible to match.
+        """
+        self._providers = dict(providers)
+        invalid_names = sorted(k for k in self._providers if _SEPARATOR in k)
+        if invalid_names:
+            raise ValueError("Invalid provider key(s): %r" % invalid_names)
+        self._unprefixed_provider = unprefixed_provider
+
+    def _parse_eid(self, experiment_id):
+        """Parse an experiment ID into prefix, sub-ID, and sub-provider.
+
+        The returned prefix may be `None` if this instance has an
+        unprefixed data provider registered. If the experiment ID is
+        invalid, this method may raise an `errors.NotFoundError`.
+        """
+        parts = experiment_id.split(_SEPARATOR, 1)
+        if len(parts) == 1:
+            if self._unprefixed_provider is None:
+                raise errors.NotFoundError(
+                    "No data provider found for unprefixed experiment ID: %r"
+                    % experiment_id
+                )
+            return (None, experiment_id, self._unprefixed_provider)
+        (prefix, sub_eid) = parts
+        sub_provider = self._providers.get(prefix)
+        if sub_provider is None:
+            raise errors.NotFoundError(
+                "Unknown prefix in experiment ID: %r" % experiment_id
+            )
+        return (prefix, sub_eid, sub_provider)
+
+    def _simple_delegate(get_method):
+        """Dispatch on experiment ID, forwarding args and result unchanged."""
+
+        def wrapper(self, *args, experiment_id, **kwargs):
+            (_, sub_eid, sub_provider) = self._parse_eid(experiment_id)
+            method = get_method(sub_provider)
+            return method(*args, experiment_id=sub_eid, **kwargs)
+
+        return wrapper
+
+    data_location = _simple_delegate(lambda p: p.data_location)
+    experiment_metadata = _simple_delegate(lambda p: p.experiment_metadata)
+    list_plugins = _simple_delegate(lambda p: p.list_plugins)
+    list_runs = _simple_delegate(lambda p: p.list_runs)
+    list_scalars = _simple_delegate(lambda p: p.list_scalars)
+    read_scalars = _simple_delegate(lambda p: p.read_scalars)
+    list_tensors = _simple_delegate(lambda p: p.list_tensors)
+    read_tensors = _simple_delegate(lambda p: p.read_tensors)
+    list_blob_sequences = _simple_delegate(lambda p: p.list_blob_sequences)
+
+    def read_blob_sequences(self, *args, experiment_id, **kwargs):
+        (prefix, sub_eid, sub_provider) = self._parse_eid(experiment_id)
+        result = sub_provider.read_blob_sequences(
+            *args, experiment_id=sub_eid, **kwargs
+        )
+        for tag_to_data in result.values():
+            for (tag, old_data) in tag_to_data.items():
+                new_data = [
+                    provider.BlobSequenceDatum(
+                        step=d.step,
+                        wall_time=d.wall_time,
+                        values=_convert_blob_references(prefix, d.values),
+                    )
+                    for d in old_data
+                ]
+                tag_to_data[tag] = new_data
+        return result
+
+    def read_blob(self, ctx, blob_key):
+        (prefix, sub_key) = _decode_blob_key(blob_key)
+        if prefix is None:
+            if self._unprefixed_provider is None:
+                raise errors.NotFoundError(
+                    "Invalid blob key: no unprefixed provider"
+                )
+            return self._unprefixed_provider.read_blob(ctx, blob_key=sub_key)
+        sub_provider = self._providers.get(prefix)
+        if sub_provider is None:
+            raise errors.NotFoundError(
+                "Invalid blob key: no such provider: %r; have: %r"
+                % (prefix, sorted(self._providers))
+            )
+        return sub_provider.read_blob(ctx, blob_key=sub_key)
+
+
+def _convert_blob_references(prefix, references):
+    """Encode all blob keys in a list of blob references.
+
+    Args:
+      prefix: The prefix of the sub-provider that generated the sub-key,
+        or `None` if this was generated by the unprefixed provider.
+      references: A list of `provider.BlobReference`s emitted by a
+        sub-provider.
+
+    Returns:
+      A new list of `provider.BlobReference`s whose blob keys have been
+      encoded per `_encode_blob_key`.
+    """
+    return [
+        provider.BlobReference(
+            blob_key=_encode_blob_key(prefix, r.blob_key), url=r.url,
+        )
+        for r in references
+    ]
+
+
+def _encode_blob_key(prefix, sub_key):
+    """Encode a blob key from prefix (optional) and sub-key.
+
+    Args:
+      prefix: The prefix of the sub-provider that generated the sub-key,
+        or `None` if this was generated by the unprefixed provider.
+      sub_key: The opaque key from the sub-provider.
+
+    Returns:
+      A string encoding `prefix` and `sub_key` injectively.
+    """
+    payload = [prefix, sub_key]
+    json_str = json.dumps(payload, separators=(",", ":"))
+    b64_str = base64.urlsafe_b64encode(json_str.encode("ascii")).decode("ascii")
+    return b64_str.rstrip("=")
+
+
+def _decode_blob_key(key):
+    """Decode a prefix (optional) and sub-key from a blob key.
+
+    Left inverse of `_encode_blob_key`.
+
+    Args:
+      key: A blob key in the form returned by `_encode_blob_key`.
+
+    Returns;
+      A tuple `(prefix, sub_key)`, where `prefix` is either `None` or a
+      sub-provider prefix, and `sub_key` is an opaque key from a
+      sub-provider.
+
+    Raises:
+      errors.NotFoundError: If `key` is invalid and has no preimage.
+    """
+    failure = errors.NotFoundError("Invalid blob key: %r" % key)
+
+    b64_str = key + "=="  # ensure adequate padding (overpadding is okay)
+    json_str = base64.urlsafe_b64decode(b64_str).decode("ascii")
+    payload = json.loads(json_str)
+    if not isinstance(payload, list) or len(payload) != 2:
+        raise failure
+    (prefix, sub_key) = payload
+    if not (prefix is None or isinstance(prefix, str)):
+        raise failure
+    if not isinstance(sub_key, str):
+        raise failure
+    return (prefix, sub_key)
diff --git a/tensorboard/data/dispatching_provider_test.py b/tensorboard/data/dispatching_provider_test.py
