wehub-resource-sync
141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Warm-prefetch job table for the MP server.
|
|
|
|
A *warm* prefetch loads a caller-supplied set of keys from L2 into L1 and leaves
|
|
them resident, retained, and unpinned, so a subsequent real lookup hits L1. It
|
|
loads every requested key not already in L1, and pins nothing -- there is no
|
|
downstream reader to pin for.
|
|
|
|
The table tracks in-flight warm prefetches: ``submit`` starts one and returns an
|
|
opaque request id; ``poll`` reports its status (pending, then completed) and
|
|
drops it once the load finishes. Status is observed reactively by the caller --
|
|
there is no background polling -- and these calls do not block. A job whose
|
|
status is never polled to completion simply lingers until a later cleanup drops
|
|
it; since nothing is pinned, no L1 is held.
|
|
"""
|
|
|
|
# Standard
|
|
from dataclasses import dataclass
|
|
import threading
|
|
import uuid
|
|
|
|
# First Party
|
|
from lmcache.logging import init_logger
|
|
from lmcache.v1.distributed.api import (
|
|
MemoryLayoutDesc,
|
|
ObjectKey,
|
|
PrefetchHandle,
|
|
PrefetchMode,
|
|
TrimPolicy,
|
|
)
|
|
from lmcache.v1.distributed.storage_manager import StorageManager
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# Job states reported by :meth:`WarmPrefetchJobs.poll`.
|
|
PENDING = "pending"
|
|
COMPLETED = "completed"
|
|
UNKNOWN = "unknown"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WarmStatus:
|
|
"""Outcome of a status poll.
|
|
|
|
Attributes:
|
|
state: ``"pending"`` (load still running), ``"completed"`` (loaded), or
|
|
``"unknown"`` (no such request id -- already completed-and-consumed,
|
|
or never submitted).
|
|
found_keys: Keys loaded into L1 (only meaningful when ``completed``).
|
|
total_keys: Keys originally requested (only meaningful when
|
|
``completed``).
|
|
"""
|
|
|
|
state: str
|
|
found_keys: int = 0
|
|
total_keys: int = 0
|
|
|
|
|
|
class WarmPrefetchJobs:
|
|
"""Thread-safe table of in-flight warm-prefetch jobs (``request_id ->
|
|
handle``).
|
|
|
|
``submit`` starts a no-lock retain prefetch and returns an opaque request
|
|
id; ``poll`` reports progress and, on completion, drops the job
|
|
(exactly-once: a subsequent poll for the same id returns ``UNKNOWN``).
|
|
Since the warm holds no read lock, there is nothing to release on
|
|
completion.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize an empty table."""
|
|
self._lock = threading.Lock()
|
|
self._jobs: dict[str, PrefetchHandle] = {}
|
|
|
|
def submit(
|
|
self,
|
|
storage_manager: StorageManager,
|
|
keys: list[ObjectKey],
|
|
layout_desc: MemoryLayoutDesc,
|
|
) -> str:
|
|
"""Start a no-lock retain prefetch and register its handle.
|
|
|
|
Args:
|
|
storage_manager: The MP server's storage manager.
|
|
keys: Object keys to load from L2 into L1.
|
|
layout_desc: Memory layout for the L1 write buffers.
|
|
|
|
Returns:
|
|
An opaque request id to pass to :meth:`poll`.
|
|
"""
|
|
handle = storage_manager.submit_prefetch_task(
|
|
keys,
|
|
layout_desc,
|
|
mode=PrefetchMode.WARM,
|
|
policy=TrimPolicy.SPARSE,
|
|
)
|
|
request_id = uuid.uuid4().hex
|
|
with self._lock:
|
|
self._jobs[request_id] = handle
|
|
return request_id
|
|
|
|
def poll(
|
|
self,
|
|
storage_manager: StorageManager,
|
|
request_id: str,
|
|
) -> WarmStatus:
|
|
"""Report a job's status; drop it once the load completes.
|
|
|
|
Args:
|
|
storage_manager: The MP server's storage manager.
|
|
request_id: The id returned by :meth:`submit`.
|
|
|
|
Returns:
|
|
A :class:`WarmStatus`. The first poll that observes completion drops
|
|
the job (and pops the controller's result bookkeeping), so later
|
|
polls for the same id return ``UNKNOWN``.
|
|
"""
|
|
with self._lock:
|
|
handle = self._jobs.get(request_id)
|
|
if handle is None:
|
|
return WarmStatus(state=UNKNOWN)
|
|
|
|
found = storage_manager.query_prefetch_status(handle)
|
|
if found is None:
|
|
return WarmStatus(state=PENDING)
|
|
|
|
with self._lock:
|
|
self._jobs.pop(request_id, None)
|
|
found_keys = found.popcount()
|
|
logger.info(
|
|
"Warm prefetch %s completed: %d/%d keys loaded into L1",
|
|
request_id,
|
|
found_keys,
|
|
handle.total_requested_keys,
|
|
)
|
|
return WarmStatus(
|
|
state=COMPLETED,
|
|
found_keys=found_keys,
|
|
total_keys=handle.total_requested_keys,
|
|
)
|