rust: add GCS logdir support

[wchargin]

Summary:
You can now run tensorboard --load_fast --logdir gs://... to launch a
Rust-backed TensorBoard reading from GCS. This patch introduces two new
Logdir implementations: one that connects to the GCS client, and one
that delegates to either a DiskLogdir or a gcs::Logdir depending on
the form of the --logdir flag.

Test Plan:
Both //tensorboard and the

bazel run //tensorboard -- \
    --logdir gs://tensorboard-bench-logs/edge_cgan --load_fast

bazel run //tensorboard/data/server:bench -- \
    --logdir gs://tensorboard-bench-logs/edge_cgan

wchargin-branch: rust-gcs-logdir

[wchargin]

@nfelt, re: discussion: you can see from the structure of these types
and implementations how they could be easily macroed into existence—all
except DynLogdir::new, of course. It’s basically “the same thing with
more $ signs”:
https://gist.github.com/wchargin/f44c370bebb61210dbeb11037f0b2d99

I think that I have a weak preference for the explicit implementation,
but the macro’s actually not that bad. It does make new cases a one-line
change. Open to suggestions.

[nfelt]

Explicit seems clearer to me given especially that we only have 2 implementations now.

What I had in mind when talking about "it would be nice to have a more automatic way to do this" was really more about whether there was something that would generalize not just to additional implementations but also cover the full set of trait methods as well (i.e. the dispatching logic for each method would be autogenerated without even having to include the macro template for each method). So it'd be sort of a compromise position between enums and trait objects, where the downside is that you need both a closed set of types known to the compiler and the trait must be object safe, but the upside is that you have no vtable overhead and you don't have to hand-roll the dispatching code for each trait method.

I'm not even sure that would work for this case where we have the associated type as a complicating factor, but it seems like the sort of thing that would in principle be possible to autogenerate in the basic case.

But this is just hypothesizing; no changes required or expected.

[wchargin]

@nfelt: Friendly ping? This blocks #4647, and I would like to get these
merged with a bit of buffer to release a data server with GCS support
and teach tb-nightly to depend on it.

[nfelt]

Explicit seems clearer to me given especially that we only have 2 implementations now.

What I had in mind when talking about "it would be nice to have a more automatic way to do this" was really more about whether there was something that would generalize not just to additional implementations but also cover the full set of trait methods as well (i.e. the dispatching logic for each method would be autogenerated without even having to include the macro template for each method). So it'd be sort of a compromise position between enums and trait objects, where the downside is that you need both a closed set of types known to the compiler and the trait must be object safe, but the upside is that you have no vtable overhead and you don't have to hand-roll the dispatching code for each trait method.

I'm not even sure that would work for this case where we have the associated type as a complicating factor, but it seems like the sort of thing that would in principle be possible to autogenerate in the basic case.

But this is just hypothesizing; no changes required or expected.

[nfelt]

Just wondering, did we determine that the buffer is a bottleneck below this size? It's just a significant chunk of RAM when we add it up across many files, so my default inclination would probably be to start out smaller with something like 8-16 MiB. Or maybe at least worth making it tunable with an env var or something (like the TF GCS readahead buffer was).

[wchargin]

Nope, I haven’t analyzed this. I picked 128 MiB because we use that
number for some similar readahead buffers internally, but I’m happy to
tone it down to 16 MiB for now. Agreed that it should be tunable, and
I’m happy to run benchmark sweeps, too; just don’t want to introduce too
much right off the bat.

With a quick unscientific analysis, I can fully load edge_cgan from
GCS in about 18 seconds with a 128 MiB buffer or 20 seconds with 16 MiB.
The standard deviation here is on the order of seconds in each case, so
it seems fine to proceed conservatively.

[nfelt]

extra line?

[wchargin]

Fixed; thanks.

[nfelt]

I may not be reading this right but it looks like we have 2 copies here - one inside gcs::Client where we do body.to_vec() to go from Bytes to Vec<u8> and then another one here where we go from Vec<u8> to our final [u8].

I doubt it's really a big deal performance-wise but if that is indeed happening, it seems slightly nicer if we could do only one copy? E.g. if we were to have our client directly return Bytes rather than Vec<u8> it seems like we could directly call copy_from_slice(Bytes) since Bytes is Borrow<[u8]>?

[wchargin]

Indeed! I was planning to do this as an optimization in a follow-up,
since using Bytes directly requires adding a dep on the bytes crate.
But it’s well spotted, and a perfectly appropriate use of Bytes.

I don’t see an easy way to zero-copy the bytes all the way from the HTTP
response into the io::Read output: e.g., neither reqwest nor hyper
seems to expose a response.bytes_into(&mut buf).

There is a similar large copy in the ReadBlob RPC handler, which in
previous drafts I had replaced with an Arc<[u8]>, but Bytes might be
better there, too. (After all, Bytes is “Arc<[u8]> on steroids”.)

[nfelt]

Gotcha, fine by me to do it in a follow-up.

I may be misunderstanding but isn't it impossible to do a zero-copy transfer of HTTP response bytes into an &mut [u8] which is what io::Read takes? Since the signature requires that we stuff the bytes into a mutable buffer we don't control? Unless what you're saying is that we can actually change the underlying buffer pointed to by buf to literally just point to the one occupied by Bytes, but my understanding was that we can only mutate the contents of buf, not repoint it entirely.

[wchargin]

I agree that I don’t see a way to do it with the reqwest/et al. APIs,
but you could imagine that it could be possible if each of the layers of
the stack supported it. Something like:

// mod reqwest::blocking::response
impl Response {
    /// Streams some bytes from the response body directly into the
    /// given buffer. Returns the number of bytes written.
    fn bytes_into<T: BufMut>(&self, x: &mut T) -> crate::Result<usize>;
}

// mod gcs::client
impl Client {
    fn read(&self, bucket, object, buf: impl BufMut) -> self::Result<usize> {
        let res = self.http.get(...).send()?;
        res.bytes_into(&mut buf)
    }
}

// mod gcs::logdir
impl Read for File {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let n = self.gcs.read(bucket, obj, buf)?;
        self.pos += n as u64;
        Ok(n)
    }
}

Or, from a lower-level perspective: at the end of the day, we receive
network data by calling recvfrom(2) and passing the kernel a mutable
pointer to a user-space buffer into which to write the received data, so
there’s no fundamental reason that that pointer can’t point into the
buffer given to io::Read::read (at least, as far as I can tell).

but my understanding was that we can only mutate the contents of buf, not repoint it entirely.

Correct, and that is indeed not what I am alluding to.

[wchargin]

thanks!

[wchargin]

Nope, I haven’t analyzed this. I picked 128 MiB because we use that
number for some similar readahead buffers internally, but I’m happy to
tone it down to 16 MiB for now. Agreed that it should be tunable, and
I’m happy to run benchmark sweeps, too; just don’t want to introduce too
much right off the bat.

With a quick unscientific analysis, I can fully load edge_cgan from
GCS in about 18 seconds with a 128 MiB buffer or 20 seconds with 16 MiB.
The standard deviation here is on the order of seconds in each case, so
it seems fine to proceed conservatively.

[wchargin]

Fixed; thanks.

[wchargin]

Indeed! I was planning to do this as an optimization in a follow-up,
since using Bytes directly requires adding a dep on the bytes crate.
But it’s well spotted, and a perfectly appropriate use of Bytes.

I don’t see an easy way to zero-copy the bytes all the way from the HTTP
response into the io::Read output: e.g., neither reqwest nor hyper
seems to expose a response.bytes_into(&mut buf).

There is a similar large copy in the ReadBlob RPC handler, which in
previous drafts I had replaced with an Arc<[u8]>, but Bytes might be
better there, too. (After all, Bytes is “Arc<[u8]> on steroids”.)