rust: implement data provider server with fake data (#4355)

Summary:
The `//tensorboard/data/server` binary now provides a gRPC server that
implements the `TensorBoardDataProvider` service. It only supports
scalars, and it has entirely fake data. But it’s enough to be wired up
to TensorBoard and show real charts in the UI.

Test Plan:
Run the server with `bazel run -c opt //tensorboard/data/server`.
(Running with `-c opt` is not necessary, but the server is notably
faster that way.) Then, in a different shell, use `grpc_cli`:

```
$ grpc_cli --channel_creds_type=insecure \
>   --protofiles tensorboard/data/proto/data_provider.proto \
>   call localhost:6806 TensorBoardDataProvider.ListRuns '' 2>/dev/null
runs {
  name: "train"
  start_time: 1605752017
}
```

Interestingly, this takes 13.9 ± 4.9 ms on my machine, whereas the demo
server added in #4318 took only 5.2 ± 2.6 ms. Both are basically doing
no work on the server, so I suspect that the difference may be due to
`grpc_cli` having to do more work to parse our real proto files. And,
indeed, making the calls from Python instead takes only 0.8–1.1 ms.

wchargin-branch: rust-real-fake-data-provider

tensorboard/data/server/BUILD

@@ -34,6 +34,7 @@ rust_library(
         "reservoir.rs",
         "run.rs",
         "scripted_reader.rs",
+        "server.rs",
         "tf_record.rs",
         "types.rs",
         "writer.rs",
@@ -42,10 +43,12 @@ rust_library(
     deps = [
         "//third_party/rust:byteorder",
         "//third_party/rust:crc",
+        "//third_party/rust:futures_core",
         "//third_party/rust:prost",
         "//third_party/rust:rand",
         "//third_party/rust:rand_chacha",
         "//third_party/rust:thiserror",
+        "//third_party/rust:tokio",
         "//third_party/rust:tonic",
     ],
 )

tensorboard/data/server/lib.rs

@@ -15,12 +15,15 @@ limitations under the License.
 
 //! Core functionality for TensorBoard data loading.
 
+#![allow(clippy::needless_update)] // https://github.com/rust-lang/rust-clippy/issues/6323
+
 pub mod commit;
 pub mod data_compat;
 pub mod event_file;
 pub mod masked_crc;
 pub mod reservoir;
 pub mod run;
+pub mod server;
 pub mod tf_record;
 pub mod types;
 

tensorboard/data/server/main.rs

@@ -13,37 +13,17 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#![allow(clippy::needless_update)] // https://github.com/rust-lang/rust-clippy/issues/6323
+use tonic::transport::Server;
 
-use tonic::{transport::Server, Request, Response};
-
-use rustboard_core::proto::demo as pb;
-
-#[derive(Debug, Default)]
-struct DemoHandler;
-
-#[tonic::async_trait]
-impl pb::demo_server::Demo for DemoHandler {
-    async fn add(
-        &self,
-        request: Request<pb::AddRequest>,
-    ) -> Result<Response<pb::AddResponse>, tonic::Status> {
-        let request = request.into_inner();
-        let sum: i32 = request.term.into_iter().sum();
-        let response = pb::AddResponse {
-            sum,
-            ..Default::default()
-        };
-        Ok(Response::new(response))
-    }
-}
+use rustboard_core::proto::tensorboard::data::tensor_board_data_provider_server::TensorBoardDataProviderServer;
+use rustboard_core::server::DataProviderHandler;
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let addr = "[::0]:6789".parse::<std::net::SocketAddr>()?;
-    let handler = DemoHandler::default();
+    let addr = "[::0]:6806".parse::<std::net::SocketAddr>()?;
+    let handler = DataProviderHandler;
     Server::builder()
-        .add_service(pb::demo_server::DemoServer::new(handler))
+        .add_service(TensorBoardDataProviderServer::new(handler))
         .serve(addr)
         .await?;
     Ok(())

tensorboard/data/server/server.rs

@@ -0,0 +1,207 @@
+/* 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.
+==============================================================================*/
+
+use futures_core::Stream;
+use std::pin::Pin;
+use tonic::{Request, Response, Status};
+
+use crate::proto::tensorboard::data;
+use data::tensor_board_data_provider_server::TensorBoardDataProvider;
+
+/// Data provider gRPC service implementation.
+#[derive(Debug)]
+pub struct DataProviderHandler;
+
+const FAKE_START_TIME: f64 = 1605752017.0;
+
+#[tonic::async_trait]
+impl TensorBoardDataProvider for DataProviderHandler {
+    async fn list_plugins(
+        &self,
+        _request: Request<data::ListPluginsRequest>,
+    ) -> Result<Response<data::ListPluginsResponse>, Status> {
+        let mut res: data::ListPluginsResponse = Default::default();
+        res.plugins.push(data::Plugin {
+            name: "scalars".to_string(),
+            ..Default::default()
+        });
+        Ok(Response::new(res))
+    }
+
+    async fn list_runs(
+        &self,
+        _request: Request<data::ListRunsRequest>,
+    ) -> Result<Response<data::ListRunsResponse>, Status> {
+        let mut res: data::ListRunsResponse = Default::default();
+        res.runs.push(data::Run {
+            name: "train".to_string(),
+            start_time: FAKE_START_TIME,
+            ..Default::default()
+        });
+        Ok(Response::new(res))
+    }
+
+    async fn list_scalars(
+        &self,
+        _request: Request<data::ListScalarsRequest>,
+    ) -> Result<Response<data::ListScalarsResponse>, Status> {
+        let mut res: data::ListScalarsResponse = Default::default();
+        let mut run: data::list_scalars_response::RunEntry = Default::default();
+        run.run_name = "train".to_string();
+        run.tags.push(data::list_scalars_response::TagEntry {
+            tag_name: "accuracy".to_string(),
+            metadata: Some(data::ScalarMetadata {
+                max_step: 5,
+                max_wall_time: FAKE_START_TIME + 5.0,
+                ..Default::default()
+            }),
+            ..Default::default()
+        });
+        res.runs.push(run);
+        Ok(Response::new(res))
+    }
+
+    async fn read_scalars(
+        &self,
+        _request: Request<data::ReadScalarsRequest>,
+    ) -> Result<Response<data::ReadScalarsResponse>, Status> {
+        let mut res: data::ReadScalarsResponse = Default::default();
+        let mut run: data::read_scalars_response::RunEntry = Default::default();
+        run.run_name = "train".to_string();
+        run.tags.push(data::read_scalars_response::TagEntry {
+            tag_name: "accuracy".to_string(),
+            data: Some(data::ScalarData {
+                step: vec![0, 1, 2, 3, 4],
+                wall_time: (0..=4)
+                    .map(|i| FAKE_START_TIME + ((i + 1) as f64))
+                    .collect(),
+                value: vec![0.1, 0.5, 0.8, 0.9, 0.95],
+                ..Default::default()
+            }),
+            ..Default::default()
+        });
+        res.runs.push(run);
+        Ok(Response::new(res))
+    }
+
+    async fn list_tensors(
+        &self,
+        _request: Request<data::ListTensorsRequest>,
+    ) -> Result<Response<data::ListTensorsResponse>, Status> {
+        Err(Status::unimplemented("not yet implemented"))
+    }
+
+    async fn read_tensors(
+        &self,
+        _request: Request<data::ReadTensorsRequest>,
+    ) -> Result<Response<data::ReadTensorsResponse>, Status> {
+        Err(Status::unimplemented("not yet implemented"))
+    }
+
+    async fn list_blob_sequences(
+        &self,
+        _request: Request<data::ListBlobSequencesRequest>,
+    ) -> Result<Response<data::ListBlobSequencesResponse>, Status> {
+        Err(Status::unimplemented("not yet implemented"))
+    }
+
+    async fn read_blob_sequences(
+        &self,
+        _request: Request<data::ReadBlobSequencesRequest>,
+    ) -> Result<Response<data::ReadBlobSequencesResponse>, Status> {
+        Err(Status::unimplemented("not yet implemented"))
+    }
+
+    type ReadBlobStream =
+        Pin<Box<dyn Stream<Item = Result<data::ReadBlobResponse, Status>> + Send + Sync + 'static>>;
+
+    async fn read_blob(
+        &self,
+        _request: Request<data::ReadBlobRequest>,
+    ) -> Result<Response<Self::ReadBlobStream>, Status> {
+        Err(Status::unimplemented("not yet implemented"))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_list_plugins() {
+        let handler = DataProviderHandler;
+        let req = Request::new(data::ListPluginsRequest {
+            experiment_id: "123".to_string(),
+            ..Default::default()
+        });
+        let res = handler.list_plugins(req).await.unwrap().into_inner();
+        assert_eq!(
+            res.plugins.into_iter().map(|p| p.name).collect::<Vec<_>>(),
+            vec!["scalars"]
+        );
+    }
+
+    #[tokio::test]
+    async fn test_list_runs() {
+        let handler = DataProviderHandler;
+        let req = Request::new(data::ListRunsRequest {
+            experiment_id: "123".to_string(),
+            ..Default::default()
+        });
+        let res = handler.list_runs(req).await.unwrap().into_inner();
+        assert_eq!(res.runs.len(), 1);
+        let run = &res.runs[0];
+        assert_eq!(run.start_time, FAKE_START_TIME);
+        assert_eq!(run.name, "train");
+    }
+
+    #[tokio::test]
+    async fn test_list_scalars() {
+        let handler = DataProviderHandler;
+        let req = Request::new(data::ListScalarsRequest {
+            experiment_id: "123".to_string(),
+            plugin_filter: Some(data::PluginFilter {
+                plugin_name: "scalars".to_string(),
+                ..Default::default()
+            }),
+            ..Default::default()
+        });
+        let res = handler.list_scalars(req).await.unwrap().into_inner();
+        assert_eq!(res.runs.len(), 1);
+        assert_eq!(res.runs[0].tags.len(), 1);
+        // fake data; don't bother checking the contents
+    }
+
+    #[tokio::test]
+    async fn test_read_scalars() {
+        let handler = DataProviderHandler;
+        let req = Request::new(data::ReadScalarsRequest {
+            experiment_id: "123".to_string(),
+            plugin_filter: Some(data::PluginFilter {
+                plugin_name: "scalars".to_string(),
+                ..Default::default()
+            }),
+            downsample: Some(data::Downsample {
+                num_points: 1000,
+                ..Default::default()
+            }),
+            ..Default::default()
+        });
+        let res = handler.read_scalars(req).await.unwrap().into_inner();
+        assert_eq!(res.runs.len(), 1);
+        assert_eq!(res.runs[0].tags.len(), 1);
+        // fake data; don't bother checking the contents
+    }
+}