Buffer Design¶

{
    "dl_config": {
        "use_legacy_buffering": true
    }
}

The new buffering system transitions from a monolithic design to a decoupled, policy-driven architecture. By separating how data is stored from how it is sampled or evicted, the system becomes highly extensible and easier to test.

Policy-based Approach¶

• Readiness Policy: Determines if the buffer has enough data to start training (e.g., waiting for a threshold).
• Sampling Policy: Defines how items are selected for the trainer (e.g., FIFO, Random, or Importance Sampling).
• Eviction Policy: Defines which items to remove when the buffer is full (e.g., drop oldest, drop random, or block).

Components¶

Storage (storage.py)¶

Storage classes are containers responsible only for the physical layout of data in memory.

• BaseStorage: An abstract interface requiring __len__ and append.

• DequeStorage: For standard FIFO operations (fast pops from the left).

• ListStorage: Designed for Reservoir-like sampling, maintaining separate lists for seen and not_seen samples to track data usage. This can also act as a single list for general-purpose usage.

Policies (policy.py)¶

• ReadinessPolicy: Manages the transition between the Reception Phase (simulation running) and the Draining Phase (simulation finished).

• SamplingPolicy: This allows a single sampler to act as a "Reservoir" (keep data) during reception and a "Queue" (remove data) during the final flush.

• EvictionPolicy: Handles overflow logic separately from sampling logic when _on_full is called. (e.g, Eviction on write in case of Reservoir, Do nothing (block) in case of FIFO, FIRO)

The Queue (base.py)¶

BaseQueue handles the thread synchronization (threading.Condition), while PolicyQueue acts as the orchestrator.

• signal_reception_over(): This method propagates a signal through the policies, switching the buffer from "accumulation mode" to "drain mode."

• _post_reception_error_handling(): A mechanism that monitors buffer size during the drain phase. If the buffer stops shrinking after the reception is over, it raises a FatalError to prevent infinite training loops.

The "Flush/Drain Phase" Lifecycle¶

1. Reception: readiness.reception_over is False. The SamplingPolicy yields data but keeps it in Storage.

2. Signal: signal_reception_over() is called.

3. Draining: The ReadinessPolicy now ignores thresholds (allowing the last few items to be read). The SamplingPolicy receives evict=True and begins removing items from Storage.

4. Validation: The PolicyQueue verifies that the storage size decreases with every get().

Adding a Custom Buffering Scheme¶

To implement a custom buffering scheme, you must define the Storage (how data is held), the SamplingPolicy (how data is retrieved and/or evicted at the end), and the EvictionPolicy (how to handle a full buffer). The core PolicyQueue then orchestrates these components, handling thread safety and the transition from data reception to the final draining phase.

To implement a priority-based scheme, you need a Sampling Policy that evaluates the "priority" of each item in the storage and a Storage backend that supports searching or filtering.

Example Implementation: Priority Reservoir¶

In this example, we use ListStorage and a placeholder priority_criterion function. This demonstrates how to identify the "best" item to sample and, more importantly, how to ensure that same logic handles eviction on read during the drain phase. Unlike the regular Reservoir that randomly samples items from the storage, this version balances exploration (random sampling) with exploitation (priority sampling) using a priority_ratio.

It also ensures that the _get method correctly bridges the internal state of the buffer (the "drain" status) with the logic required by the sampler.

import random
import numpy as np
from typing_extensions import override

from melissa.server.deep_learning.metric_logger import BaseLogger
from melissa.server.deep_learning.buffer import (
    PolicyQueue,
    SamplingPolicy,
    ThresholdReadiness,
    ListStorage,
    Sample,
)
from melissa.server.deep_learning.buffer.reservoir import BasicReservoirEviction


def calculate_priority(sample: Sample) -> float:
    """Example function to determine the importance of a sample.
    Lower 'seen' count results in higher priority to ensure 
    data variety during training."""
    return 1.0 / (sample.seen + 1) 


class PrioritySampler(SamplingPolicy[ListStorage]):
    @override
    def sample(self, storage: ListStorage, evict: bool, **kwargs) -> Sample:

        # probability of picking the 'best' item vs a random item
        ratio = kwargs.get("priority_ratio", 0.5)

        if np.random.rand() > ratio:
            # exploration
            idx = random.randrange(len(storage))
        else:
            # exploitation
            idx = max(
                range(len(storage)), 
                key=lambda i: calculate_priority(storage[i])
            )

        # evict while draining after reception is over
        if evict:
            return storage.pop(idx)

        # return the item but keep it in storage for reuse
        return storage[idx]


class PriorityReservoir(PolicyQueue[ListStorage]):
    def __init__(
        self,
        maxsize: int,
        threshold: int,
        metric_logger: BaseLogger | None = None,
        priority_ratio: float = 0.6,
    ):
        super().__init__(
            maxsize=maxsize,
            readiness=ThresholdReadiness(threshold), # begin sampling after threshold is met
            sampler=PrioritySampler(),               # priority sampling defined above
            evictor=BasicReservoirEviction(),        # random eviction of seen items
            storage_factory=ListStorage              # unified list storage of seen and unseen lists
            metric_logger=metric_logger,
        )
        self.priority_ratio = priority_ratio

    @override
    def _get(self) -> Sample:
        """Overrides the internal get logic to pass the priority."""
        return self.sampler.sample(
            self.queue,
            evict=self.readiness.reception_over,
            priority_ratio=self.priority_ratio,
        )

Finally, update types.py and __init__.py from the melissa/server/deep_learning/buffer module to add the custom implementation.