[OutstanderWang]

Problem

Workloads that read one large payload from many keys store it once per key. In GRPO, for example, each prompt video is read by its n responses, so putting the video into every response key multiplies its memory and put traffic by n.

The same workloads also suffer from unbalanced storage units. SimpleStorage stores global_index on unit global_index % num_data_storage_units, and indexes were handed out globally in sequence. When puts alternate between two partitions on two units, all four rows of one partition land on unit 0 ([0, 0, 0, 0]). Large shared payloads make this imbalance much worse.

Change

1. Balance each partition's rows across SimpleStorage units

• PartitionIndexManager counts each partition's rows per storage unit. Each new index goes to the unit where that partition has the fewest rows. On a tie it prefers a reusable index, then the counter's next index. The example above now spreads as [0, 0, 1, 1].
• Indexes skipped to reach a unit wait in that unit's reuse pool. Each pool is FIFO.
• The number of units comes from SimpleStorage.num_data_storage_units in store_config. Other backends, and SimpleStorage with one unit, keep the previous FIFO allocation.
• Controller checkpoints keep the flat reusable_indexes list, which is regrouped by unit on load.

2. Share payloads through __public__ and kv_link

• tq.PUBLIC_PARTITION ("__public__") is an ordinary partition that holds shared payloads. It is written and read with the usual KV APIs, and change 1 balances its rows like any other partition's.
• tq.kv_link(keys, partition_id, public_key, fields) makes fields of keys read the payload of public_key without copying it. public_key is one key for all keys, or a list aligned with keys to link a batch of prompts at once. Missing keys are created. The link is a single controller request that checks everything first: either every key is linked or none is, and a refused call creates no key.
• Linked fields take the public key's dtype, shape, and readiness. A linked public key is read-only: kv_put, kv_update and clears on it are refused until no key links to it, so linking keys never see their shared payload change.
• tq.kv_clear_unlinked_public(stale_keys=None) clears every public key that keys linked to before but none links to now, and returns the cleared keys. A public key that was put but never linked is left alone, so a producer between its kv_put and kv_link is safe. Pass such keys in stale_keys once their producer is given up on, and they are cleared too if nothing links to them.
• tq.kv_exists(keys, partition_id) maps each existing key to its ready fields, linked ones included, and leaves missing keys out without logging. It is the probe for "has another worker already put this shared prompt?", where kv_describe would raise and kv_batch_get would warn.

import transfer_queue as tq

# One copy of the video, shared by the n responses of prompt 0.
tq.kv_put(key="video_0", partition_id=tq.PUBLIC_PARTITION, fields={"video": video})
keys = [f"traj_0_{i}" for i in range(n)]
tq.kv_link(keys, partition_id="train", public_key="video_0", fields="video")
tq.kv_batch_put(keys=keys, partition_id="train", fields=TensorDict({"response": responses}, batch_size=n))

batch = tq.kv_batch_get(keys=keys, partition_id="train")  # "video" and "response" for every key

# After the step: clear the trajectories, then free the videos nobody reads any more.
tq.kv_clear(keys, partition_id="train")
tq.kv_clear_unlinked_public()  # ["video_0"]

3. Link, write and clear through BatchMeta

Rows without keys can share payloads too. TransferQueueClient.link(metadata, public_metadata, fields) (and async_link) links the rows of metadata to the aligned rows of public_metadata and returns a BatchMeta whose _links point at the public rows. It refuses what kv_link refuses: linking into __public__, misaligned rows, a public field not produced yet, and rows that already store the field.

client = tq.get_client()
public_meta = client.put(TensorDict({"video": video.unsqueeze(0)}, batch_size=1), partition_id=tq.PUBLIC_PARTITION)
rows_meta = client.put(TensorDict({"response": responses}, batch_size=n), partition_id="train")

linked = client.link(rows_meta, public_meta.select_samples([0] * n), "video")
batch = client.get_data(linked)  # "video" from the public row, "response" from each row

client.clear_samples(linked)
tq.kv_clear_unlinked_public()  # frees the public row although it has no key

• A put through BatchMeta into __public__ asks the controller first and is refused while other rows link to the target rows, so a linked public row is read-only on both write paths.
• kv_clear_unlinked_public sweeps public rows by index, so keyless rows are freed as well. It still returns only the keys it cleared.

4. data_parser on linked fields

kv_batch_get(..., data_parser=..., parser_args=...) used to refuse linked fields. A linked field is now parsed on the storage unit that holds its public row, so only the parsed part crosses the network, e.g. the frames each trajectory needs instead of the whole video. The public row is read and parsed once per distinct parser_args entry among the keys linking to it: keys with equal plain entries (scalars, and lists or tuples of them, compared by type and value) share one parse, and any other entry, such as a tensor, is parsed once per key. Without parser_args, each public row is parsed once.

How links are implemented:

• The controller records links in DataPartitionStatus.links ({global_index: {field: public global_index}}) and indexes them by source in linked_by, so clear checks scale with the size of the clear.
• BatchMeta carries the links to the storage manager. AsyncSimpleStorageManager.get_data reads each linked public row once, however many keys in the batch link to it, or with data_parser once per distinct parser_args entry.

How other APIs behave with links

API	Behavior
kv_batch_get, get_data, samplers, StreamingDataLoader	Read linked fields from the public row
kv_put / put through BatchMeta on a linked field	Refused before anything is written; leave linked fields out of a read-modify-write, or kv_clear(fields=...) first to store a private value
kv_update on a linked field	Refused before anything is written
kv_put / kv_update / put through BatchMeta on a linked public key or row	Refused before anything is written
kv_clear of a linking key or field	Drops only the link
kv_clear, clear_samples, clear_partition on a public key that is still linked	Refused; the error names the linking keys
kv_link on a field the key already stores itself	Refused; kv_clear(fields=...) first, so no unread private copy stays in storage
kv_describe	fields lists every readable field, linked ones included; unfinalized its not-ready subset; links maps each linked field to its public key
dump_data_by_key / load_data_by_key	Links are dumped by public key name in format version 5, which older readers reject. Dump the public keys too and load them first; the load checks that their fields are ready, relinks, then writes the payloads
save_checkpoint / load_checkpoint	Links are restored with the controller state, including the clear guard
data_parser	Linked fields are parsed on the public row's storage unit, once per distinct parser_args entry
kv_exists	Lists linked fields among each key's ready fields
Backends other than SimpleStorage	kv_link, TransferQueueClient.kv_link and TransferQueueClient.link raise NotImplementedError

Concurrency

• kv_link refuses a public key that is being cleared (mark_clearing has zeroed its production status). A destination key being cleared cannot be told apart from one not produced yet, so a link to it is accepted; the clear then drops that link.
• kv_clear_unlinked_public selects and marks the rows in one controller request, so a racing kv_link is refused for the keys it clears.
• A kv_put during a public key's clear can make the key ready again and let a kv_link through. When the clear then releases the row, the controller drops that link, so the field becomes unwritten instead of pointing at an index that another partition may reuse.
• kv_link validates field types before it creates keys, so a dtype conflict leaves no empty keys behind.

Limitations

• Units are balanced by row count, not bytes. Public payloads of very different sizes can still leave some units fuller than others.
• Links name public keys, so a public key cleared and put again under the same name serves the new payload to keys relinked to it later, including keys restored from a dump. __public__ is shared by every job on one TransferQueue; prefix public keys with a job id.
• Linking is SimpleStorage-only, and kv_update cannot modify linked fields.
• Parse results are shared within one kv_batch_get call only; nothing is cached across calls. Keys sharing a parse may receive the same Python object for a non-tensor field, so modifying one row in place changes the others. As before, a data_parser must not modify the stored values it receives: one public row requested with different parser_args hands the same stored object to each parse.

Tests

• tests/e2e/test_kv_link_e2e.py (new) covers:
  • balanced spread of public and private rows;
  • reads through links, single and batched kv_link;
  • argument checks, including empty keys, fields the key stores itself and the client-side backend check;
  • the clear guard on every clear path, and its error naming the linking keys;
  • refusal of puts on linked fields, writes to linked public keys, and kv_update;
  • kv_describe on linked fields;
  • client.link reads, its refusals, the read-only guard on BatchMeta puts into __public__, and keyless public rows freed by kv_clear_unlinked_public;
  • data_parser slicing and selecting elements of a linked field, with different parser_args per key;
  • dump/load and checkpoint round trips, and a load into an incompatible partition writing nothing;
  • kv_clear_unlinked_public leaving never-linked keys alone and clearing stale_keys;
  • kv_exists on present, missing, linked and being-cleared keys;
  • the concurrency cases above.
• tests/e2e/test_kv_interface_e2e.py: repeated keys keep their own parser_args.
• tests/e2e/test_data_dump_e2e.py: version 3 and 4 dumps still load.
• tests/test_controller_data_partitions.py:
  • the allocator balances each partition and stays FIFO with one unit;
  • the link index follows put and clear;
  • a checkpoint saved before links existed still loads.
• tests/test_metadata.py: add_fields and union handle links.
• tests/test_async_simple_storage_manager.py: a linked field is requested from its public row's unit once per distinct parser_args entry, and which entries count as equal.
• python -m pytest -q tests: 899 passed, 12 skipped. The 8 tests in tests/test_yuanrong_storage_client_e2e.py error locally because they need a Yuanrong environment; they error the same way without this PR.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-19T05:25:05.995929Z	5837879	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5837879910

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Bump the dump format before writing links

When a link-bearing dump is restored by the pre-change build, dump_info.json still advertises format 4, which that reader accepts. That reader ignores the new links member while fields deliberately omits linked payloads, so the restore succeeds with those fields silently lost. Allocate a new dump format version (while retaining v3/v4 read compatibility) so older readers reject link-aware dumps.

Useful? React with 👍 / 👎.

[OutstanderWang]

Confirmed: main's reader loads a link-bearing dump without error and the linked field is gone. Fixed in 8fdf3ff: dumps are written as format 5, and this build still reads 3 and 4 (a test rewrites a fresh dump as each older version and restores it).

[chatgpt-codex-connector]

Preserve duplicate samples through parser grouping

When kv_batch_get receives repeated keys/global indexes together with data_parser and per-key parser_args, this dictionary coalesces occurrences before routing. The parser runs only once with one occurrence's argument and the reordered output duplicates that result, so another duplicate ignores its own argument; the previous routing path kept duplicate positions. Preserve occurrences for parser-backed reads or reject duplicate read keys.

Useful? React with 👍 / 👎.

[OutstanderWang]

Confirmed, and it was wider than duplicates: batch_positions indexed the deduplicated row list, so every key after a repeat got its neighbour's argument ([a, a, b, c] with [1, 2, 3, 4] read [11, 11, 22, 33]; main reads [11, 12, 23, 34]). Fixed in 4c441ca: a sample's own fields are read once per occurrence, and only linked public rows are deduplicated. Covered in test_kv_interface_e2e.py.

[chatgpt-codex-connector]

Delete private cells when linking existing fields

When a destination key already stores any requested field, kv_link only changes controller link maps; no client or storage-manager call removes the destination cell. Reads switch to the public source, but the old tensor or SSD file remains allocated until an explicit overwrite or clear, so relinking an existing dataset retains all N copies and can still exhaust memory. Reject already-produced destinations or clear their private cells through the data plane when linking.

AGENTS.md reference: AGENTS.md:L13-L14

Useful? React with 👍 / 👎.

[OutstanderWang]

Confirmed. Fixed in 78134b5 by refusing instead of deleting through the data plane: kv_link raises when a key already stores one of the fields itself, before creating any key, and asks for kv_clear(fields=...) first.

[chatgpt-codex-connector]

Enforce the backend check in the low-level link client

Under Mooncake or YuanRong, callers of the exported get_client().kv_link/TransferQueueClient.kv_link reach this RPC directly and bypass the backend guard that exists only in interface.async_kv_link. The controller returns success and marks fields linked, but those storage managers do not resolve BatchMeta._links, so subsequent reads use missing or stale destination cells. Put the backend check in this low-level client path as well.

AGENTS.md reference: AGENTS.md:L17-L18

Useful? React with 👍 / 👎.

[OutstanderWang]

Confirmed. Fixed in ae87d25: the SimpleStorage check moved into TransferQueueClient.async_kv_link, which the sync client and tq.kv_link both go through; the interface-level copy is gone.

[chatgpt-codex-connector]

Reject empty destination lists before recording links

kv_link([], ...) is accepted here: the controller creates a destination partition and link_fields records linked_by[source] = set() despite having no rows. mark_unlinked_public_clearing considers the reverse-map keys rather than nonempty sets, so that public source is skipped by every bulk-clear call until the empty destination partition is removed. Reject the empty list before the RPC or avoid creating an empty reverse entry.

Useful? React with 👍 / 👎.

[OutstanderWang]

Confirmed (kv_link([]) left an empty reverse entry and kv_clear_unlinked_public() returned []). Fixed in 5162c4f: the controller's kv_link rejects an empty key list, so direct client calls are covered too.

[chatgpt-codex-connector]

Validate linked field types before loading payloads

When loading a linked dump into an existing partition that declares the linked field with an incompatible dtype or tensor kind, this preflight checks only public readiness because linked fields are absent from row_index["field_schema"]. load_rows_by_key then registers keys and may write shards, and the subsequent client.kv_link is the first type validation and fails, leaving a partial restore. Validate the destination against the public fields' schemas before calling load_rows_by_key.

Useful? React with 👍 / 👎.

[OutstanderWang]

Confirmed (the payloads of the linking keys were written before kv_link failed). Fixed in b0d2acb: load relinks first. kv_link checks field types before creating any key, so an incompatible destination now fails with nothing written; there is a test for it.

Chase-Rong

[Chase-Rong]

[AI comments] Links reference a mutable key name, not a snapshot. A linked public row can still be rewritten, and L1271-L1279 propagates the new value and readiness to every linking key. If a later step reuses a public key name (video_0) while older trajectories still link it, they silently read the new payload.

Dump/load has the same gap. Repro: link 4 keys to v0 (values 0..5), dump both, clear them, kv_put("v0", <all 9s>), then load_data_by_key(link_dir). The load succeeds and the restored keys read [9, 9, ...]. The preflight only checks that the public field is ready.

Suggestion: refuse writes to a public row while it is linked, symmetric to the clear guard; this also removes the need for the cascade loop. Record a fingerprint or dump id of the public rows in the link dump so load can verify it, or offer an option to dereference links at dump time.

[Chase-Rong]

kv_clear_unlinked_public() sweeps every unlinked key, including keys another producer has kv_put but not yet kv_linked. The documented precondition ("no writer between kv_put and kv_link") rarely holds in async/off-policy GRPO, where the rollout of step k+1 overlaps the trainer cleanup of step k. The producer kv_link then fails ("must have produced" / "not found") and its payload is gone. __public__ is also one global namespace, so two jobs sharing a TransferQueue can collide on key names and sweep each other's payloads.

[OutstanderWang]

Agreed, the precondition was too strong for overlapping steps. 2e812ec: the controller records which public rows have ever been linked, and kv_clear_unlinked_public frees only those that no key links to any more. A key that is put but not linked yet is never swept, so a producer between kv_put and kv_link is safe; a payload whose producer died before linking needs an explicit kv_clear. The record survives controller checkpoints. For the shared namespace, the docstring now says to prefix public keys with a job id.

[Chase-Rong]

Each call to kv_link requires a public_key, so processing B prompts requires B RPC calls, and the entire process is not atomic: if a failure occurs midway, it will result in an earlier set of prompts being linked while a later set remains unlinked. I suggest adding an optional batch processing feature, eg. accept public_keys: list[str] aligned with keys or a {key: public_key} mapping

[OutstanderWang]

Done in 1166b53: public_key also accepts a list aligned with keys. The controller checks every public key, field and destination before creating or linking anything, so a batch is linked by one atomic request. load_data_by_key uses it too, with one call per linked field set.

[Chase-Rong]

[AI comments] A plain read-modify-write silently turns every link into a private copy. Repro: data = kv_batch_get(keys, P), data["resp"] += 100, kv_batch_put(keys, P, data). Afterwards links is {} for every key, video is stored n times, and kv_clear_unlinked_public() frees v0 as unreferenced. client.put(data, metadata=meta) behaves the same way through add_fields. This is the most common RL update pattern, so dedup is lost without any signal.

Suggestion: refuse a put that lands on a linked field and require an explicit unlink (kv_clear(fields=...)), consistent with how kv_update already refuses. At minimum, log a warning.

[Chase-Rong]

[AI comments] kv_describe disagrees with the read path on linked fields. fields omits video, while kv_retrieve_meta and kv_batch_get include it. After kv_update(public_key, ..., finalize=False), kv_batch_get(select_fields="video") raises "not ready", but kv_describe(...)["unfinalized"] is [], which breaks its documented contract. There is also no reverse lookup: when clearing a public key is refused, the error names only the public key, and nothing reports which keys or partitions link to it.

Suggestion: keep fields meaning "readable fields", include linked fields in unfinalized when their source is not ready, and keep links as extra information. Dump can then use fields - links.keys() to pick the payloads to read. Include the linking partitions and keys in the refusal message.

[Chase-Rong]

The PR description says kv_link refuses keys that are being cleared, but only the public source's production status is checked. I conduct a simple test: kv_put("t0", P), mark_clearing([t0]), then kv_link("t0", P, "v0", "video") is accepted, and video reads as ready during the clear window. The following clear_meta drops the link, so the state converges and the impact is small. A destination being cleared cannot be told apart from one not produced yet by production status alone, so correcting the description is likely simpler than adding the check.

[OutstanderWang]

Confirmed with the same repro. I agree that correcting the description is simpler. The PR description now says a destination being cleared is accepted and that the clear then drops the link.

[OutstanderWang]

Re: #200 (review)

Agreed on the write side. bba58db makes a linked public key read-only: kv_put, kv_update and tag writes through the KV path are refused until no key links to it, symmetric to the clear guard, and the error names the linking keys.

I kept the propagation loop. A put through a BatchMeta the sampler fetched from __public__ writes storage before the controller hears about it, so the controller cannot refuse it up front; for that path the loop keeps linking keys consistent with the new payload's readiness and shape. The KV paths no longer reach it.

For dump/load, links stay by name, and the docstring and docs/data_dump.md now say that load relinks to whatever the names hold then. A fingerprint would mean hashing payloads at dump and load. Dereferencing at dump time would bring back the copies the feature removes. I'd rather add one of them if a workload needs it.

[OutstanderWang]

Re: #200 (comment)

Confirmed. Fixed in 11d6a38: a put whose BatchMeta marks a written field as linked is refused before anything is written, for kv_put/kv_batch_put and client.put(data, metadata=meta) alike. A read-modify-write must leave linked fields out (data.select(...)), and kv_clear(fields=...) is the explicit way to store a private value, consistent with kv_update.

[OutstanderWang]

Re: #200 (comment)

Confirmed on both points. 514ee42: fields now lists readable fields, linked ones included; unfinalized is computed over all of them, since a linking row's readiness already follows its source; links stays as extra information. The dump strips linked fields itself, so its row index still records only stored payloads. A refused clear or write on a public key now names the linking keys, e.g. (3 keys, e.g. ['train:traj_0_0', ...]).

[0oshowero0]

[AI Review]

Two reproducible P2 findings in the BatchMeta link and linked-field parser paths.

[0oshowero0]

[AI Review] [P2] Preserve joint parser input for fields on the same public source row and document cross-source limitations

This suggestion applies only when multiple selected fields link to the same public source row. They have the same actual global index and reside on the same StorageUnit, so one GET can read them together and pass them to the parser.

The current implementation still splits these fields into single-field requests. As a result, a parser that works when reading the public record directly fails when reading the same fields through links.

For example, public record video_0 stores both a padded video and its valid length. Trajectory traj_0 links both fields to video_0:

import torch
import transfer_queue as tq

# Assume TransferQueue has been initialized with SimpleStorage.
tq.kv_put(
    key="video_0",
    partition_id=tq.PUBLIC_PARTITION,
    fields={
        "video": torch.tensor([10, 20, 0, 0]),
        "length": torch.tensor([2]),
    },
)

tq.kv_put(
    key="traj_0",
    partition_id="train",
    fields={"response": torch.tensor([99])},
)

tq.kv_link(
    keys=["traj_0"],
    partition_id="train",
    public_key="video_0",
    fields=["video", "length"],
)

def crop_video(data, args):
    return {
        "video": [
            video[:int(length.item())]
            for video, length in zip(data["video"], data["length"])
        ],
        "length": data["length"],
    }

# Reading the public record directly succeeds: video becomes [10, 20].
direct = tq.kv_batch_get(
    keys=["video_0"],
    partition_id=tq.PUBLIC_PARTITION,
    select_fields=["video", "length"],
    data_parser=crop_video,
)

# Reading through links fails with RuntimeError reporting a missing field.
via_link = tq.kv_batch_get(
    keys=["traj_0"],
    partition_id="train",
    select_fields=["video", "length"],
    data_parser=crop_video,
)

The direct read sends one request:

source row: the public row for video_0
fields: ["video", "length"]

The linked read instead produces:

request 1: the same public source row, fields=["video"]
request 2: the same public source row, fields=["length"]

The parser therefore receives only video or only length in each invocation and raises KeyError inside the StorageUnit when accessing the other field. The SDK exposes this as a RuntimeError reporting the missing-field error.

Please group fields from the same public source row with the same parser-argument identity into one read request, and add a regression test that executes the cross-field parser above. This grouping should retain the existing _parser_arg_identity and deduplication policy, without introducing value comparisons for tensor arguments. Reading and parsing still happen on the source StorageUnit; this does not require combining data across storage locations.

For reads that mix public and non-public fields, please explicitly document the limitation in the kv_batch_get and get_data / async_get_data docstrings:

data_parser executes per storage read group. Each invocation receives only the fields in its current group. When selected fields of one logical sample come from a public source row and a regular row, they may be parsed in separate requests. Callers must not assume that one parser invocation receives all selected fields of that logical sample, or depend on fields from different source rows within that sample being joined for parsing.

This comment does not request joining fields that resolve to different source rows or storage locations within a single logical sample. A parser invocation may still process multiple source rows as a batch. Even when all fields belong to the public partition, joining fields from different source rows within one logical sample is outside the scope of this suggestion.

[0oshowero0]

[AI Review] [P2] Preserve caller-local metadata when returning linked BatchMeta

The new BatchMeta returned by link() does not retain the input batch's extra_info or local custom_meta changes not yet synchronized to the controller. The original input object remains unchanged, and this does not delete tags already persisted in the controller. The concern is the information missing from the returned object used for subsequent processing or writes.

Repro with one destination row whose persisted custom metadata is empty, and one public row:

rows_meta.extra_info["epoch"] = 3
rows_meta.update_custom_meta([{"group": "g0"}])

linked = client.link(rows_meta, public_meta, "prompt")

assert linked.extra_info == rows_meta.extra_info   # fails: {} vs {"epoch": 3}
assert linked.custom_meta == rows_meta.custom_meta # [{}] vs [{"group": "g0"}]

The returned metadata is missing both values, although they remain present in rows_meta. This can also remove per-sample information used to construct get-parser arguments.

Please distinguish the handling of these two kinds of metadata:

• Preserve the caller's extra_info in the returned batch.
• Define how unsynchronized local custom_meta is reconciled with the controller's current custom metadata, including precedence when both contain a value for the same key. Simply copying the old input metadata over the controller response could overwrite newer controller values. The controller's updated links, field schema, and readiness should remain authoritative.

This concerns metadata carried by the returned object; it does not require link() to persist local tags automatically. Please cover extra_info and unsynchronized custom_meta in regression tests, and make any conflict-handling policy explicit.

[0oshowero0]

[AI Review] Suggestion: document the write-back contract for linked fields

A concrete GRPO-style case: one prompt video is stored once in PUBLIC_PARTITION, and n response rows link to it. A worker reads both the video and responses for computation, updates only the responses, and selects response in the metadata before writing back, expecting that selection to restrict the written fields:

linked = client.link(
    rows_meta, public_meta.select_samples([0] * n), "video"
)
batch = client.get_data(linked)  # contains both video and response
batch["response"] = updated_responses

write_meta = linked.select_fields(["response"])
client.put(batch, metadata=write_meta)

However, batch still contains video. The put path writes data.keys() to the destination row indexes, while the selected metadata no longer carries the excluded field's link information. This can store private video values under the trajectory rows and remove their original links, even though the worker only changed response. The public source itself is not modified.

For now, I suggest adding a Warning to the put / async_put docstrings explaining that:

• metadata.field_names is not a write-field whitelist; the fields to write come from data.keys().
• BatchMeta.select_fields() only filters local metadata. It neither filters the data being written nor unlinks fields in the controller.
• For a BatchMeta-based put, the client-side check against writing linked destination fields depends on the link information retained in the supplied metadata. When writing back a batch containing shared fields, explicitly select the private fields to write.

This observation concerns linked destination fields in a BatchMeta-based put. The guard against overwriting the public source row is a separate controller check.

A safe-write example would make the intended usage clear:

batch = client.get_data(linked)
batch["response"] = updated_responses
client.put(batch.select("response"), metadata=linked)

A short note in BatchMeta.select_fields() pointing out that data must be selected separately for write-back would also help.

[0oshowero0]

[AI Review] Suggestion: document that a public-row overwrite must not overlap with linking that row

For now, I suggest adding a caller-side concurrency warning to the put / async_put and link / async_link docstrings, without changing the implementation:

Overwriting an existing row in PUBLIC_PARTITION must not run concurrently with operations that establish links to that same row, including operations from other workers or clients. The write check observes the current link relationships; it does not reserve the row until the storage write completes. Remove existing links before overwriting a public row, wait for the overwrite to complete, and then establish new links.

The following SDK-level case explains the restriction.

Initial state

Assume TransferQueue is initialized with SimpleStorage. Store one public video and one trajectory response:

import torch
from tensordict import TensorDict
import transfer_queue as tq

client = tq.get_client()
public_meta = client.put(
    TensorDict({"video": torch.tensor([[1, 2, 3]])}, batch_size=1),
    partition_id=tq.PUBLIC_PARTITION,
)
traj_meta = client.put(
    TensorDict({"response": torch.tensor([[0]])}, batch_size=1),
    partition_id="train",
)

Call the allocated public row p and the trajectory row t. At this point, p.video = [1, 2, 3], its production status is ready, and no row links to p. The trajectory stores only response.

Pass the metadata to two workers connected to the same TransferQueue. The following calls overlap; the snippets alone do not force a particular scheduling order.

Worker A: overwrite the existing public row

client_a = tq.get_client()

def prepare(data):
    return {name: values + 10 for name, values in data.items()}

client_a.put(
    TensorDict({"video": torch.tensor([[7, 8, 9]])}, batch_size=1),
    metadata=public_meta,
    data_parser=prepare,
)

Supplying metadata=public_meta targets the existing row p. The parser will turn the replacement value into [17, 18, 19].

Worker B: establish a link and read through it

client_b = tq.get_client()
linked = client_b.link(traj_meta, public_meta, "video")
first = client_b.get_data(linked)
# Further processing occurs before the next read.
second = client_b.get_data(linked)

One permitted interleaving

1. A asks the controller whether public row p can be written. The SDK sends CHECK_PUBLIC_WRITE with global_indexes=[p]. The controller checks the reverse link indexes of its partitions: for example, train.linked_by[p] = {t} would mean that trajectory t references this public row. There are no such entries yet, so the check passes. This check does not mark p as being written or change its ready status.
2. A has not yet sent/performed the storage write. There is an asynchronous boundary between the successful controller check and the storage-manager put.
3. B's link() reaches the controller. The rows exist, the public video field still has production status 1, the destination does not store its own video, and the schema is compatible. The link is accepted and the controller records t.video -> p.video. No pending-write state prevents this.
4. B's first get reads [1, 2, 3]. The storage manager resolves the link to public row p, which still holds the old value. The link records a row location, not a payload snapshot.
5. A's storage write executes. The StorageUnit runs prepare and replaces p.video with [17, 18, 19]. The write was already approved in step 1. This storage operation does not recheck the newly established controller link.
6. A notifies the controller that the write completed. The current update path records the field state and propagates readiness/shape to linking rows; it does not reject the completed overwrite.
7. B's second get, using the same linked metadata, reads [17, 18, 19]. The public payload changed after B successfully established its link.

The deterministic reproduction used two asynchronous SDK operations sharing one client in the same process, with events pausing A after its controller check and before its storage-manager put, allowing B to link and read, then resuming A. The two-worker SDK snippets above illustrate the corresponding possible interleaving; they do not force it by themselves. The scheduling hook only fixes the interleaving; parsing and payload reads/writes still use real SimpleStorage units. The same concurrency restriction applies without a put parser; this parser simply makes the replacement value easy to distinguish.

A write-once public-row workflow, with the initial put completed before linking, avoids this overwrite/link overlap. Please make the same-row serialization requirement explicit in the docstrings so callers do not interpret the current write check as protection against concurrent link creation.