replay-rec 0.20.3rc0__py3-none-any.whl → 0.21.0__py3-none-any.whl

replay/data/nn/parquet/impl/indexing.py

@@ -0,0 +1,123 @@
+from typing import Union
+
+import torch
+
+
+def raw_get_offsets(lengths: torch.LongTensor) -> torch.LongTensor:
+    """
+    Performs offset calculation, defined simply as a cumulative sum of
+    the provided lengths tensor.
+
+    :param lengths: A tensor containing lengths of each individual row in a dataset's column.
+    :return: A tensor of offsets for each row.
+    """
+    zero = torch.zeros((1,), device=lengths.device, dtype=torch.int64)
+    cumsum = torch.cumsum(lengths, dim=-1)
+    return torch.cat([zero, cumsum])
+
+
+def get_offsets(lengths: torch.LongTensor) -> torch.LongTensor:
+    """
+    Sanitizes row lengths, then calculates offsets for each row.
+    The calculation itself is performed via the ``raw_get_offsets`` method.
+
+    :param lengths: A tensor containing lengths of each individual row in a dataset's column.
+    :raises ValueError: If the lengths tensor is of invalid shape or contains negative values.
+
+    :return: A tensor of offsets for each row.
+    """
+    if lengths.ndim != 1:
+        msg = f"Lengths must be strictly 1D. Got {lengths.ndim}D."
+        raise ValueError(msg)
+    min_length = torch.min(lengths.detach()).cpu().item()
+    if min_length < 0:
+        msg = f"There is a negative length. Got {min_length}."
+        raise ValueError(msg)
+    return raw_get_offsets(lengths)
+
+
+LengthType = Union[int, torch.LongTensor]
+
+
+def raw_get_mask(
+    indices: torch.LongTensor,
+    offsets: torch.LongTensor,
+    length: LengthType,
+) -> tuple[torch.BoolTensor, torch.LongTensor]:
+    """
+    Performs mask construction.
+    Given the data itself, its offsets and the expected sequence length, returns two tensors.
+
+    The first tensor is the padding mask, where ``False`` represents a padded value that was not present in the data,
+    and ``True`` represents a real element from the dataset.
+
+    The second tensor is the data itself, left-padded with a 0 to the desired length.
+
+    :param indices: A tensor of indices to be sampled from the dataset.
+    :param offsets: A tensor containing individual offsets for each of the column's rows.
+    :param length: THe total number of elements in a dataset's column.
+
+    :return: Constructed mask.
+    """
+    length = torch.asarray(length, dtype=torch.int64, device=indices.device)
+
+    # For every "line", start element index matches the offset, while end is the offset of the next line
+    last = offsets[indices + 1]
+    first = offsets[indices + 0]
+
+    per_line = length - (last - first)
+
+    arange = torch.arange(length, dtype=torch.int64, device=offsets.device)
+    raw_indices = (first[:, None] - per_line[:, None]) + arange[None, :]
+    mask = (first[:, None] <= raw_indices) & (raw_indices < last[:, None])
+
+    assert torch.all(torch.sum(mask, dim=-1, dtype=torch.int64) == torch.minimum(last - first, length)).cpu().item()
+
+    output_indices = torch.where(mask, raw_indices, 0)
+    assert torch.all((torch.max(output_indices, dim=-1).values < last) | (last == first)).cpu().item()
+    return (mask, output_indices)
+
+
+def get_mask(
+    indices: torch.LongTensor,
+    offsets: torch.LongTensor,
+    length: LengthType,
+) -> tuple[torch.BoolTensor, torch.LongTensor]:
+    """
+    Perform input sanity checks, then contructs a mask from inputs.
+    The mask calculation itself is performed via the ``raw_get_mask`` method.
+
+    :param indices: A tensor of indices to be sampled from the dataset.
+    :param offsets: A tensor containing individual offsets for each of the column's rows.
+    :param length: THe total number of elements in a dataset's column.
+
+    :raises ValueError: When mishaped or otherwise invalid arguments are provided.
+    :raises IndexError: When sampling indices missing from dataset or none at all.
+    :raises RuntimeError: When provided tensors are not on the same device.
+
+    :return: Constructed mask.
+    """
+    if torch.asarray(length).cpu().item() < 1:
+        msg = f"Length must be a positive number. Got {length}"
+        raise ValueError(msg)
+    if torch.numel(indices) < 1:
+        msg = f"Indices must be non-empty. Got {torch.numel(indices)}."
+        raise IndexError(msg)
+    if indices.device != offsets.device:  # pragma: no cover
+        msg = f"Devices must match. Got {indices.device} vs {offsets.device}"
+        raise RuntimeError(msg)
+    if offsets.ndim != 1:
+        msg = f"Offsets must be strictly 1D. Got {offsets.ndim}D."
+        raise ValueError(msg)
+    min_index = torch.min(indices.detach()).cpu().item()
+    if min_index < 0:
+        msg = f"Index is too small. Got {min_index}."
+        raise IndexError(msg)
+    max_index = torch.max(indices.detach()).cpu().item()
+    if torch.numel(offsets) < max_index:
+        msg = f"Index is too large. Got {max_index}."
+        raise IndexError(msg)
+    if not torch.all(offsets[:-1] <= offsets[1:]).cpu().item():
+        msg = "Offset sequence is not monothonous."
+        raise ValueError(msg)
+    return raw_get_mask(indices, offsets, length)

replay/data/nn/parquet/impl/masking.py

@@ -0,0 +1,20 @@
+from typing import Callable
+
+from replay.data.nn.parquet.collate import general_collate
+from replay.data.nn.parquet.constants.batches import GeneralCollateFn
+from replay.data.nn.parquet.info.replicas import ReplicasInfo, ReplicasInfoProtocol
+
+DEFAULT_COLLATE_FN: GeneralCollateFn = general_collate
+
+DEFAULT_MASK_POSTFIX: str = "_mask"
+
+
+def default_make_mask_name(postfix: str) -> Callable[[str], str]:
+    def function(name: str) -> str:
+        return f"{name}{postfix}"
+
+    return function
+
+
+DEFAULT_MAKE_MASK_NAME = default_make_mask_name(DEFAULT_MASK_POSTFIX)
+DEFAULT_REPLICAS_INFO: ReplicasInfoProtocol = ReplicasInfo()

replay/data/nn/parquet/impl/named_columns.py

@@ -0,0 +1,100 @@
+from collections.abc import Sequence
+from typing import Callable
+
+import torch
+
+from replay.data.nn.parquet.impl.masking import DEFAULT_MAKE_MASK_NAME
+
+from .column_protocol import ColumnProtocol
+
+Batch = dict[str, torch.Tensor]
+
+
+def deduce_device(columns: Sequence[ColumnProtocol]) -> torch.device:
+    """
+    Sanity check for matching devices on all of dataset's columns.
+
+    :param columns: A list of dataset's column data.
+    :raises RuntimeError: If any of the columns have mismatching devices.
+    :return: The determined columns' device.
+    """
+    assert len(columns) > 0
+    device = columns[0].device
+
+    def is_correct_device(column: ColumnProtocol) -> bool:
+        return column.device == device
+
+    if not all(map(is_correct_device, columns)):  # pragma: no cover
+        msg = "Columns must be all on the same device."
+        raise RuntimeError(msg)
+    return device
+
+
+def deduce_length(columns: Sequence[ColumnProtocol]) -> int:
+    """
+    Sanity check for matching lengths on all of dataset's columns.
+
+    :param columns: A list of dataset's column data.
+    :raises RuntimeError: If any of the columns has less rows than others.
+    :return: The determined columns' length.
+    """
+    assert len(columns) > 0
+    length = columns[0].length
+
+    def is_correct_length(column: ColumnProtocol) -> bool:
+        return column.length == length
+
+    if not all(map(is_correct_length, columns)):
+        msg = "Columns must have the same lengths."
+        raise RuntimeError(msg)
+    assert length > 0
+    return length
+
+
+def deduce_length_device(columns: dict[str, ColumnProtocol]) -> tuple[int, torch.device]:
+    """A combination check for both matching devices and lengths."""
+    raw = [*columns.values()]
+    columns_length = deduce_length(raw)
+    columns_device = deduce_device(raw)
+    del raw
+    return (columns_length, columns_device)
+
+
+class NamedColumns:
+    """
+    Representation of a data batch read from the filesystem.
+    This representation contains all of the columns read into memory, as well as
+    metadata such as their length and current device.
+    """
+
+    def __init__(
+        self,
+        columns: dict[str, ColumnProtocol],
+        make_mask_name: Callable[[str], str] = DEFAULT_MAKE_MASK_NAME,
+    ) -> None:
+        """
+        :param columns: Column data read from the filesystem.
+        :param make_mask_name: A function generating matching mask names for each column.
+        """
+        self.columns_length, self.columns_device = deduce_length_device(columns)
+
+        self.columns = columns
+        self.make_mask_name = make_mask_name
+
+    @property
+    def length(self) -> int:
+        return self.columns_length
+
+    @property
+    def device(self) -> torch.device:
+        return self.columns_device
+
+    def __len__(self) -> int:
+        return self.columns_length
+
+    def __getitem__(self, indices: torch.LongTensor) -> Batch:
+        indices = indices.to(device=self.device)
+        result = {}
+        for name, column in self.columns.items():
+            result[self.make_mask_name(name)], result[name] = column[indices]
+        return result

replay/data/nn/parquet/impl/numeric_column.py

@@ -0,0 +1,110 @@
+from typing import Any, Optional
+
+import pyarrow as pa
+import torch
+
+from replay.data.nn.parquet.constants.device import DEFAULT_DEVICE
+from replay.data.nn.parquet.constants.metadata import DEFAULT_PADDING
+from replay.data.nn.parquet.metadata import Metadata, get_numeric_columns
+from replay.data.utils.typing.dtype import pyarrow_to_torch
+
+from .column_protocol import OutputType
+from .utils import ensure_mutable
+
+
+class NumericColumn:
+    """A representation of a numeric column, containing a single number in each of its rows."""
+
+    def __init__(
+        self,
+        data: torch.Tensor,
+        mask: Optional[torch.BoolTensor] = None,
+        padding: Any = DEFAULT_PADDING,
+    ) -> None:
+        """
+        :param data: A tensor containing column data.
+        :param mask: A mask tensor to differentiate real values from paddings. Default: ``None``.
+        :param padding: Padding to use for future indexing of non-existent data. Default: value of ``DEFAULT_PADDING``.
+        """
+        self.padding: Any = padding
+        self.data = data
+        self.mask = mask
+
+    @property
+    def length(self) -> int:
+        result = torch.numel(self.data)
+        if self.mask is not None:
+            assert result == torch.numel(self.mask)
+        return result
+
+    def __len__(self) -> int:
+        return self.length
+
+    @property
+    def device(self) -> torch.device:
+        result = self.data.device
+        if self.mask is not None:
+            assert result == self.mask.device
+        return result
+
+    def _get_mask(self, indices: torch.LongTensor) -> torch.BoolTensor:
+        mask = torch.ones_like(indices, dtype=torch.bool) if self.mask is None else self.mask[indices]
+        return mask
+
+    def __getitem__(self, indices: torch.LongTensor) -> OutputType:
+        indices = indices.to(device=self.device)
+        mask = self._get_mask(indices)
+        output = torch.where(mask, self.data[indices], self.padding)
+        return (mask, output)
+
+
+def to_torch(array: pa.Array, device: torch.device = DEFAULT_DEVICE, padding: Any = DEFAULT_PADDING) -> OutputType:
+    """
+    Converts a PyArrow array into a PyTorch tensor.
+
+    :param array: Original PyArrow array.
+    :param device: Target device to send the resulting tensor to. Default: value of ``DEFAULT_DEVICE``.
+    :param padding: Value to fill null values with. Default: value of to ``DEFAULT_PADDING``.
+
+    :return: A PyTorch tensor obtained from original array.
+    """
+    dtype = pyarrow_to_torch(array.type)
+
+    mask_torch = None
+    if array.null_count > 0:
+        mask_torch = torch.asarray(
+            ensure_mutable(array.is_valid().to_numpy(zero_copy_only=False)),
+            device=device,
+            dtype=torch.bool,
+        )
+
+    array_torch = torch.asarray(
+        ensure_mutable(array.fill_null(padding).to_numpy()),
+        device=device,
+        dtype=dtype,
+    )
+    return (mask_torch, array_torch)
+
+
+def to_numeric_columns(
+    data: pa.RecordBatch,
+    metadata: Metadata,
+    device: torch.device = DEFAULT_DEVICE,
+    padding: Any = DEFAULT_PADDING,
+) -> dict[str, NumericColumn]:
+    """
+    Converts a PyArrow batch of data to a set of ``NumericColumn``s.
+    This function filters only those columns matching its format from the full batch.
+
+    :param data: A PyArrow batch of column data.
+    :param metadata: Metadata containing information about columns' formats.
+    :param device: Target device to send column tensors to. Default: value of ``DEFAULT_DEVICE``
+    :param padding: Padding to use for future indexing of non-existent data. Default: value of ``DEFAULT_PADDING``.
+
+    :return: A dict of tensors containing dataset's numeric columns.
+    """
+    result = {}
+    for column_name in get_numeric_columns(metadata):
+        mask, torch_array = to_torch(data.column(column_name), device, padding)
+        result[column_name] = NumericColumn(data=torch_array, mask=mask, padding=padding)
+    return result

replay/data/nn/parquet/impl/utils.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+WRITEABLE_FLAG: str = "WRITEABLE"
+
+
+def ensure_mutable(array: np.array) -> np.array:
+    """
+    Ensures the resulting NumPy array is mutable by making a copy if it's not.
+
+    :param array: Array to be checked for mutability.
+    :return: Mutable copy of `array`.
+    """
+    if not array.flags[WRITEABLE_FLAG]:
+        result = array.copy()
+        assert result.flags[WRITEABLE_FLAG]
+        return result
+    return array

replay/data/nn/parquet/info/distributed_info.py

@@ -0,0 +1,40 @@
+from typing import Protocol
+
+import torch.distributed as dist
+
+
+class DistributedInfo:
+    """Wrapper class for Torch's distibuted environment metadata."""
+
+    def __iter__(self):
+        yield self.rank
+        yield self.world_size
+
+    @property
+    def is_distributed(self) -> bool:
+        if dist.is_available():
+            return dist.is_initialized()
+        return False
+
+    @property
+    def rank(self) -> int:
+        if self.is_distributed:
+            return dist.get_rank()
+        return 0
+
+    @property
+    def world_size(self) -> int:
+        if self.is_distributed:
+            return dist.get_world_size()
+        return 1
+
+
+class DistributedInfoProtocol(Protocol):
+    @property
+    def rank(self) -> int: ...
+
+    @property
+    def world_size(self) -> int: ...
+
+
+DEFAULT_DISTRIBUTED_INFO: DistributedInfo = DistributedInfo()

replay/data/nn/parquet/info/partitioning.py

@@ -0,0 +1,132 @@
+from functools import lru_cache
+from math import ceil
+from typing import Optional, Union
+
+import torch
+
+from replay.data.nn.parquet.constants.device import DEFAULT_DEVICE
+
+
+def validate_length(length: int) -> int:
+    if length < 1:
+        msg = f"Length is invalid. Got {length}."
+        raise ValueError(msg)
+    return length
+
+
+def validate_num_replicas(num_replicas: int) -> int:
+    if num_replicas < 1:
+        msg = f"Num Replicas is invalid. Got {num_replicas}."
+        raise ValueError(msg)
+    return num_replicas
+
+
+def validate_curr_replica(curr_replica: int, num_replicas: int) -> int:
+    num_replicas = validate_num_replicas(num_replicas)
+    if (curr_replica < 0) or (num_replicas <= curr_replica):
+        msg = f"Curr Replicas is invalid. Got {curr_replica}."
+        raise ValueError(msg)
+    return curr_replica
+
+
+@lru_cache
+def _partitioning_length(length: int, num_replicas: int) -> int:
+    length = validate_length(length)
+    num_replicas = validate_num_replicas(num_replicas)
+
+    result = length
+    if length % num_replicas != 0:
+        raw_per_replica = length / num_replicas
+        per_replica = ceil(raw_per_replica)
+        new_length = per_replica * num_replicas
+        assert (new_length - length) < num_replicas
+        result = new_length
+    assert result % num_replicas == 0
+    assert length <= result
+    return result
+
+
+def partitioning_length(length: int, num_replicas: int) -> int:
+    return _partitioning_length(length, num_replicas)
+
+
+@lru_cache
+def _partitioning_per_replica(length: int, num_replicas: int) -> int:
+    full_length = partitioning_length(length, num_replicas)
+    result = full_length // num_replicas
+    assert result <= length
+    assert result > 0
+    return result
+
+
+def partitioning_per_replica(length: int, num_replicas: int) -> int:
+    return _partitioning_per_replica(length, num_replicas)
+
+
+class Partitioning:
+    """Utility class for calculating valid indices across multiple replicas."""
+
+    def __init__(
+        self,
+        curr_replica: int,
+        num_replicas: int,
+        device: Union[torch.device, str] = DEFAULT_DEVICE,
+        generator: Optional[torch.Generator] = None,
+    ) -> None:
+        """
+        :param curr_replica: Id of the curreent replica.
+        :param num_replicas: Total number of active replicas.
+        :param device: Target device to send the indices tensor to.
+            Default: value of ``DEFAULT_DEVICE``.
+        :param generator: A pseudo-random number generator for index shuffling. Default: ``None``.
+        """
+        self.device = torch.device(device)
+        self.generator = generator
+        self.num_replicas = validate_num_replicas(num_replicas)
+        self.curr_replica = validate_curr_replica(curr_replica, self.num_replicas)
+
+    def generate_raw_indices(self, length: int) -> torch.LongTensor:
+        full_length = partitioning_length(length, self.num_replicas)
+
+        if self.generator is None:
+            raw_indices = torch.arange(full_length, dtype=torch.int64, device=self.device)
+        else:
+            raw_indices = torch.randperm(full_length, dtype=torch.int64, generator=self.generator)
+            raw_indices = raw_indices.to(device=self.device)
+
+        assert torch.max(raw_indices).cpu().item() < full_length
+        assert torch.numel(raw_indices) == full_length
+        assert raw_indices.device == self.device
+
+        return raw_indices
+
+    def replica_indices(self, raw_indices: torch.LongTensor) -> torch.LongTensor:
+        full_length = torch.numel(raw_indices)
+        slc = slice(self.curr_replica, full_length, self.num_replicas)
+        replica_indices = raw_indices[slc].clone()
+
+        assert torch.max(replica_indices).cpu().item() < full_length
+
+        return replica_indices
+
+    def generate(self, length: int) -> torch.LongTensor:
+        raw_indices = self.generate_raw_indices(length)
+        full_length = partitioning_length(length, self.num_replicas)
+
+        assert torch.numel(raw_indices) == full_length
+
+        replica_indices = self.replica_indices(raw_indices)
+        per_replica = partitioning_per_replica(length, self.num_replicas)
+
+        assert torch.numel(replica_indices) == per_replica
+
+        indices = torch.remainder(replica_indices, length)
+
+        assert torch.max(indices).cpu().item() < length
+        assert torch.numel(indices) == per_replica
+        assert indices.device == self.device
+
+        return indices
+
+    def __call__(self, length: int) -> torch.LongTensor:
+        return self.generate(length)

replay/data/nn/parquet/info/replicas.py

@@ -0,0 +1,67 @@
+from typing import Protocol
+
+from .distributed_info import DEFAULT_DISTRIBUTED_INFO, DistributedInfoProtocol
+from .worker_info import DEFAULT_WORKER_INFO, WorkerInfoProtocol
+
+
+def num_replicas(
+    worker_info: WorkerInfoProtocol = DEFAULT_WORKER_INFO,
+    distributed_info: DistributedInfoProtocol = DEFAULT_DISTRIBUTED_INFO,
+) -> int:
+    return worker_info.num_workers * distributed_info.world_size
+
+
+def curr_replica(
+    worker_info: WorkerInfoProtocol = DEFAULT_WORKER_INFO,
+    distributed_info: DistributedInfoProtocol = DEFAULT_DISTRIBUTED_INFO,
+) -> int:
+    result = worker_info.id + worker_info.num_workers * distributed_info.rank
+    assert result < num_replicas(worker_info, distributed_info)
+    return result
+
+
+class ReplicasInfoProtocol(Protocol):
+    @property
+    def num_replicas(self) -> int: ...
+
+    @property
+    def curr_replica(self) -> int: ...
+
+
+class ReplicasInfo:
+    """
+    A replica metadata geneartor.
+
+    By default, assumes standard Torch DDP training/inference procedure,
+    where each replica (a distinct worker on a specific device) is expected to process
+    a separate chunk of the dataset.
+
+    This behavior can be modified by providing custom ``worker_info`` and ``distributed_info`` objects
+    able to provide infor about local worker count and world size/rank respectively.
+    """
+
+    def __init__(
+        self,
+        worker_info: WorkerInfoProtocol = DEFAULT_WORKER_INFO,
+        distributed_info: DistributedInfoProtocol = DEFAULT_DISTRIBUTED_INFO,
+    ) -> None:
+        """
+        :param worker_info: An object adhering to the ``WorkerInfoProtocol`` and used to obtain local worker count.
+            Default: value of ``DEFAULT_WORKER_INFO`` - an implementation using ``torch.utils.data.get_worker_info()``.
+        :param distributed_info: An object adhering to the ``DistributedInfoProtocol`` and used to obtain
+            world size and rank. Default: value of ``DEFAULT_WORKER_INFO`` - an implementation using the
+            ``torch.distributed`` module.
+        """
+        self.worker_info = worker_info
+        self.distributed_info = distributed_info
+
+    @property
+    def num_replicas(self) -> int:
+        return num_replicas(worker_info=self.worker_info, distributed_info=self.distributed_info)
+
+    @property
+    def curr_replica(self) -> int:
+        return curr_replica(worker_info=self.worker_info, distributed_info=self.distributed_info)
+
+
+DEFAULT_REPLICAS_INFO: ReplicasInfoProtocol = ReplicasInfo()

replay/data/nn/parquet/info/worker_info.py

@@ -0,0 +1,43 @@
+from typing import Any, Iterator, Optional, Protocol
+
+import torch.utils.data as data
+
+
+class WorkerInfoProtocol(Protocol):
+    @property
+    def id(self) -> int: ...
+
+    @property
+    def num_workers(self) -> int: ...
+
+
+class WorkerInfo:
+    """Wrapper class for Torch's worker metadata."""
+
+    def __iter__(self) -> Iterator[int]:
+        yield self.id
+
+    @property
+    def worker_info(self) -> Optional[Any]:
+        return data.get_worker_info()
+
+    @property
+    def is_parallel(self) -> bool:
+        return self.worker_info is not None
+
+    @property
+    def id(self) -> int:
+        wi: Optional[data.WorkerInfo] = self.worker_info
+        if wi is not None:
+            return wi.id
+        return 0
+
+    @property
+    def num_workers(self) -> int:
+        wi: Optional[data.WorkerInfo] = self.worker_info
+        if wi is not None:
+            return wi.num_workers
+        return 1
+
+
+DEFAULT_WORKER_INFO: WorkerInfo = WorkerInfo()