Repository navigation
[WIP] Feat/kv share - #200
OutstanderWang wants to merge 20 commits into
Conversation
GRPO-style rollouts put one prompt and link its n responses to it, so the prompt payload is stored once. Links are controller metadata only; the SimpleStorage read path resolves them and reads each source row once. A put on a linked field replaces the link, clearing a linked key drops only the link, and a source cannot be cleared while other keys still link to it. dump_data_by_key records links and load_data_by_key relinks.
kv_link checked field metadata, which stays until the clear finishes, so a link made between mark_clearing and the storage delete pointed at a payload about to vanish. Check production status, which mark_clearing zeroes first.
A repeated missing key was registered twice, leaving a keyless row that held a link and could never be cleared. Reject duplicates before any RPC, as the batch writes do.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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".
| <dump_dir>/ | ||
| dump_info.json # version, partition, counts | ||
| row_index.pt # key -> {global_index, fields, unfinalized, tag} | ||
| row_index.pt # key -> {global_index, fields, unfinalized, links, tag} |
There was a problem hiding this comment.
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 👍 / 👎.
| fields_by_row: dict[int, set[str]] = defaultdict(set) | ||
| for global_index, links in zip(metadata.global_indexes, metadata._links, strict=True): | ||
| for field in metadata.field_names: | ||
| fields_by_row[links.get(field, global_index)].add(field) |
There was a problem hiding this comment.
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 👍 / 👎.
| metadata = self.kv_retrieve_meta(keys, partition_id, create=True) | ||
| partition = self._get_partition(partition_id) | ||
| assert partition is not None | ||
| partition.link_fields(metadata.global_indexes, public, source, fields) |
There was a problem hiding this comment.
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 👍 / 👎.
| await self._request_controller( | ||
| socket=socket, | ||
| request_type=ZMQRequestType.KV_LINK, | ||
| response_type=ZMQRequestType.KV_LINK_RESPONSE, | ||
| body={"keys": keys, "partition_id": partition_id, "public_key": public_key, "fields": fields}, |
There was a problem hiding this comment.
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 👍 / 👎.
| keys = [keys] if isinstance(keys, str) else keys | ||
| _check_unique_keys(keys) |
There was a problem hiding this comment.
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 👍 / 👎.
| key | ||
| for row in rows.values() | ||
| for name, key in row.get("links", {}).items() | ||
| if name not in public_rows[key]["fields"] or name in public_rows[key]["unfinalized"] |
There was a problem hiding this comment.
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 👍 / 👎.
Chase-Rong
left a comment
Chase-Rong
left a comment
There was a problem hiding this comment.
[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.
| assert partition is not None | ||
| partition.link_fields(metadata.global_indexes, public, source, fields) | ||
|
|
||
| def mark_unlinked_public_clearing(self) -> tuple[BatchMeta, list[str]]: |
There was a problem hiding this comment.
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.
| tq_client._run_coroutine(async_kv_clear(keys=keys, partition_id=partition_id, fields=fields)) | ||
|
|
||
|
|
||
| def kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: list[str] | str) -> None: |
There was a problem hiding this comment.
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
|
[AI comments] A plain read-modify-write silently turns every link into a private copy. Repro: Suggestion: refuse a put that lands on a linked field and require an explicit unlink ( |
|
[AI comments] Suggestion: keep |
| raise KeyError(f"key {public_key!r} not found in partition {PUBLIC_PARTITION!r}") | ||
| source = public.keys_mapping[public_key] | ||
| # Production status, not field metadata: mark_clearing zeroes it before storage deletes the payload. | ||
| missing = [ | ||
| name | ||
| for name in fields | ||
| if name not in public.field_name_mapping | ||
| or public.production_status[source, public.field_name_mapping[name]] != 1 | ||
| ] | ||
| if not fields or missing: | ||
| raise ValueError(f"public key {public_key!r} must have produced fields {missing or fields}") |
There was a problem hiding this comment.
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.
2 participants
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
nresponses, so putting the video into every response key multiplies its memory and put traffic byn.The same workloads also suffer from unbalanced storage units. SimpleStorage stores
global_indexon unitglobal_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
PartitionIndexManagercounts 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].SimpleStorage.num_data_storage_unitsinstore_config. Other backends, and SimpleStorage with one unit, keep the previous FIFO allocation.reusable_indexeslist, which is regrouped by unit on load.2. Share payloads through
__public__andkv_linktq.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)makesfieldsofkeysread the payload ofpublic_keywithout copying it. Missing keys are created. The link is a single controller request: either every key is linked or none is.tq.kv_clear_unlinked_public()clears every public key that no key links to, and returns the cleared keys.How links are implemented:
DataPartitionStatus.links({global_index: {field: public global_index}}) and indexes them by source inlinked_by, so clear checks scale with the size of the clear.BatchMetacarries the links to the storage manager.AsyncSimpleStorageManager.get_datareads each linked public row once, however many keys in the batch link to it.How other APIs behave with links
kv_batch_get,get_data, samplers,StreamingDataLoaderkv_put/ put throughBatchMetaon a linked fieldkv_updateon a linked fieldkv_putreplaces the link insteadkv_clearof a linking key or fieldkv_clear,clear_samples,clear_partitionon a public key that is still linkedkv_describelinks:{field: public key}dump_data_by_key/load_data_by_keysave_checkpoint/load_checkpointdata_parserNotImplementedErrorwhen a selected field is linkedkv_linkraisesNotImplementedErrorConcurrency
kv_linkrefuses a public key that is being cleared (mark_clearinghas zeroed its production status), and keys that are being cleared.kv_clear_unlinked_publicselects and marks the unlinked rows in one controller request, so a racingkv_linkis refused for the keys it clears.kv_putduring a public key's clear can make the key ready again and let akv_linkthrough. 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_linkvalidates field types before it creates keys, so a dtype conflict leaves no empty keys behind.Limitations
kv_clear_unlinked_publicalso clears a public key that was put but not linked yet. Call it when no writer is between itskv_putandkv_link, for example after a step's trajectories are cleared.kv_linkis SimpleStorage-only. It does not supportdata_parser, andkv_updatecannot modify linked fields.Tests
tests/e2e/test_kv_link_e2e.py(new) covers:kv_updateanddata_parser;kv_clear_unlinked_public;tests/test_controller_data_partitions.py:tests/test_metadata.py:add_fieldsandunionhandle links.python -m pytest -q: 883 passed, 12 skipped.tests/test_yuanrong_storage_client_e2e.pywas deselected locally because it needs a Yuanrong environment.