Propagate failed data notifications to producers

[rawsh]

Propagate missing, invalid, negative and error data-update responses to the producer instead of completing put successfully.
Require a successful ACK from the configured controller and discard pooled sockets after failed or cancelled waits.
Proof: allocated job 140220 executes 8/8 candidate service controls with zero skips; exact-main controls expose 7 failures and 1 positive pass; Ruff and compileall pass. Full suite and runtime-image adoption remain pending.

[0oshowero0]

[AI Review]

Thanks for the PR. The direction is right: on main, a rejected or failed data-update notification only gets logged, put still returns success, and consumers never see the samples. This diff is small and fixes real problems. However, I don't think it can merge as-is, for three reasons:

1. Regression: the new sender check makes every put fail after restoring a checkpoint in a fresh process (reproduced below).
2. Conflict with #182: #182 changed this same code to return bool, while this PR raises. A rebase has to settle on one contract.
3. Unstated behavior change: a notify timeout now fails put instead of being logged. That deserves an explicit decision and documentation.

What this PR fixes well

• Rejected updates now fail put. For example, if the partition was cleared, the controller replies success=False. Before this PR, that was only logged and put returned success.
• REQUEST_ERROR surfaces immediately. The old loop only accepted NOTIFY_DATA_UPDATE_ACK, so an error reply was skipped and the caller waited the full TQ_DATA_UPDATE_RESPONSE_TIMEOUT (30s by default).
• Cancellation no longer poisons the socket pool. except Exception does not catch CancelledError, so a cancelled wait used to return the socket to the pool with an unread ACK, which the next borrower would then read. except BaseException + close() fixes this.
• Dropping the receive loop is a safe simplification. The socket is leased exclusively and closed on any failure, so there is no stale traffic left to skip.

Blocking issues

1. The sender check breaks put after tq.load_checkpoint in a new process

if response_msg.sender_id != self.controller_info.id:
    raise RuntimeError(f"Unexpected data status update sender ...")

TransferQueueController.load_checkpoint overwrites self.controller_id with the id stored in the checkpoint (controller.py#L2239). However, zmq_server_info.id, which clients receive as controller_info.id, keeps the id from startup (controller.py#L1760-L1762). After a restart-and-restore, every ACK has a different sender_id, so every put fails.

Repro (SimpleStorage, 1 storage unit, Ray 2.56.0):

tq.init(conf)
tq.get_client().put(data(1), partition_id="p")
tq.save_checkpoint(ckpt)
tq.close()

tq.init(conf)            # new controller, new random controller_id
tq.load_checkpoint(ckpt)
tq.get_client().put(data(2), partition_id="p")

this PR: RuntimeError: Unexpected data status update sender TQ_CONTROLLER_4ad186c4; expected TQ_CONTROLLER_3d9eb5c1
base:    OK

The check also guards against something the transport already rules out. A pooled socket is leased per controller address and connected to exactly one controller, so the reply cannot come from anywhere else. I'd suggest removing the check. The controller overwriting its own id on restore is a separate latent issue that could be fixed in its own PR.

2. Merge conflict with #182: the error contracts disagree

#182 made notify_data_update return bool. update_data (kv_update) depends on that value to raise "Storage update committed, but the controller did not accept its metadata; do not retry this merge blindly". That hint matters because retrying a merge folds the same value in twice. Applying this PR's raise-based approach directly would make the if not published branch dead code and lose that hint.

I'd suggest keeping the raise-based contract: one error path that carries the reason, and it lets us delete the bool return. update_data can preserve the hint by wrapping the exception:

try:
    await self.notify_data_update(...)
except Exception as e:
    raise RuntimeError(
        "Storage update committed, but the controller did not accept its metadata; "
        "do not retry this merge blindly"
    ) from e

The tests on main that assert on the return value also need updating: the _notify_and_wait(...) is accepted case and the return_value=False rejection case in tests/test_storage_request_retry.py. It would also be worth confirming the intended behavior of the concurrent asyncio.gather over notifications in load_rows_by_index once they raise.

3. Timeout now fails put. Please make this an explicit decision

The removed comment documented the old design intent: "Logged rather than raised, so a slow controller does not fail the put." This PR reverses that, but the description doesn't mention it.

When the timeout fires, the payload is already in storage, and the controller may still apply the update later, so the outcome is unknown:

• If the producer retries with partition_id, it gets new indexes, and a late-applied update produces duplicate samples.
• If the producer does not retry, the payload is orphaned in storage.

The safe retry is to reuse the same metadata: the storage write overwrites, and a repeated notify is idempotent.

Suggestions:

• Say this in the error message, e.g. "payload stored, publication outcome unknown; retry with the same metadata".
• Add it to the Raises section of the put docstring.
• Call out the new meaning of TQ_DATA_UPDATE_RESPONSE_TIMEOUT in the PR description.

Non-blocking suggestions

• The controller_info null branch is unreachable. __init__ → _connect_to_controller already raises ValueError for anything that isn't a ZMQServerInfo. I'd delete the branch rather than turn it into a raise, and drop test_missing_controller_configuration_fails_public_put, which sets the field to None by hand.
• Rejection messages are uninformative. The controller's negative ACK has no message key, so users always see "success is not true". Including the partition_id, plus a hint that the partition may have been cleared, would help.
• Tests:
  • Subclassing TransferQueueController.__ray_metadata__.modified_class relies on a private Ray attribute and may break on Ray upgrades.
  • Three of the four parametrized modes (wrong_sender, missing_success, unexpected_type) build replies the real controller never sends. wrong_sender effectively encodes the bug above as expected behavior.
  • The cancellation fix, arguably the most subtle change, has no test. In the timeout test, the outer wait_for(8s) is longer than the internal 2s timeout, so cancellation never happens.
  • Every test starts a fresh controller and storage actor. The rejection, REQUEST_ERROR, and cancellation paths could use a fake socket instead, like the _SocketLease helper in tests/test_storage_request_retry.py, which runs in milliseconds.
  • There is a leftover debug print(...).
  • The manual clear_data(meta) cleanup shows that a rejected put leaves its payload in storage. main has the same silent leak, so this could be a follow-up.
• PR description: please drop internal job ids. It also notes that the full suite is still pending, and no CI has run on this PR yet.

What I verified (on bca3a31)

• python -m compileall -q transfer_queue tutorial tests passes.
• The new tests plus test_async_simple_storage_manager.py, test_kv_storage_manager.py, and tests/e2e/test_checkpoint_e2e.py all pass: 56 tests.
• The checkpoint-restart repro above fails on this PR and passes on the base commit.
• git merge-tree against current main reports a content conflict in transfer_queue/storage/managers/base.py.