rust: implement ListScalars and ReadScalars

[wchargin]

Summary:
The Rust gRPC server now implements the RPCs required to serve scalar
data to the TensorBoard frontend.

There is some repetition in the implementation: ListScalars and
ReadScalars have very similar bodies. There will be more repetition
with the tensor RPCs, and to a lesser degree with those for blob
sequences. I plan to consider refactoring that later, once all the raw
materials are there, but I’m open to feedback.

Test Plan:
Unit tests included. For an end-to-end test, cherry-pick #4356, launch
the data server (bazel run //tensorboard/data/server -- LOGDIR), and
concurrently run TensorBoard with --grpc_data_provider localhost:6806
and --generic_data true flags.

wchargin-branch: rust-listscalars-readscalars

[nfelt]

This could also be pf.unwrap_or_default().plugin_name right? (akin to parse_rtf below). That seems somewhat more straightforward and if I have Done The Thing correctly it looks like it optimizes to the same code: https://rust.godbolt.org/z/YWsrf3

[nfelt]

So I get why we have this: to abstract over the various run-tag-indexed response types and avoid having to do something like res.runs.iter().find(|r| r.run_name == "train").unwrap() inside the test, which would certainly be ugly. That said, I'm a bit lukewarm about the map conversion approach too; it just feels a bit clunky to have to pass in the conversion specifier closures each time in order to unwrap each response, and it's a little hard for me to look at the callsite of unroll_map() and immediately understand what the parameters are doing.

What if we instead we tucked the cumbersome indexing expression inside a tests-only Index impl for the protobuf response types?

impl Index<&str> for data::ListScalarsResponse {
    type Output = data::list_scalars_response::RunEntry;
    fn index(&self, key: &str) -> &Self::Output {
        self.runs.iter().find(|r| r.run_name == key).unwrap()
    }
}

This would allow us in the test itself to just go let train_run = &rsp.runs["train"].

It's true that we would need to repeat this boilerplate once per Run + Tag for each response type, but for test code I don't mind boilerplate so much, and we would only need to write the impls once and then no matter how many test methods we have they all get the nice concise lookup syntax.

Thoughts?

[wchargin]

That’s quite clever. I will also note that it’s also “barely” allowed:
the orphan rules say that you can only impl Trait for Type if either
Trait or Type is defined in your crate, to ensure trait coherence
(i.e., that there aren’t ever conflicting impls). And the Prost bindings
are in our crate, but I don’t view that as set in stone. It’s true
that that’s the formulation that typical prost-build/tonic-build
users will use, but a more Bazel-idiomatic workflow would probably put
each proto_library in its own crate with proper dependencies among
them. And indeed rules_rust provides rust_proto_library of
this form, and bazelbuild/rules_rust#479 implements similar
functionality for Prost/Tonic (written by a Googler).

So I’m a little hesitant to depend on this. We could newtype it, but
that starts getting less ergonomic.

Here’s an approach that avoids having to pass in the conversion
specifier closures, and removes the type parameters entirely. We take
advantage of the shared run_name/tag_name structure via a macro.
Very Pythonic. What do you think?

[nfelt]

Heh, I'll take a "clever" :) Yeah, I realize it depends on the orphan rule; originally I was thinking it would just be a new test-only trait with fn get() but then I realized Index could be used and makes it even pithier.

But I think the macro version is quite nice and addresses the main things I was lukewarm about with unroll_map so I'm perfectly happy to keep it the way you have it now.

[wchargin]

Yeah, I would also be totally fine with a test-only extension trait (no
orphan rule issues since the trait is in-crate). And I would probably
use a macro to implement that trait for each response type.

I still might switch to that. It has type information earlier, so it’s
easier to return a HashMap<Run, HashMap<Tag, V>> given a mapping
function FnMut(Self::TagEntry) -> V. And that would be nice to avoid
the newly introduced .as_ref() dances. You can do it with the macro,
but it’s a bit annoying because of inference details. (Basically: you
can’t write (|x| x.foo())(some_x) without explicitly specifying the
type for x, but you can work around this by capturing x’s type, Java
style: std::iter::once(some_x).map(|x| x.foo()).nth(0).unwrap() is
one awkward way to do so.)

[nfelt]

Do we not want to actually assert anything on the response? Seems like it might not hurt to check that there are indeed no points.

[wchargin]

Well, there are points, because we don’t actually downsample yet. :-)

We can assert on the keys, though, and I’ll add a draft assertion to
enable once downsampling is implemented.

[nfelt]

FWIW, iterating over the filter instead of the keyset would seem likely to save effort on net (assuming single-run lookups against a big experiment are fairly plausible, while sending a filter with many non-existent runs is unlikely; I guess most optimal would be to iterate over the shorter one). Thoughts? It probably does make the code a little messier so maybe better to wait.

[wchargin]

Agreed. It does complicate the code. My first attempt is:

// call site: nice and easy, modulo `&*` (for the `RwLockReadGuard`)
for (run, data) in run_filter.pluck(&*runs) {
    // ...
}

// implementation
impl<T: Hash + Eq> Filter<T> {
    fn pluck<'f, 'm: 'f, V: 'm>(
        &'f self,
        map: &'m HashMap<T, V>,
    ) -> impl Iterator<Item = (&'m T, &'m V)> + 'f
    where
        T: Eq + Hash + 'm,
    {
        match self {
            Filter::All => Either::Left(map.iter()),
            Filter::Just(which) => {
                let it = which.iter().filter_map(move |k| map.get_key_value(k));
                Either::Right(it)
            }
        }
    }
}

enum Either<A, B> {
    Left(A),
    Right(B),
}

impl<A, B, I> Iterator for Either<A, B>
where
    A: Iterator<Item = I>,
    B: Iterator<Item = I>,
{
    type Item = I;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Either::Left(i) => i.next(),
            Either::Right(i) => i.next(),
        }
    }
}

…where the Either layer is needed to form the output type of pluck,
since the iterators in the two branches are of different types: one is
a hash_map::Iter<'m, T, V>, and the other is something like
a iter::FilterMap<hash_set::Iter<'f, T>, C> where C is the
anonymous type of the closure.

(edited: in an earlier version of this comment, I had the right
structure, but forgot to actually make the important change and iterate
over which rather than just iterating over the input map. The main
point stands.)

Another approach involves not using Iterator::filter_map and instead
just implementing our own iterator struct explicitly. This might be
simpler (I haven’t written it out). But it’s still not trivial.

Another approach uses Box<dyn Iterator<Item = (&K, &V)>> in place of
the Either layer, pushing more dispatching to runtime.

I suspect that generators would make this a lot easier, but they’re
not yet stable (and not really close, afaict).

It’s hard for me to write that and then claim with a straight face that
Rust reads like Python. Given that this only affects RPC times (not load
throughput) and that even “large” experiments probably won’t have more
than a few thousand entries in these hash maps, I’m inclined to punt, at
least until we can measure it.

[nfelt]

Thanks for the detailed thoughts! Didn't mean to make this too much of a digression, totally fine to leave as just a potential TODO for now. Maybe if I have any particularly good ideas I'll take a stab at this myself :)

[wchargin]

Yep, TODO included. And, sure, feel free to grab it. :-)

[nfelt]

Maybe worth a TODO to track this on TimeSeries? Would be nice not to have to rescan every single scalar wall time each time we hit ListScalars. I guess it's hard to get this entirely for "free" on TimeSeries due to preemptions, but still seems overall cheaper to compute it there (where we could take a cheap path if no preemption happened).

[wchargin]

Done (as TODO). Agreed on all counts.

[wchargin]

PTAL re: run_tag_map!; the rest is mostly just done.

[nfelt]

Heh, I'll take a "clever" :) Yeah, I realize it depends on the orphan rule; originally I was thinking it would just be a new test-only trait with fn get() but then I realized Index could be used and makes it even pithier.

But I think the macro version is quite nice and addresses the main things I was lukewarm about with unroll_map so I'm perfectly happy to keep it the way you have it now.

[nfelt]

Thanks for the detailed thoughts! Didn't mean to make this too much of a digression, totally fine to leave as just a potential TODO for now. Maybe if I have any particularly good ideas I'll take a stab at this myself :)