[CW/R2] Tensorstore `Malformed StorageGeneration`

[ravwojdyla-agent]

Describe the bug

The cw-ci-test integration test fails with a tensorstore Malformed StorageGeneration error when reading tokenized zarr3 data from R2 (S3-compatible) storage. The failure is in the quickstart-tests/tokenized_2fdf2156 step of the marin-itest pipeline.

This failure is unrelated to the PR under test (#4370, K8s stream_logs OOM fix) — it appears to be a flaky interaction between tensorstore and R2.

To Reproduce

1. Push any PR that triggers the cw-ci-test workflow.
2. The marin-itest pipeline submits a quickstart tokenization job.
3. A subsequent step tries to open the tokenized zarr3 array and fails with Malformed StorageGeneration.

Expected behavior

The tokenized zarr3 array should be readable after the tokenization step completes.

Additional context

Error from CI run:

ValueError: Error opening "zarr3" driver: Error reading
  "temp/ci/marin-itest-22712054/quickstart-tests/tokenized-20383a/train/part-00000-of-00001/input_ids/offsets/zarr.json":
  Malformed StorageGeneration

The tensorstore spec uses "recheck_cached_data": false, which may cause stale cache entries to collide with freshly written data. The kvstore driver is s3 pointed at the R2 endpoint (74981a43be0de7712369306c7b19133d.r2.cloudflarestorage.com, bucket marin-na).

Possible causes:

• R2 returning an unexpected ETag format that tensorstore cannot parse as a StorageGeneration.
• A race between the writer finishing and the reader opening the array when caching is aggressive (recheck_cached_data: false).
• Transient R2 API inconsistency during the CI window.

[ravwojdyla]

This is definitely not a CI fluke, this caused 4 step failures on nemotron tokenization (out of 6), here's one example:

E20260402 08:07:01 123469092288320 marin.execution.step_runner Step failed: tokenized/nemotron_cc/low_synth_8fac7840
Traceback (most recent call last):
  File "/app/lib/marin/src/marin/execution/step_runner.py", line 153, in _harvest
    status = handle.wait(raise_on_failure=True)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/fray/src/fray/v2/local_backend.py", line 65, in wait
    self._future.result(timeout=timeout)
  File "/uv/cache/python/cpython-3.11.14-linux-x86_64-gnu/lib/python3.11/concurrent/futures/_base.py", line 449, in result
    return self.__get_result()
           ^^^^^^^^^^^^^^^^^^^
  File "/uv/cache/python/cpython-3.11.14-linux-x86_64-gnu/lib/python3.11/concurrent/futures/_base.py", line 401, in __get_result
    raise self._exception
  File "/uv/cache/python/cpython-3.11.14-linux-x86_64-gnu/lib/python3.11/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/marin/src/marin/execution/step_runner.py", line 281, in worker_fn
    run_step(step)
  File "/app/lib/marin/src/marin/execution/step_runner.py", line 328, in run_step
    result = step.fn(output_path)  # pyrefly: ignore[not-callable]
             ^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/marin/src/marin/execution/executor.py", line 700, in resolved_fn
    return captured_fn(captured_config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/marin/src/marin/processing/tokenize/tokenize.py", line 448, in tokenize
    run_pipeline(ctx, train_groups, "train")
  File "/app/lib/marin/src/marin/processing/tokenize/tokenize.py", line 414, in run_pipeline
    ledger = consolidate_shard_caches(
             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/levanter/src/levanter/store/cache.py", line 446, in consolidate_shard_caches
    per_shard_sizes = list(executor.map(_get_data_sizes, shard_cache_paths))
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/uv/cache/python/cpython-3.11.14-linux-x86_64-gnu/lib/python3.11/concurrent/futures/_base.py", line 619, in result_iterator
    yield _result_or_cancel(fs.pop())
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/uv/cache/python/cpython-3.11.14-linux-x86_64-gnu/lib/python3.11/concurrent/futures/_base.py", line 317, in _result_or_cancel
    return fut.result(timeout)
           ^^^^^^^^^^^^^^^^^^^
  File "/uv/cache/python/cpython-3.11.14-linux-x86_64-gnu/lib/python3.11/concurrent/futures/_base.py", line 456, in result
    return self.__get_result()
           ^^^^^^^^^^^^^^^^^^^
  File "/uv/cache/python/cpython-3.11.14-linux-x86_64-gnu/lib/python3.11/concurrent/futures/_base.py", line 401, in __get_result
    raise self._exception
  File "/uv/cache/python/cpython-3.11.14-linux-x86_64-gnu/lib/python3.11/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/levanter/src/levanter/store/cache.py", line 442, in _get_data_sizes
    store = TreeStore.open(exemplar, shard_path, mode="r", cache_metadata=True)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/levanter/src/levanter/store/tree_store.py", line 62, in open
    tree = _construct_builder_tree(exemplar, path, mode, cache_metadata)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/levanter/src/levanter/store/tree_store.py", line 193, in _construct_builder_tree
    return jtu.tree_map_with_path(open_builder, exemplar, is_leaf=heuristic_is_leaf)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/.venv/lib/python3.11/site-packages/jax/_src/tree_util.py", line 1239, in tree_map_with_path
    return treedef.unflatten(f(*xs) for xs in zip(*all_keypath_leaves))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/.venv/lib/python3.11/site-packages/jax/_src/tree_util.py", line 1239, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_keypath_leaves))
                             ^^^^^^
  File "/app/lib/levanter/src/levanter/store/tree_store.py", line 185, in open_builder
    return JaggedArrayStore.open(
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/levanter/src/levanter/store/jagged_array.py", line 177, in open
    offsets = _ts_open_sync(offset_path, jnp.int64, [1], mode=mode)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/lib/levanter/src/levanter/store/jagged_array.py", line 624, in _ts_open_sync
    raise e
  File "/app/lib/levanter/src/levanter/store/jagged_array.py", line 619, in _ts_open_sync
    ).result()
      ^^^^^^^^
ValueError: Error opening "zarr3" driver: Error reading "marin/rav/tokenized/nemotron_cc/low_synth/train/part-00000-of-01788/input_ids/offsets/zarr.json": Malformed StorageGeneration [source locations='tensorstore/internal/cache/kvs_backed_cache.h:220\ntensorstore/driver/kvs_backed_chunk_driver.cc:1124\ntensorstore/driver/driver.cc:116'] [tensorstore_spec='{\"context\":{\"aws_credentials\":{\"type\":\"default\"},\"cache_pool\":{\"total_bytes_limit\":1000000000},\"data_copy_concurrency\":{},\"s3_request_concurrency\":{},\"s3_request_retries\":{}},\"driver\":\"zarr3\",\"dtype\":\"int64\",\"kvstore\":{\"aws_region\":\"us-east-1\",\"bucket\":\"marin-na\",\"driver\":\"s3\",\"endpoint\":\"<TRUNCATED>\",\"path\":\"marin/rav/tokenized/nemotron_cc/low_synth/train/part-00000-of-01788/input_ids/offsets/\"},\"metadata\":{\"chunk_grid\":{\"configuration\":{\"chunk_shape\":[134217728]},\"name\":\"regular\"},\"codecs\":[{\"configuration\":{\"chunk_shape\":[262144],\"codecs\":[{\"configuration\":{\"clevel\":5},\"name\":\"blosc\"}]},\"name\":\"sharding_indexed\"}],\"node_type\":\"array\"},\"recheck_cached_data\":\"open\",\"schema\":{\"domain\":{\"exclusive_max\":[18014398509481984],\"inclusive_min\":[0]}},\"transform\":{\"input_exclusive_max\":[[18014398509481984]],\"input_inclusive_min\":[0]}}']

[ravwojdyla]

Ok - so I can confirm that I can take the tensorstore spec, open and read some of that data, so this must be some kind of race condition?

[ravwojdyla]

Just a random/unrelated pearl from https://google.github.io/tensorstore/kvstore/s3/index.html#json-kvstore/s3:

This non-atomicity can lead to unexpected behavior when writing to an S3-backed TensorStore. For example, writing to a zarr array can in some cases lead to a delete rather than a write (if it happens to match the fill value)

[hammer]

Related? #4433

[ravwojdyla]

@hammer yep, thanks for linking that. Will push a general fix and revert #4436. A general solution here is to NOT use our process-level cache for levanter consolidation during tokenization, the cache on R2 gets stale and causes the Malformed StorageGeneration.