Concurrent kv_put(tag=...) on one key loses tag updates #188

Problem

Concurrent tag updates on the same key lose updates. Two writers that set different tag names on one key can leave it with only one of them.

kv_put(key, partition_id, tag={...}) (and kv_batch_put(..., tags=[...]), plus the async versions) is a read-modify-write split between client and controller:

1. The client fetches the key's metadata, including its current tags.
2. It merges the new tag into that copy locally (BatchMeta.update_custom_meta).
3. It sends the whole merged dict back, and the controller stores it with custom_meta.update(...) (DataPartitionStatus.set_custom_meta), which replaces each sample's dict.

When two writers interleave between steps 1 and 3, the second one writes a dict that was read before the first one's tag landed, so the first tag is dropped.

Repro (on main)

import threading
import torch
import transfer_queue as tq

tq.init()
P = "tag_race"
keys = [f"k{i}" for i in range(300)]
for key in keys:
    tq.kv_put(key, P, fields={"x": torch.tensor([0])})
barriers = [threading.Barrier(2) for _ in keys]

def writer(tag_name):
    for key, barrier in zip(keys, barriers):
        barrier.wait()
        tq.kv_put(key, P, tag={tag_name: 1})

threads = [threading.Thread(target=writer, args=(name,)) for name in ("a", "b")]
for t in threads:
    t.start()
for t in threads:
    t.join()
tags = tq.kv_list(P)[P]
print(sum(not {"a", "b"} <= set(tags[k]) for k in keys))  # 300: every key lost one tag

Run sequentially, both tags survive.

Why it matters

Retry or bookkeeping logic that keeps state in tags (for example an attempt id written by one worker while another writes a score) silently loses one of the writes. Callers can work around it today by holding a lock around the update (the KV locks in #186), but a tag-only update should be safe on its own.

Suggested fix

Merge per key in the controller instead of on the client: send only the new tag entries, and have the controller apply them to the sample's existing dict (custom_meta.setdefault(idx, {}).update(new_tags)) in its single request loop. That makes each tag update atomic without any lock.

Found while reviewing #186.