[feat] add kv_update and field-level kv_clear

[OutstanderWang]

Summary

Adds two operations on existing KV keys:

• kv_update / kv_batch_update merge new values into fields that are already stored, using merge_fn(old, new). The merge runs on the SimpleStorage unit that holds each sample, so the stored payload never makes a round trip through the client.
• kv_clear(keys, partition_id, fields=...) clears only the given fields and keeps the keys, for example to release pixel_values after rollout while training still reads the rest of the sample. It works on every storage backend. Without fields, kv_clear behaves as before.

To make multi-field updates all-or-nothing on a unit, SSD-offloaded put_data is now two-phase, which also affects plain multi-field put (see "SSD offload").

Usage

kv_update: append the response to the prompt

import torch
import transfer_queue as tq

tq.kv_put(key="s0", partition_id="train", fields={
    "input_ids": torch.tensor([10, 11, 12]),
    "attention_mask": torch.tensor([1, 1, 1]),
})

tq.kv_update(
    key="s0",
    partition_id="train",
    fields={"input_ids": torch.tensor([20, 21]), "attention_mask": torch.tensor([1, 1])},
    merge_fn=lambda old, new: torch.cat([old, new]),
)
# input_ids      -> tensor([10, 11, 12, 20, 21])
# attention_mask -> tensor([1, 1, 1, 1, 1])

For a batch, kv_batch_update(keys, partition_id, fields: TensorDict, merge_fn) takes a TensorDict whose batch size matches keys (nested tensors are allowed here). Both have async_ variants.

kv_clear(fields=...): release pixel values after rollout

In a multimodal RL step, the rollout engine needs the raw pixel_values, while training only needs the processed tensors. Once rollout is done, the pixel values only take up storage:

tq.kv_batch_put(keys=keys, partition_id="train", fields=TensorDict({
    "input_ids": input_ids,
    "pixel_values": pixel_values,  # large; only rollout reads it
    "mm_tensors": mm_tensors,      # what training reads
}, batch_size=len(keys)))

# ... rollout reads pixel_values and writes the responses back ...

tq.kv_clear(keys=keys, partition_id="train", fields="pixel_values")

# Training reads input_ids / responses / mm_tensors exactly as before.

After a field-level clear:

• The storage backend no longer holds the values. With SSD offload on, their SSD files are deleted.
• The field reads as never written for those keys: it is not ready, and the controller no longer describes its old shape or dtype for them.
• The keys, their other fields, tags and consumption status are unchanged, so samples are not consumed again.
• A later kv_put to the field writes it normally and marks it ready again.
• Missing keys and fields are ignored, matching whole-key kv_clear. Clearing every field still keeps the key in kv_list; removing the key itself is still kv_clear without fields.

Design

kv_update (data plane, SimpleStorage only)

interface.async_kv_batch_update → client.async_update → AsyncSimpleStorageManager.update_data → new ZMQ request UPDATE_DATA → SimpleStorageUnit._handle_update → StorageUnitData.apply_update.

• Preconditions. Every key must exist and every field must already be produced for every key; otherwise ValueError. Nothing is created.
• Compute, validate, then write once. Each unit computes every merge result and checks that each field keeps its tensor/non-tensor type and tensor dtype, then commits with a single put_data. A raising merge_fn or an incompatible result leaves that unit unchanged.
• Schema. Each unit reports the shapes it stored; the manager assembles one batch-level field_schema in global_indexes order (the controller reads per_sample_shapes by position) and notifies the controller. The client merges it into its BatchMeta with the new BatchMeta.apply_field_schema (add_fields now uses the same helper).
• Controller ACK is checked. notify_data_update now returns the ACK's success. On the update path, a rejected or missing ACK raises RuntimeError stating that storage may already have committed and the merge must not be retried blindly.
• No replays. Updates use max_attempts=1, because a replay after a lost reply would fold the same value in twice. A timeout therefore has an unknown outcome: the unit may commit after the caller stops waiting. This is documented on the API.
• Atomicity is per storage-unit request. A batch spanning several units is not a distributed transaction.
• Other backends. The base StorageManager.update_data raises NotImplementedError("kv_update is not supported by <ManagerClass>; it requires the SimpleStorage backend"). KV backends store each sample-field under its own key and have no hook to run merge_fn where the value lives.

Field-level kv_clear (control plane first, then storage)

interface.async_kv_clear(fields=...) → client.async_kv_clear_fields → new controller request KV_CLEAR_FIELDS → storage_manager.clear_data(metadata, fields=...).

• Controller first. TransferQueueController.kv_clear_fields resolves the keys and captures the metadata of the affected cells, including custom backend metadata that KV backends need for deletion. Then DataPartitionStatus.clear_fields zeroes those production-status cells, calls the existing FieldMeta.remove_samples (which already falls back from nested to regular), and drops the per-field backend metadata. It does not touch key mappings, tags, consumption status or indexes, so there is no index-reuse race and no MARK_CLEARING step is needed.
• Then storage. clear_data gains an optional fields filter. KV backends (Mooncake, Yuanrong, Ray) narrow the metadata with select_fields. SimpleStorage forwards fields in CLEAR_DATA, and units drop only those fields, unlinking their SSD files; the sample stays counted against unit capacity.
• The controller resolves fields per cell rather than relying on kv_retrieve_meta, which only reports fields that every requested key has, so a field present on only some keys is still cleared on those keys.

SSD offload: all-or-nothing multi-field put (#162)

HybridStorageUnitData.put_data used to write, commit and clean up one field at a time, so a failure on a later field (for example, a full disk) left earlier fields replaced. It is now two-phase: write every field's SSD files first (on failure, unlink only the files this call wrote and re-raise), then swap all references and unlink superseded files. The base StorageUnitData.put_data checks every field's length before writing any.

This also affects plain multi-field put with SSD offload on. The trade-off is a higher temporary peak: the old files plus all replacement files exist at once. docs/ssd_offload.md documents it.

Shared backing memory

SimpleStorage slices dense batches without copying, so samples written by one batched put share an allocation. Replacing or clearing some of them does not return memory until every sibling is replaced or cleared. docs/ssd_offload.md documents this alongside the existing data_parser note.

Docs and tutorial

• README: (async_)kv_update / (async_)kv_batch_update added, kv_clear notes the fields option, and the StorageManager.clear_data signature is updated.
• tutorial/02_kv_interface.py: a new step that concatenates a field with kv_update and then clears it with kv_clear(fields=...).
• docs/ssd_offload.md: shared backing memory and SSD headroom for atomic multi-field writes.

Testing

• python -m compileall -q transfer_queue tutorial tests and pre-commit (ruff, ruff-format, mypy) are clean.
• tests/e2e/test_kv_interface_e2e.py on SimpleStorage with 2 units, sync and async: 64 passed, 2 skipped (the skips are KV-backend-only rejection tests).
• Full local suite, python -m pytest -q: 699 passed, 12 skipped, 8 errors. The errors are all in tests/test_yuanrong_storage_client_e2e.py at fixture setup, because the Yuanrong datasystem package is not installed locally.
• Not run: the Mooncake and Yuanrong e2e backends.

New and updated tests:

• tests/test_kv_update.py: merge, atomicity on merge error or incompatible result, unproduced field rejection, decoded SSD-offloaded old values, atomicity when a later field fails to reach SSD, cross-unit schema assembly, KV-backend rejection naming the class.
• tests/test_storage_request_retry.py: the controller ACK status is returned; a rejected ACK after a storage commit raises; an update timeout is neither retried nor published.
• tests/test_controller_data_partitions.py: field-level clear keeps keys, tags, consumption status and other fields; the cleared field's schema for the remaining samples is unchanged; a later nested write to the same field succeeds; FieldMeta is dropped once no sample holds the field.
• tests/test_simple_storage_unit.py: field-level clear in memory and on SSD (only that field's file is unlinked); all-or-nothing multi-field SSD put.
• tests/test_kv_storage_manager.py: clear_data(fields=...) deletes only those keys with their backend metadata.
• tests/e2e/test_kv_interface_e2e.py: TestKVUpdateE2E (update then clear a field, then nested writes; batch update across units; rejected arguments) and TestKVClearE2E::test_kv_clear_fields_keeps_keys_and_other_fields.

Out of scope

• kv_batch_get(select_fields=...) on a field that is not ready returns an empty result rather than raising. This is existing behavior.

[0oshowero0]

[AI Review]

Summary: I'd ask for changes before merging. The feature works and the PR's own tests pass, but there is a controller-metadata regression that the PR's main use case (releasing pixel_values in multimodal RL) can trigger. Also, the empty mode duplicates what kv_put already does, so the PR could be considerably smaller.

Permalinks below point to the branch head caf3753.

Reliability

1. (High) kv_empty on one key rewrites the field's metadata for the whole partition

FieldMeta is kept per (partition, field), not per key. The new non-tensor branch marks the whole column as non-tensor and drops every sample's shape:

TransferQueue/transfer_queue/controller.py

Lines 256 to 263 in caf3753

	elif new_is_non_tensor:
	# Some samples now hold plain objects (e.g. kv_update empty stores None), so no
	# single dtype/shape describes the column any more. dtype is kept so a later
	# tensor write still has to agree with the one this field was created with.
	self.is_non_tensor = True
	self.is_nested = False
	self.shape = None
	self.per_sample_shapes.clear()

Controller-only repro:

import torch
from transfer_queue.controller import DataPartitionStatus

t4 = {"dtype": torch.float32, "shape": (4,), "is_nested": False, "is_non_tensor": False}
empty = {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True}
nested = {"dtype": torch.float32, "shape": None, "is_nested": True, "is_non_tensor": False,
          "per_sample_shapes": {5: (3,), 6: (7,)}}

p = DataPartitionStatus(partition_id="train")
p.update_production_status([0, 1, 2, 3], ["pv"], {"pv": t4})
p.update_production_status([0], ["pv"], {"pv": empty})            # kv_empty on one key
print(p.field_metadata["pv"].to_batch_schema([1, 2, 3]))           # keys 1..3: is_non_tensor=True, shape=None
print(p.update_production_status([5, 6], ["pv"], {"pv": nested}))  # False: AssertionError inside FieldMeta.update

End to end on SimpleStorage (2 units): put (4,) tensors for keys a/b/c, release key a, then kv_batch_put jagged nested tensors for keys d/e into the same field:

this PR (kv_empty):
  after releasing a   schema(b,c) = {'shape': None, 'is_non_tensor': True}   # b, c still hold (4,) tensors
  nested kv_batch_put for d,e returns without error
  after nested put    schema(d,e) = {'shape': None, 'is_non_tensor': True}   # no per_sample_shapes
main (kv_put(fields={"pv": None})):
  after releasing a   schema(b,c) = {'shape': [4], 'is_non_tensor': False}
  after nested put    schema(d,e) = {'is_nested': True, 'per_sample_shapes': [[3], [7]]}

This has three consequences:

• Other keys get wrong metadata. Samples that still hold tensors are described as non-tensor.
• It never recovers. A later regular-tensor write takes the tensor branch, which does nothing once is_non_tensor is True. The field only resets after every sample in it is removed.
• Later nested writes fail silently. The empty branch sets is_nested=False and shape=None, so a later nested write hits assert self.shape is not None in FieldMeta.update. update_production_status catches it and returns False, and that result is ignored (see item 2). The put reports success while the controller metadata is incomplete. The same sequence works on main, so this is a regression introduced here.

This is the multimodal RL pattern the PR targets: the previous batch's pixel_values have been emptied while the next batch's variable-size pixel_values are being written into the same partition.

The branch isn't specific to kv_update either. Any backend that writes a non-tensor value into a tensor field now takes it, and KV backends size their read buffers from metadata.get_shapes() / get_dtypes():

TransferQueue/transfer_queue/storage/managers/base.py

Lines 736 to 738 in caf3753

	for field_name in sorted(metadata.field_names):
	field_shape = metadata.get_shapes(field_name)
	field_dtype = metadata.get_dtypes(field_name)

So "no change to put/get semantics for KV backends" doesn't strictly hold.

Suggestion: don't express a per-key release through the column-level schema. One option is to treat a release as clearing the field for those samples:

• reset their production-status bit for that field;
• call the existing FieldMeta.remove_samples, which already handles falling back from nested to regular.

This doesn't block training. scan_data_status only checks the requested fields, so training that reads input_ids / mm_tensors is unaffected:

TransferQueue/transfer_queue/controller.py

Lines 720 to 763 in caf3753

	def scan_data_status(self, field_names: list[str], task_name: str) -> list[int]:
	"""
	Scan data status to find samples ready for consumption.
	This replaces the original _scan_data_status functionality.
	Args:
	field_names: List of required field names
	task_name: Name of the consumer task
	Returns:
	List of sample indices that are ready for consumption
	"""
	# Check if all requested fields are registered
	for field_name in field_names:
	if field_name not in self.field_name_mapping:
	return []
	row_mask = torch.ones(self.allocated_samples_num, dtype=torch.bool)
	# Apply consumption filter (exclude already consumed samples)
	_, consumption_status = self.get_consumption_status(task_name, mask=False)
	if consumption_status is not None:
	unconsumed_mask = consumption_status == 0
	row_mask &= unconsumed_mask
	# Create column mask for requested fields
	col_mask = torch.zeros(self.allocated_fields_num, dtype=torch.bool)
	field_indices = [self.field_name_mapping[field] for field in field_names]
	if field_indices:
	col_mask[field_indices] = True
	# Filter production status by masks
	relevant_status = self.production_status[row_mask][:, col_mask]
	# Check if all required fields are ready for each sample
	all_fields_ready = torch.all(relevant_status, dim=1)
	ready_indices_in_filtered = torch.nonzero(all_fields_ready, as_tuple=False).flatten()
	# Map back to original sample indices
	all_indices = torch.where(row_mask)[0]
	ready_sample_indices = all_indices[ready_indices_in_filtered].tolist()
	return ready_sample_indices

It also means a released field can't be confused with a stored None. If keeping the field ready and returning None is the intended behavior, the controller's column metadata should at least be left untouched.

2. (Medium) Storage commits first, and a failed controller update is silently ignored

The controller ACK carries success, but _notify_and_wait only checks the message type:

TransferQueue/transfer_queue/storage/managers/base.py

Lines 300 to 305 in caf3753

	if response_msg.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE_ACK: # type: ignore[arg-type]
	logger.debug(
	f"[{self.storage_manager_id}]: Get data status update ACK response "
	f"from controller id #{response_msg.sender_id} successfully."
	)
	return

This predates the PR, but kv_update makes it much easier to hit. Besides the nested case above, a parser that changes dtype (e.g. torch.cat promoting int64 + float) gets rejected by the controller's dtype check. The storage unit, the controller, and the caller's BatchMeta (already merged via apply_field_schema) then all disagree, while the caller sees success. I'd suggest at least raising on success=False in the update path.

3. (Medium) A timeout doesn't mean the update wasn't applied

The parser runs on the storage unit's request thread, the same thread that serves put, get, and the diagnostic probe. As a result:

• a slow parser blocks every request to that unit;
• past TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, the caller gets StorageUnitTimeout, but the unit still commits later.

max_attempts=1 prevents automatic replay. But a user who sees a timeout will naturally retry, and with a concat parser that appends twice. data_parser behaves the same way. The kv_update docstring should say that the outcome is unknown after a timeout.

4. (Low) kv_empty only frees memory once every sibling sample is released

Storage units decode tensors zero-copy from the received ZMQ frame:

TransferQueue/transfer_queue/utils/serial_utils.py

Line 309 in caf3753

arr = torch.frombuffer(buffer, dtype=torch.uint8)

So samples sent to one unit in a single put are views of one buffer. Emptying part of a kv_batch_put batch frees nothing until the whole batch on that unit has been emptied.

The PR example uses per-key kv_put, so it isn't affected, but batch puts are the common path. The tests only assert that the value reads back as None. docs/ssd_offload.md already documents the same limitation for data_parser; it would be worth documenting here too.

5. (Low) The SSD two-phase put looks correct but belongs in its own PR

I didn't find a problem in the logic: on failure only the files written by this call are removed, and the duplicate-index retention still holds across fields. Two notes:

• It changes plain multi-field put. That is a separate concern from kv_update; AGENTS.md asks for one minimal diff per issue.
• Peak SSD usage goes up. Before, only one field had old and new files on disk at the same time; now every field does. That makes ENOSPC more likely when the disk is nearly full. Worth stating as a trade-off.

Simplicity and maintainability

6. The empty mode duplicates kv_put(fields={f: None}). On main, kv_put(key, partition_id, fields={"pv": None}) already reads back [None], and HybridStorageUnitData.put_data deletes the replaced SSD files. The only thing kv_empty adds is the check that the key exists.

Today empty threads through six layers: interface, client.update, update_data, the ZMQ body, _handle_update, and apply_update. Each layer has its own validation, and the manager adds a dedicated retry rule for it. Without the flag:

• kv_empty becomes kv_retrieve_meta(create=False) plus one put in interface.py;
• the update path only has to handle the parser case.

7. The same argument validation is written four times: in _normalize_kv_update_args, update_data, _handle_update, and apply_update.

TransferQueue/transfer_queue/storage/managers/simple_storage_manager.py

Lines 544 to 553 in caf3753

	if empty:
	if values is not None or parser is not None:
	raise ValueError("empty update must not include values or a parser")
	else:
	if values is None or parser is None:
	raise ValueError("update requires values and a parser unless empty is set")
	if values.batch_size[0] != metadata.size:
	raise ValueError(
	f"Batch size of values ({values.batch_size[0]}) does not match metadata size ({metadata.size})"
	)

Validating at the public entry point and at the RPC boundary is enough. _normalize_kv_update_args also returns parser and empty unchanged, and its dict→TensorDict block is copied from async_kv_put.

8. KVStorageManager.update_data only rewords the error. The base StorageManager.update_data already raises NotImplementedError with the class name, and the KV override only changes the message:

TransferQueue/transfer_queue/storage/managers/base.py

Lines 809 to 825 in caf3753

	async def update_data(
	self,
	metadata: BatchMeta,
	field_names: list[str],
	values: TensorDict | None = None,
	parser: Callable[[Any, Any], Any] | None = None,
	empty: bool = False,
	) -> dict[str, dict[str, Any]]:
	"""kv_update is only implemented for SimpleStorage.
	A KV backend stores each sample-field under its own key and offers no hook to run
	parser(old, new) where the value lives, so read-modify-write cannot be made atomic.
	"""
	raise NotImplementedError(
	f"kv_update is not supported by {type(self).__name__}; it requires the SimpleStorage backend."
	)

Moving "requires the SimpleStorage backend" into the base message removes the override and four tests that only check its wording (Mooncake, Yuanrong, Ray, base KV).

9. The update RPC path mostly copies put. update_data / _update_to_single_storage_unit mirror put_data / _put_to_single_storage_unit: the same routing, gather, error logging, timeout handling, and heavy-op logging. On top of that come three new ZMQ message types and a new metrics label. One option is to reuse PUT_DATA with an optional merge function in the body, and have units return the stored shapes only when a merge function is present.

10. recipe/kv_update_e2e.py (389 lines) is really a test. It runs 10 assertion checks and refuses to attach to an existing deployment, while recipe/ holds demos. The e2e config already uses 2 storage units, so the cross-unit per_sample_shapes check and the remote-worker check can move into tests/e2e/test_kv_interface_e2e.py, and the recipe can be dropped.

Interface clarity

11. fields and values have to be kept in sync. A single field takes a bare value. Multiple fields need a dict whose keys exactly match fields, which is why the missing/extra checks exist. The dict keys already name the fields, and kv_put takes fields: dict. kv_update(key, partition_id, fields: dict[str, Any], <merge fn>) would remove that validation entirely.

12. Two entry points do the same thing. kv_update(..., empty=True) is kv_empty. Also, empty=True makes values and parser illegal, which is a mode switch. Suggest keeping only kv_empty.

13. Only a single key is supported. kv_clear accepts keys: list[str] | str, but kv_update / kv_empty take one key. The motivating case releases a whole batch of pixel_values, so per-key calls cost N × 3 RPCs (retrieve meta, storage unit, controller notify). client.update already handles batches; the single-key limit exists only in the interface wrapper.

14. The name parser. The existing data_parser takes a batch dict and resolves references into content. Here, parser is a per-sample merge(old, new). Same name, different contract; a merge-style name would be clearer.

Suggested split

• kv_empty: accept a batch of keys, build it on the existing put or on per-field production-status clearing, and fix item 1.
• SSD two-phase multi-field put: its own PR.
• kv_update with a merge function: it would help to name the concrete consumer that needs storage-side concat, then slim it down per items 6, 9 and 11.

---

Verification at caf3753:

• python -m compileall -q transfer_queue tutorial tests recipe is clean.
• tests/test_kv_update.py, tests/test_simple_storage_unit.py, and tests/test_controller_data_partitions.py pass (96 passed).
• The comparisons above were run on a local Ray cluster with SimpleStorage, against both this branch and main.
• Not run: the full suite and the Mooncake/Yuanrong backends.

[OutstanderWang]

Thanks for the detailed review. I reproduced the controller metadata regression, ignored controller ACK, unknown timeout outcome, and shared-allocation behavior before making the changes.

Reliability

1. kv_empty no longer uses the update RPC or changes partition-level FieldMeta. It now retrieves existing metadata and performs an ordinary batched put of None. The field remains ready and reads back as None, while the controller's tensor metadata stays unchanged. A regression test covers emptying one key followed by nested writes to the same field.

2. kv_update now only accepts fields already produced for every requested key. Each storage unit computes all merge results and validates tensor/non-tensor type and tensor dtype before its single commit. The update path also checks the controller ACK; a rejected or missing ACK raises an explicit error stating that storage may already have committed and the merge must not be retried blindly.

3. Merge updates remain non-retryable (max_attempts=1). The public documentation now states that a timeout has an unknown outcome because the storage unit may commit after the caller stops waiting.

4. The shared-backing-allocation behavior is now documented as the zero-copy trade-off: partially emptying a dense batch may retain the allocation until all sibling views are replaced or cleared.

5. I kept the SSD two-phase write in this PR because multi-field merges need per-unit all-or-nothing storage behavior. It remains isolated in its own commit, and the documentation now calls out the temporary peak requirement for the old files plus all replacement files.

API and implementation simplification

• Removed empty=True from the update path; kv_empty is now the only emptying API.
• Changed the API to kv_update(key, partition_id, fields: dict, merge_fn) and added kv_batch_update(keys, partition_id, fields: TensorDict, merge_fn).
• kv_empty now accepts one key or a batch of keys.
• Renamed parser to merge_fn.
• Removed the redundant KV backend override and moved the SimpleStorage requirement into the base error.
• Kept a distinct UPDATE_DATA operation because merges are non-idempotent, have different retry semantics, return stored shape information, and need separate observability from ordinary puts.
• Deleted the assertion-heavy recipe. Sync/async and cross-storage-unit coverage now lives in the existing KV E2E suite; the remote-worker-only scenario was removed.

Update atomicity remains per storage-unit request; a batch spanning multiple units is not a distributed transaction, and this is now documented.

Verification

• compileall and pre-commit pass.
• Focused update tests: 42 passed.
• Sync/async two-unit E2E: 6 passed, 2 backend-specific skips.
• Full local run: 692 passed, 12 skipped; the remaining Yuanrong optional-backend fixture setup errors reproduce unchanged on main in this environment and are unrelated to this PR.

[0oshowero0]

On kv_empty: it's hard to tell apart from kv_clear, and I think the naming problem comes from the semantics. I'd suggest turning it into a real field-level clear and folding it into kv_clear:

tq.kv_clear(keys, partition_id)                         # unchanged: remove the whole key
tq.kv_clear(keys, partition_id, fields="pixel_values")  # new: clear only these fields, keep the key

Why change the semantics, not just the name

Today kv_empty overwrites the field with None. The field stays ready, reads back as None, and the controller keeps its old shape/dtype. That is effectively kv_put(fields={"pixel_values": None}) plus a check that the key exists. Calling that "clear" would be misleading, and if we keep these semantics I don't think it needs a new API at all.

A field-level clear would instead do the following:

• Storage: delete the values for these (key, field) pairs, including their SSD files.
• Controller: reset production status for these (sample, field) cells and remove the samples from that field's FieldMeta.
• Kept as is: the key, its other fields, tags, consumption status, and the global index.

Benefits:

1. Clearer reads. A cleared field behaves like a field that was never written for that key. It can't be confused with a stored None, and a consumer that requests the field won't see the sample as ready and get None.
2. Accurate controller metadata. FieldMeta.remove_samples already handles the fallback from nested to regular, so the stale shape after emptying goes away.
3. More backends. kv_empty is SimpleStorage-only, but the KV backends' clear_data already deletes one key per (field, sample) taken from metadata.field_names. Passing metadata.select_fields(fields) therefore clears only those fields, with little or no change on the storage side.
4. Smaller API. One verb with an optional scope, consistent with kv_batch_get(select_fields=...). kv_empty / async_kv_empty can be removed.

What needs to change

• SimpleStorage: CLEAR_DATA only carries global indexes, and StorageUnitData.clear drops whole rows. Add an optional fields filter, and have HybridStorageUnitData.clear unlink only those fields' SSD files.
• Controller: add a partition method that clears fields for the given samples. It zeroes the production-status cells, calls remove_samples on those fields, and drops the matching field_custom_backend_meta. It leaves key mappings, tags, and consumption status alone and does not release indexes.
• Client / interface: when fields is given, async_kv_clear narrows the metadata with select_fields and takes the field-clear path.
• Order: mark the cells as not produced in the controller first, then delete from storage. Neither the key nor the index is released, so there is no index-reuse race. That makes this simpler than the three-step whole-key clear.

Edge cases to pin down

• Clearing every field of a key leaves the key in place: it is still listed by kv_list and keeps its tags. Removing the key itself is still kv_clear without fields.
• Keys that don't have the field should be ignored, matching the current idempotent kv_clear.
• A later kv_put to a cleared field writes normally and marks the field produced again.
• Consumption status is not reset, so tasks that already consumed the sample won't consume it again.

A separate kv_clear_fields with a required fields argument would also work. Its only advantage is that forgetting fields can't clear the whole key. I slightly prefer the fields= parameter, because omitting it gives exactly today's behavior.

Either way, the shared-buffer caveat in docs/ssd_offload.md still applies: samples written by one batched put share an allocation, so memory is only returned once all of them are cleared.

[OutstanderWang]

Thanks — agreed that the semantics were the problem, not just the name. kv_empty is gone, and its job is now a real field-level clear folded into kv_clear (cc28484):

tq.kv_clear(keys, partition_id)                         # unchanged: remove the whole key
tq.kv_clear(keys, partition_id, fields="pixel_values")  # clear only these fields, keep the key

Controller first. A new KV_CLEAR_FIELDS request resolves the keys and captures the affected cells' metadata, including custom backend metadata. Then DataPartitionStatus.clear_fields:

• zeroes those production-status cells;
• calls FieldMeta.remove_samples;
• drops the per-field backend metadata.

Key mappings, tags, consumption status and indexes are left alone.

Then storage. clear_data(metadata, fields=None) gains an optional filter:

• KV backends narrow with select_fields, so this now works on Mooncake, Yuanrong and Ray too.
• SimpleStorage forwards fields in CLEAR_DATA. Units drop only those fields and unlink their SSD files.

One deviation from your sketch: the controller resolves fields from its own per-cell state instead of narrowing a kv_retrieve_meta result. kv_retrieve_meta only reports fields that every requested key has, so a field present on only some keys would have been skipped. Doing it in the controller also makes it one controller RPC instead of two.

Edge cases, as you listed:

• Missing keys and fields are ignored.
• Clearing every field keeps the key in kv_list with its tags.
• A later kv_put writes the field normally and marks it ready again.
• Consumption status is not reset.

The shared-buffer caveat in docs/ssd_offload.md now says "replacing or clearing".

Tests:

• Your controller repro: release one key's field, then nested-write the same field. The remaining samples keep shape=(4,), and the nested write succeeds with correct per_sample_shapes.
• SimpleStorage memory and SSD clears.
• KV-backend clear_data(fields=...).
• An e2e test on 2 units, sync and async.

Results:

• KV e2e: 64 passed, 2 skipped.
• Full suite: 699 passed, 12 skipped. The 8 errors are the Yuanrong client module, whose datasystem package isn't installed locally.
• The Mooncake and Yuanrong e2e backends were not run.

I also rewrote the PR description to match the current API.