wehub-resource-sync
655 lines
25 KiB
Python
655 lines
25 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
SerdeL2AdapterWrapper: wraps an inner L2 adapter with a SerdeProcessor
|
|
so controllers see a plain ``L2AdapterInterface`` while data is
|
|
transparently serialized on store and deserialized on load.
|
|
|
|
Threading: the wrapper owns an internal poll thread that reacts to
|
|
inner-adapter and serde event notifiers and chains
|
|
|
|
store : caller → serialize → inner.store → signal store_efd
|
|
load : caller → inner.load → deserialize → signal load_efd
|
|
|
|
Lookup / unlock / delete / eviction pass straight through to the inner
|
|
adapter (no serde transform involved).
|
|
|
|
Temp buffer lifecycle: temp byte buffers come from the injected
|
|
``L1Manager`` so the extra memory shows up in L1 accounting just like
|
|
the non-serde path's temporary KV buffers. For a store, the temp holds
|
|
serialized bytes; for a load, the temp catches the bytes L2 reads
|
|
before deserialize copies them into the caller-provided KV buffer.
|
|
|
|
Failure policy: all-or-nothing per submit. A partial temp-alloc
|
|
failure fails the whole task (``success=False`` for store, all-zeros
|
|
bitmap for load). This preserves the coarse-grained success semantic
|
|
of ``L2AdapterInterface`` and means the caller's lock / lifecycle
|
|
invariants don't need to change.
|
|
"""
|
|
|
|
# Future
|
|
from __future__ import annotations
|
|
|
|
# Standard
|
|
from dataclasses import dataclass, field
|
|
import enum
|
|
import select
|
|
import threading
|
|
|
|
# First Party
|
|
from lmcache.logging import init_logger
|
|
from lmcache.native_storage_ops import Bitmap
|
|
from lmcache.v1.distributed.api import KeyListPage, MemoryLayoutDesc, ObjectKey
|
|
from lmcache.v1.distributed.error import L1Error
|
|
from lmcache.v1.distributed.internal_api import L2AdapterListener, L2StoreResult
|
|
from lmcache.v1.distributed.l1_manager import L1Manager
|
|
from lmcache.v1.distributed.l2_adapters.base import (
|
|
AdapterUsage,
|
|
L2AdapterInterface,
|
|
L2TaskId,
|
|
)
|
|
from lmcache.v1.distributed.serde import (
|
|
SerdeProcessor,
|
|
SerdeTaskId,
|
|
make_temp_key,
|
|
serialized_layout_desc,
|
|
)
|
|
from lmcache.v1.memory_management import MemoryObj
|
|
from lmcache.v1.platform import consume_fd, create_event_notifier
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_POLL_TIMEOUT_MS = 500
|
|
|
|
|
|
class _StorePhase(enum.Enum):
|
|
SERIALIZE = enum.auto()
|
|
INNER_STORE = enum.auto()
|
|
|
|
|
|
@dataclass
|
|
class _StoreTaskState:
|
|
wrapped_id: L2TaskId
|
|
keys: list[ObjectKey]
|
|
temp_keys: list[ObjectKey]
|
|
temp_objs: list[MemoryObj]
|
|
phase: _StorePhase
|
|
"""SERIALIZE while temps are write-locked; INNER_STORE after the
|
|
serialize→store transition. Only read on shutdown to pick the right
|
|
lock-release path; assignment is done under ``self._lock``."""
|
|
|
|
|
|
@dataclass
|
|
class _LoadTaskState:
|
|
wrapped_id: L2TaskId
|
|
keys: list[ObjectKey]
|
|
dst_objs: list[MemoryObj]
|
|
temp_keys: list[ObjectKey]
|
|
temp_objs: list[MemoryObj]
|
|
load_bitmap: Bitmap = field(default_factory=lambda: Bitmap(0))
|
|
"""Inner adapter's per-key load bitmap; populated in
|
|
``_drain_inner_load`` before the task transitions to the deserialize
|
|
stage. ``Bitmap(0)`` means "not populated yet" — by the time
|
|
``_drain_deserialize`` reads it, this placeholder has been
|
|
overwritten with the real bitmap."""
|
|
|
|
|
|
class SerdeL2AdapterWrapper(L2AdapterInterface):
|
|
"""L2 adapter that adds transparent serde on top of an inner adapter.
|
|
|
|
Args:
|
|
inner: The wrapped L2 adapter doing the actual storage.
|
|
serde: The SerdeProcessor used to (de)serialize KV data.
|
|
l1_manager: L1 manager used to allocate temp byte buffers.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
inner: L2AdapterInterface,
|
|
serde: SerdeProcessor,
|
|
l1_manager: L1Manager,
|
|
) -> None:
|
|
super().__init__()
|
|
self._inner = inner
|
|
self._serde = serde
|
|
self._l1_manager = l1_manager
|
|
|
|
# Our own notifiers for store/load completion. Lookup passes the
|
|
# inner adapter's fd straight through (no chaining needed there).
|
|
self._store_efd = create_event_notifier()
|
|
self._load_efd = create_event_notifier()
|
|
|
|
# Task-id space separate from inner's. Reverse maps let the
|
|
# internal thread pair inner / serde completions back to our
|
|
# wrapped task id.
|
|
self._lock = threading.Lock()
|
|
self._next_task_id: L2TaskId = 0
|
|
self._store_tasks: dict[L2TaskId, _StoreTaskState] = {}
|
|
self._load_tasks: dict[L2TaskId, _LoadTaskState] = {}
|
|
self._serde_to_store: dict[SerdeTaskId, L2TaskId] = {}
|
|
self._inner_to_store: dict[L2TaskId, L2TaskId] = {}
|
|
self._inner_to_load: dict[L2TaskId, L2TaskId] = {}
|
|
self._serde_to_load: dict[SerdeTaskId, L2TaskId] = {}
|
|
|
|
# User-visible completion queues (drained by controller polls).
|
|
self._completed_store: dict[L2TaskId, L2StoreResult] = {}
|
|
self._completed_load: dict[L2TaskId, Bitmap] = {}
|
|
|
|
self._stop_flag = threading.Event()
|
|
self._thread = threading.Thread(
|
|
target=self._loop,
|
|
name="serde-l2-wrapper",
|
|
daemon=True,
|
|
)
|
|
self._thread.start()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Event fds
|
|
# ------------------------------------------------------------------
|
|
|
|
def get_store_event_fd(self) -> int:
|
|
return self._store_efd.fileno()
|
|
|
|
def get_load_event_fd(self) -> int:
|
|
return self._load_efd.fileno()
|
|
|
|
def get_lookup_and_lock_event_fd(self) -> int:
|
|
# Lookup doesn't touch serde; passing through the inner adapter's
|
|
# fd avoids a useless thread-hop per lookup.
|
|
return self._inner.get_lookup_and_lock_event_fd()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Store
|
|
# ------------------------------------------------------------------
|
|
|
|
def submit_store_task(
|
|
self,
|
|
keys: list[ObjectKey],
|
|
objects: list[MemoryObj],
|
|
) -> L2TaskId:
|
|
"""Submit a wrapped store (serialize → inner.store).
|
|
|
|
All-or-nothing: if temp alloc fails for any key or serialize
|
|
submission raises, the whole task is marked failed and the
|
|
caller's next ``pop_completed_store_tasks`` call sees it.
|
|
"""
|
|
with self._lock:
|
|
wrapped_id = self._next_task_id
|
|
self._next_task_id += 1
|
|
|
|
temp_keys, temp_objs = self._alloc_temp_buffers(keys, objects)
|
|
if temp_objs is None:
|
|
logger.warning(
|
|
"Serde wrapper: temp alloc failed for store task %d",
|
|
wrapped_id,
|
|
)
|
|
self._finalize_store(wrapped_id, success=False)
|
|
return wrapped_id
|
|
|
|
# Hold the wrapper lock across submit + reverse-map registration
|
|
# so the internal drain thread cannot observe a half-state where
|
|
# the serde already signaled completion but ``_serde_to_store``
|
|
# has no entry — which would leave the wrapped task hanging.
|
|
state = _StoreTaskState(
|
|
wrapped_id=wrapped_id,
|
|
keys=list(keys),
|
|
temp_keys=temp_keys,
|
|
temp_objs=temp_objs,
|
|
phase=_StorePhase.SERIALIZE,
|
|
)
|
|
try:
|
|
with self._lock:
|
|
self._store_tasks[wrapped_id] = state
|
|
serde_task_id = self._serde.submit_serialize(objects, temp_objs)
|
|
self._serde_to_store[serde_task_id] = wrapped_id
|
|
except Exception:
|
|
logger.exception(
|
|
"Serde wrapper: submit_serialize raised for store task %d",
|
|
wrapped_id,
|
|
)
|
|
with self._lock:
|
|
self._store_tasks.pop(wrapped_id, None)
|
|
self._release_write_temps(temp_keys)
|
|
self._finalize_store(wrapped_id, success=False)
|
|
return wrapped_id
|
|
return wrapped_id
|
|
|
|
def pop_completed_store_tasks(self) -> dict[L2TaskId, L2StoreResult]:
|
|
with self._lock:
|
|
result = self._completed_store
|
|
self._completed_store = {}
|
|
return result
|
|
|
|
# ------------------------------------------------------------------
|
|
# Lookup / unlock (pure delegation)
|
|
# ------------------------------------------------------------------
|
|
|
|
def submit_lookup_and_lock_task(
|
|
self, keys: list[ObjectKey], layout_desc: MemoryLayoutDesc
|
|
) -> L2TaskId:
|
|
return self._inner.submit_lookup_and_lock_task(keys, layout_desc)
|
|
|
|
def query_lookup_and_lock_result(self, task_id: L2TaskId) -> Bitmap | None:
|
|
return self._inner.query_lookup_and_lock_result(task_id)
|
|
|
|
def submit_unlock(self, keys: list[ObjectKey]) -> None:
|
|
self._inner.submit_unlock(keys)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Load
|
|
# ------------------------------------------------------------------
|
|
|
|
def submit_load_task(
|
|
self,
|
|
keys: list[ObjectKey],
|
|
objects: list[MemoryObj],
|
|
) -> L2TaskId:
|
|
"""Submit a wrapped load (inner.load → deserialize).
|
|
|
|
All-or-nothing: if temp alloc or inner submission fails, the
|
|
caller gets an all-zeros bitmap on next ``query_load_result``.
|
|
"""
|
|
with self._lock:
|
|
wrapped_id = self._next_task_id
|
|
self._next_task_id += 1
|
|
|
|
temp_keys, temp_objs = self._alloc_temp_buffers(keys, objects)
|
|
if temp_objs is None:
|
|
logger.warning(
|
|
"Serde wrapper: temp alloc failed for load task %d",
|
|
wrapped_id,
|
|
)
|
|
self._finalize_load(wrapped_id, Bitmap(len(keys)))
|
|
return wrapped_id
|
|
|
|
# Hold the wrapper lock across submit + reverse-map registration
|
|
# so the internal drain thread cannot observe a half-state where
|
|
# the inner already signaled completion but ``_inner_to_load``
|
|
# has no entry.
|
|
state = _LoadTaskState(
|
|
wrapped_id=wrapped_id,
|
|
keys=list(keys),
|
|
dst_objs=list(objects),
|
|
temp_keys=temp_keys,
|
|
temp_objs=temp_objs,
|
|
)
|
|
try:
|
|
with self._lock:
|
|
self._load_tasks[wrapped_id] = state
|
|
inner_task_id = self._inner.submit_load_task(keys, temp_objs)
|
|
self._inner_to_load[inner_task_id] = wrapped_id
|
|
except Exception:
|
|
logger.exception(
|
|
"Serde wrapper: inner.submit_load_task raised for task %d",
|
|
wrapped_id,
|
|
)
|
|
with self._lock:
|
|
self._load_tasks.pop(wrapped_id, None)
|
|
self._release_write_temps(temp_keys)
|
|
self._finalize_load(wrapped_id, Bitmap(len(keys)))
|
|
return wrapped_id
|
|
return wrapped_id
|
|
|
|
def query_load_result(self, task_id: L2TaskId) -> Bitmap | None:
|
|
with self._lock:
|
|
return self._completed_load.pop(task_id, None)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Eviction / metadata / listeners (delegate to inner)
|
|
# ------------------------------------------------------------------
|
|
|
|
@property
|
|
def inner_adapter(self) -> L2AdapterInterface:
|
|
"""Return the wrapped L2 adapter."""
|
|
return self._inner
|
|
|
|
@property
|
|
def supports_global_eviction(self) -> bool:
|
|
return self._inner.supports_global_eviction
|
|
|
|
def get_usage(self) -> AdapterUsage:
|
|
return self._inner.get_usage()
|
|
|
|
def delete(self, keys: list[ObjectKey]) -> None:
|
|
self._inner.delete(keys)
|
|
|
|
def list_l2_keys(
|
|
self,
|
|
model_name: str | None = None,
|
|
page_size: int = 500,
|
|
cursor: str | None = None,
|
|
) -> KeyListPage:
|
|
return self._inner.list_l2_keys(
|
|
model_name=model_name,
|
|
page_size=page_size,
|
|
cursor=cursor,
|
|
)
|
|
|
|
def register_listener(self, listener: L2AdapterListener) -> None:
|
|
# Listeners track what's actually stored — which is inner's job.
|
|
self._inner.register_listener(listener)
|
|
|
|
def report_status(self) -> dict:
|
|
inner_status = self._inner.report_status()
|
|
return {**inner_status, "serde_wrapped": True}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Lifecycle
|
|
# ------------------------------------------------------------------
|
|
|
|
def close(self) -> None:
|
|
self._stop_flag.set()
|
|
self._thread.join()
|
|
|
|
# Shut down the inner adapter and serde processor BEFORE
|
|
# releasing temp buffers. Both ``close()`` calls block until
|
|
# their in-flight reads / writes against the temp MemoryObjs
|
|
# finish; releasing the temps first would let L1 reclaim memory
|
|
# that the inner adapter or serde thread pool is still touching
|
|
# (use-after-free).
|
|
self._inner.close()
|
|
self._serde.close()
|
|
|
|
# Now safe to release leftover temp buffers. Store tasks in
|
|
# SERIALIZE phase hold write locks on their temps; tasks in
|
|
# INNER_STORE phase hold read locks (transitioned after
|
|
# serialize). Load tasks always hold write locks.
|
|
with self._lock:
|
|
write_locked: list[ObjectKey] = []
|
|
read_locked: list[ObjectKey] = []
|
|
for s in self._store_tasks.values():
|
|
if s.phase is _StorePhase.SERIALIZE:
|
|
write_locked.extend(s.temp_keys)
|
|
else:
|
|
read_locked.extend(s.temp_keys)
|
|
for load in self._load_tasks.values():
|
|
write_locked.extend(load.temp_keys)
|
|
self._store_tasks.clear()
|
|
self._load_tasks.clear()
|
|
self._serde_to_store.clear()
|
|
self._inner_to_store.clear()
|
|
self._inner_to_load.clear()
|
|
self._serde_to_load.clear()
|
|
|
|
if write_locked:
|
|
try:
|
|
self._l1_manager.finish_write(write_locked)
|
|
self._l1_manager.delete(write_locked)
|
|
except Exception:
|
|
logger.exception(
|
|
"Serde wrapper: error releasing write-locked leftover temps"
|
|
)
|
|
if read_locked:
|
|
try:
|
|
self._l1_manager.finish_read(read_locked)
|
|
except Exception:
|
|
logger.exception(
|
|
"Serde wrapper: error releasing read-locked leftover temps"
|
|
)
|
|
|
|
self._store_efd.close()
|
|
self._load_efd.close()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Internal loop
|
|
# ------------------------------------------------------------------
|
|
|
|
def _loop(self) -> None:
|
|
poller = select.poll()
|
|
inner_store_efd = self._inner.get_store_event_fd()
|
|
inner_load_efd = self._inner.get_load_event_fd()
|
|
serialize_efd = self._serde.get_serialize_event_fd()
|
|
deserialize_efd = self._serde.get_deserialize_event_fd()
|
|
|
|
poller.register(inner_store_efd, select.POLLIN)
|
|
poller.register(inner_load_efd, select.POLLIN)
|
|
poller.register(serialize_efd, select.POLLIN)
|
|
poller.register(deserialize_efd, select.POLLIN)
|
|
|
|
while not self._stop_flag.is_set():
|
|
ready = poller.poll(_POLL_TIMEOUT_MS)
|
|
for fd, events in ready:
|
|
if not (events & select.POLLIN):
|
|
continue
|
|
try:
|
|
consume_fd(fd)
|
|
except OSError:
|
|
pass
|
|
try:
|
|
if fd == serialize_efd:
|
|
self._drain_serialize()
|
|
elif fd == inner_store_efd:
|
|
self._drain_inner_store()
|
|
elif fd == inner_load_efd:
|
|
self._drain_inner_load()
|
|
elif fd == deserialize_efd:
|
|
self._drain_deserialize()
|
|
except Exception:
|
|
logger.exception("Serde wrapper: internal loop error on fd %d", fd)
|
|
|
|
def _drain_serialize(self) -> None:
|
|
"""Poll pending serialize tasks; on success submit inner store."""
|
|
with self._lock:
|
|
pending = list(self._serde_to_store.keys())
|
|
for serde_id in pending:
|
|
result = self._serde.query_serialize_result(serde_id)
|
|
if result is None:
|
|
continue
|
|
with self._lock:
|
|
wrapped_id = self._serde_to_store.pop(serde_id, None)
|
|
state = (
|
|
self._store_tasks.get(wrapped_id)
|
|
if wrapped_id is not None
|
|
else None
|
|
)
|
|
if wrapped_id is None or state is None:
|
|
continue
|
|
|
|
if not result:
|
|
self._release_write_temps(state.temp_keys)
|
|
self._finalize_store(wrapped_id, success=False)
|
|
continue
|
|
|
|
# Serialize succeeded — transition temps write → read so inner
|
|
# can safely read them during the store.
|
|
self._l1_manager.finish_write_and_reserve_read(state.temp_keys)
|
|
try:
|
|
inner_id = self._inner.submit_store_task(state.keys, state.temp_objs)
|
|
except Exception:
|
|
logger.exception(
|
|
"Serde wrapper: inner.submit_store_task raised for task %d",
|
|
wrapped_id,
|
|
)
|
|
# Temps are now read-locked and temporary — finish_read
|
|
# is enough; the entries auto-delete.
|
|
self._l1_manager.finish_read(state.temp_keys)
|
|
self._finalize_store(wrapped_id, success=False)
|
|
continue
|
|
|
|
# Phase flip and reverse-map insert happen under the same
|
|
# lock so ``close()``'s cleanup can't observe a half-way
|
|
# transition.
|
|
with self._lock:
|
|
state.phase = _StorePhase.INNER_STORE
|
|
self._inner_to_store[inner_id] = wrapped_id
|
|
|
|
def _drain_inner_store(self) -> None:
|
|
"""Drain inner store completions; release temp read locks (auto-
|
|
delete) and finalize the wrapped tasks."""
|
|
completed = self._inner.pop_completed_store_tasks()
|
|
for inner_id, result in completed.items():
|
|
with self._lock:
|
|
wrapped_id = self._inner_to_store.pop(inner_id, None)
|
|
state = (
|
|
self._store_tasks.get(wrapped_id)
|
|
if wrapped_id is not None
|
|
else None
|
|
)
|
|
if wrapped_id is None:
|
|
logger.warning(
|
|
"Serde wrapper: inner store task %d has no wrapped id",
|
|
inner_id,
|
|
)
|
|
continue
|
|
if state is not None:
|
|
self._l1_manager.finish_read(state.temp_keys)
|
|
self._finalize_store(
|
|
wrapped_id, result.is_successful(), result.bytes_transferred()
|
|
)
|
|
|
|
def _drain_inner_load(self) -> None:
|
|
"""Drain inner load completions; on per-key success submit
|
|
deserialize, otherwise fail the keys immediately."""
|
|
with self._lock:
|
|
pending = list(self._inner_to_load.keys())
|
|
for inner_id in pending:
|
|
bitmap = self._inner.query_load_result(inner_id)
|
|
if bitmap is None:
|
|
continue
|
|
with self._lock:
|
|
wrapped_id = self._inner_to_load.pop(inner_id, None)
|
|
state = (
|
|
self._load_tasks.get(wrapped_id) if wrapped_id is not None else None
|
|
)
|
|
if wrapped_id is None or state is None:
|
|
continue
|
|
|
|
src_objs: list[MemoryObj] = []
|
|
dst_objs: list[MemoryObj] = []
|
|
for i in range(len(state.keys)):
|
|
if bitmap.test(i):
|
|
src_objs.append(state.temp_objs[i])
|
|
dst_objs.append(state.dst_objs[i])
|
|
|
|
if not src_objs:
|
|
# Inner loaded nothing — skip deserialize, finalize.
|
|
self._release_write_temps(state.temp_keys)
|
|
self._finalize_load(wrapped_id, bitmap)
|
|
continue
|
|
|
|
state.load_bitmap = bitmap
|
|
try:
|
|
serde_id = self._serde.submit_deserialize(src_objs, dst_objs)
|
|
except Exception:
|
|
logger.exception(
|
|
"Serde wrapper: submit_deserialize raised for task %d",
|
|
wrapped_id,
|
|
)
|
|
self._release_write_temps(state.temp_keys)
|
|
self._finalize_load(wrapped_id, Bitmap(len(state.keys)))
|
|
continue
|
|
with self._lock:
|
|
self._serde_to_load[serde_id] = wrapped_id
|
|
|
|
def _drain_deserialize(self) -> None:
|
|
"""Drain deserialize completions; report inner's load bitmap on
|
|
success, all-zeros on deserialize failure."""
|
|
with self._lock:
|
|
pending = list(self._serde_to_load.keys())
|
|
for serde_id in pending:
|
|
result = self._serde.query_deserialize_result(serde_id)
|
|
if result is None:
|
|
continue
|
|
with self._lock:
|
|
wrapped_id = self._serde_to_load.pop(serde_id, None)
|
|
state = (
|
|
self._load_tasks.get(wrapped_id) if wrapped_id is not None else None
|
|
)
|
|
if wrapped_id is None or state is None:
|
|
continue
|
|
|
|
if result:
|
|
# ``load_bitmap`` was populated by _drain_inner_load before
|
|
# the task was registered in ``_serde_to_load``.
|
|
final_bitmap = state.load_bitmap
|
|
else:
|
|
logger.warning(
|
|
"Serde wrapper: deserialize failed for task %d; "
|
|
"reporting all keys as failed",
|
|
wrapped_id,
|
|
)
|
|
final_bitmap = Bitmap(len(state.keys))
|
|
|
|
self._release_write_temps(state.temp_keys)
|
|
self._finalize_load(wrapped_id, final_bitmap)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _alloc_temp_buffers(
|
|
self,
|
|
keys: list[ObjectKey],
|
|
objects: list[MemoryObj],
|
|
) -> tuple[list[ObjectKey], list[MemoryObj] | None]:
|
|
"""Reserve one temp byte buffer per input key. All-or-nothing:
|
|
any single failure releases the partial successes and returns
|
|
``(temp_keys, None)``.
|
|
|
|
Args:
|
|
keys: Original logical keys; used only to derive temp keys.
|
|
objects: Source (store) or destination (load) MemoryObjs.
|
|
All entries must share a single ``(shape, dtype)`` — the
|
|
caller (store/prefetch controller) is responsible for
|
|
shape-grouping before submission.
|
|
"""
|
|
shape_0 = objects[0].get_shapes()
|
|
dtype_0 = objects[0].get_dtypes()
|
|
temp_keys = [make_temp_key(k) for k in keys]
|
|
layout = serialized_layout_desc(
|
|
MemoryLayoutDesc(shapes=shape_0, dtypes=dtype_0), self._serde
|
|
)
|
|
results = self._l1_manager.reserve_write(
|
|
keys=temp_keys,
|
|
is_temporary=[True] * len(temp_keys),
|
|
layout_desc=layout,
|
|
mode="new",
|
|
)
|
|
# First pass: collect every key whose reserve_write succeeded.
|
|
# We must scan the full list (not bail on the first failure)
|
|
# so a mixed-success result still releases all reserved keys.
|
|
successful_temp_keys: list[ObjectKey] = []
|
|
for temp_key in temp_keys:
|
|
r = results.get(temp_key)
|
|
if r is not None and r[0] == L1Error.SUCCESS:
|
|
successful_temp_keys.append(temp_key)
|
|
if len(successful_temp_keys) != len(temp_keys):
|
|
self._release_write_temps(successful_temp_keys)
|
|
return temp_keys, None
|
|
temp_objs = [results[tk][1] for tk in temp_keys]
|
|
return temp_keys, temp_objs
|
|
|
|
def _release_write_temps(self, temp_keys: list[ObjectKey]) -> None:
|
|
"""Release write-locked temps and delete them. No-op on empty."""
|
|
if not temp_keys:
|
|
return
|
|
try:
|
|
self._l1_manager.finish_write(temp_keys)
|
|
self._l1_manager.delete(temp_keys)
|
|
except Exception:
|
|
logger.exception("Serde wrapper: failed releasing write-locked temps")
|
|
|
|
def _finalize_store(
|
|
self,
|
|
wrapped_id: L2TaskId,
|
|
success: bool,
|
|
bytes_transferred: int = 0,
|
|
) -> None:
|
|
with self._lock:
|
|
self._store_tasks.pop(wrapped_id, None)
|
|
self._completed_store[wrapped_id] = L2StoreResult(
|
|
success, bytes_transferred
|
|
)
|
|
try:
|
|
self._store_efd.notify()
|
|
except OSError:
|
|
logger.exception("Serde wrapper: failed to signal store notifier")
|
|
|
|
def _finalize_load(self, wrapped_id: L2TaskId, bitmap: Bitmap) -> None:
|
|
with self._lock:
|
|
self._load_tasks.pop(wrapped_id, None)
|
|
self._completed_load[wrapped_id] = bitmap
|
|
try:
|
|
self._load_efd.notify()
|
|
except OSError:
|
|
logger.exception("Serde wrapper: failed to signal load notifier")
|