tsgrpc_kvstore driver - how to signal `generation_if_equal` mismatch on write?

[Mononofu]

Hi,

I'm trying to implement a gRPC service for the tsgrpc_kvstore driver, and trying to understand how to correctly support the generation_if_equal option for WriteRequest. In particular, how to handle the case where the generation does not match - the comment in the proto says "the write is aborted" on a mismatch, what exactly does this mean?

Both returning an error from the RPC and returning an error in the WriteResponse.status field leads to the TensorStore write to fail immediately, which leads me to believe this is not what I should do?

Looking at PerformWriteback I see that the transaction is re-tried if the returned generation is "unknown" (== empty). However if I try this, TensorStore just sends writes with generation_if_equal=StorageGeneration::NoValue() in a loop, instead of what I would expect:

1. Attempt a write with an old generation, receive StorageGeneration::Unknown
2. Issue a read to obtain the current value of the file
3. Issue a new write, with the update generation and value

thank you!

PS: I'm also a bit unclear what I should use for the timestamp value of GenerationAndTimestamp - is this intended to be the current time of the gRPC server? The last modified time of the file?

[laramiel]

The easiest references are likely the memory and file kvstore drivers.

https://github.com/google/tensorstore/blob/master/tensorstore/kvstore/memory/memory_key_value_store.cc#L401

tensorstore/tensorstore/kvstore/file/file_key_value_store.cc

Line 602 in 933712f

struct WriteTask {

The if_equal condition does not mach, then the return value should be { timestamp, Unknown() }.
The timestamp is essentially a point in time at which the server has read the generation.

As for the transaction behavior, what does your testcase look like? You should probably start with a testcase like the tsgrpc/kvstore_server_test, but that has not, I think, been updated to include all of the newer tests.

[Mononofu]

Thanks for the quick response!

Unless I missed something, then it seems I'm already returning the correct value:

// StorageGeneration::Unknown(), used to indicate writes that failed due to
// generation mismatch.
// https://github.com/google/tensorstore/blob/933712f0628f9170ed00eafdd561402eb6fb77fe/tensorstore/kvstore/generation.h#L199
const GENERATION_UNKNOWN: &str = "";


WriteOutput::EtagMismatch(_) => WriteResponse {
    generation_and_timestamp: Some(GenerationAndTimestamp {
        // GENERATION_UNKNOWN is used to indicate that the write should be retried.
        // https://github.com/google/tensorstore/blob/933712f0628f9170ed00eafdd561402eb6fb77fe/tensorstore/kvstore/transaction.cc#L135
        generation: Bytes::from_owner(GENERATION_UNKNOWN.to_string()),
        timestamp: Some(prost_types::Timestamp::from(SystemTime::now())),
    }),
    ..Default::default()
}

I'll have a look at the test cases you mentioned!

[Mononofu]

The test code in question is very simple:

dataset = await ts_open(ts_port, path, create=True)
# this write succeeds
dataset[0, 0:4].write([8, 7, 9, 10]).result()

# this write hangs forever; dumping the RPC logs 
# I observe the infinite write retry loop described above
dataset[0, 4:8].write([1, 2, 3, 4]).result()

[laramiel]

You can try your test with the provided tsgrpc/kvstore_server and see what happens.

One note: you picked a timestamp after you checked the generation, when you should have picked a timestamp before you checked the generation.

[Mononofu]

I think that's the same as the memory example you linked though? It uses the current time?

[Mononofu]

Ok I think I found the issues!

1. It seems TensorStore expects the first byte of the generation to be 0, but the gRPC driver just directly copies whatever value we return instead of handling this invariant.
2. TensorStore expects generations to be different even if the same file content is later written back to the file, i.e. hash(file_contents) is not a valid generation.

Not sure how easy it is to detect violations of (2), but at least for (1) it would be useful to either handle this in the gRPC driver or raise a clear error message.

[jbms]

Ok I think I found the issues!

1. It seems TensorStore expects the first byte of the generation to be 0, but the gRPC driver just directly copies whatever value we return instead of handling this invariant.

We probably should indeed just handle this in tensorstore.

2. TensorStore expects generations to be different even if the same file content is later written back to the file, i.e. hash(file_contents) is not a valid generation.

Hash(filecontents) is okay actually.

Not sure how easy it is to detect violations of (2), but at least for (1) it would be useful to either handle this in the gRPC driver or raise a clear error message.

[Mononofu]

Hash(filecontents) is okay actually.

Oh, that would be great! The current comment implies that this is not allowed, is the comment out of date?

[laramiel]

I think that's the same as the memory example you linked though? It uses the current time?

Somewhat. the timestamp indicates a point at which the version does not need to be rechecked.
In memory kvstore there is a lock, so any time before the function returns is equivalent.

[laramiel]

Ok I think I found the issues!

1. It seems TensorStore expects the first byte of the generation to be 0, but the gRPC driver just directly copies whatever value we return instead of handling this invariant.

The StorageGeneration has internal structure and it is generally expected that the value is constructed using one named constructors such as FromString, etc..

Hash(filecontents) is okay actually.

Oh, that would be great! The current comment implies that this is not allowed, is the comment out of date?

I don't think that we actually have tests to verify whether using the content-hash will work properly. There are a few things going on there since we try to elide identical writes.

Conceptually it should work, as the potentially problematic case of an observed sequence of generation hashes {H1, H2, H3, H1' }, any retries expecting writes expecting H1 and seeing H1' should be ok since either they are part of a transaction which only modifies that block, or all the blocks have the proper hash, or other attempted writes in the transaction would fail.

[Mononofu]

The named constructors make sense for C++, but unfortunately I can't easily use them as I'm writing the gRPC service in Rust.

With the adjustments as discussed above I now see correct behaviour for simple uncontended writes and reads. However I see data loss for concurrent but non-overlapping writes to the same tensor:

def simulated_learner(learner_i: int, path: str, ports: Sequence[int]):
    with make_kv_service(ports) as ts_port:
        config = ts_config(ts_port, path, create=False)
        dataset = ts.open(config).result()
        bs, seq_len = dataset.shape


        def do_write(i: int):
            config = ts_config(ts_port, path, create=False)
            dataset = ts.open(config).result()
            dataset[0, i].write(i).result()

        # Concurrently write to non-overlapping sections of the tensor.
        indices = [i for i in range(seq_len) if i % n_writers == learner_i]
        with concurrent.futures.ThreadPool() as pool:
            pool.map(do_write, indices)

If I read the resulting array afterwards I observe incorrect 0 entries where the writes got lost, while the expected value is np.arange(seq_len).

While I double check that my implementation is correct, could I confirm that concurrent writes to distinct regions of the same tensor are supported for the tsgrpc_kvstore driver? I do see correct behaviour with the local file driver.

[Mononofu]

Looking carefully at the logs of the gRPC requests that my server receives from the tsgrpc_kvstore driver, I see the following pattern:

• incoming ReadRequest { key: b"bar_tensor/c/0/0", generation_if_not_equal: b"", generation_if_equal: b"", staleness_bound: Some(Timestamp { seconds: -9223372036854775808, nanos: 0 }), byte_range: None }
• file does not yet exist, so my server responds with ReadResponse { state: State::Missing }
• incoming WriteRequest { key: b"bar_tensor/c/0/0", generation_if_equal: b"" }, i.e. generation_if_equal=StorageGeneration::Unknown()

Isn't this incorrect? I would expect the write request to specify generation_if_equal=StorageGeneration::NoValue(), i.e. generation_if_equal=b"\2". Otherwise if we get two concurrent writes they will both succeed, overwriting each other.

[Mononofu]

I noticed that I was not setting any generation in the ReadResponse; if I change this to also specify StorageGeneration::NoValue() then any subsequent write requests will include generation_if_equal=b"\0\2".

Note the unexpected extra 0 at the beginning. However at least I can work around this by checking for both b"\0" and b"\0\2" and treating both of them as StorageGeneration::NoValue().

With this I no longer see any data loss, all writes seem to get retried correctly if there are concurrent modifications!

Thank you for all the help!