Propagate shard clear failures before releasing ownership

[rawsh]

Surface shard clear failures after every shard settles, before the client releases controller metadata and allocator ownership.
Keep uncleared payloads owned and permit an idempotent retry after the shard recovers.
Proof: allocated job140204 executes 6/6 candidate service controls with zero skips; exact-main controls expose4 failures and2 positive passes; Ruff and compileall pass. Full suite and runtime-image adoption remain pending.

[0oshowero0]

[AI Review]

Thanks for tracking this down. The direction is right: today clear_data logs and swallows shard failures, so the client goes on to release controller metadata and indexes while the failed unit still holds the payload (which keeps counting against that unit's _active_keys capacity and leaves SSD-offload files behind). Raising after every shard settles is a small, correct fix, and it aligns SimpleStorage with the KV backends, whose clear_data already propagates errors. A few things should change before merging.

Verification

• The 6 new tests pass. With the source change reverted, the 4 failure-path cases fail and the 2 success-path cases pass, so they do exercise the fix.
• Full unit suite (pytest tests --ignore=tests/e2e): one failure introduced by this PR (see 1 below). The Yuanrong errors are environment-only (datasystem not installed). e2e tests were not run.
• The branch merges cleanly into current main.

Must fix

1. Breaks an existing test. test_async_storage_manager_error_handling asserts that clear_data does not raise. It should become with pytest.raises(RuntimeError, match="Mock CLEAR error"), which is also the natural unit test for this change.
2. The new test is order-dependent. The module fixture asserts not ray.is_initialized(), but tests/test_simple_storage_unit.py has a session-scoped ray.init that stays alive until the session ends. pytest tests/test_simple_storage_unit.py::test_clear_data tests/test_clear_failures.py errors in all 4 failure-path cases; the default run only passes because clear sorts before simple. Use the if not ray.is_initialized(): ray.init(...) pattern from the e2e fixtures, or see 3.

Suggested

3. The test is heavy for a 6-line fix and coupled to internals: the pickled controller checkpoint layout (state["index_manager"]["allocated_indexes"]), the private _group_by_hash / _get_from_single_storage_unit driven through asyncio.run on a second event loop next to the client's own loop, and an unused controller handle. The invariant comes down to two properties that the existing fixtures can cover in a few lines (untested sketch):

   # test_async_simple_storage_manager.py: every unit settles before the error surfaces
   attempted = []
   
   async def fake_clear(global_indexes, target_storage_unit=None, **kwargs):
       attempted.append(target_storage_unit)
       if target_storage_unit == "storage_0":
           raise RuntimeError("Mock CLEAR error")
   
   manager._clear_single_storage_unit = fake_clear
   with pytest.raises(RuntimeError, match="Mock CLEAR error"):
       await manager.clear_data(batch_meta)
   assert sorted(attempted) == ["storage_0", "storage_1"]
   
   # test_client.py (client_setup): a storage failure must not release controller metadata
   client.storage_manager.clear_data = AsyncMock(side_effect=RuntimeError("unit down"))
   client._clear_meta_in_controller = AsyncMock()
   with pytest.raises(RuntimeError, match="unit down"):
       client.clear_samples(meta)
   client._clear_meta_in_controller.assert_not_called()

   If an end-to-end check is still wanted, tests/e2e/test_e2e_lifecycle_consistency.py already has suitable fixtures.

4. Document the contract in docstrings instead of the inline comment. "Keep controller ownership until every physical clear has succeeded" is the caller's invariant. Please state in StorageManager.clear_data and the SimpleStorage implementation that it raises if any unit fails, that every unit is attempted first, and that retrying is safe.

5. Nit: the error is wrapped three times (Failed to clear storage unit X: Failed to clear storage X: ..., then Error in clear_samples operation: ... in the client). _clear_single_storage_unit already names the unit and logs, so raise result is enough.

Reliability notes

Checked and fine: surviving units always finish before the error surfaces; retry is idempotent (StorageUnitData.clear uses pop(key, None), mark_clearing is idempotent, clear_partition re-fetches metadata); the socket pool discards a failed lease, so a retry gets a fresh socket; zip(routing, results, strict=True) matches the task order.

Worth addressing or at least noting:

• Post-failure state. After a failed clear, the samples are invisible to consumers (production status zeroed by MARK_CLEARING) but keep their indexes and the failed unit's payload until the caller retries. That is the right trade-off, but the client docstrings should tell callers to retry; a caller that swallows the exception now leaks controller indexes instead of storage.
• Transient timeouts now reach users. _clear_single_storage_unit has no retry. On main, put/get/update go through _request_with_retry, which retries StorageUnitTimeout; clear is idempotent and a good fit for it. Fine as a follow-up PR.
• Field-level clear after rebase. [feat] add kv_update and field-level kv_clear #182 added kv_clear(fields=...), which updates the controller before deleting from storage. With this PR that path starts raising too, but a retry is not guaranteed to recover: kv_clear_fields filters by partition.field_metadata, and clear_fields drops a field once no sample holds it, so the retry can get empty metadata and never re-delete the leftover payload. Please rebase and either call this out as out of scope or open a follow-up.

PR description

Please rewrite it for external reviewers: the problem, the behavior change (clear failures now reach the caller), how to recover (retry the same call), and what was verified. Internal job ids and terms like "candidate service controls" are hard to follow.