ReMe/deploy/benchmark_docker/_cleanup.py
xyf2020 48585f7523
feat(docker): Add Docker benchmark sandbox and pooled runner (#608)
* Add Docker benchmark sandbox and pooled runner

* Fix benchmark worker reuse after Docker command timeout

* Simplify benchmark worker requests and preserve artifacts

* Harden benchmark sandbox reuse and isolate query state

* Deduplicate benchmark sandbox cleanup and artifact checks

* Report benchmark cleanup failures and preserve diagnostic artifacts

* Handle optional Docker dependencies in benchmark unit tests

* Keep benchmark Docker dependencies optional in CI

* Fix cancellation of thread-pool background jobs

* Finish background worker cleanup before propagating cancellation

* Revert background worker cancellation changes

* Fix benchmark worker plugin defaults before config validation

* Stop interrupted benchmark containers before collecting diagnostics

* Fix candidate exclusions and benchmark upload serialization

* Fix late Docker creation cleanup and preserve query artifacts

* Stream benchmark workspace archives and extract plan validation

---------

Co-authored-by: jinli.yl <jinli.yl@alibaba-inc.com>
2026-10-19 16:26:38 +08:00

52 lines
2.1 KiB
Python

"""Finish owned-resource cleanup before propagating caller cancellation."""
from __future__ import annotations
import asyncio
from collections.abc import Awaitable, Iterable
def format_error(error: BaseException) -> str:
"""Include cleanup notes when an exception becomes a result string."""
return "; ".join([f"{type(error).__name__}: {error}", *getattr(error, "__notes__", ())])
async def finish_cleanup(
operations: Iterable[Awaitable[None]],
) -> tuple[list[None | BaseException], asyncio.CancelledError | None]:
"""Collect all failures, shielding teardown from repeated cancellation.
The caller must propagate the returned cancellation after reporting errors.
Individual cleanup operations may still fail or cancel themselves.
"""
async def collect(operation: Awaitable[None]) -> None | BaseException:
try:
await operation
except BaseException as exc:
# gather(return_exceptions=True) replaces CancelledError and loses
# its notes. Keep the original exception for cleanup diagnostics.
return exc
return None
cleanup = asyncio.gather(*(collect(operation) for operation in operations))
cancelled = None
while not cleanup.done():
try:
await asyncio.shield(cleanup)
except asyncio.CancelledError as exc:
cancelled = exc
return cleanup.result(), cancelled
async def cleanup_on_error(operations: Iterable[tuple[str, Awaitable[None]]], error: BaseException) -> None:
"""Attach teardown failures to the primary error before propagating cancellation."""
operations = list(operations)
outcomes, cancelled = await finish_cleanup(operation for _, operation in operations)
for (name, _), outcome in zip(operations, outcomes):
if isinstance(outcome, BaseException):
error.add_note(f"{name} close failed: {format_error(outcome)}")
if cancelled is not None and not isinstance(error, asyncio.CancelledError):
for note in getattr(error, "__notes__", ()):
cancelled.add_note(note)
raise cancelled from error