_named_data_store.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

import hashlib
import math
from dataclasses import dataclass

# from dataclasses import dataclass
from typing import Dict, List, Optional


@dataclass
class BufferEntry:
    """A class to hold the buffer entries for serialization.

    Attributes:
        buffer: The buffer bytes.
        alignment: The alignment of the buffer.
    """

    buffer: bytes
    alignment: int


@dataclass
class NamedDataStoreOutput:
    """
    Holds named data for serialization.

    Attributes:
        buffers: A list of unique buffer entries.
        pte_data: Contains data that is stored inside the PTE file. A mapping from
            {key: buffer_index}.
        external_data: Contains data that is stored external to the PTE. A mapping
            from {filename: {key: buffer_index}}.
    """

    buffers: List[BufferEntry]
    pte_data: Dict[str, int]
    external_data: Dict[str, Dict[str, int]]


class NamedDataStore:
    """
    NamedDataStore manages the data that delegates want to share. Backends add
    bytes to the store under a unique key. These bytes can be retrieved at
    runtime using the same key with the NamedDataMap.

    Note:
    - Keys are unique in the data store, regardless of whether they are stored
        in the PTE or externally.
    - Multiple keys can point to the same buffer entry.
    - The same data can be added multiple times and all keys will point to one
        buffer. If a duplicate blob is added with a different alignment, the
        lcm of the current and new alignment is taken for that blob.
    """

    # List of unique blobs.
    buffers: List[BufferEntry]
    # Named data stored inside the PTE file. Map of {key: buffer_index}.
    pte_data: Dict[str, int]
    # Named data stored outside of the PTE file.
    # Map of {filename: {key: buffer_index}}.
    external_data: Dict[str, Dict[str, int]]

    # Cache of the data hash for deduplication.
    # Use a hash instead of the data as a key because a sha256 collision is
    # unlikely, and the data may be large.
    data_hash_to_buffer_idx: Dict[bytes, int]
    # Cache of the key to buffer idx to ensure uniqueness.
    # If a key is added multiple times, check the buffer idx to ensure that the
    # data is identical too.
    key_to_buffer_idx: Dict[str, int]

    def __init__(self) -> None:
        """
        Initializes a new NamedDataStore.
        """
        self.buffers = []
        self.pte_data = {}
        self.external_data = {}

        self.data_hash_to_buffer_idx = {}
        self.key_to_buffer_idx = {}

    def _add_named_data_to_map(
        self,
        key: str,
        data: bytes,
        alignment: int,
        local_key_to_buffer_idx: Dict[str, int],
    ) -> None:
        """
        Add data to a map and update the alignment. Ensure that the key-data
        pair is unique.
        - If the key exists, the data must be identical.
        - If multiple unique keys exist for the same data, those keys should
            point to the same buffer.

        Args:
            key (str): key associated with the data.
            data (bytes): Bytes being requested to be serialized.
            alignment (int): alignment for bytes to be serialized with.
            local_key_to_buffer_idx (Dict[str, int]): map to add the data to.
        Raises:
            ValueError: when the key exists in the store, and corresponding data
                is different.
        """
        # Get data hash.
        hashed = hashlib.sha256(data).digest()

        # Check if the key exists.
        buffer_idx = self.key_to_buffer_idx.get(key, -1)
        if buffer_idx != -1:
            # If the key exists, the corresponding data must be identical.
            if self.data_hash_to_buffer_idx.get(hashed, -1) != buffer_idx:
                raise ValueError(
                    f"Duplicate key {key} with different data. "
                    f"Existing data: {self.buffers[buffer_idx].buffer}. "
                    f"New data: {data}."
                )
            self.buffers[buffer_idx].alignment = math.lcm(
                self.buffers[buffer_idx].alignment, alignment
            )
        else:
            # Key doesn't exist; check if the data exists.
            buffer_idx = self.data_hash_to_buffer_idx.get(hashed, -1)
            if buffer_idx != -1:
                # The data exists; update the alignment.
                self.buffers[buffer_idx].alignment = math.lcm(
                    self.buffers[buffer_idx].alignment, alignment
                )
            else:
                # The data doesn't exist; add it to the data store.
                buffer_idx = len(self.buffers)
                self.buffers.append(BufferEntry(data, alignment))
                self.data_hash_to_buffer_idx[hashed] = buffer_idx

            # Add key to the map and the key cache.
            local_key_to_buffer_idx[key] = buffer_idx
            self.key_to_buffer_idx[key] = buffer_idx

    def add_named_data(
        self,
        key: str,
        data: bytes,
        alignment: Optional[int] = 1,
        external_tag: Optional[str] = None,
    ) -> None:
        """
        Adds a named blob to the NamedDataStore.
        Args:
            key (str): key associated with the data.
            data (bytes): Bytes being requested to be serialized.
            alignment (int): alignment for bytes to be serialized with.
            external (Optional[str]): the external filename that this data is saved to.
        Raises:
            ValueError: when the key exists in the store, and corresponding data
                is different.
        """

        # Set default alignment.
        if alignment is None:
            alignment = 1
        if alignment <= 0:
            raise ValueError(f"Alignment must be greater than 0, received {alignment}.")

        if external_tag is None:
            self._add_named_data_to_map(key, data, alignment, self.pte_data)
        else:
            self._add_named_data_to_map(
                key, data, alignment, self.external_data.setdefault(external_tag, {})
            )

    def get_named_data_store_output(self) -> NamedDataStoreOutput:
        # Clean up empty maps inside self.external_data
        self.external_data = {k: v for k, v in self.external_data.items() if len(v) > 0}
        return NamedDataStoreOutput(self.buffers, self.pte_data, self.external_data)

    def merge_named_data_store(self, other: NamedDataStoreOutput) -> None:
        """
        Merge another NamedDataStore into this one.
        Args:
            other (NamedDataStore): the other NamedDataStore to merge.
        Raises:
            ValueError: when the key exists in both stores, and corresponding
                data is different between them.
        """
        # Merge the pte_data.
        for key, buffer_idx in other.pte_data.items():
            self.add_named_data(
                key,
                other.buffers[buffer_idx].buffer,
                other.buffers[buffer_idx].alignment,
            )

        # Merge the external_data.
        for filename, key_to_buffer_idx in other.external_data.items():
            for key, buffer_idx in key_to_buffer_idx.items():
                self.add_named_data(
                    key,
                    other.buffers[buffer_idx].buffer,
                    other.buffers[buffer_idx].alignment,
                    external_tag=filename,
                )