data: check liveness before blessing data server (#4851)

Summary:
Issue #4844 shows that there are circumstances when communication
between TensorBoard and the local data server process cannot be
established. As of this patch, we send a trivial RPC to the data server
and wait for its response before committing to use the new loading
paths. This adds about 1–5ms total time to the happy path on my machine,
with a worst case penalty of 5s due to the timeout. If the server is not
reachable, we print a warning and fall back to the legacy paths.

Test Plan:

  - Test with normal working TensorBoard.

  - Simulate failed data server connection by changing the definition of
    `addr` to `"localhost:%d" % (port + 777)` on line 175. This should
    print the `UNAVAILABLE`/“failed to connect” message from #4844.

  - Simulate slow data server by adding the following to `cli.rs`:

    ```rust
    tokio::time::sleep(Duration::from_secs(3)).await;
    ```

    Add this right before the `Server::builder().(...).await?` call at
    the end of `main`: i.e., after we write the port file, but before we
    actually respond to requests. Note that TensorBoard still works with
    a 3-second delay, but that it actually delays printing the startup
    message for those 3 seconds as it determines which data provider to
    use.

  - Simulate extra-slow data provider as above but waiting 6 seconds,
    and note that TensorBoard prints a `DEADLINE_EXCEEDED` error, falls
    back to the legacy paths, and shows valid data.

wchargin-branch: data-liveness-check
wchargin-source: 7fd916ebcf426d92babaeb57c38789d47446df6e

Author: wchargin

tensorboard/data/BUILD

@@ -95,6 +95,7 @@ py_library(
         ":ingester",
         "//tensorboard:expect_grpc_installed",
         "//tensorboard:expect_pkg_resources_installed",
+        "//tensorboard/data/proto:protos_all_py_pb2",
         "//tensorboard/util:tb_logging",
     ],
 )

tensorboard/data/server_ingester.py

@@ -26,6 +26,7 @@
 
 from tensorboard.data import grpc_provider
 from tensorboard.data import ingester
+from tensorboard.data.proto import data_provider_pb2
 from tensorboard.util import tb_logging
 
 
@@ -47,7 +48,8 @@ def __init__(self, address, *, channel_creds_type):
           channel_creds_type: `grpc_util.ChannelCredsType`, as passed to
             `--grpc_creds_type`.
         """
-        self._data_provider = _make_provider(address, channel_creds_type)
+        stub = _make_stub(address, channel_creds_type)
+        self._data_provider = grpc_provider.GrpcDataProvider(address, stub)
 
     @property
     def data_provider(self):
@@ -170,13 +172,23 @@ def start(self):
             )
 
         addr = "localhost:%d" % port
-        self._data_provider = _make_provider(addr, self._channel_creds_type)
+        stub = _make_stub(addr, self._channel_creds_type)
         logger.info(
-            "Established connection to data server at pid %d via %s",
+            "Opened channel to data server at pid %d via %s",
             popen.pid,
             addr,
         )
 
+        req = data_provider_pb2.GetExperimentRequest()
+        try:
+            stub.GetExperiment(req, timeout=5)  # should be near-instant
+        except grpc.RpcError as e:
+            msg = "Failed to communicate with data server at %s: %s" % (addr, e)
+            logging.warning("%s", msg)
+            raise DataServerStartupError(msg) from e
+        logger.info("Got valid response from data server")
+        self._data_provider = grpc_provider.GrpcDataProvider(addr, stub)
+
 
 def _maybe_read_file(path):
     """Read a file, or return `None` on ENOENT specifically."""
@@ -189,12 +201,11 @@ def _maybe_read_file(path):
         raise
 
 
-def _make_provider(addr, channel_creds_type):
+def _make_stub(addr, channel_creds_type):
     (creds, options) = channel_creds_type.channel_config()
     options.append(("grpc.max_receive_message_length", 1024 * 1024 * 256))
     channel = grpc.secure_channel(addr, creds, options=options)
-    stub = grpc_provider.make_stub(channel)
-    return grpc_provider.GrpcDataProvider(addr, stub)
+    return grpc_provider.make_stub(channel)
 
 
 class NoDataServerError(RuntimeError):