class CPUOffloadingManager(OffloadingManager):
"""An OffloadingManager with a pluggable CachePolicy, resolved by name via
CachePolicyFactory (built in: "lru", "arc"; external policies can either
register their own or be loaded out-of-tree via cache_policy_module_path).
The manager owns all shared logic: ref-counting, event emission,
chunk pool management, and the prepare_store/complete_store skeletons.
Policy-specific chunk organization and eviction decisions are delegated
to the CachePolicy implementation.
"""
def __init__(
self,
num_chunks: int,
cache_policy: str = "lru",
cache_policy_module_path: str | None = None,
enable_events: bool = False,
store_threshold: int = 1,
max_tracker_size: int = 64_000,
):
self.medium: Medium = Medium.CPU
self._num_chunks: int = num_chunks
self._num_allocated_chunks: int = 0
self._free_list: list[int] = []
self.events: list[OffloadingEvent] | None = [] if enable_events else None
policy_cls = CachePolicyFactory.get_cache_policy_cls(
cache_policy, cache_policy_module_path
)
self._policy: CachePolicy = policy_cls(cache_capacity=num_chunks)
# Track the number of chunks in the cache that are evictable. i.e. ref_cnt 0.
self._num_evictable_cache_chunks: int = 0
# Track chunks with an in-flight store (ref_cnt -1, not yet completed).
self._num_write_pending_chunks: int = 0
self.store_threshold: int = store_threshold
self.max_tracker_size: int = max_tracker_size
self.stores_skipped_in_current_batch: int = 0
self.allocation_sizes_in_current_batch: list[int] = []
self._cache_generation = 0
# Number of chunk references. It is ordered so can evict the LRU entry in O(1).
self.counts: OrderedDict[OffloadKey, int] | None = (
OrderedDict() if store_threshold >= 2 else None
)
# --- chunk pool ---
def _get_num_free_chunks(self) -> int:
return len(self._free_list) + self._num_chunks - self._num_allocated_chunks
def _allocate_chunks(self, keys: list[OffloadKey]) -> list[ChunkStatus]:
num_fresh = min(len(keys), self._num_chunks - self._num_allocated_chunks)
num_reused = len(keys) - num_fresh
assert len(self._free_list) >= num_reused
# allocate fresh chunks
chunks: list[ChunkStatus] = []
for _ in range(num_fresh):
chunks.append(ChunkStatus(self._num_allocated_chunks))
self._num_allocated_chunks += 1
# allocate reused chunks
for _ in range(num_reused):
chunks.append(ChunkStatus(self._free_list.pop()))
return chunks
def _free_chunk(self, chunk: ChunkStatus) -> None:
self._free_list.append(chunk.chunk_id)
def _get_load_store_spec(
self,
keys: Iterable[OffloadKey],
chunks: Iterable[ChunkStatus],
) -> CPULoadStoreSpec:
return CPULoadStoreSpec([chunk.chunk_id for chunk in chunks])
def _record_accesses(self, keys: Collection[OffloadKey]) -> None:
"""Record an offer without evicting its tracked candidates."""
assert self.counts is not None
protected: set[OffloadKey] = set()
for key in keys:
if key in self.counts:
self.counts.move_to_end(key)
self.counts[key] += 1
protected.add(key)
num_unprotected = len(self.counts) - len(protected)
for key in keys:
if key in self.counts:
continue
if len(self.counts) >= self.max_tracker_size:
if num_unprotected == 0:
continue
self.counts.popitem(last=False)
num_unprotected -= 1
self.counts[key] = 1
def _get_request_cache_access(self, req_context: ReqContext) -> _RequestCacheAccess:
state = req_context.get_state(_RequestCacheAccess)
if (
state is None
or state.owner is not self
or state.cache_generation != self._cache_generation
):
state = _RequestCacheAccess(
owner=self, cache_generation=self._cache_generation
)
req_context.set_state(state)
return state
def _record_request_cache_access(
self,
keys: Iterable[OffloadKey],
req_context: ReqContext,
inserted_keys: Iterable[OffloadKey] = (),
reused_keys: Iterable[OffloadKey] = (),
) -> None:
state = self._get_request_cache_access(req_context)
for key in keys:
if key in state.seen_keys:
continue
assert not state.finished, (
"New cache keys observed after request finalization"
)
group_idx = get_offload_group_idx(key)
state.key_groups.setdefault(group_idx, []).append(key)
state.seen_keys.add(key)
state.inserted_keys.update(inserted_keys)
# Re-reading a chunk inserted by this request is an internal transfer
# (for example, a tiering cascade), not a second cache access.
state.reused_keys.update(
key for key in reused_keys if key not in state.inserted_keys
)
# --- OffloadingManager interface ---
@override
def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
self._get_request_cache_access(req_context)
return RequestOffloadingContext()
@override
def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
chunk = self._policy.get(key)
if chunk is None:
return LookupResult.MISS
if not chunk.is_ready:
return LookupResult.HIT_PENDING
return LookupResult.HIT
@override
def prepare_load(
self,
keys: Collection[OffloadKey],
req_context: ReqContext,
) -> LoadStoreSpec:
return self._prepare_load(keys, req_context, record_access=True)
def _prepare_load(
self,
keys: Collection[OffloadKey],
req_context: ReqContext,
*,
record_access: bool,
) -> LoadStoreSpec:
chunks = []
for key in keys:
chunk = self._policy.get(key)
assert chunk is not None, f"Chunk {key!r} not found in cache"
assert chunk.is_ready, f"Chunk {key!r} is not ready for reading"
if chunk.ref_cnt == 0:
self._policy.mark_non_evictable(key)
self._num_evictable_cache_chunks -= 1 # ref_cnt 0 -> 1
assert self._num_evictable_cache_chunks >= 0
chunk.ref_cnt += 1
chunks.append(chunk)
if record_access:
self._record_request_cache_access(keys, req_context, reused_keys=keys)
return self._get_load_store_spec(keys, chunks)
@override
def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
self._policy.touch(keys, req_context)
@override
def complete_load(
self, keys: Collection[OffloadKey], req_context: ReqContext
) -> None:
for key in keys:
chunk = self._policy.get(key)
assert chunk is not None, f"Chunk {key!r} not found"
assert chunk.ref_cnt > 0, f"Chunk {key!r} ref_cnt is already 0"
chunk.ref_cnt -= 1
if chunk.ref_cnt == 0:
self._num_evictable_cache_chunks += 1 # ref_cnt 1 -> 0
self._policy.mark_evictable(key)
@override
def prepare_store(
self,
keys: Collection[OffloadKey],
req_context: ReqContext,
) -> PrepareStoreOutput | None:
keys = list(keys)
if self.counts is not None:
self._record_accesses(keys)
# Partition the original offer before admission filtering. Ready
# resident chunks are request reuses even if their tracker entry was
# aged out; the threshold applies only to new store candidates.
keys_to_store: list[OffloadKey] = []
ready_existing_keys: list[OffloadKey] = []
for key in keys:
chunk = self._policy.get(key)
if chunk is None:
keys_to_store.append(key)
else:
if chunk.is_ready:
ready_existing_keys.append(key)
if self.counts is not None:
num_store_candidates = len(keys_to_store)
keys_to_store = [
key
for key in keys_to_store
if self.counts.get(key, 0) >= self.store_threshold
]
self.stores_skipped_in_current_batch += num_store_candidates - len(
keys_to_store
)
state = self._get_request_cache_access(req_context)
new_store_misses = [
key for key in keys_to_store if key not in state.store_miss_keys
]
if new_store_misses:
# ARC learns from B1/B2 before insert() removes the ghost entry.
# Deduplication makes this one policy observation per request.
self._policy.on_store_miss(new_store_misses, req_context)
state.store_miss_keys.update(new_store_misses)
recorded_keys = set(keys_to_store)
recorded_keys.update(ready_existing_keys)
# Record once before any allocation-related early return. Store
# candidates are classified as insertions only after allocation succeeds.
self._record_request_cache_access(
(key for key in keys if key in recorded_keys),
req_context,
reused_keys=ready_existing_keys,
)
if not keys_to_store:
return PrepareStoreOutput(
keys_to_store=[],
store_spec=self._get_load_store_spec([], []),
evicted_keys=[],
)
self.allocation_sizes_in_current_batch.append(len(keys_to_store))
num_chunks_to_evict = len(keys_to_store) - self._get_num_free_chunks()
to_evict: list[OffloadKey] = []
if num_chunks_to_evict > 0:
if num_chunks_to_evict > self._num_evictable_cache_chunks:
# Eviction will fail.
return None
# There is a still a chance for eviction failure as some of the
# idle chunks might be in the protected list.
# Chunks from the original input are excluded from eviction candidates:
# a chunk that was already stored must remain in the cache after this call.
protected = set(keys)
evicted = self._policy.evict(num_chunks_to_evict, protected)
if evicted is None:
return None
# cache-policy removes only idle chunks.
self._num_evictable_cache_chunks -= len(evicted)
assert self._num_evictable_cache_chunks >= 0
for key, chunk in evicted:
self._free_chunk(chunk)
to_evict.append(key)
if to_evict and self.events is not None:
self.events.append(
OffloadingEvent(
keys=to_evict,
medium=self.medium,
removed=True,
)
)
chunks = self._allocate_chunks(keys_to_store)
assert len(chunks) == len(keys_to_store), (
"Chunk pool did not allocate the expected number of chunks"
)
for key, chunk in zip(keys_to_store, chunks):
self._policy.insert(key, chunk)
self._num_write_pending_chunks += len(keys_to_store)
state.inserted_keys.update(keys_to_store)
# build store specs for allocated chunks
store_spec = self._get_load_store_spec(keys_to_store, chunks)
return PrepareStoreOutput(
keys_to_store=keys_to_store,
store_spec=store_spec,
evicted_keys=to_evict,
)
@override
def complete_store(
self,
keys: Collection[OffloadKey],
req_context: ReqContext,
success: bool = True,
) -> None:
stored_keys: list[OffloadKey] = []
if success:
for key in keys:
chunk = self._policy.get(key)
if chunk is not None and not chunk.is_ready:
chunk.ref_cnt = 0
self._num_write_pending_chunks -= 1
self._num_evictable_cache_chunks += 1
self._policy.mark_evictable(key)
stored_keys.append(key)
else:
for key in keys:
chunk = self._policy.get(key)
if chunk is not None and not chunk.is_ready:
self._num_write_pending_chunks -= 1
self._policy.remove(key)
self._free_chunk(chunk)
if stored_keys and self.events is not None:
self.events.append(
OffloadingEvent(
keys=stored_keys,
medium=self.medium,
removed=False,
)
)
@override
def on_request_finished(self, req_context: ReqContext) -> None:
state = req_context.get_state(_RequestCacheAccess)
if (
state is None
or state.owner is not self
or state.cache_generation != self._cache_generation
or state.finished
):
return
state.finished = True
key_groups = []
for group_idx in sorted(state.key_groups):
keys = state.key_groups[group_idx]
positions = {
key: position
for key in keys
if (position := req_context.get_offload_key_position(key)) is not None
}
if len(positions) == len(keys):
keys = sorted(keys, key=positions.__getitem__)
key_groups.append(tuple(keys))
self._policy.on_request_finished(
tuple(key_groups),
state.inserted_keys - state.reused_keys,
state.reused_keys,
req_context,
)
@override
def reset_cache(self) -> None:
# Clear ALL chunks unconditionally. The scheduler's _stale_job_threshold
# guarantees that complete_load / complete_store are never called for
# pre-reset jobs, so no lazy cleanup is needed. The scheduler also
# flushes in-flight load job IDs to the workers before any new stores
# can begin, preventing a cross-direction data race on reused offload chunk IDs.
self._policy.clear()
self._cache_generation += 1
self._num_evictable_cache_chunks = 0
self._num_write_pending_chunks = 0
self._free_list.clear()
self._num_allocated_chunks = 0
@override
def take_events(self) -> Iterable[OffloadingEvent]:
if self.events is not None:
yield from self.events
self.events.clear()
def get_stats(self) -> OffloadingConnectorStats | None:
stats = OffloadingConnectorStats()
# Compute cache usage.
num_used = (
self._num_allocated_chunks
- len(self._free_list)
- self._num_evictable_cache_chunks
)
usage = num_used / self._num_chunks if self._num_chunks > 0 else 0.0
stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC, usage)
for allocation_size in self.allocation_sizes_in_current_batch:
stats.observe_histogram(
CPUOffloadingMetrics.CPU_ALLOCATION_SIZE, allocation_size
)
self.allocation_sizes_in_current_batch.clear()
write_usage = (
self._num_write_pending_chunks / self._num_chunks
if self._num_chunks > 0
else 0.0
)
read_usage = max(usage - write_usage, 0.0)
stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC, write_usage)
stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC, read_usage)
if self.store_threshold >= 2:
stats.increase_counter(
CPUOffloadingMetrics.STORES_SKIPPED,
self.stores_skipped_in_current_batch,
)
self.stores_skipped_in_current_batch = 0
return stats