rust: implement RPC-time downsampling (#4403)

Summary:
A new `downsample` function performs TensorBoard-style downsampling of a
vector: random subsequence, relative order preserved, last element
always kept. The RPC server uses this to honor the downsampling
specifications for `ReadScalars` RPCs.

Test Plan:
Unit tests included for both the utility function and the RPC.

wchargin-branch: rust-downsample

Author: wchargin

tensorboard/data/server/BUILD

@@ -29,6 +29,7 @@ rust_library(
         "lib.rs",
         "commit.rs",
         "data_compat.rs",
+        "downsample.rs",
         "event_file.rs",
         "logdir.rs",
         "masked_crc.rs",

tensorboard/data/server/downsample.rs

@@ -0,0 +1,109 @@
+/* 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.
+==============================================================================*/
+
+//! Downsampling for in-memory sequences.
+
+use rand::SeedableRng;
+use rand_chacha::ChaCha20Rng;
+
+/// Downsamples `xs` in place to contain at most `k` elements, always including the last element.
+///
+/// If `k == 0`, then `xs` is cleared. If `k >= xs.len()`, then `xs` is returned unchanged.
+/// Otherwise, a uniformly random sample of `k - 1` of the first `n - 1` elements of `xs` is chosen
+/// and retained, and the final element of `xs` is also retained. The relative order of elements of
+/// `xs` is unchanged.
+///
+/// More declaratively: among all subsequences of `xs` of length `min(k, xs.len())` that include
+/// the last element, one is selected uniformly at random, and `xs` is updated in place to
+/// represent that subsequence.
+///
+/// The random number generator is initialized with a fixed seed, so this function is
+/// deterministic.
+pub fn downsample<T>(xs: &mut Vec<T>, k: usize) {
+    let n = xs.len();
+    if k == 0 {
+        xs.clear();
+        return;
+    }
+    if k >= n {
+        return;
+    }
+
+    let mut rng = ChaCha20Rng::seed_from_u64(0);
+
+    // Choose `k - 1` of the `n - 1` indices to keep, and move their elements into place. Then,
+    // move the last element into place and drop extra elements.
+    let mut indices = rand::seq::index::sample(&mut rng, n - 1, k - 1).into_vec();
+    indices.sort_unstable();
+    for (dst, src) in indices.into_iter().enumerate() {
+        xs.swap(dst, src);
+    }
+    xs.swap(k - 1, n - 1);
+    xs.truncate(k);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Clones `xs` and [`downsample`]s the result to `k` elements.
+    fn downsample_cloned<T: Clone>(xs: &Vec<T>, k: usize) -> Vec<T> {
+        let mut ys = xs.clone();
+        downsample(&mut ys, k);
+        ys
+    }
+
+    #[test]
+    fn test_deterministic() {
+        let xs: Vec<char> = "abcdefg".chars().collect();
+        let expected = downsample_cloned(&xs, 4);
+        assert_eq!(expected.len(), 4);
+        for _ in 0..100 {
+            assert_eq!(downsample_cloned(&xs, 4), expected);
+        }
+    }
+
+    #[test]
+    fn test_ok_when_k_greater_than_n() {
+        let xs: Vec<char> = "abcdefg".chars().collect();
+        assert_eq!(downsample_cloned(&xs, 10), xs);
+        assert_eq!(downsample_cloned(&xs, usize::MAX), xs);
+    }
+
+    #[test]
+    fn test_inorder_plus_last() {
+        let xs: Vec<u32> = downsample_cloned(&(0..10000).collect(), 100);
+        let mut ys = xs.clone();
+        ys.sort();
+        assert_eq!(xs, ys);
+        assert_eq!(xs.last(), Some(&9999));
+    }
+
+    #[test]
+    fn test_zero_k() {
+        for n in 0..3 {
+            let xs: Vec<u32> = (0..n).collect();
+            assert_eq!(downsample_cloned(&xs, 0), Vec::<u32>::new());
+        }
+    }
+
+    #[test]
+    fn test_zero_n() {
+        let xs: Vec<u32> = vec![];
+        for k in 0..3 {
+            assert_eq!(downsample_cloned(&xs, k), Vec::<u32>::new());
+        }
+    }
+}

tensorboard/data/server/lib.rs

@@ -19,6 +19,7 @@ limitations under the License.
 
 pub mod commit;
 pub mod data_compat;
+pub mod downsample;
 pub mod event_file;
 pub mod logdir;
 pub mod masked_crc;

tensorboard/data/server/server.rs

@@ -17,12 +17,14 @@ use futures_core::Stream;
 use std::borrow::Borrow;
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::convert::TryInto;
 use std::hash::Hash;
 use std::pin::Pin;
 use std::sync::{RwLock, RwLockReadGuard};
 use tonic::{Request, Response, Status};
 
 use crate::commit::{self, Commit};
+use crate::downsample;
 use crate::proto::tensorboard::data;
 use crate::types::{Run, Tag, WallTime};
 use data::tensor_board_data_provider_server::TensorBoardDataProvider;
@@ -158,7 +160,7 @@ impl TensorBoardDataProvider for DataProviderHandler {
         let req = req.into_inner();
         let want_plugin = parse_plugin_filter(req.plugin_filter)?;
         let (run_filter, tag_filter) = parse_rtf(req.run_tag_filter);
-        let _downsample = parse_downsample(req.downsample)?; // TODO(@wchargin): Use `downsample`.
+        let num_points = parse_downsample(req.downsample)?;
         let runs = self.read_runs()?;
 
         let mut res: data::ReadScalarsResponse = Default::default();
@@ -183,11 +185,13 @@ impl TensorBoardDataProvider for DataProviderHandler {
                     continue;
                 }
 
-                let n = ts.valid_values().count();
+                let mut points = ts.valid_values().collect::<Vec<_>>();
+                downsample::downsample(&mut points, num_points);
+                let n = points.len();
                 let mut steps = Vec::with_capacity(n);
                 let mut wall_times = Vec::with_capacity(n);
                 let mut values = Vec::with_capacity(n);
-                for (step, wall_time, &commit::ScalarValue(value)) in ts.valid_values() {
+                for (step, wall_time, &commit::ScalarValue(value)) in points {
                     steps.push(step.into());
                     wall_times.push(wall_time.into());
                     values.push(value);
@@ -278,7 +282,7 @@ fn parse_rtf(rtf: Option<data::RunTagFilter>) -> (Filter<Run>, Filter<Tag>) {
 }
 
 /// Parses `Downsample.num_points` from a request, failing if it's not given or invalid.
-fn parse_downsample(downsample: Option<data::Downsample>) -> Result<i64, Status> {
+fn parse_downsample(downsample: Option<data::Downsample>) -> Result<usize, Status> {
     let num_points = downsample
         .ok_or_else(|| Status::invalid_argument("must specify downsample"))?
         .num_points;
@@ -288,7 +292,13 @@ fn parse_downsample(downsample: Option<data::Downsample>) -> Result<i64, Status>
             num_points
         )));
     }
-    Ok(num_points)
+    num_points.try_into().map_err(|_| {
+        Status::out_of_range(format!(
+            "num_points ({}) is too large for this system; max: {}",
+            num_points,
+            usize::MAX
+        ))
+    })
 }
 
 /// A predicate that accepts either all values or just an explicit set of values.
@@ -575,8 +585,6 @@ mod tests {
         let map = run_tag_map!(res.runs);
         let train_run = &map[&Run("train".to_string())];
         let xent_data = &train_run[&Tag("xent".to_string())].data.as_ref().unwrap();
-        // TODO(@wchargin): Enable once downsampling is implemented.
-        // assert_eq!(xent_data.value.len(), 0);
-        let _ = xent_data;
+        assert_eq!(xent_data.value, Vec::new());
     }
 }