program: add gRPC data ingester (#4411)

Summary:
TensorBoard can now start RustBoard if `--load_fast` is given at runtime
and `--define=link_data_server=true` was given at build time. The
default configuration still has no Rust code, so our Pip packages are
still portable. When we actually deploy this, we can distribute the Rust
binary in a separate Pip package that TensorBoard imports, but the UX
can stay the same: “just add `--load_fast`.”

Test Plan:
Test in the following configurations:

  - with just a `--logdir`, everything works as normal;
  - with `--define=link_data_server=true -- --logdir ... --load_fast`,
    the subprocess is spawned and cleaned up on a successful exit or a
    SIGKILL to TensorBoard, and INFO logs are shown iff `--verbosity 0`
    is specified;
  - with `--logdir ... --load_fast` but no `--define`, TensorBoard fails
    to start the server and prints a message before exiting; and
  - with `--grpc_data_provider localhost:6806`, TensorBoard connects to
    an existing server without needing `--logdir` or `--define`.

To test the “data server died” case, comment out the `--logdir=%s` flag,
which will cause the server to fail with a usage message. That message
should appear in the logs.

To test the polling, add `thread::sleep(Duration::from_secs(3))` before
the server writes its port file, and run with `--verbosity 0` to note
the “Polling for data server port” messages.

This also works after syncing into Google, in all relevant
configurations; see <http://cl/344900410> and <http://cl/344955833>.

wchargin-branch: grpc-ingester

Author: wchargin

tensorboard/BUILD

@@ -216,11 +216,10 @@ py_library(
         ":version",
         "//tensorboard:expect_absl_flags_installed",
         "//tensorboard:expect_absl_logging_installed",
-        "//tensorboard:expect_grpc_installed",
         "//tensorboard/backend:application",
         "//tensorboard/backend/event_processing:data_ingester",
         "//tensorboard/backend/event_processing:event_file_inspector",
-        "//tensorboard/data:grpc_provider",
+        "//tensorboard/data:server_ingester",
         "//tensorboard/util:argparse_util",
         "@org_pocoo_werkzeug",
         "@org_pythonhosted_six",

tensorboard/backend/event_processing/BUILD

@@ -37,6 +37,7 @@ py_library(
         ":data_provider",
         ":event_multiplexer",
         ":tag_types",
+        "//tensorboard/data:ingester",
         "//tensorboard/plugins/audio:metadata",
         "//tensorboard/plugins/histogram:metadata",
         "//tensorboard/plugins/image:metadata",

tensorboard/backend/event_processing/data_ingester.py

@@ -24,6 +24,7 @@
 from tensorboard.backend.event_processing import data_provider
 from tensorboard.backend.event_processing import plugin_event_multiplexer
 from tensorboard.backend.event_processing import tag_types
+from tensorboard.data import ingester
 from tensorboard.plugins.audio import metadata as audio_metadata
 from tensorboard.plugins.histogram import metadata as histogram_metadata
 from tensorboard.plugins.image import metadata as image_metadata
@@ -48,7 +49,7 @@
 logger = tb_logging.get_logger()
 
 
-class LocalDataIngester(object):
+class LocalDataIngester(ingester.DataIngester):
     """Data ingestion implementation to use when running locally."""
 
     def __init__(self, flags):

tensorboard/data/BUILD

@@ -64,6 +64,59 @@ py_test(
     ],
 )
 
+py_library(
+    name = "ingester",
+    srcs = ["ingester.py"],
+    srcs_version = "PY3",
+)
+
+py_test(
+    name = "ingester_test",
+    size = "small",
+    srcs = ["ingester_test.py"],
+    srcs_version = "PY3",
+    tags = ["support_notf"],
+    deps = [
+        ":ingester",
+    ],
+)
+
+config_setting(
+    name = "link_data_server",
+    define_values = {"link_data_server": "true"},
+)
+
+py_library(
+    name = "server_ingester",
+    srcs = ["server_ingester.py"],
+    data = select({
+        ":link_data_server": ["//tensorboard/data/server"],
+        "//conditions:default": [],
+    }),
+    srcs_version = "PY3",
+    deps = [
+        ":grpc_provider",
+        ":ingester",
+        "//tensorboard:expect_grpc_installed",
+        "//tensorboard/util:tb_logging",
+    ],
+)
+
+py_test(
+    name = "server_ingester_test",
+    size = "medium",  # time.sleep
+    timeout = "short",
+    srcs = ["server_ingester_test.py"],
+    srcs_version = "PY3",
+    tags = ["support_notf"],
+    deps = [
+        ":grpc_provider",
+        ":server_ingester",
+        "//tensorboard:expect_grpc_installed",
+        "//tensorboard:test",
+    ],
+)
+
 py_library(
     name = "grpc_provider",
     srcs = ["grpc_provider.py"],

tensorboard/data/ingester.py

@@ -0,0 +1,46 @@
+# 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.
+# ==============================================================================
+"""Abstraction for data ingestion logic."""
+
+import abc
+
+
+class DataIngester(metaclass=abc.ABCMeta):
+    """Link between a data source and a data provider.
+
+    A data ingester starts a reload operation in the background and
+    provides a data provider as a view.
+    """
+
+    @property
+    @abc.abstractmethod
+    def data_provider(self):
+        """Returns a `DataProvider`.
+
+        It may be an error to dereference this before `start` is called.
+        """
+        pass
+
+    @abc.abstractmethod
+    def start(self):
+        """Starts ingesting data.
+
+        This may start a background thread or process, and will return
+        once communication with that task is established. It won't block
+        forever as data is reloaded.
+
+        Must only be called once.
+        """
+        pass

tensorboard/data/ingester_test.py

@@ -0,0 +1,21 @@
+# 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.
+# ==============================================================================
+"""Unit tests for `tensorboard.ingester`."""
+
+from tensorboard.data import ingester
+
+# This is a pure abstract class. There's really nothing to test other than that
+# it executes successfully.
+del ingester

tensorboard/data/server/DEVELOPMENT.md

@@ -36,6 +36,45 @@ Rust source files in your editor. For editor setup, consult
 
 [ra-docs]: https://rust-analyzer.github.io/
 
+## Running under TensorBoard
+
+You can point TensorBoard at a data server in two ways: start the server
+yourself and give TensorBoard an address, or tell TensorBoard to start the
+server as a subprocess.
+
+To connect to an existing server, pass `--grpc_data_provider ADDRESS`, where the
+address is like `localhost:6806`. Thus:
+
+```
+bazel run -c opt //tensorboard -- --grpc_data_provider localhost:6806
+```
+
+You don’t have to pass a `--logdir` if you do this, but you do have to
+concurrently run `//tensorboard/data/server` (say, in the background, or in a
+separate shell). You can also swap out the data server whenever you want without
+restarting TensorBoard; new RPCs will transparently reconnect. The server
+doesn’t have to be running when TensorBoard starts.
+
+To tell TensorBoard to start the server as a subprocess, build with
+`--define=link_data_server=true` and pass `--load_fast` to TensorBoard along
+with a normal `--logdir`. Thus:
+
+```
+bazel run -c opt --define=link_data_server=true //tensorboard -- \
+    --load_fast --logdir ~/tensorboard_data/mnist/ --bind_all --verbosity 0
+```
+
+This is an easier one-shot solution, but requires a `--define` flag, offers less
+flexibility over the flags to the data server, and requires restarting
+TensorBoard if you want to restart the data server (though that’s not usually a
+big deal). The data server will automatically shut down when TensorBoard exits
+for any reason.
+
+As an alternative to `--define=link_data_server=true`, you can set the
+`TENSORBOARD_DATA_SERVER_BINARY` environment variable to the path to a data
+server binary, and pass `--load_fast`. If running with `bazel run`, this should
+be an absolute path.
+
 ## Adding third-party dependencies
 
 Rust dependencies are usually hosted on [crates.io]. We use [`cargo-raze`][raze]
@@ -51,8 +90,8 @@ dependency:
     package on <https://crates.io/>.
 3.  Change into the `tensorboard/data/server/` directory.
 4.  Run `cargo fetch` to update `Cargo.lock`. Running this before `cargo raze`
-    ensures that the `http_archive` workspace rules in the generated build
-    files will have `sha256` checksums.
+    ensures that the `http_archive` workspace rules in the generated build files
+    will have `sha256` checksums.
 5.  Run `cargo raze` to update `third_party/rust/...`.
 
 This will add a new target like `//third_party/rust:rand`. Manually build it

tensorboard/data/server/main.rs

@@ -14,10 +14,11 @@ limitations under the License.
 ==============================================================================*/
 
 use clap::Clap;
-use log::{debug, info, LevelFilter};
-use std::io::Read;
+use log::{debug, error, info, LevelFilter};
+use std::fs::File;
+use std::io::{Read, Write};
 use std::net::{IpAddr, SocketAddr};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::str::FromStr;
 use std::thread;
 use std::time::{Duration, Instant};
@@ -74,6 +75,15 @@ struct Opts {
     /// sense) but not killed.
     #[clap(long)]
     die_after_stdin: bool,
+
+    /// Write bound port to this file
+    ///
+    /// Once a server socket is opened, write the port on which it's listening to the file at this
+    /// path. Useful with `--port 0`. Port will be written as ASCII decimal followed by a newline
+    /// (e.g., "6806\n"). If the server fails to start, this file may not be written at all. If the
+    /// port file is specified but cannot be written, the server will die.
+    #[clap(long)]
+    port_file: Option<PathBuf>,
 }
 
 /// A duration in seconds.
@@ -113,21 +123,39 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let bound = listener.local_addr()?;
     eprintln!("listening on {:?}", bound);
 
+    if let Some(port_file) = opts.port_file {
+        let port = bound.port();
+        if let Err(e) = write_port_file(&port_file, port) {
+            error!(
+                "Failed to write port \"{}\" to {}: {}",
+                port,
+                port_file.display(),
+                e
+            );
+            std::process::exit(1);
+        }
+        info!("Wrote port \"{}\" to {}", port, port_file.display());
+    }
+
     // Leak the commit object, since the Tonic server must have only 'static references. This only
     // leaks the outer commit structure (of constant size), not the pointers to the actual data.
     let commit: &'static Commit = Box::leak(Box::new(Commit::new()));
 
     thread::Builder::new()
         .name("Reloader".to_string())
-        .spawn(move || {
-            let mut loader = LogdirLoader::new(commit, opts.logdir);
-            loop {
-                info!("Starting load cycle");
-                let start = Instant::now();
-                loader.reload();
-                let end = Instant::now();
-                info!("Finished load cycle ({:?})", end - start);
-                thread::sleep(opts.reload_interval.duration());
+        .spawn({
+            let logdir = opts.logdir;
+            let reload_interval = opts.reload_interval;
+            move || {
+                let mut loader = LogdirLoader::new(commit, logdir);
+                loop {
+                    info!("Starting load cycle");
+                    let start = Instant::now();
+                    loader.reload();
+                    let end = Instant::now();
+                    info!("Finished load cycle ({:?})", end - start);
+                    thread::sleep(reload_interval.duration());
+                }
             }
         })
         .expect("failed to spawn reloader thread");
@@ -156,3 +184,11 @@ fn die_after_stdin() {
     info!("Stdin closed; exiting");
     std::process::exit(0);
 }
+
+/// Writes `port` to file `path` as an ASCII decimal followed by newline.
+fn write_port_file(path: &Path, port: u16) -> std::io::Result<()> {
+    let mut f = File::create(path)?;
+    writeln!(f, "{}", port)?;
+    f.sync_all()?;
+    Ok(())
+}