From 2f2e009506e20624472a25d2feaa1e7e784bd572 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 11:58:01 -0700 Subject: [PATCH 1/8] Add torch and lsdb dependencies and device helpers Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 2 ++ src/lsdb_ml/device.py | 77 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 src/lsdb_ml/device.py diff --git a/pyproject.toml b/pyproject.toml index 799fdb2..8f999d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,8 @@ classifiers = [ dynamic = ["version"] requires-python = ">=3.11" dependencies = [ + "lsdb>=0.11.0", # Brings nested-pandas, pyarrow and dask[complete] (incl. distributed) + "torch>=2.4", ] [project.urls] diff --git a/src/lsdb_ml/device.py b/src/lsdb_ml/device.py new file mode 100644 index 0000000..3dc0c55 --- /dev/null +++ b/src/lsdb_ml/device.py @@ -0,0 +1,77 @@ +"""Device resolution and batch movement helpers.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import torch + +DeviceLike = str | torch.device | None + + +def resolve_device(device: DeviceLike = "cpu") -> torch.device: + """Turn a device request into a concrete ``torch.device``. + + Parameters + ---------- + device : str, torch.device or None, default "cpu" + ``"cpu"`` (or ``None``) for the CPU, ``"auto"`` to pick CUDA if available, then + MPS, then CPU. Any other value (e.g. ``"cuda"``, ``"cuda:1"``, ``"mps"``, or a + ``torch.device``) is passed to ``torch.device`` and checked for availability. + + Returns + ------- + torch.device + The resolved device. + + Raises + ------ + RuntimeError + If an explicitly requested accelerator is not available. + + Examples + -------- + >>> resolve_device("cpu") + device(type='cpu') + >>> resolve_device(None) + device(type='cpu') + """ + if device is None: + return torch.device("cpu") + if isinstance(device, str) and device == "auto": + if torch.cuda.is_available(): + return torch.device("cuda") + if torch.backends.mps.is_available(): + return torch.device("mps") + return torch.device("cpu") + + resolved = torch.device(device) + if resolved.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"Device {resolved} was requested, but CUDA is not available.") + if resolved.type == "mps" and not torch.backends.mps.is_available(): + raise RuntimeError(f"Device {resolved} was requested, but MPS is not available.") + return resolved + + +def move_batch(batch: Mapping[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]: + """Move every tensor in a batch dict to ``device``. + + Non-tensor values are passed through unchanged. + + Parameters + ---------- + batch : Mapping[str, torch.Tensor] + A collated batch, as produced by :meth:`lsdb_ml.ResolvedSpec.collate`. + device : torch.device + The target device. + + Returns + ------- + dict[str, torch.Tensor] + A new dict with the tensors on ``device``. + """ + non_blocking = device.type == "cuda" + return { + key: value.to(device, non_blocking=non_blocking) if isinstance(value, torch.Tensor) else value + for key, value in batch.items() + } From e2668758c97080909f75ab7bdf782d6c14f6efc8 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 11:58:01 -0700 Subject: [PATCH 2/8] Add ColumnSpec: meta validation and padded nested collation ColumnSpec maps flat and nested catalog columns to tensors. resolve() validates against catalog meta (nested sub-column dtypes come from NestedDtype.column_dtypes) without computing data, and reports every problem at once. Nested columns are padded to the batch maximum with a mask and lengths, fully vectorized via nest.len() and nest.to_flat(). Co-Authored-By: Claude Opus 5.5 --- src/lsdb_ml/__init__.py | 17 +- src/lsdb_ml/collate.py | 167 ++++++++++++ src/lsdb_ml/spec.py | 470 ++++++++++++++++++++++++++++++++++ tests/lsdb_ml/conftest.py | 57 +++++ tests/lsdb_ml/test_collate.py | 169 ++++++++++++ tests/lsdb_ml/test_device.py | 26 ++ tests/lsdb_ml/test_spec.py | 120 +++++++++ 7 files changed, 1025 insertions(+), 1 deletion(-) create mode 100644 src/lsdb_ml/collate.py create mode 100644 src/lsdb_ml/spec.py create mode 100644 tests/lsdb_ml/test_collate.py create mode 100644 tests/lsdb_ml/test_device.py create mode 100644 tests/lsdb_ml/test_spec.py diff --git a/src/lsdb_ml/__init__.py b/src/lsdb_ml/__init__.py index e42fde6..6502887 100644 --- a/src/lsdb_ml/__init__.py +++ b/src/lsdb_ml/__init__.py @@ -1,4 +1,19 @@ from ._version import __version__ +from .device import move_batch, resolve_device from .example_module import greetings, meaning +from .spec import ColumnSpec, Field, Flat, Nested, ResolvedSpec, SpecError, TensorInfo -__all__ = ["greetings", "meaning", "__version__"] +__all__ = [ + "ColumnSpec", + "Field", + "Flat", + "Nested", + "ResolvedSpec", + "SpecError", + "TensorInfo", + "move_batch", + "resolve_device", + "greetings", + "meaning", + "__version__", +] diff --git a/src/lsdb_ml/collate.py b/src/lsdb_ml/collate.py new file mode 100644 index 0000000..5c2b239 --- /dev/null +++ b/src/lsdb_ml/collate.py @@ -0,0 +1,167 @@ +"""Vectorized conversion of (nested) dataframe columns into padded numpy arrays. + +These helpers do the numeric work behind :class:`lsdb_ml.Flat` and :class:`lsdb_ml.Nested`. +Nothing here loops over rows in Python. +""" + +from __future__ import annotations + +from typing import Literal + +import numpy as np +import pandas as pd +import torch + +NullPolicy = Literal["raise"] | float +TruncatePolicy = Literal["raise", "head", "tail"] + + +def numpy_dtype(dtype: torch.dtype) -> np.dtype: + """Return the numpy dtype that corresponds to a torch dtype. + + Parameters + ---------- + dtype : torch.dtype + A torch dtype with a numpy equivalent. + + Returns + ------- + np.dtype + The matching numpy dtype. + + Raises + ------ + TypeError + If the torch dtype has no numpy equivalent (e.g. ``torch.bfloat16``). + + Examples + -------- + >>> numpy_dtype(torch.float32) + dtype('float32') + """ + try: + return torch.empty(0, dtype=dtype).numpy().dtype + except TypeError as err: + raise TypeError(f"torch dtype {dtype} has no numpy equivalent and cannot be collated.") from err + + +def frame_to_array(frame: pd.DataFrame, dtype: np.dtype, null: NullPolicy, what: str) -> np.ndarray: + """Convert the columns of ``frame`` into a writable ``(N, F)`` array. + + Parameters + ---------- + frame : pd.DataFrame + The columns to convert; may be arrow- or numpy-backed. + dtype : np.dtype + Output dtype. + null : "raise" or float + ``"raise"`` to fail on missing values, otherwise the value to fill them with. + what : str + Description of the data used in error messages. + + Returns + ------- + np.ndarray + A C-contiguous, writable array of shape ``(len(frame), len(frame.columns))``. + """ + missing = frame.isna().to_numpy() + if missing.any(): + if isinstance(null, str): + n_missing = int(missing.any(axis=1).sum()) + raise ValueError( + f"{what} contains {n_missing} row(s) with missing values. Pass null= " + "to fill them, or drop those rows before collation." + ) + frame = frame.fillna(null) + # np.array (not np.asarray) always copies: arrow-backed columns can convert to read-only + # views, which torch.from_numpy would warn about and share. + return np.array(frame.to_numpy(dtype=dtype), dtype=dtype, order="C", copy=True) + + +def pad_ragged( + values: np.ndarray, + lengths: np.ndarray, + *, + max_length: int | None = None, + truncate: TruncatePolicy = "raise", + pad_to: int | None = None, + pad_value: float = 0, + what: str = "nested column", +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Scatter flat ragged values into a padded ``(B, L, F)`` array with a mask. + + Parameters + ---------- + values : np.ndarray + Flat values of shape ``(sum(lengths), F)``, rows in order. + lengths : np.ndarray + Number of values in each of the ``B`` rows. + max_length : int or None, default None + Longest sequence to keep. Longer rows are handled by ``truncate``. + truncate : {"raise", "head", "tail"}, default "raise" + For rows over the limit: raise, keep the first values, or keep the last values. + pad_to : int or None, default None + Fixed padded length. When None, pad to the longest (kept) row in the batch. Rows + longer than ``pad_to`` are handled by ``truncate``. + pad_value : float, default 0 + Value written into padded positions. + what : str, default "nested column" + Description of the data used in error messages. + + Returns + ------- + padded : np.ndarray + Shape ``(B, L, F)``. + mask : np.ndarray + Boolean, shape ``(B, L)``; True where ``padded`` holds real data. + kept_lengths : np.ndarray + int64, shape ``(B,)``; number of real values per row after truncation. + + Examples + -------- + >>> padded, mask, kept = pad_ragged(np.arange(5.0)[:, None], np.array([2, 0, 3])) + >>> padded[..., 0] + array([[0., 1., 0.], + [0., 0., 0.], + [2., 3., 4.]]) + >>> mask.astype(int) + array([[1, 1, 0], + [0, 0, 0], + [1, 1, 1]]) + """ + lengths = np.asarray(lengths, dtype=np.int64) + n_rows = len(lengths) + n_features = values.shape[1] + if int(lengths.sum()) != len(values): + raise ValueError(f"{what}: lengths sum to {int(lengths.sum())} but there are {len(values)} values.") + + starts = np.cumsum(lengths) - lengths + limits = [limit for limit in (max_length, pad_to) if limit is not None] + limit = min(limits) if limits else None + + kept = lengths + source_starts = starts + if limit is not None and n_rows and int(lengths.max()) > limit: + if truncate == "raise": + raise ValueError( + f"{what} has rows of length up to {int(lengths.max())}, over the limit of {limit}. " + "Pass truncate='head' or truncate='tail' to cut them." + ) + kept = np.minimum(lengths, limit) + if truncate == "tail": + source_starts = starts + lengths - kept + elif truncate != "head": + raise ValueError(f"truncate must be 'raise', 'head' or 'tail', got {truncate!r}.") + + width = pad_to if pad_to is not None else (int(kept.max()) if n_rows else 0) + padded = np.full((n_rows, width, n_features), pad_value, dtype=values.dtype) + mask = np.zeros((n_rows, width), dtype=bool) + + total = int(kept.sum()) + if total: + row = np.repeat(np.arange(n_rows), kept) + position = np.arange(total) - np.repeat(np.cumsum(kept) - kept, kept) + source = np.repeat(source_starts, kept) + position + padded[row, position] = values[source] + mask[row, position] = True + return padded, mask, kept.astype(np.int64) diff --git a/src/lsdb_ml/spec.py b/src/lsdb_ml/spec.py new file mode 100644 index 0000000..fb76eda --- /dev/null +++ b/src/lsdb_ml/spec.py @@ -0,0 +1,470 @@ +"""Column specifications: map flat and nested catalog columns to tensors. + +A :class:`ColumnSpec` names the tensors a model needs and where each comes from. It is +checked against a catalog's ``meta`` (an empty ``NestedFrame`` that carries every dtype, +including nested sub-column dtypes) before any data is computed, and the resulting +:class:`ResolvedSpec` collates dataframe rows into a ``dict[str, torch.Tensor]``. + +>>> import lsdb_ml +>>> spec = lsdb_ml.ColumnSpec({"x": lsdb_ml.Flat(["ra", "dec"]), "lc": lsdb_ml.Nested("lc", ["flux"])}) +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass + +import numpy as np +import pandas as pd +import pyarrow as pa +import torch +from nested_pandas.series.dtype import NestedDtype + +from lsdb_ml.collate import NullPolicy, TruncatePolicy, frame_to_array, numpy_dtype, pad_ragged + + +class SpecError(ValueError): + """Raised when a column specification does not match a catalog's schema.""" + + +@dataclass(frozen=True) +class TensorInfo: + """Description of one tensor a resolved spec produces. + + Attributes + ---------- + dtype : torch.dtype + The tensor dtype. + shape : tuple + The shape; ``None`` marks a dimension that varies by batch (always the first one). + """ + + dtype: torch.dtype + shape: tuple[int | None, ...] + + +def _as_list(columns: str | Sequence[str]) -> list[str]: + return [columns] if isinstance(columns, str) else list(columns) + + +def _is_numeric(dtype) -> bool: + """True for pandas/numpy/pyarrow integer, float, and bool dtypes.""" + if isinstance(dtype, pa.DataType): + return pa.types.is_integer(dtype) or pa.types.is_floating(dtype) or pa.types.is_boolean(dtype) + if isinstance(dtype, pd.ArrowDtype): + return _is_numeric(dtype.pyarrow_dtype) + if isinstance(dtype, NestedDtype): + return False + try: + return pd.api.types.is_numeric_dtype(dtype) or pd.api.types.is_bool_dtype(dtype) + except TypeError: + return False + + +class Field(ABC): + """Base class for one entry of a :class:`ColumnSpec`. + + Subclasses implement :meth:`resolve`, which validates the field against catalog meta + and returns a :class:`ResolvedField` that does the actual collation. New column + representations (e.g. ragged tensors) are added as new ``Field`` types. + """ + + @abstractmethod + def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None: + """Validate against ``meta`` and build the collator. + + Parameters + ---------- + key : str + The output key this field was registered under. + meta : pd.DataFrame + The catalog meta (an empty frame with full dtypes). + problems : list[str] + Append a message for every problem found, instead of raising, so that all + problems in a spec are reported together. + + Returns + ------- + ResolvedField or None + None if any problem was found. + """ + + +class ResolvedField(ABC): + """A field validated against a schema, ready to collate frames.""" + + #: Top-level columns this field reads. + columns: list[str] + + @abstractmethod + def describe(self) -> dict[str, TensorInfo]: + """Return the tensors this field produces, keyed by output key.""" + + @abstractmethod + def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]: + """Convert the rows of ``frame`` to tensors.""" + + +@dataclass +class Flat(Field): + """One or more flat (non-nested) columns, stacked into a single tensor. + + Parameters + ---------- + columns : str or sequence of str + A single column name produces a ``(B,)`` tensor; a list produces ``(B, F)``. + dtype : torch.dtype, default torch.float32 + Output dtype. + null : "raise" or float, default "raise" + Missing values raise an error, or are replaced by this fill value. + transform : callable, optional + ``transform(frame[columns]) -> array-like`` of shape ``(B,)`` or ``(B, F)``. Use it + for non-numeric columns or custom encodings. When given, the dtype check is skipped + and the result is cast to ``dtype``. It must be picklable to be used from Dask + tasks or DataLoader workers. + """ + + columns: str | Sequence[str] + dtype: torch.dtype = torch.float32 + null: NullPolicy = "raise" + transform: Callable[[pd.DataFrame], np.ndarray] | None = None + + def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None: + """Validate the columns against ``meta``; see :meth:`Field.resolve`.""" + columns = _as_list(self.columns) + start = len(problems) + if not columns: + problems.append(f"'{key}': Flat needs at least one column.") + for column in columns: + if column not in meta.columns: + problems.append(f"'{key}': column '{column}' is not in the catalog.") + elif isinstance(meta.dtypes[column], NestedDtype): + problems.append(f"'{key}': column '{column}' is nested; use Nested('{column}', [...]).") + elif self.transform is None and not _is_numeric(meta.dtypes[column]): + problems.append( + f"'{key}': column '{column}' has non-numeric dtype {meta.dtypes[column]}; " + "pass transform= to encode it." + ) + try: + np_dtype = numpy_dtype(self.dtype) + except TypeError as err: + problems.append(f"'{key}': {err}") + if len(problems) > start: + return None + return _ResolvedFlat(key, columns, isinstance(self.columns, str), self, np_dtype) + + +class _ResolvedFlat(ResolvedField): + def __init__(self, key: str, columns: list[str], scalar: bool, spec: Flat, np_dtype: np.dtype): + self.key = key + self.columns = columns + self.scalar = scalar + self.spec = spec + self.np_dtype = np_dtype + + def describe(self) -> dict[str, TensorInfo]: + if self.spec.transform is not None: + shape: tuple[int | None, ...] = (None,) if self.scalar else (None, None) + else: + shape = (None,) if self.scalar else (None, len(self.columns)) + return {self.key: TensorInfo(self.spec.dtype, shape)} + + def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]: + what = f"'{self.key}' (columns {self.columns})" + if self.spec.transform is not None: + array = np.array(self.spec.transform(frame[self.columns]), dtype=self.np_dtype, copy=True) + if len(array) != len(frame): + raise ValueError(f"{what}: transform returned {len(array)} rows for {len(frame)} input rows.") + else: + array = frame_to_array(frame[self.columns], self.np_dtype, self.spec.null, what) + if self.scalar: + array = array[:, 0].copy() + return {self.key: torch.from_numpy(array)} + + +@dataclass +class Nested(Field): + """Sub-columns of a nested column, padded to a rectangular tensor with a mask. + + For output key ``k`` this produces: + + - ``k``: values, ``(B, L)`` for a single sub-column or ``(B, L, F)`` for a list; + - ``k_mask``: bool ``(B, L)``, True where a value is real rather than padding; + - ``k_lengths``: int64 ``(B,)``, the number of real values per row. + + ``L`` is the longest row in the batch (after truncation), or ``pad_to`` if set. Rows + whose nested value is null or empty have length 0. The padding value carries no + meaning; always use the mask. + + Parameters + ---------- + column : str + The nested column. + fields : str or sequence of str + Sub-column name(s) to read. + dtype : torch.dtype, default torch.float32 + Output dtype of the values. + max_length : int, optional + Longest sequence to keep; see ``truncate``. + truncate : {"raise", "head", "tail"}, default "raise" + What to do with rows longer than ``max_length`` (or ``pad_to``): raise, keep the + first values, or keep the last values. + pad_to : int, optional + Pad every batch to this fixed length instead of the batch maximum. + pad_value : float, default 0 + Value written into padded positions. + null : "raise" or float, default "raise" + Missing values *inside* a list raise an error, or are replaced by this value. + transform : callable, optional + ``transform(flat_values) -> array-like`` with one row per flat value, of shape + ``(N,)`` or ``(N, F)``, where ``flat_values`` is a DataFrame of the requested + sub-columns. When given, the dtype check is skipped. + """ + + column: str + fields: str | Sequence[str] + dtype: torch.dtype = torch.float32 + max_length: int | None = None + truncate: TruncatePolicy = "raise" + pad_to: int | None = None + pad_value: float = 0 + null: NullPolicy = "raise" + transform: Callable[[pd.DataFrame], np.ndarray] | None = None + + def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None: + """Validate the nested column and sub-columns against ``meta``; see :meth:`Field.resolve`.""" + subcolumns = _as_list(self.fields) + start = len(problems) + if not subcolumns: + problems.append(f"'{key}': Nested needs at least one sub-column.") + if self.truncate not in ("raise", "head", "tail"): + problems.append(f"'{key}': truncate must be 'raise', 'head' or 'tail', got {self.truncate!r}.") + for name, limit in (("max_length", self.max_length), ("pad_to", self.pad_to)): + if limit is not None and (isinstance(limit, bool) or not isinstance(limit, int) or limit < 1): + problems.append(f"'{key}': {name} must be a positive integer, got {limit!r}.") + if self.column not in meta.columns: + problems.append(f"'{key}': column '{self.column}' is not in the catalog.") + elif not isinstance(meta.dtypes[self.column], NestedDtype): + problems.append(f"'{key}': column '{self.column}' is not nested; use Flat('{self.column}').") + else: + available = meta.dtypes[self.column].column_dtypes + for sub in subcolumns: + if sub not in available: + problems.append( + f"'{key}': nested column '{self.column}' has no sub-column '{sub}' " + f"(available: {sorted(available)})." + ) + elif self.transform is None and not _is_numeric(available[sub]): + problems.append( + f"'{key}': sub-column '{self.column}.{sub}' has non-numeric type {available[sub]}; " + "pass transform= to encode it." + ) + try: + np_dtype = numpy_dtype(self.dtype) + except TypeError as err: + problems.append(f"'{key}': {err}") + if len(problems) > start: + return None + return _ResolvedNested(key, subcolumns, isinstance(self.fields, str), self, np_dtype) + + +class _ResolvedNested(ResolvedField): + def __init__(self, key: str, subcolumns: list[str], scalar: bool, spec: Nested, np_dtype: np.dtype): + self.key = key + self.columns = [spec.column] + self.subcolumns = subcolumns + self.scalar = scalar + self.spec = spec + self.np_dtype = np_dtype + + def describe(self) -> dict[str, TensorInfo]: + length = self.spec.pad_to + if self.scalar: + values_shape: tuple[int | None, ...] = (None, length) + else: + n_features = None if self.spec.transform is not None else len(self.subcolumns) + values_shape = (None, length, n_features) + return { + self.key: TensorInfo(self.spec.dtype, values_shape), + f"{self.key}_mask": TensorInfo(torch.bool, (None, length)), + f"{self.key}_lengths": TensorInfo(torch.int64, (None,)), + } + + def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]: + what = f"'{self.key}' (nested column '{self.spec.column}', fields {self.subcolumns})" + series = frame[self.spec.column] + # nest.len() counts null rows as 0, and to_flat() returns the values of all rows + # concatenated in row order; both are consistent across iloc slices and concat. + lengths = np.asarray(series.nest.len(), dtype=np.int64) + flat = series.nest.to_flat(self.subcolumns) + if self.spec.transform is not None: + values = np.array(self.spec.transform(flat), dtype=self.np_dtype, copy=True) + if len(values) != len(flat): + raise ValueError(f"{what}: transform returned {len(values)} rows for {len(flat)} values.") + if values.ndim == 1: + values = values[:, None] + else: + values = frame_to_array(flat, self.np_dtype, self.spec.null, what) + padded, mask, kept = pad_ragged( + values, + lengths, + max_length=self.spec.max_length, + truncate=self.spec.truncate, + pad_to=self.spec.pad_to, + pad_value=self.spec.pad_value, + what=what, + ) + if self.scalar and padded.shape[2] == 1: + padded = padded[..., 0] + return { + self.key: torch.from_numpy(np.ascontiguousarray(padded)), + f"{self.key}_mask": torch.from_numpy(mask), + f"{self.key}_lengths": torch.from_numpy(kept), + } + + +@dataclass +class ColumnSpec: + """Mapping from output keys to catalog columns. + + Parameters + ---------- + fields : Mapping[str, Field] + Output key to field, e.g. ``{"x": Flat(["ra", "dec"]), "lc": Nested("lc", "flux")}``. + index_key : str or None, default "_index" + Every batch also includes the frame index (``_healpix_29`` for LSDB catalogs) as an + int64 ``(B,)`` tensor under this key, so batches can be traced back to rows. None + disables it. + + Examples + -------- + >>> import nested_pandas as npd + >>> nf = npd.NestedFrame({"a": [1.0, 2.0]}, index=[10, 11]) + >>> nf = nf.join_nested(npd.NestedFrame({"f": [1.0, 2.0, 3.0]}, index=[10, 10, 11]), "lc") + >>> resolved = ColumnSpec({"a": Flat("a"), "lc": Nested("lc", "f")}).resolve(nf) + >>> batch = resolved.collate(nf) + >>> batch["lc"] + tensor([[1., 2.], + [3., 0.]]) + >>> batch["lc_mask"] + tensor([[ True, True], + [ True, False]]) + >>> batch["_index"] + tensor([10, 11]) + """ + + fields: Mapping[str, Field] + index_key: str | None = "_index" + + def resolve(self, schema) -> ResolvedSpec: + """Validate the spec against a catalog's schema without computing any data. + + Parameters + ---------- + schema : lsdb.Catalog or pd.DataFrame + A catalog (its ``meta`` is used) or any dataframe with the same columns and + dtypes, such as ``catalog.meta`` or a computed partition. + + Returns + ------- + ResolvedSpec + The validated spec, which can collate frames. + + Raises + ------ + SpecError + Listing every problem found. + """ + meta = schema.meta if hasattr(schema, "meta") and not isinstance(schema, pd.DataFrame) else schema + if not isinstance(meta, pd.DataFrame): + raise TypeError(f"Expected an lsdb catalog or a DataFrame, got {type(schema).__name__}.") + problems: list[str] = [] + if not self.fields: + problems.append("The spec has no fields.") + resolved: dict[str, ResolvedField] = {} + for key, spec_field in self.fields.items(): + if not isinstance(spec_field, Field): + problems.append(f"'{key}': expected a Field such as Flat or Nested, got {type(spec_field)}.") + continue + result = spec_field.resolve(key, meta, problems) + if result is not None: + resolved[key] = result + + seen: dict[str, str] = {} + if self.index_key is not None: + seen[self.index_key] = "the index" + if not pd.api.types.is_integer_dtype(meta.index.dtype): + problems.append( + f"index_key={self.index_key!r} needs an integer index, but the index has dtype " + f"{meta.index.dtype}; pass index_key=None to disable it." + ) + for key, result in resolved.items(): + for output_key in result.describe(): + if output_key in seen: + problems.append( + f"Output key '{output_key}' from '{key}' collides with {seen[output_key]}." + ) + seen[output_key] = f"'{key}'" + + if problems: + hint = "" + missing = [p for p in problems if "is not in the catalog" in p] + all_columns = getattr(schema, "all_columns", None) + if missing and all_columns is not None: + hint = ( + "\nSome columns may exist on disk but were not loaded; the catalog was opened " + f"with columns {list(meta.columns)}." + ) + raise SpecError("Column spec does not match the catalog:\n - " + "\n - ".join(problems) + hint) + return ResolvedSpec(resolved, self.index_key) + + +class ResolvedSpec: + """A :class:`ColumnSpec` validated against a schema; collates frames into tensors. + + Create it with :meth:`ColumnSpec.resolve`. It is picklable as long as any + ``transform`` callables are, so it can be shipped to Dask tasks and DataLoader workers. + """ + + def __init__(self, fields: dict[str, ResolvedField], index_key: str | None): + self.fields = fields + self.index_key = index_key + + @property + def columns(self) -> list[str]: + """Top-level catalog columns needed for collation, in first-use order.""" + columns: list[str] = [] + for resolved in self.fields.values(): + columns.extend(c for c in resolved.columns if c not in columns) + return columns + + def describe(self) -> dict[str, TensorInfo]: + """Return the dtype and shape of every tensor a batch contains, keyed by output key.""" + info: dict[str, TensorInfo] = {} + for resolved in self.fields.values(): + info.update(resolved.describe()) + if self.index_key is not None: + info[self.index_key] = TensorInfo(torch.int64, (None,)) + return info + + def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]: + """Convert the rows of ``frame`` to a batch of tensors. + + Parameters + ---------- + frame : pd.DataFrame + Rows to convert, e.g. a slice of a catalog partition. Row order is preserved: + row ``i`` of every output tensor comes from row ``i`` of ``frame``. + + Returns + ------- + dict[str, torch.Tensor] + One entry per key in :meth:`describe`. + """ + batch: dict[str, torch.Tensor] = {} + for resolved in self.fields.values(): + batch.update(resolved.collate(frame)) + if self.index_key is not None: + batch[self.index_key] = torch.from_numpy(np.array(frame.index.to_numpy(), dtype=np.int64)) + return batch diff --git a/tests/lsdb_ml/conftest.py b/tests/lsdb_ml/conftest.py index e69de29..ec295e2 100644 --- a/tests/lsdb_ml/conftest.py +++ b/tests/lsdb_ml/conftest.py @@ -0,0 +1,57 @@ +import lsdb +import nested_pandas as npd +import numpy as np +import pandas as pd +import pytest + + +def make_nested_frame(n_rows: int = 500, seed: int = 0, max_length: int = 6) -> npd.NestedFrame: + """Synthetic object table with a ragged nested light curve column. + + Light curve lengths range from 0 to ``max_length`` - 1, so some rows have empty + (null) nested values. + """ + rng = np.random.default_rng(seed) + base = pd.DataFrame( + { + "id": np.arange(n_rows, dtype=np.int64), + "ra": rng.uniform(0, 360, n_rows), + "dec": rng.uniform(-60, 60, n_rows), + "mag": rng.normal(20, 1, n_rows).astype(np.float32), + "flag": rng.integers(0, 2, n_rows).astype(bool), + "name": [f"obj{i}" for i in range(n_rows)], + } + ) + lengths = rng.integers(0, max_length, n_rows) + n_flat = int(lengths.sum()) + flat = pd.DataFrame( + { + "t": rng.uniform(0, 100, n_flat), + "flux": rng.normal(0, 1, n_flat).astype(np.float32), + "band": rng.choice(["g", "r"], n_flat), + }, + index=np.repeat(np.arange(n_rows), lengths), + ) + return npd.NestedFrame(base).join_nested(flat, "lc") + + +@pytest.fixture(scope="session") +def nested_frame() -> npd.NestedFrame: + """The in-memory frame behind ``nested_catalog``, indexed by row number.""" + return make_nested_frame() + + +@pytest.fixture(scope="session") +def nested_catalog(nested_frame) -> lsdb.Catalog: + """A lazy catalog with many small partitions and a ragged nested column.""" + return lsdb.from_dataframe(nested_frame, partition_rows=40, catalog_name="toy") + + +@pytest.fixture +def unloadable_catalog(nested_catalog) -> lsdb.Catalog: + """Same schema as ``nested_catalog``, but computing any partition raises.""" + + def fail(_): + raise AssertionError("A partition was computed, but no data should be read here.") + + return nested_catalog.map_partitions(fail, meta=nested_catalog.meta) diff --git a/tests/lsdb_ml/test_collate.py b/tests/lsdb_ml/test_collate.py new file mode 100644 index 0000000..2369eb2 --- /dev/null +++ b/tests/lsdb_ml/test_collate.py @@ -0,0 +1,169 @@ +import nested_pandas as npd +import numpy as np +import pandas as pd +import pytest +import torch + +from lsdb_ml import ColumnSpec, Flat, Nested +from lsdb_ml.collate import pad_ragged + + +def reference_rows(frame, column, fields): + """Per-row python reference: a list of (length, n_fields) arrays.""" + rows = [] + for value in frame[column].array: + if value is None or not isinstance(value, pd.DataFrame): + rows.append(np.zeros((0, len(fields)))) + else: + rows.append(value[fields].to_numpy(dtype=np.float64)) + return rows + + +def reference_pad(rows, width, keep="head"): + """Pad reference rows to ``width`` the slow way.""" + padded = np.zeros((len(rows), width, rows[0].shape[1] if rows else 0)) + mask = np.zeros((len(rows), width), dtype=bool) + for i, row in enumerate(rows): + row = row[:width] if keep == "head" else row[max(len(row) - width, 0) :] + padded[i, : len(row)] = row + mask[i, : len(row)] = True + return padded, mask + + +@pytest.fixture +def frame(nested_catalog): + """A frame mixing a sliced partition, a whole partition, and a concatenation.""" + # A computed partition, plus a slice and a concatenation, all exercise the + # flat-values/offsets bookkeeping on non-trivial arrow layouts. + parts = [nested_catalog.partitions[i].compute(progress_bar=False) for i in range(3)] + return pd.concat([parts[0].iloc[3:], parts[1], parts[2].iloc[:5]]) + + +def test_flat_columns(frame): + """Flat columns stack into (B,) or (B, F) tensors with the requested dtype.""" + resolved = ColumnSpec( + {"x": Flat(["ra", "dec", "mag"]), "mag": Flat("mag"), "flag": Flat("flag", dtype=torch.bool)} + ).resolve(frame) + batch = resolved.collate(frame) + assert batch["x"].shape == (len(frame), 3) and batch["x"].dtype == torch.float32 + np.testing.assert_allclose(batch["x"].numpy(), frame[["ra", "dec", "mag"]].to_numpy(dtype=np.float32)) + assert batch["mag"].shape == (len(frame),) + assert batch["flag"].dtype == torch.bool + np.testing.assert_array_equal(batch["flag"].numpy(), frame["flag"].to_numpy(dtype=bool)) + np.testing.assert_array_equal(batch["_index"].numpy(), frame.index.to_numpy()) + assert all(t.numpy().flags.writeable for t in batch.values()) + + +@pytest.mark.parametrize("fields", [["t", "flux"], ["flux"]]) +def test_nested_matches_reference(frame, fields): + """Vectorized nested padding matches a per-row reference, including empty rows.""" + resolved = ColumnSpec({"lc": Nested("lc", fields, dtype=torch.float64)}).resolve(frame) + batch = resolved.collate(frame) + rows = reference_rows(frame, "lc", fields) + lengths = np.array([len(r) for r in rows]) + assert (lengths == 0).any(), "fixture should include empty / null nested rows" + expected, expected_mask = reference_pad(rows, lengths.max()) + np.testing.assert_allclose(batch["lc"].numpy(), expected) + np.testing.assert_array_equal(batch["lc_mask"].numpy(), expected_mask) + np.testing.assert_array_equal(batch["lc_lengths"].numpy(), lengths) + + +def test_nested_single_field_is_2d(frame): + """A single nested sub-column produces (B, L) values.""" + batch = ColumnSpec({"f": Nested("lc", "flux")}).resolve(frame).collate(frame) + assert batch["f"].ndim == 2 + assert batch["f"].shape == batch["f_mask"].shape + + +@pytest.mark.parametrize("keep", ["head", "tail"]) +def test_truncation(frame, keep): + """Rows over max_length keep their first or last values.""" + resolved = ColumnSpec({"lc": Nested("lc", ["t"], max_length=2, truncate=keep, dtype=torch.float64)}) + batch = resolved.resolve(frame).collate(frame) + rows = reference_rows(frame, "lc", ["t"]) + expected, expected_mask = reference_pad(rows, 2, keep=keep) + np.testing.assert_allclose(batch["lc"].numpy(), expected) + np.testing.assert_array_equal(batch["lc_mask"].numpy(), expected_mask) + np.testing.assert_array_equal(batch["lc_lengths"].numpy(), np.minimum([len(r) for r in rows], 2)) + + +def test_truncation_raises_by_default(frame): + """Rows over max_length raise unless a truncate policy is given.""" + with pytest.raises(ValueError, match="over the limit of 2"): + ColumnSpec({"lc": Nested("lc", "t", max_length=2)}).resolve(frame).collate(frame) + + +def test_pad_to_fixed_width(frame): + """pad_to fixes the padded width and fills with pad_value.""" + batch = ColumnSpec({"lc": Nested("lc", "t", pad_to=9, pad_value=-1)}).resolve(frame).collate(frame) + assert batch["lc"].shape == (len(frame), 9) + assert (batch["lc"][~batch["lc_mask"]] == -1).all() + + +def test_transforms(frame): + """Transforms encode non-numeric flat and nested columns.""" + resolved = ColumnSpec( + { + "name_len": Flat("name", transform=lambda df: df["name"].str.len().to_numpy()), + "is_g": Nested("lc", "band", transform=lambda df: (df["band"] == "g").to_numpy()), + } + ).resolve(frame) + batch = resolved.collate(frame) + np.testing.assert_array_equal(batch["name_len"].numpy(), frame["name"].str.len().to_numpy()) + first = next(i for i, v in enumerate(frame["lc"].array) if isinstance(v, pd.DataFrame)) + expected = (frame["lc"].array[first]["band"] == "g").to_numpy() + np.testing.assert_array_equal(batch["is_g"][first, : len(expected)].numpy(), expected) + + +def test_empty_frame(frame): + """Zero-row frames collate to zero-row tensors.""" + resolved = ColumnSpec({"x": Flat(["ra", "dec"]), "lc": Nested("lc", ["t", "flux"])}).resolve(frame) + batch = resolved.collate(frame.iloc[:0]) + assert batch["x"].shape == (0, 2) + assert batch["lc"].shape == (0, 0, 2) + assert batch["lc_mask"].shape == (0, 0) + assert batch["_index"].shape == (0,) + + +def test_all_rows_empty_nested(frame): + """A batch where every nested row is empty pads to width 0.""" + lengths = np.asarray(frame["lc"].nest.len()) + empty_rows = frame[lengths == 0] + batch = ColumnSpec({"lc": Nested("lc", "flux")}).resolve(frame).collate(empty_rows) + assert batch["lc"].shape == (len(empty_rows), 0) + assert not batch["lc_mask"].any() + + +def test_null_values(): + """Missing values raise by default and are filled when a fill value is given.""" + nf = npd.NestedFrame({"a": [1.0, None, 3.0]}, index=[0, 1, 2]) + nf = nf.join_nested(npd.NestedFrame({"f": [1.0, None, 2.0]}, index=[0, 0, 2]), "lc") + with pytest.raises(ValueError, match="1 row\\(s\\) with missing values"): + ColumnSpec({"a": Flat("a")}).resolve(nf).collate(nf) + with pytest.raises(ValueError, match="missing values"): + ColumnSpec({"lc": Nested("lc", "f")}).resolve(nf).collate(nf) + batch = ColumnSpec({"a": Flat("a", null=-9.0), "lc": Nested("lc", "f", null=0.5)}).resolve(nf).collate(nf) + np.testing.assert_array_equal(batch["a"].numpy(), [1, -9, 3]) + np.testing.assert_array_equal(batch["lc"].numpy(), [[1, 0.5], [0, 0], [2, 0]]) + np.testing.assert_array_equal(batch["lc_lengths"].numpy(), [2, 0, 1]) + + +def test_transform_length_mismatch(frame): + """A transform that changes the row count is rejected.""" + resolved = ColumnSpec({"x": Flat("ra", transform=lambda df: [1.0])}).resolve(frame) + with pytest.raises(ValueError, match="transform returned 1 rows"): + resolved.collate(frame) + + +def test_pad_ragged_length_mismatch(): + """Inconsistent lengths and values are rejected.""" + with pytest.raises(ValueError, match="lengths sum to 3"): + pad_ragged(np.zeros((2, 1)), np.array([1, 2])) + + +def test_row_order_preserved(frame): + """Row i of every tensor comes from row i of the frame, even for a shuffled frame.""" + shuffled = frame.sample(frac=1, random_state=1) + batch = ColumnSpec({"id": Flat("id", dtype=torch.int64)}).resolve(shuffled).collate(shuffled) + np.testing.assert_array_equal(batch["id"].numpy(), shuffled["id"].to_numpy()) + np.testing.assert_array_equal(batch["_index"].numpy(), shuffled.index.to_numpy()) diff --git a/tests/lsdb_ml/test_device.py b/tests/lsdb_ml/test_device.py new file mode 100644 index 0000000..cb90019 --- /dev/null +++ b/tests/lsdb_ml/test_device.py @@ -0,0 +1,26 @@ +import pytest +import torch + +from lsdb_ml import move_batch, resolve_device + + +def test_resolve_device(): + """Known device requests resolve to torch devices.""" + assert resolve_device("cpu") == torch.device("cpu") + assert resolve_device(torch.device("cpu")) == torch.device("cpu") + assert resolve_device("auto").type in {"cpu", "cuda", "mps"} + + +@pytest.mark.skipif(torch.cuda.is_available(), reason="checks the error when CUDA is missing") +def test_resolve_missing_cuda(): + """Requesting CUDA without CUDA fails clearly.""" + with pytest.raises(RuntimeError, match="CUDA is not available"): + resolve_device("cuda") + + +def test_move_batch(): + """Tensors move to the device and other values pass through.""" + batch = {"x": torch.zeros(2), "meta": "keep"} + moved = move_batch(batch, torch.device("cpu")) + assert moved["x"].device.type == "cpu" + assert moved["meta"] == "keep" diff --git a/tests/lsdb_ml/test_spec.py b/tests/lsdb_ml/test_spec.py new file mode 100644 index 0000000..73e2896 --- /dev/null +++ b/tests/lsdb_ml/test_spec.py @@ -0,0 +1,120 @@ +import pytest +import torch + +from lsdb_ml import ColumnSpec, Flat, Nested, SpecError, TensorInfo + + +def test_resolve_valid_spec_from_catalog_without_compute(unloadable_catalog): + """A valid spec resolves against a catalog without computing any partition.""" + spec = ColumnSpec( + { + "x": Flat(["ra", "dec", "mag"]), + "flag": Flat("flag", dtype=torch.bool), + "lc": Nested("lc", ["t", "flux"], max_length=4, truncate="tail"), + "flux": Nested("lc", "flux", pad_to=8, truncate="head"), + } + ) + resolved = spec.resolve(unloadable_catalog) + assert resolved.columns == ["ra", "dec", "mag", "flag", "lc"] + assert resolved.describe() == { + "x": TensorInfo(torch.float32, (None, 3)), + "flag": TensorInfo(torch.bool, (None,)), + "lc": TensorInfo(torch.float32, (None, None, 2)), + "lc_mask": TensorInfo(torch.bool, (None, None)), + "lc_lengths": TensorInfo(torch.int64, (None,)), + "flux": TensorInfo(torch.float32, (None, 8)), + "flux_mask": TensorInfo(torch.bool, (None, 8)), + "flux_lengths": TensorInfo(torch.int64, (None,)), + "_index": TensorInfo(torch.int64, (None,)), + } + + +def test_unloadable_catalog_fixture_raises(unloadable_catalog): + """Guard: the no-compute fixture really fails when a partition is computed.""" + with pytest.raises(Exception, match="no data should be read"): + unloadable_catalog.partitions[0].compute(progress_bar=False) + + +def test_resolve_accepts_meta_frame(nested_catalog): + """resolve() also accepts the meta frame directly.""" + resolved = ColumnSpec({"x": Flat("ra")}).resolve(nested_catalog.meta) + assert resolved.columns == ["ra"] + + +def test_all_problems_reported_together(unloadable_catalog): + """Every problem in a spec is reported in one error.""" + spec = ColumnSpec( + { + "missing": Flat("nope"), + "nested_as_flat": Flat("lc"), + "string": Flat("name"), + "flat_as_nested": Nested("ra", "x"), + "missing_sub": Nested("lc", ["flux", "nope"]), + "string_sub": Nested("lc", "band"), + "bad_trunc": Nested("lc", "flux", truncate="middle"), + "bad_len": Nested("lc", "flux", max_length=0), + } + ) + with pytest.raises(SpecError) as err: + spec.resolve(unloadable_catalog) + message = str(err.value) + for expected in [ + "column 'nope' is not in the catalog", + "column 'lc' is nested", + "'name' has non-numeric dtype", + "column 'ra' is not nested", + "has no sub-column 'nope'", + "'lc.band' has non-numeric type", + "truncate must be", + "max_length must be a positive integer", + ]: + assert expected in message + assert "not loaded" in message # hint for missing columns on a catalog + + +def test_transform_allows_non_numeric(unloadable_catalog): + """A transform lifts the numeric dtype requirement.""" + spec = ColumnSpec( + { + "name_len": Flat("name", transform=lambda df: df["name"].str.len()), + "band": Nested("lc", "band", transform=lambda df: (df["band"] == "g").to_numpy()), + } + ) + spec.resolve(unloadable_catalog) + + +def test_key_collisions(unloadable_catalog): + """Output keys may not collide with each other or the index key.""" + with pytest.raises(SpecError, match="'lc_mask' from 'lc' collides"): + ColumnSpec({"lc_mask": Flat("ra"), "lc": Nested("lc", "flux")}).resolve(unloadable_catalog) + with pytest.raises(SpecError, match="'_index' from '_index' collides with the index"): + ColumnSpec({"_index": Flat("ra")}).resolve(unloadable_catalog) + ColumnSpec({"_index": Flat("ra")}, index_key=None).resolve(unloadable_catalog) + + +def test_non_field_and_empty_spec(unloadable_catalog): + """Non-Field values and empty specs are rejected.""" + with pytest.raises(SpecError, match="expected a Field"): + ColumnSpec({"x": "ra"}).resolve(unloadable_catalog) + with pytest.raises(SpecError, match="no fields"): + ColumnSpec({}).resolve(unloadable_catalog) + + +def test_unsupported_dtype(unloadable_catalog): + """Torch dtypes without a numpy equivalent are rejected.""" + with pytest.raises(SpecError, match="no numpy equivalent"): + ColumnSpec({"x": Flat("ra", dtype=torch.bfloat16)}).resolve(unloadable_catalog) + + +def test_non_integer_index(nested_frame): + """The index key requires an integer index, and can be disabled.""" + frame = nested_frame.reset_index(drop=True).set_index("name") + with pytest.raises(SpecError, match="needs an integer index"): + ColumnSpec({"x": Flat("ra")}).resolve(frame) + ColumnSpec({"x": Flat("ra")}, index_key=None).resolve(frame) + + +def test_bad_schema_type(): + """resolve() rejects objects that are neither catalogs nor frames.""" + with pytest.raises(TypeError, match="Expected an lsdb catalog"): + ColumnSpec({"x": Flat("ra")}).resolve(42) From 2186d0bdbf77f06c7a4213bacb24e3edf5e246f0 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 12:20:16 -0700 Subject: [PATCH 3/8] Add CatalogIterableDataset with partition-level sharding Partitions are sharded across DDP ranks and DataLoader workers from a shared per-epoch permutation, so every row is read exactly once per epoch for any worker count and world size. Each shard streams its partitions through lsdb's CatalogStream; batches are cut within a chunk and never span chunks. Supports fixed epochs and open-ended streaming. Co-Authored-By: Claude Opus 5.5 --- src/lsdb_ml/data/__init__.py | 14 ++ src/lsdb_ml/data/dataset.py | 193 ++++++++++++++++++++++++++++ src/lsdb_ml/data/sharding.py | 116 +++++++++++++++++ src/lsdb_ml/data/source.py | 148 +++++++++++++++++++++ tests/lsdb_ml/data/test_dataset.py | 162 +++++++++++++++++++++++ tests/lsdb_ml/data/test_sharding.py | 32 +++++ 6 files changed, 665 insertions(+) create mode 100644 src/lsdb_ml/data/__init__.py create mode 100644 src/lsdb_ml/data/dataset.py create mode 100644 src/lsdb_ml/data/sharding.py create mode 100644 src/lsdb_ml/data/source.py create mode 100644 tests/lsdb_ml/data/test_dataset.py create mode 100644 tests/lsdb_ml/data/test_sharding.py diff --git a/src/lsdb_ml/data/__init__.py b/src/lsdb_ml/data/__init__.py new file mode 100644 index 0000000..f50496b --- /dev/null +++ b/src/lsdb_ml/data/__init__.py @@ -0,0 +1,14 @@ +"""Streaming LSDB catalogs into PyTorch.""" + +from .dataset import CatalogIterableDataset, iter_slices +from .sharding import ShardInfo, assign_partitions +from .source import CatalogStreamSource, PartitionSource + +__all__ = [ + "CatalogIterableDataset", + "CatalogStreamSource", + "PartitionSource", + "ShardInfo", + "assign_partitions", + "iter_slices", +] diff --git a/src/lsdb_ml/data/dataset.py b/src/lsdb_ml/data/dataset.py new file mode 100644 index 0000000..04b3868 --- /dev/null +++ b/src/lsdb_ml/data/dataset.py @@ -0,0 +1,193 @@ +"""An ``IterableDataset`` that streams collated batches from an LSDB catalog.""" + +from __future__ import annotations + +import itertools +import logging +import warnings +from collections.abc import Iterator + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import IterableDataset, get_worker_info + +from lsdb_ml.data.sharding import ShardInfo, assign_partitions, distributed_rank_and_world_size, shard_seed +from lsdb_ml.data.source import CatalogStreamSource, PartitionSource +from lsdb_ml.spec import ColumnSpec, ResolvedSpec + +logger = logging.getLogger(__name__) + + +def iter_slices(frame: pd.DataFrame, batch_size: int | None) -> Iterator[pd.DataFrame]: + """Cut one chunk into consecutive row slices; slices never span chunks. + + Parameters + ---------- + frame : pd.DataFrame + One chunk. + batch_size : int or None + Maximum rows per slice. None yields the whole chunk as one slice. + + Yields + ------ + pd.DataFrame + Non-empty slices, in row order, covering every row of ``frame`` once. + + Examples + -------- + >>> [len(s) for s in iter_slices(pd.DataFrame({"a": range(7)}), 3)] + [3, 3, 1] + >>> [len(s) for s in iter_slices(pd.DataFrame({"a": []}), None)] + [] + """ + if len(frame) == 0: + return + if batch_size is None: + yield frame + return + for start in range(0, len(frame), batch_size): + yield frame.iloc[start : start + batch_size] + + +class CatalogIterableDataset(IterableDataset): + """Stream batches of tensors from an LSDB catalog for training. + + Each catalog partition is read by exactly one consumer: partitions are sharded across + DDP ranks and DataLoader workers, so with any ``num_workers`` and world size every row + appears exactly once per epoch. Batches are cut from the chunks the stream produces + (one or more stacked partitions) and never span two chunks, so batch sizes are ragged. + Use it with ``DataLoader(dataset, batch_size=None, num_workers=...)``: batches arrive + already collated. + + Parameters + ---------- + catalog : lsdb.Catalog + The catalog to stream. It is projected to the columns the spec needs. + spec : ColumnSpec or ResolvedSpec + Which columns become which tensors. Validated against ``catalog.meta`` here, + without computing data. + batch_size : int or None, default None + Maximum rows per batch. None yields each chunk as one batch. + shuffle : bool, default True + Shuffle partition order each epoch and rows within each chunk. Mixing is limited + to rows within one chunk; raise ``partitions_per_chunk`` for more. + seed : int, optional + Base seed. If None, one is drawn at construction, which DataLoader workers then + share. Required when ``shuffle`` is used with ``world_size > 1``. + epochs : int or None, default 1 + Passes over the catalog per iteration. Each pass is reshuffled reproducibly. None + streams forever. + partitions_per_chunk : int, default 1 + Partitions stacked per chunk. Bigger chunks use more memory but give more Dask + parallelism per fetch and better row mixing. + rank, world_size : int, optional + DDP rank and world size. Default to ``torch.distributed`` values when it is + initialized, otherwise ``0`` and ``1``. They are read here rather than during + iteration because ``torch.distributed`` is not initialized in DataLoader workers. + use_dask_client, scheduler_address, worker_scheduler + Passed to :class:`~lsdb_ml.data.source.CatalogStreamSource`. + source : PartitionSource, optional + Replace the default ``CatalogStreamSource``. + + Notes + ----- + Every yielded batch is a ``dict[str, torch.Tensor]`` with the keys described by + ``dataset.spec.describe()``, including the row index under ``"_index"`` by default. + + :meth:`set_epoch` takes effect at the next ``iter()``. With + ``DataLoader(persistent_workers=True)`` workers keep their own copy of the dataset, so + ``set_epoch`` does not reach them; use ``epochs`` instead. + """ + + def __init__( + self, + catalog, + spec: ColumnSpec | ResolvedSpec, + *, + batch_size: int | None = None, + shuffle: bool = True, + seed: int | None = None, + epochs: int | None = 1, + partitions_per_chunk: int = 1, + rank: int | None = None, + world_size: int | None = None, + use_dask_client: bool = True, + scheduler_address: str | None = None, + worker_scheduler: str = "synchronous", + source: PartitionSource | None = None, + ): + super().__init__() + self.spec = spec if isinstance(spec, ResolvedSpec) else spec.resolve(catalog) + if batch_size is not None and (isinstance(batch_size, bool) or int(batch_size) < 1): + raise ValueError(f"batch_size must be a positive integer or None, got {batch_size!r}.") + if epochs is not None and (isinstance(epochs, bool) or int(epochs) < 1): + raise ValueError(f"epochs must be a positive integer or None, got {epochs!r}.") + self.batch_size = batch_size + self.shuffle = shuffle + self.epochs = epochs + self.epoch = 0 + + default_rank, default_world_size = distributed_rank_and_world_size() + self.rank = default_rank if rank is None else rank + self.world_size = default_world_size if world_size is None else world_size + if not 0 <= self.rank < self.world_size: + raise ValueError( + f"rank must be in [0, world_size), got rank={self.rank}, world_size={self.world_size}." + ) + + if seed is None: + if shuffle and self.world_size > 1: + raise ValueError( + "seed is required when shuffle=True and world_size > 1: every rank must " + "compute the same partition permutation, or shards would overlap." + ) + seed = int(np.random.default_rng().integers(2**31)) + self.seed = int(seed) + + if source is None: + source = CatalogStreamSource( + catalog[self.spec.columns], + partitions_per_chunk=partitions_per_chunk, + use_dask_client=use_dask_client, + scheduler_address=scheduler_address, + worker_scheduler=worker_scheduler, + ) + self.source = source + + def set_epoch(self, epoch: int) -> None: + """Set the epoch the next iteration starts from (like ``DistributedSampler.set_epoch``).""" + self.epoch = int(epoch) + + def shard(self) -> ShardInfo: + """Return the shard for the current process (call during iteration in workers).""" + info = get_worker_info() + worker_id, num_workers = (0, 1) if info is None else (info.id, info.num_workers) + return ShardInfo.from_parts(self.rank, self.world_size, worker_id, num_workers) + + def __iter__(self) -> Iterator[dict[str, torch.Tensor]]: + shard = self.shard() + npartitions = self.source.npartitions + if shard.num_shards > npartitions: + warnings.warn( + f"There are {shard.num_shards} shards (ranks x DataLoader workers) but only " + f"{npartitions} partitions, so some shards get no data in an epoch. Use fewer " + "workers or a catalog with more partitions.", + stacklevel=2, + ) + epochs = ( + itertools.count(self.epoch) + if self.epochs is None + else range(self.epoch, self.epoch + self.epochs) + ) + for epoch in epochs: + partitions = assign_partitions(npartitions, shard, epoch, self.seed, self.shuffle) + if len(partitions) == 0 and self.epochs is None and not self.shuffle: + return # this shard never gets data; don't spin forever + logger.debug("Shard %s epoch %d reads partitions %s", shard, epoch, partitions) + chunks = self.source.iter_chunks( + partitions, shuffle=self.shuffle, seed=shard_seed(self.seed, epoch, shard.shard_id) + ) + for chunk in chunks: + for rows in iter_slices(chunk, self.batch_size): + yield self.spec.collate(rows) diff --git a/src/lsdb_ml/data/sharding.py b/src/lsdb_ml/data/sharding.py new file mode 100644 index 0000000..e843054 --- /dev/null +++ b/src/lsdb_ml/data/sharding.py @@ -0,0 +1,116 @@ +"""Partition-level sharding across DDP ranks and DataLoader workers. + +Every shard computes the same per-epoch permutation of partition indices from a shared +seed and takes a strided slice of it, so shards are disjoint and together cover every +partition exactly once per epoch. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch.distributed as dist + + +@dataclass(frozen=True) +class ShardInfo: + """Which slice of the partitions a consumer reads. + + Attributes + ---------- + shard_id : int + This consumer's shard, ``rank * num_workers + worker_id``. + num_shards : int + Total number of shards, ``world_size * num_workers``. + """ + + shard_id: int + num_shards: int + + @classmethod + def from_parts(cls, rank: int, world_size: int, worker_id: int, num_workers: int) -> ShardInfo: + """Build the shard for a DDP rank and a DataLoader worker. + + Parameters + ---------- + rank : int + The DDP rank (0 when not distributed). + world_size : int + The DDP world size (1 when not distributed). + worker_id : int + The DataLoader worker id (0 in the main process). + num_workers : int + The number of DataLoader workers (1 in the main process). + + Returns + ------- + ShardInfo + The shard. + + Examples + -------- + >>> ShardInfo.from_parts(rank=1, world_size=2, worker_id=0, num_workers=3) + ShardInfo(shard_id=3, num_shards=6) + """ + return cls(shard_id=rank * num_workers + worker_id, num_shards=world_size * num_workers) + + +def distributed_rank_and_world_size() -> tuple[int, int]: + """Return ``(rank, world_size)`` from ``torch.distributed``, or ``(0, 1)`` if not initialized.""" + if dist.is_available() and dist.is_initialized(): + return dist.get_rank(), dist.get_world_size() + return 0, 1 + + +def assign_partitions(npartitions: int, shard: ShardInfo, epoch: int, seed: int, shuffle: bool) -> np.ndarray: + """Return the partition indices a shard reads in a given epoch. + + Parameters + ---------- + npartitions : int + Number of partitions in the catalog. + shard : ShardInfo + The shard. + epoch : int + The epoch; with ``shuffle``, each epoch uses a different permutation. + seed : int + Base seed shared by every shard. + shuffle : bool + Whether to permute partitions before striding. + + Returns + ------- + np.ndarray + Partition indices for this shard. + + Examples + -------- + >>> assign_partitions(5, ShardInfo(0, 2), epoch=0, seed=0, shuffle=False) + array([0, 2, 4]) + """ + if shuffle: + order = np.random.default_rng(np.random.SeedSequence([seed, epoch])).permutation(npartitions) + else: + order = np.arange(npartitions) + return order[shard.shard_id :: shard.num_shards] + + +def shard_seed(seed: int, epoch: int, shard_id: int) -> int: + """Derive the seed one shard uses for its own ordering in one epoch. + + Parameters + ---------- + seed : int + Base seed. + epoch : int + The epoch. + shard_id : int + The shard. + + Returns + ------- + int + A 32-bit seed. + """ + return int(np.random.SeedSequence([seed, epoch, shard_id]).generate_state(1)[0]) diff --git a/src/lsdb_ml/data/source.py b/src/lsdb_ml/data/source.py new file mode 100644 index 0000000..40292a2 --- /dev/null +++ b/src/lsdb_ml/data/source.py @@ -0,0 +1,148 @@ +"""Partition sources: the Dask-side producer that feeds the torch-side consumer.""" + +from __future__ import annotations + +import contextlib +import logging +from collections.abc import Iterator, Sequence +from typing import Protocol, runtime_checkable + +import dask +import numpy as np +import pandas as pd +from torch.utils.data import get_worker_info + +logger = logging.getLogger(__name__) + + +@runtime_checkable +class PartitionSource(Protocol): + """Streams chunks of rows for a given set of partitions. + + This is the producer/consumer seam: everything downstream (slicing, collation, the + training session) only sees the frames yielded here. + """ + + @property + def npartitions(self) -> int: + """Number of partitions available to stream.""" + + def iter_chunks(self, partitions: Sequence[int], *, shuffle: bool, seed: int) -> Iterator[pd.DataFrame]: + """Yield frames covering exactly the rows of ``partitions``, each row once.""" + + +class CatalogStreamSource: + """A :class:`PartitionSource` backed by lsdb's ``CatalogStream``. + + The requested partitions are selected as a sub-catalog and streamed with + ``CatalogStream``, which keeps lsdb's one-chunk prefetch and, with ``shuffle``, its + partition-order and within-chunk row shuffling. + + Parameters + ---------- + catalog : lsdb.Catalog + The catalog to stream. + partitions_per_chunk : int, default 1 + Partitions stacked into each chunk. + use_dask_client : bool, default True + In the main process, stream through the active ``distributed`` client if there is + one, so the next chunk is computed while the current one is consumed. + scheduler_address : str, optional + Connect to this Dask scheduler instead, from the main process and from every + DataLoader worker process. + worker_scheduler : str, default "synchronous" + Dask scheduler used inside DataLoader worker processes when no + ``scheduler_address`` is given. Workers run in parallel already, so the default + avoids oversubscribing cores with a thread pool per worker. + + Notes + ----- + Without ``shuffle``, ``CatalogStream`` streams partitions in reverse pixel order. + """ + + def __init__( + self, + catalog, + partitions_per_chunk: int = 1, + use_dask_client: bool = True, + scheduler_address: str | None = None, + worker_scheduler: str = "synchronous", + ): + if isinstance(partitions_per_chunk, bool) or int(partitions_per_chunk) < 1: + raise ValueError( + f"partitions_per_chunk must be a positive integer, got {partitions_per_chunk!r}." + ) + self.catalog = catalog + self.partitions_per_chunk = int(partitions_per_chunk) + self.use_dask_client = use_dask_client + self.scheduler_address = scheduler_address + self.worker_scheduler = worker_scheduler + self._owned_client = None + + @property + def npartitions(self) -> int: + """Number of partitions in the catalog.""" + return self.catalog.npartitions + + def __getstate__(self): + # A Client cannot cross into DataLoader worker processes; each one reconnects. + state = self.__dict__.copy() + state["_owned_client"] = None + return state + + def _client(self): + if self.scheduler_address is not None: + if self._owned_client is None: + from distributed import Client + + self._owned_client = Client(self.scheduler_address, set_as_default=False) + return self._owned_client + if get_worker_info() is not None or not self.use_dask_client: + return None + try: + from distributed import get_client + + return get_client() + except (ImportError, ValueError): + return None + + def iter_chunks(self, partitions: Sequence[int], *, shuffle: bool, seed: int) -> Iterator[pd.DataFrame]: + """Stream the rows of ``partitions`` in chunks of ``partitions_per_chunk`` partitions. + + Parameters + ---------- + partitions : Sequence[int] + Partition indices to stream. + shuffle : bool + Shuffle partition order and rows within each chunk. + seed : int + Seed for the shuffle. + + Yields + ------ + pd.DataFrame + One chunk at a time (a ``NestedFrame`` for nested catalogs). Chunks may be + empty when partitions are empty. + """ + from lsdb.streams import CatalogStream + + partitions = np.asarray(partitions, dtype=np.int64) + if len(partitions) == 0: + return + sub_catalog = self.catalog.partitions[partitions.tolist()] + client = self._client() + + in_worker = get_worker_info() is not None + config = ( + dask.config.set(scheduler=self.worker_scheduler) + if in_worker and client is None + else contextlib.nullcontext() + ) + with config: + yield from CatalogStream( + sub_catalog, + client=client, + partitions_per_chunk=self.partitions_per_chunk, + shuffle=shuffle, + seed=seed, + ) diff --git a/tests/lsdb_ml/data/test_dataset.py b/tests/lsdb_ml/data/test_dataset.py new file mode 100644 index 0000000..1f0d28a --- /dev/null +++ b/tests/lsdb_ml/data/test_dataset.py @@ -0,0 +1,162 @@ +import itertools +import math +from collections import Counter + +import numpy as np +import pytest +import torch +from distributed import Client +from torch.utils.data import DataLoader + +from lsdb_ml import CatalogIterableDataset, ColumnSpec, Flat, Nested +from lsdb_ml.data import iter_slices + +SPEC = ColumnSpec({"id": Flat("id", dtype=torch.int64), "lc": Nested("lc", ["t", "flux"])}) + + +def collect_ids(dataset, num_workers=0): + """Iterate a dataset through a DataLoader and return the ids of every batch.""" + loader = DataLoader(dataset, batch_size=None, num_workers=num_workers) + return [batch["id"].tolist() for batch in loader] + + +def flatten(batches): + """Concatenate per-batch id lists.""" + return [i for batch in batches for i in batch] + + +# Spawned DataLoader workers cost a few seconds each to start, so the multi-worker cases +# are a subset of the full grid. +@pytest.mark.parametrize( + "num_workers,world_size,shuffle", + [(0, 1, True), (0, 2, True), (0, 1, False), (0, 2, False), (2, 1, True), (2, 2, False)], +) +def test_every_row_exactly_once(nested_catalog, num_workers, world_size, shuffle): + """Across DataLoader workers and (simulated) DDP ranks, no row is duplicated or dropped.""" + ids = [] + for rank in range(world_size): + dataset = CatalogIterableDataset( + nested_catalog, SPEC, batch_size=16, shuffle=shuffle, seed=3, rank=rank, world_size=world_size + ) + ids += flatten(collect_ids(dataset, num_workers)) + assert sorted(ids) == list(range(500)) + + +def test_index_key_matches_catalog(nested_catalog): + """The _index tensor carries the catalog's _healpix_29 values, once each.""" + dataset = CatalogIterableDataset(nested_catalog, SPEC, batch_size=50, seed=0) + index = np.concatenate([batch["_index"].numpy() for batch in dataset]) + expected = nested_catalog.compute(progress_bar=False).index.to_numpy() + assert sorted(index.tolist()) == sorted(expected.tolist()) + + +def test_epochs_reshuffle_reproducibly(nested_catalog): + """Each epoch covers every row; order changes per epoch and is fixed by the seed.""" + + def run(seed): + dataset = CatalogIterableDataset(nested_catalog, SPEC, batch_size=16, seed=seed, epochs=2) + return flatten(collect_ids(dataset)) + + ids = run(seed=7) + assert Counter(ids) == Counter({i: 2 for i in range(500)}) + first, second = ids[:500], ids[500:] + assert sorted(first) == list(range(500)) and first != second + assert run(seed=7) == ids + assert run(seed=8) != ids + + +def test_set_epoch(nested_catalog): + """set_epoch selects the epoch's permutation.""" + dataset = CatalogIterableDataset(nested_catalog, SPEC, batch_size=None, seed=1, epochs=2) + two_epochs = flatten(collect_ids(dataset)) + dataset = CatalogIterableDataset(nested_catalog, SPEC, batch_size=None, seed=1, epochs=1) + dataset.set_epoch(1) + assert flatten(collect_ids(dataset)) == two_epochs[500:] + + +def test_batches_never_span_chunks(nested_catalog): + """Batches are cut within one partition: sizes follow each partition's row count.""" + partition_of = {} + sizes = [] + for i, partition in enumerate(nested_catalog.partitions): + ids = partition.compute(progress_bar=False)["id"].tolist() + partition_of.update(dict.fromkeys(ids, i)) + sizes.append(len(ids)) + dataset = CatalogIterableDataset(nested_catalog, SPEC, batch_size=7, seed=2, partitions_per_chunk=1) + batches = collect_ids(dataset) + assert all(len({partition_of[i] for i in batch}) == 1 for batch in batches) + assert len(batches) == sum(math.ceil(n / 7) for n in sizes) + + +@pytest.mark.parametrize("partitions_per_chunk", [1, 3]) +def test_batch_size_none_yields_whole_chunks(nested_catalog, partitions_per_chunk): + """With batch_size=None every chunk is one batch.""" + dataset = CatalogIterableDataset( + nested_catalog, SPEC, batch_size=None, seed=2, partitions_per_chunk=partitions_per_chunk + ) + batches = collect_ids(dataset) + assert len(batches) == math.ceil(nested_catalog.npartitions / partitions_per_chunk) + assert sorted(flatten(batches)) == list(range(500)) + + +def test_open_ended_stream(nested_catalog): + """epochs=None keeps streaming; every pass covers every row once.""" + dataset = CatalogIterableDataset(nested_catalog, SPEC, batch_size=None, seed=4, epochs=None) + per_epoch = nested_catalog.npartitions + batches = [b["id"].tolist() for b in itertools.islice(iter(dataset), 3 * per_epoch)] + for epoch in range(3): + assert sorted(flatten(batches[epoch * per_epoch : (epoch + 1) * per_epoch])) == list(range(500)) + + +def test_empty_partitions(nested_catalog): + """Partitions emptied by a filter are skipped, and every remaining row arrives once.""" + filtered = nested_catalog.query("mag > 21") + expected = sorted(filtered.compute(progress_bar=False)["id"].tolist()) + sizes = [len(p.compute(progress_bar=False)) for p in filtered.partitions] + assert 0 in sizes, "the filter should empty some partitions" + batches = collect_ids(CatalogIterableDataset(filtered, SPEC, batch_size=4, seed=0)) + assert all(batches) + assert sorted(flatten(batches)) == expected + + +def test_more_shards_than_partitions(nested_catalog): + """With more shards than partitions some shards are empty; coverage stays exact.""" + small = nested_catalog.partitions[[0, 1, 2]] + expected = sorted(small.compute(progress_bar=False)["id"].tolist()) + ids = [] + with pytest.warns(UserWarning, match="some shards get no data"): + for rank in range(5): + dataset = CatalogIterableDataset(small, SPEC, batch_size=8, seed=0, rank=rank, world_size=5) + ids += flatten(collect_ids(dataset)) + assert sorted(ids) == expected + + +def test_with_distributed_client(nested_catalog): + """In the main process the stream uses the active distributed client.""" + with Client(processes=False, n_workers=2, threads_per_worker=1, dashboard_address=None): + dataset = CatalogIterableDataset(nested_catalog, SPEC, batch_size=32, seed=0, partitions_per_chunk=4) + assert sorted(flatten(collect_ids(dataset))) == list(range(500)) + + +def test_validation(nested_catalog, unloadable_catalog): + """Bad arguments and specs fail at construction, before any data is read.""" + with pytest.raises(ValueError, match="seed is required"): + CatalogIterableDataset(unloadable_catalog, SPEC, rank=0, world_size=2) + CatalogIterableDataset(unloadable_catalog, SPEC, rank=0, world_size=2, shuffle=False) + with pytest.raises(ValueError, match="rank must be in"): + CatalogIterableDataset(unloadable_catalog, SPEC, seed=0, rank=2, world_size=2) + with pytest.raises(ValueError, match="batch_size"): + CatalogIterableDataset(unloadable_catalog, SPEC, batch_size=0) + with pytest.raises(ValueError, match="epochs"): + CatalogIterableDataset(unloadable_catalog, SPEC, epochs=0) + with pytest.raises(ValueError, match="partitions_per_chunk"): + CatalogIterableDataset(unloadable_catalog, SPEC, partitions_per_chunk=0) + with pytest.raises(ValueError, match="not in the catalog"): + CatalogIterableDataset(unloadable_catalog, ColumnSpec({"x": Flat("nope")})) + + +def test_iter_slices_cover_frame(nested_frame): + """Slices are in order and cover every row once.""" + slices = list(iter_slices(nested_frame, 64)) + assert [len(s) for s in slices] == [64] * 7 + [52] + assert np.array_equal(np.concatenate([s["id"].to_numpy() for s in slices]), nested_frame["id"].to_numpy()) diff --git a/tests/lsdb_ml/data/test_sharding.py b/tests/lsdb_ml/data/test_sharding.py new file mode 100644 index 0000000..16d6219 --- /dev/null +++ b/tests/lsdb_ml/data/test_sharding.py @@ -0,0 +1,32 @@ +import numpy as np +import pytest + +from lsdb_ml.data.sharding import ShardInfo, assign_partitions, shard_seed + + +@pytest.mark.parametrize("shuffle", [True, False]) +@pytest.mark.parametrize("num_shards", [1, 2, 3, 7, 40]) +def test_shards_are_disjoint_and_complete(shuffle, num_shards): + """Every partition goes to exactly one shard, for any shard count.""" + for epoch in range(3): + parts = [ + assign_partitions(33, ShardInfo(i, num_shards), epoch, seed=5, shuffle=shuffle) + for i in range(num_shards) + ] + combined = np.concatenate(parts) + assert sorted(combined.tolist()) == list(range(33)) + + +def test_shuffle_changes_by_epoch_and_is_reproducible(): + """Each epoch has its own permutation, determined by the seed.""" + shard = ShardInfo(0, 1) + first = assign_partitions(33, shard, 0, seed=1, shuffle=True) + assert np.array_equal(first, assign_partitions(33, shard, 0, seed=1, shuffle=True)) + assert not np.array_equal(first, assign_partitions(33, shard, 1, seed=1, shuffle=True)) + assert not np.array_equal(first, assign_partitions(33, shard, 0, seed=2, shuffle=True)) + + +def test_shard_seed_varies(): + """Shards and epochs get different derived seeds.""" + seeds = {shard_seed(0, epoch, shard) for epoch in range(3) for shard in range(3)} + assert len(seeds) == 9 From 7d3ab5ed2e5934939c526a210a7ecae4af7716b1 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 12:20:16 -0700 Subject: [PATCH 4/8] Add TrainSession, SupervisedStep and checkpointing Context-manager session modeled on Hyrax's train_stream: iterate (batch, metrics), stop(), close() returns the model. The optimization step is a user callable (SupervisedStep covers the common case). Periodic and final atomic checkpoints, rank-0 only; DDP runs under model.join() so uneven batch counts across ranks do not hang. Co-Authored-By: Claude Opus 5.5 --- src/lsdb_ml/inputs.py | 41 +++++ src/lsdb_ml/training/__init__.py | 7 + src/lsdb_ml/training/checkpoint.py | 88 +++++++++ src/lsdb_ml/training/session.py | 235 +++++++++++++++++++++++++ src/lsdb_ml/training/steps.py | 88 +++++++++ tests/lsdb_ml/training/test_ddp.py | 62 +++++++ tests/lsdb_ml/training/test_session.py | 139 +++++++++++++++ 7 files changed, 660 insertions(+) create mode 100644 src/lsdb_ml/inputs.py create mode 100644 src/lsdb_ml/training/__init__.py create mode 100644 src/lsdb_ml/training/checkpoint.py create mode 100644 src/lsdb_ml/training/session.py create mode 100644 src/lsdb_ml/training/steps.py create mode 100644 tests/lsdb_ml/training/test_ddp.py create mode 100644 tests/lsdb_ml/training/test_session.py diff --git a/src/lsdb_ml/inputs.py b/src/lsdb_ml/inputs.py new file mode 100644 index 0000000..6e18329 --- /dev/null +++ b/src/lsdb_ml/inputs.py @@ -0,0 +1,41 @@ +"""Selecting model inputs from a collated batch.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from typing import Any + +import torch + +#: How to build the model's positional arguments from a batch: one key, a list of keys, +#: or a callable ``batch -> argument`` (a tuple return value is unpacked into arguments). +InputsLike = str | Sequence[str] | Callable[[Mapping[str, torch.Tensor]], Any] + + +def model_args(batch: Mapping[str, torch.Tensor], inputs: InputsLike) -> tuple: + """Return the positional arguments to call a model with. + + Parameters + ---------- + batch : Mapping[str, torch.Tensor] + A collated batch. + inputs : str, sequence of str, or callable + A key, keys (passed positionally in order), or ``callable(batch)``. + + Returns + ------- + tuple + Positional arguments. + + Examples + -------- + >>> batch = {"x": 1, "y": 2} + >>> model_args(batch, "x"), model_args(batch, ["y", "x"]), model_args(batch, lambda b: (b["x"], 5)) + ((1,), (2, 1), (1, 5)) + """ + if isinstance(inputs, str): + return (batch[inputs],) + if callable(inputs): + result = inputs(batch) + return result if isinstance(result, tuple) else (result,) + return tuple(batch[key] for key in inputs) diff --git a/src/lsdb_ml/training/__init__.py b/src/lsdb_ml/training/__init__.py new file mode 100644 index 0000000..c6fa0ee --- /dev/null +++ b/src/lsdb_ml/training/__init__.py @@ -0,0 +1,7 @@ +"""Training on streamed LSDB catalogs.""" + +from .checkpoint import load_checkpoint, save_checkpoint +from .session import TrainSession +from .steps import SupervisedStep, TrainStep + +__all__ = ["SupervisedStep", "TrainSession", "TrainStep", "load_checkpoint", "save_checkpoint"] diff --git a/src/lsdb_ml/training/checkpoint.py b/src/lsdb_ml/training/checkpoint.py new file mode 100644 index 0000000..360a412 --- /dev/null +++ b/src/lsdb_ml/training/checkpoint.py @@ -0,0 +1,88 @@ +"""Saving and loading training checkpoints.""" + +from __future__ import annotations + +import os +import tempfile +from pathlib import Path +from typing import Any + +import torch +from torch import nn + + +def unwrap_model(model: nn.Module) -> nn.Module: + """Return the underlying module of a ``DistributedDataParallel``/``DataParallel`` wrapper.""" + return ( + model.module if isinstance(model, (nn.parallel.DistributedDataParallel, nn.DataParallel)) else model + ) + + +def save_checkpoint(path: str | Path, model: nn.Module, step: Any = None, **counters: Any) -> Path: + """Atomically write a checkpoint. + + The file is a ``torch.save`` dict with ``"model"`` (the state dict), ``"step"`` (the + step's ``state_dict()`` if it has one) and any extra counters. + + Parameters + ---------- + path : str or Path + Destination file. Written to a temporary file first, then renamed, so a crash never + leaves a truncated checkpoint. + model : nn.Module + The model; DDP wrappers are unwrapped. + step : optional + The training step; saved if it defines ``state_dict()``. + **counters + Extra values to store, e.g. ``batch_count=100``. + + Returns + ------- + Path + The checkpoint path. + """ + from lsdb_ml import __version__ + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + payload = {"model": unwrap_model(model).state_dict(), "lsdb_ml_version": __version__, **counters} + if step is not None and hasattr(step, "state_dict"): + payload["step"] = step.state_dict() + fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp") + try: + with os.fdopen(fd, "wb") as handle: + torch.save(payload, handle) + os.replace(tmp_name, path) + except BaseException: + Path(tmp_name).unlink(missing_ok=True) + raise + return path + + +def load_checkpoint( + path: str | Path, model: nn.Module, step: Any = None, map_location="cpu" +) -> dict[str, Any]: + """Load a checkpoint written by :func:`save_checkpoint` into ``model`` (and ``step``). + + Parameters + ---------- + path : str or Path + Checkpoint file. + model : nn.Module + Model to load weights into; DDP wrappers are unwrapped. + step : optional + Training step to restore, if it defines ``load_state_dict()``. + map_location : default "cpu" + Passed to ``torch.load``. + + Returns + ------- + dict + The remaining entries (counters such as ``batch_count``). + """ + payload = torch.load(path, map_location=map_location, weights_only=True) + unwrap_model(model).load_state_dict(payload.pop("model")) + step_state = payload.pop("step", None) + if step_state is not None and step is not None and hasattr(step, "load_state_dict"): + step.load_state_dict(step_state) + return payload diff --git a/src/lsdb_ml/training/session.py b/src/lsdb_ml/training/session.py new file mode 100644 index 0000000..32084ed --- /dev/null +++ b/src/lsdb_ml/training/session.py @@ -0,0 +1,235 @@ +"""A context-manager training session that drives the per-batch loop.""" + +from __future__ import annotations + +import contextlib +import logging +from collections.abc import Callable, Iterable, Iterator, Mapping +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from torch import nn + +from lsdb_ml.device import DeviceLike, move_batch, resolve_device +from lsdb_ml.training.checkpoint import load_checkpoint, save_checkpoint, unwrap_model +from lsdb_ml.training.steps import TrainStep + +logger = logging.getLogger(__name__) + +Callback = Callable[["TrainSession", Mapping[str, torch.Tensor], Mapping[str, Any]], None] + + +def batch_rows(batch: Mapping[str, Any]) -> int: + """Number of rows in a collated batch (the first dimension of its first tensor).""" + for value in batch.values(): + if isinstance(value, torch.Tensor): + return int(value.shape[0]) if value.dim() else 1 + return 0 + + +class TrainSession: + """Train a model on a stream of batches, one batch at a time. + + Streams have no length and, in open-ended mode, no epochs, so instead of an epoch + loop the session hands control back after every batch:: + + with TrainSession(model, step, data=loader, checkpoint_dir="ckpt") as session: + for batch, metrics in session: # metrics is None for a skipped batch + if metrics and metrics["loss"] < 1e-3: + session.stop() + model = session.close() # already closed by the with block + + Parameters + ---------- + model : nn.Module + The model; moved to ``device``. May be wrapped in ``DistributedDataParallel``, in + which case iteration runs under ``model.join()`` so ranks with fewer batches do + not hang the others. + step : TrainStep + Callable ``step(model, batch) -> metrics`` that performs one optimization step, + e.g. :class:`~lsdb_ml.training.SupervisedStep`. + data : iterable of batches, optional + Usually a ``DataLoader`` over a :class:`~lsdb_ml.data.CatalogIterableDataset`. + Without it, drive the session manually with :meth:`train_batch`. + device : str or torch.device, default "cpu" + Where the model and batches go; ``"auto"`` picks CUDA, then MPS, then CPU. + checkpoint_dir : str or Path, optional + Where checkpoints are written. None disables checkpointing. + checkpoint_every : int, optional + Write ``checkpoint-.pt`` every this many trained batches. + keep_last : int, optional + Keep only this many periodic checkpoints. The final checkpoint is always kept. + min_batch_size : int, optional + Skip batches with fewer rows, e.g. 2 for models with BatchNorm. Empty batches are + always skipped. + callbacks : list of callables, optional + Each is called as ``callback(session, batch, metrics)`` after every trained batch; + use them for logging (MLflow, TensorBoard, ...). + resume_from : str or Path, optional + Checkpoint to restore model weights, step state and ``batch_count`` from. The data + position is not restored. + + Notes + ----- + Every trained batch is one optimizer step, whatever its size, so small batches (e.g. + the last slice of a chunk) weigh as much as full ones. Use ``min_batch_size`` to skip + them. Under DDP only rank 0 writes checkpoints. + """ + + def __init__( + self, + model: nn.Module, + step: TrainStep, + data: Iterable[Mapping[str, torch.Tensor]] | None = None, + *, + device: DeviceLike = "cpu", + checkpoint_dir: str | Path | None = None, + checkpoint_every: int | None = None, + keep_last: int | None = None, + min_batch_size: int | None = None, + callbacks: Iterable[Callback] = (), + resume_from: str | Path | None = None, + ): + if checkpoint_every is not None and checkpoint_dir is None: + raise ValueError("checkpoint_every needs a checkpoint_dir.") + for name, value in ( + ("checkpoint_every", checkpoint_every), + ("keep_last", keep_last), + ("min_batch_size", min_batch_size), + ): + if value is not None and (isinstance(value, bool) or int(value) < 1): + raise ValueError(f"{name} must be a positive integer or None, got {value!r}.") + self.device = resolve_device(device) + self.model = model.to(self.device) + self.step = step + self.data = data + self.checkpoint_dir = None if checkpoint_dir is None else Path(checkpoint_dir) + self.checkpoint_every = checkpoint_every + self.keep_last = keep_last + self.min_batch_size = min_batch_size + self.callbacks = list(callbacks) + self.batch_count = 0 + self.rows_seen = 0 + self._stop_requested = False + self._closed = False + self._periodic_checkpoints: list[Path] = [] + if resume_from is not None: + counters = load_checkpoint(resume_from, self.model, self.step, map_location=self.device) + self.batch_count = int(counters.get("batch_count", 0)) + self.rows_seen = int(counters.get("rows_seen", 0)) + + def __iter__(self) -> Iterator[tuple[Mapping[str, torch.Tensor], Mapping[str, Any] | None]]: + """Train on every batch from ``data``, yielding ``(batch, metrics)`` after each.""" + if self.data is None: + raise RuntimeError("This session has no data source; call train_batch(batch) instead.") + self._check_open() + self._stop_requested = False + with contextlib.ExitStack() as stack: + if isinstance(self.model, nn.parallel.DistributedDataParallel): + stack.enter_context(self.model.join()) + iterator = iter(self.data) + # Closing the iterator shuts down DataLoader workers and pending fetches. + stack.callback(lambda: getattr(iterator, "close", lambda: None)()) + for batch in iterator: + metrics = self.train_batch(batch) + yield batch, metrics + if self._stop_requested: + break + + def train_batch(self, batch: Mapping[str, torch.Tensor]) -> Mapping[str, Any] | None: + """Train on one batch. + + Parameters + ---------- + batch : Mapping[str, torch.Tensor] + A collated batch. + + Returns + ------- + Mapping or None + The step's metrics, or None if the batch was skipped (empty, or smaller than + ``min_batch_size``). + + Raises + ------ + RuntimeError + If the session is closed. + """ + self._check_open() + rows = batch_rows(batch) + if rows == 0: + logger.debug("Skipping an empty batch.") + return None + if self.min_batch_size is not None and rows < self.min_batch_size: + logger.debug("Skipping a batch of %d rows (min_batch_size=%d).", rows, self.min_batch_size) + return None + + self.model.train() + metrics = self.step(self.model, move_batch(batch, self.device)) + self.batch_count += 1 + self.rows_seen += rows + for callback in self.callbacks: + callback(self, batch, metrics) + if self.checkpoint_every is not None and self.batch_count % self.checkpoint_every == 0: + self._save_periodic() + return metrics + + def stop(self) -> None: + """End the current iteration after the batch being processed.""" + self._stop_requested = True + + def save_checkpoint(self, name: str) -> Path | None: + """Write a checkpoint named ``name`` into ``checkpoint_dir``. + + Returns the path, or None when checkpointing is disabled or this is not rank 0. + """ + if self.checkpoint_dir is None or not _is_rank_zero(): + return None + return save_checkpoint( + self.checkpoint_dir / name, + self.model, + self.step, + batch_count=self.batch_count, + rows_seen=self.rows_seen, + ) + + def _save_periodic(self) -> None: + path = self.save_checkpoint(f"checkpoint-{self.batch_count:08d}.pt") + if path is None: + return + self._periodic_checkpoints.append(path) + if self.keep_last is not None: + while len(self._periodic_checkpoints) > self.keep_last: + self._periodic_checkpoints.pop(0).unlink(missing_ok=True) + + def close(self) -> nn.Module: + """Stop, write the final checkpoint (``checkpoint-final.pt``) and return the model. + + Safe to call more than once. + + Returns + ------- + nn.Module + The trained model, unwrapped from any DDP wrapper. + """ + if not self._closed: + self.stop() + self.save_checkpoint("checkpoint-final.pt") + self._closed = True + return unwrap_model(self.model) + + def _check_open(self) -> None: + if self._closed: + raise RuntimeError("This TrainSession is closed.") + + def __enter__(self) -> TrainSession: + return self + + def __exit__(self, *exc_info) -> None: + self.close() + + +def _is_rank_zero() -> bool: + return not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0 diff --git a/src/lsdb_ml/training/steps.py b/src/lsdb_ml/training/steps.py new file mode 100644 index 0000000..7dea14f --- /dev/null +++ b/src/lsdb_ml/training/steps.py @@ -0,0 +1,88 @@ +"""Training steps: the seam between the session and the model's optimization logic.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any, Protocol, runtime_checkable + +import torch +from torch import nn + +from lsdb_ml.inputs import InputsLike, model_args + + +@runtime_checkable +class TrainStep(Protocol): + """One optimization step on one batch. + + The step owns everything between receiving a batch and returning metrics: zeroing + gradients, the forward pass, the loss, backward, and the optimizer step. The session + owns everything around it (device placement, skipping, counting, checkpoints). + + Steps may also define ``state_dict()`` and ``load_state_dict()``; the session then + includes their state (e.g. the optimizer's) in checkpoints. + """ + + def __call__(self, model: nn.Module, batch: Mapping[str, torch.Tensor]) -> Mapping[str, Any]: + """Train on ``batch`` and return metrics such as ``{"loss": 0.12}``.""" + + +class SupervisedStep: + """A standard supervised step: ``loss_fn(model(*inputs), batch[target])``. + + Parameters + ---------- + loss_fn : callable + ``loss_fn(prediction, target) -> scalar tensor``, e.g. ``nn.MSELoss()``. + optimizer : torch.optim.Optimizer + Optimizer over the model's parameters. + inputs : str, sequence of str, or callable + Batch key(s) passed positionally to the model, or ``callable(batch)``. + target : str + Batch key of the target. + scheduler : optional + Learning-rate scheduler, stepped after every optimizer step. + + Notes + ----- + Every call steps the optimizer once, whatever the batch size: a small batch weighs as + much as a full one. + """ + + def __init__( + self, + loss_fn: Callable[[Any, torch.Tensor], torch.Tensor], + optimizer: torch.optim.Optimizer, + inputs: InputsLike, + target: str, + scheduler: Any | None = None, + ): + self.loss_fn = loss_fn + self.optimizer = optimizer + self.inputs = inputs + self.target = target + self.scheduler = scheduler + + def __call__(self, model: nn.Module, batch: Mapping[str, torch.Tensor]) -> dict[str, float]: + """Run one optimization step and return ``{"loss": float}``.""" + self.optimizer.zero_grad(set_to_none=True) + prediction = model(*model_args(batch, self.inputs)) + loss = self.loss_fn(prediction, batch[self.target]) + loss.backward() + self.optimizer.step() + if self.scheduler is not None: + self.scheduler.step() + return {"loss": loss.item()} + + def state_dict(self) -> dict[str, Any]: + """Return optimizer (and scheduler) state for checkpoints.""" + state = {"optimizer": self.optimizer.state_dict()} + if self.scheduler is not None: + state["scheduler"] = self.scheduler.state_dict() + return state + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore state saved by :meth:`state_dict`.""" + self.optimizer.load_state_dict(state["optimizer"]) + if self.scheduler is not None and "scheduler" in state: + self.scheduler.load_state_dict(state["scheduler"]) diff --git a/tests/lsdb_ml/training/test_ddp.py b/tests/lsdb_ml/training/test_ddp.py new file mode 100644 index 0000000..37283e6 --- /dev/null +++ b/tests/lsdb_ml/training/test_ddp.py @@ -0,0 +1,62 @@ +import os +import socket + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch import nn + +from lsdb_ml import SupervisedStep, TrainSession + + +def free_port(): + """Find an unused local TCP port for the process group.""" + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def run_rank(rank, world_size, port, out_dir): + """Train a DDP model on a rank-dependent number of batches.""" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + try: + torch.manual_seed(0) + model = nn.parallel.DistributedDataParallel(nn.Linear(1, 1)) + step = SupervisedStep( + nn.MSELoss(), torch.optim.SGD(model.parameters(), lr=0.1), inputs="x", target="y" + ) + # Rank 0 gets 6 batches and rank 1 only 2: without join() rank 0 would hang. + n_batches = 6 if rank == 0 else 2 + data = [{"x": torch.full((4, 1), float(rank + 1)), "y": torch.ones(4, 1)}] * n_batches + with TrainSession(model, step, data, checkpoint_dir=out_dir) as session: + for _ in session: + pass + trained = session.close() + torch.save( + {"batches": session.batch_count, "weight": trained.weight.detach().clone()}, + os.path.join(out_dir, f"rank{rank}.pt"), + ) + finally: + dist.destroy_process_group() + + +def test_ddp_uneven_inputs_do_not_hang(tmp_path): + """Ranks with different batch counts finish, stay in sync, and only rank 0 checkpoints.""" + ctx = mp.start_processes( + run_rank, args=(2, free_port(), str(tmp_path)), nprocs=2, join=False, start_method="spawn" + ) + for process in ctx.processes: + process.join(timeout=120) + hung = [p for p in ctx.processes if p.is_alive()] + for process in hung: # don't leave stuck ranks behind to block interpreter exit + process.kill() + process.join() + assert not hung, "DDP training hung" + assert all(p.exitcode == 0 for p in ctx.processes) + + results = [torch.load(tmp_path / f"rank{r}.pt") for r in range(2)] + assert [r["batches"] for r in results] == [6, 2] + assert torch.equal(results[0]["weight"], results[1]["weight"]) + assert sorted(p.name for p in tmp_path.glob("checkpoint-*.pt")) == ["checkpoint-final.pt"] diff --git a/tests/lsdb_ml/training/test_session.py b/tests/lsdb_ml/training/test_session.py new file mode 100644 index 0000000..f4d437b --- /dev/null +++ b/tests/lsdb_ml/training/test_session.py @@ -0,0 +1,139 @@ +import numpy as np +import pytest +import torch +from torch import nn + +from lsdb_ml import CatalogIterableDataset, ColumnSpec, Flat, SupervisedStep, TrainSession +from lsdb_ml.training import load_checkpoint + + +def scaled_dec(df): + """Feature: declination scaled to [-1, 1], shaped (B, 1).""" + return (df["dec"].to_numpy(dtype=np.float32) / 60)[:, None] + + +def linear_target(df): + """Target: an exact linear function of the feature, shaped (B, 1).""" + return 3 * scaled_dec(df) + 1 + + +SPEC = ColumnSpec({"x": Flat(["dec"], transform=scaled_dec), "y": Flat(["dec"], transform=linear_target)}) + + +def make_session(catalog, tmp_path=None, epochs=1, **kwargs): + """A linear model, an SGD step and a single-process data stream.""" + torch.manual_seed(0) + model = nn.Linear(1, 1) + step = SupervisedStep(nn.MSELoss(), torch.optim.SGD(model.parameters(), lr=0.1), inputs="x", target="y") + data = CatalogIterableDataset(catalog, SPEC, batch_size=16, seed=0, epochs=epochs) + if tmp_path is not None: + kwargs.setdefault("checkpoint_dir", tmp_path) + return TrainSession(model, step, data, **kwargs), model, step + + +def test_session_trains(nested_catalog): + """Iterating the session trains the model: the loss drops and the fit is close.""" + session, model, _ = make_session(nested_catalog, epochs=3) + with session: + losses = [metrics["loss"] for _, metrics in session] + trained = session.close() + assert trained is model + assert np.mean(losses[-5:]) < 0.01 * np.mean(losses[:5]) + assert session.batch_count == len(losses) + assert session.rows_seen == 1500 + assert model.weight.item() == pytest.approx(3, abs=0.05) + assert model.bias.item() == pytest.approx(1, abs=0.05) + + +def test_periodic_and_final_checkpoints(nested_catalog, tmp_path): + """Checkpoints are written every N batches, pruned to keep_last, and reload exactly.""" + session, model, _ = make_session(nested_catalog, tmp_path, checkpoint_every=5, keep_last=2) + with session: + for _ in session: + pass + n = session.batch_count + periodic = sorted(p.name for p in tmp_path.glob("checkpoint-0*.pt")) + assert periodic == [f"checkpoint-{b:08d}.pt" for b in range(5 * (n // 5) - 5, 5 * (n // 5) + 1, 5)] + assert (tmp_path / "checkpoint-final.pt").exists() + assert not list(tmp_path.glob(".*tmp")) + + fresh = nn.Linear(1, 1) + counters = load_checkpoint(tmp_path / "checkpoint-final.pt", fresh) + assert counters["batch_count"] == n + assert torch.equal(fresh.weight, model.weight) and torch.equal(fresh.bias, model.bias) + + +def test_resume(nested_catalog, tmp_path): + """resume_from restores weights, optimizer state and counters.""" + session, model, _ = make_session(nested_catalog, tmp_path) + with session: + for _ in session: + pass + resumed, resumed_model, resumed_step = make_session( + nested_catalog, resume_from=tmp_path / "checkpoint-final.pt" + ) + assert resumed.batch_count == session.batch_count + assert torch.equal(resumed_model.weight, model.weight) + assert resumed_step.optimizer.state_dict()["param_groups"][0]["lr"] == 0.1 + + +def test_stop(nested_catalog): + """stop() ends iteration after the current batch.""" + session, _, _ = make_session(nested_catalog) + with session: + for _, _ in session: + if session.batch_count == 3: + session.stop() + assert session.batch_count == 3 + + +def test_skips_empty_and_small_batches(nested_catalog): + """Empty batches and batches under min_batch_size do not step the optimizer.""" + session, model, _ = make_session(nested_catalog, min_batch_size=4) + weight = model.weight.detach().clone() + assert session.train_batch({"x": torch.zeros(0, 1), "y": torch.zeros(0, 1)}) is None + assert session.train_batch({"x": torch.ones(3, 1), "y": torch.ones(3, 1)}) is None + assert session.batch_count == 0 and torch.equal(model.weight, weight) + assert session.train_batch({"x": torch.ones(4, 1), "y": torch.ones(4, 1)})["loss"] >= 0 + assert session.batch_count == 1 + + +def test_callbacks(nested_catalog): + """Callbacks run once per trained batch with the session and metrics.""" + seen = [] + session, _, _ = make_session(nested_catalog, callbacks=[lambda s, b, m: seen.append((s.batch_count, m))]) + with session: + for _ in session: + pass + assert [count for count, _ in seen] == list(range(1, session.batch_count + 1)) + assert all("loss" in metrics for _, metrics in seen) + + +def test_close_semantics(nested_catalog, tmp_path): + """close() is idempotent and the session cannot be used afterwards.""" + session, model, _ = make_session(nested_catalog, tmp_path) + assert session.close() is model + assert session.close() is model + with pytest.raises(RuntimeError, match="closed"): + session.train_batch({"x": torch.ones(2, 1), "y": torch.ones(2, 1)}) + with pytest.raises(RuntimeError, match="closed"): + next(iter(session)) + + +def test_manual_session_without_data(): + """A session without data is driven by train_batch and cannot be iterated.""" + model = nn.Linear(1, 1) + step = SupervisedStep(nn.MSELoss(), torch.optim.SGD(model.parameters(), lr=0.1), inputs="x", target="y") + with TrainSession(model, step) as session: + with pytest.raises(RuntimeError, match="no data source"): + next(iter(session)) + assert session.train_batch({"x": torch.ones(2, 1), "y": torch.ones(2, 1)}) is not None + + +def test_argument_validation(): + """Inconsistent checkpoint and size arguments are rejected.""" + model = nn.Linear(1, 1) + with pytest.raises(ValueError, match="needs a checkpoint_dir"): + TrainSession(model, step=None, checkpoint_every=5) + with pytest.raises(ValueError, match="min_batch_size"): + TrainSession(model, step=None, min_batch_size=0) From a500c3038b49b8b8bc4e3552cad2d8cd9ca84676 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 12:20:17 -0700 Subject: [PATCH 5/8] Export training API and add a training example script Co-Authored-By: Claude Opus 5.5 --- docs/examples/train_from_catalog.py | 87 +++++++++++++++++++++++++++++ src/lsdb_ml/__init__.py | 6 ++ 2 files changed, 93 insertions(+) create mode 100644 docs/examples/train_from_catalog.py diff --git a/docs/examples/train_from_catalog.py b/docs/examples/train_from_catalog.py new file mode 100644 index 0000000..bfaa7b9 --- /dev/null +++ b/docs/examples/train_from_catalog.py @@ -0,0 +1,87 @@ +"""Train a small model on an LSDB catalog with a nested light-curve column. + +Run with ``python docs/examples/train_from_catalog.py``. It builds a synthetic catalog, +so no data download is needed. To use a real catalog, replace ``make_catalog()`` with +``lsdb.open_catalog(...)`` and adjust the spec's column names. +""" + +import lsdb +import nested_pandas as npd +import numpy as np +import pandas as pd +import torch +from distributed import Client +from torch import nn +from torch.utils.data import DataLoader + +from lsdb_ml import CatalogIterableDataset, ColumnSpec, Flat, Nested, SupervisedStep, TrainSession + + +def make_catalog(n_objects=2000, seed=0): + """Objects whose label is the mean flux of their (ragged) light curve, plus noise.""" + rng = np.random.default_rng(seed) + lengths = rng.integers(1, 30, n_objects) + flux = rng.normal(0, 1, lengths.sum()).astype(np.float32) + owner = np.repeat(np.arange(n_objects), lengths) + mean_flux = np.bincount(owner, weights=flux) / lengths + objects = pd.DataFrame( + { + "ra": rng.uniform(0, 360, n_objects), + "dec": rng.uniform(-60, 60, n_objects), + "label": (mean_flux + rng.normal(0, 0.01, n_objects)).astype(np.float32), + } + ) + light_curves = pd.DataFrame({"mjd": rng.uniform(0, 100, lengths.sum()), "flux": flux}, index=owner) + frame = npd.NestedFrame(objects).join_nested(light_curves, "lc") + return lsdb.from_dataframe(frame, partition_rows=200) + + +class MaskedMeanModel(nn.Module): + """Predict the label from the masked mean of a learned per-observation embedding.""" + + def __init__(self): + super().__init__() + self.embed = nn.Linear(1, 8) + self.head = nn.Linear(8, 1) + + def forward(self, flux, mask): + """Pool padded observations using the mask, so padding never contributes.""" + hidden = self.embed(flux.unsqueeze(-1)) * mask.unsqueeze(-1) + pooled = hidden.sum(dim=1) / mask.sum(dim=1, keepdim=True).clamp(min=1) + return self.head(pooled).squeeze(-1) + + +def main(): + """Stream the catalog for a few epochs and train until the loss is small.""" + catalog = make_catalog() + spec = ColumnSpec({"flux": Nested("lc", "flux"), "label": Flat("label")}) + + torch.manual_seed(0) + model = MaskedMeanModel() + step = SupervisedStep( + loss_fn=nn.MSELoss(), + optimizer=torch.optim.Adam(model.parameters(), lr=1e-2), + inputs=["flux", "flux_mask"], + target="label", + ) + dataset = CatalogIterableDataset(catalog, spec, batch_size=64, seed=0, epochs=20, partitions_per_chunk=2) + loader = DataLoader(dataset, batch_size=None, num_workers=0) + + recent = [] + with TrainSession(model, step, loader, checkpoint_dir="checkpoints", checkpoint_every=100) as session: + for _, metrics in session: + if metrics is None: # skipped (empty) batch + continue + recent = (recent + [metrics["loss"]])[-10:] + if session.batch_count % 10 == 0: + print(f"batch {session.batch_count}: mean loss of last 10 batches {np.mean(recent):.5f}") + if len(recent) == 10 and np.mean(recent) < 1e-3: + print(f"converged after {session.batch_count} batches") + session.stop() + print("final checkpoint: checkpoints/checkpoint-final.pt") + + +if __name__ == "__main__": + # A distributed client lets the stream compute the next chunk while the model trains. + with Client(processes=False, n_workers=2, threads_per_worker=1, dashboard_address=None): + main() diff --git a/src/lsdb_ml/__init__.py b/src/lsdb_ml/__init__.py index 6502887..d514376 100644 --- a/src/lsdb_ml/__init__.py +++ b/src/lsdb_ml/__init__.py @@ -1,9 +1,15 @@ from ._version import __version__ +from .data import CatalogIterableDataset from .device import move_batch, resolve_device from .example_module import greetings, meaning from .spec import ColumnSpec, Field, Flat, Nested, ResolvedSpec, SpecError, TensorInfo +from .training import SupervisedStep, TrainSession, TrainStep __all__ = [ + "CatalogIterableDataset", + "SupervisedStep", + "TrainSession", + "TrainStep", "ColumnSpec", "Field", "Flat", From b9c09ba6f2f6713b38bb259c4d8fa6239cae612d Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 12:27:27 -0700 Subject: [PATCH 6/8] Add OutputSpec and ModelPlugin for per-worker model installation ModelSpec ships only a factory and weights path; each worker builds the model once in its WorkerPlugin setup, loads weights (plain state dicts or TrainSession checkpoints), switches to eval and moves to the device resolved on the worker. OutputSpec builds output meta and validates every batch's output shape against it. Co-Authored-By: Claude Opus 5.5 --- src/lsdb_ml/inference/outputs.py | 164 ++++++++++++++++++++++++++++++ src/lsdb_ml/inference/plugin.py | 165 +++++++++++++++++++++++++++++++ 2 files changed, 329 insertions(+) create mode 100644 src/lsdb_ml/inference/outputs.py create mode 100644 src/lsdb_ml/inference/plugin.py diff --git a/src/lsdb_ml/inference/outputs.py b/src/lsdb_ml/inference/outputs.py new file mode 100644 index 0000000..93b56c1 --- /dev/null +++ b/src/lsdb_ml/inference/outputs.py @@ -0,0 +1,164 @@ +"""Mapping model outputs to catalog columns.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence + +import nested_pandas as npd +import numpy as np +import pandas as pd +import torch + + +class OutputSpec: + """Names and dtypes of the columns a model's output becomes. + + The model's output for a batch of ``B`` rows may be: + + - a tensor of shape ``(B,)`` or ``(B, 1)``, for a single column; + - a tensor of shape ``(B, k)``, for ``k`` columns, in order; + - a dict of tensors of shape ``(B,)`` or ``(B, 1)``, keyed by column name. + + Parameters + ---------- + columns : Mapping[str, dtype] or sequence of str + Column name to pandas/numpy dtype (e.g. ``{"z": "float32"}``), or names that all + use ``dtype``. + dtype : default "float32" + The dtype for columns given as a sequence of names. + + Examples + -------- + >>> spec = OutputSpec(["p_star", "p_galaxy"]) + >>> spec.columns + {'p_star': dtype('float32'), 'p_galaxy': dtype('float32')} + >>> spec.to_columns(torch.tensor([[0.1, 0.9], [0.8, 0.2]]), n_rows=2)["p_galaxy"] + array([0.9, 0.2], dtype=float32) + """ + + def __init__(self, columns: Mapping[str, object] | Sequence[str], dtype="float32"): + if isinstance(columns, str): + columns = [columns] + if not isinstance(columns, Mapping): + columns = dict.fromkeys(columns, dtype) + if not columns: + raise ValueError("OutputSpec needs at least one column.") + self.columns = {name: pd.api.types.pandas_dtype(d) for name, d in columns.items()} + + def meta(self, input_meta: pd.DataFrame, append: bool) -> npd.NestedFrame: + """Return the empty output frame (Dask ``meta``) for an input catalog's meta. + + Parameters + ---------- + input_meta : pd.DataFrame + The input catalog's meta. + append : bool + Keep the input columns before the new ones. + + Returns + ------- + npd.NestedFrame + Empty frame with the output columns and dtypes. + """ + empty = {name: pd.Series(dtype=d).array for name, d in self.columns.items()} + if append: + result = input_meta.copy() + for name, values in empty.items(): + result[name] = values + return result + return npd.NestedFrame(empty, index=input_meta.index[:0]) + + def to_columns(self, output, n_rows: int) -> dict[str, np.ndarray]: + """Convert one batch of model output to numpy columns, validating the shape. + + Parameters + ---------- + output : torch.Tensor or Mapping[str, torch.Tensor] + Model output for a batch. + n_rows : int + Rows in the batch; every column must have exactly this many values. + + Returns + ------- + dict[str, np.ndarray] + One 1-d array per output column. + + Raises + ------ + ValueError + If the output does not match the spec. + """ + names = list(self.columns) + if isinstance(output, Mapping): + if set(output) != set(names): + raise ValueError( + f"Model returned keys {sorted(output)}, but the outputs are {sorted(names)}." + ) + arrays = {name: _to_numpy(output[name], name) for name in names} + for name, array in arrays.items(): + if array.ndim == 2 and array.shape[1] == 1: + arrays[name] = array[:, 0] + elif array.ndim != 1: + raise ValueError(f"Output '{name}' has shape {array.shape}; expected ({n_rows},).") + else: + array = _to_numpy(output, "the model output") + if array.ndim == 1: + array = array[:, None] + if array.ndim != 2 or array.shape[1] != len(names): + raise ValueError( + f"Model returned shape {tuple(array.shape)}, but {len(names)} output column(s) " + f"{names} were declared; expected ({n_rows}, {len(names)})." + ) + arrays = {name: array[:, i] for i, name in enumerate(names)} + for name, array in arrays.items(): + if len(array) != n_rows: + raise ValueError( + f"Output '{name}' has {len(array)} values for a batch of {n_rows} rows; " + "the model must return exactly one prediction per input row." + ) + return arrays + + def build_frame( + self, partition: pd.DataFrame, columns: dict[str, np.ndarray], append: bool + ) -> pd.DataFrame: + """Attach prediction columns to their partition, by position. + + Parameters + ---------- + partition : pd.DataFrame + The input partition the predictions were computed from, in the same row order. + columns : dict[str, np.ndarray] + One array per output column, ``len(partition)`` values each. + append : bool + Keep the input columns. + + Returns + ------- + pd.DataFrame + The output partition, with ``partition``'s index. + """ + for name, values in columns.items(): + if len(values) != len(partition): + raise ValueError( + f"Output '{name}' has {len(values)} values for a partition of {len(partition)} rows." + ) + # Assigning arrays (not Series) is positional: no index alignment, which would be + # wrong for a non-unique _healpix_29 index. + typed = { + name: pd.Series(columns[name], dtype=d, copy=False).array for name, d in self.columns.items() + } + if append: + result = partition.copy() + for name, values in typed.items(): + result[name] = values + return result + result = npd.NestedFrame(index=partition.index) + for name, values in typed.items(): + result[name] = values + return result + + +def _to_numpy(value, name: str) -> np.ndarray: + if isinstance(value, torch.Tensor): + return value.detach().to("cpu").numpy() + raise TypeError(f"Expected {name} to be a torch.Tensor, got {type(value).__name__}.") diff --git a/src/lsdb_ml/inference/plugin.py b/src/lsdb_ml/inference/plugin.py new file mode 100644 index 0000000..8ee3d96 --- /dev/null +++ b/src/lsdb_ml/inference/plugin.py @@ -0,0 +1,165 @@ +"""Installing a model once per Dask worker.""" + +from __future__ import annotations + +import copy +import hashlib +import logging +import threading +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path + +import cloudpickle +import torch +from distributed import WorkerPlugin +from torch import nn + +from lsdb_ml.device import DeviceLike, resolve_device + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class ModelSpec: + """How to build a model on a worker. + + Only the spec travels to the workers: each worker calls ``factory()``, loads + ``weights`` itself, switches to ``eval()`` and moves the model to ``device``. Prefer a + factory plus a weights path for large models, so the weights never pass through the + Dask scheduler. + + Parameters + ---------- + factory : callable + Picklable ``factory() -> nn.Module``. + weights : str or Path, optional + File readable from every worker: a ``state_dict``, or a checkpoint written by + :class:`~lsdb_ml.TrainSession` (its ``"model"`` entry is used). + device : str or torch.device, default "cpu" + Resolved on each worker; ``"auto"`` picks CUDA, then MPS, then CPU. + torch_threads : int, optional + Calls ``torch.set_num_threads`` on the worker, to avoid oversubscribing cores with + Dask threads times torch intra-op threads. It affects the whole worker process. + """ + + factory: Callable[[], nn.Module] + weights: str | Path | None = None + device: DeviceLike = "cpu" + torch_threads: int | None = None + + @classmethod + def from_module(cls, module: nn.Module, device: DeviceLike = "cpu") -> ModelSpec: + """Wrap an in-memory module; it is pickled into the plugin once per worker.""" + return cls(factory=_CopyFactory(copy.deepcopy(module).to("cpu")), device=device) + + def token(self) -> str: + """A short content hash, used to name the worker plugin.""" + payload = cloudpickle.dumps((self.factory, str(self.weights), str(self.device), self.torch_threads)) + return hashlib.sha256(payload).hexdigest()[:16] + + def build(self) -> tuple[nn.Module, torch.device]: + """Build the model for inference: weights loaded, ``eval()`` mode, on its device.""" + device = resolve_device(self.device) + if self.torch_threads is not None: + torch.set_num_threads(self.torch_threads) + model = self.factory() + if not isinstance(model, nn.Module): + raise TypeError(f"The model factory returned {type(model).__name__}, not a torch.nn.Module.") + if self.weights is not None: + state = torch.load(self.weights, map_location=device, weights_only=True) + if "model" in state and "lsdb_ml_version" in state: + state = state["model"] + model.load_state_dict(state) + model.eval() + return model.to(device), device + + +class _CopyFactory: + """Picklable factory returning a fresh copy of a module.""" + + def __init__(self, module: nn.Module): + self.module = module + + def __call__(self) -> nn.Module: + return copy.deepcopy(self.module) + + +class ModelPlugin(WorkerPlugin): + """Dask worker plugin that builds a model once when a worker starts (or joins). + + Tasks look the model up with :func:`get_model`. Each worker has its own instance, so + there is no state shared between workers, and tasks on one worker share only the + read-only model. + + Parameters + ---------- + spec : ModelSpec + How to build the model. + name : str + Plugin name; tasks find the model by it. + """ + + def __init__(self, spec: ModelSpec, name: str): + self.spec = spec + self.name = name + self.model = None + self.device = None + + def setup(self, worker=None): + """Build the model on this worker.""" + self.model, self.device = self.spec.build() + logger.info("Loaded model plugin %s on %s (device %s).", self.name, worker, self.device) + + def teardown(self, worker=None): + """Release the model.""" + self.model = None + if self.device is not None and self.device.type == "cuda": + torch.cuda.empty_cache() + + +# Models deployed without a distributed client, for the synchronous and threaded +# schedulers: built once in this process and shared by its threads. +_LOCAL_MODELS: dict[str, tuple[nn.Module, torch.device]] = {} +_LOCAL_LOCK = threading.Lock() + + +def deploy_locally(spec: ModelSpec, name: str) -> None: + """Build the model in this process, if it is not already there.""" + with _LOCAL_LOCK: + if name not in _LOCAL_MODELS: + _LOCAL_MODELS[name] = spec.build() + + +def undeploy_locally(name: str) -> None: + """Drop a model built by :func:`deploy_locally`.""" + with _LOCAL_LOCK: + _LOCAL_MODELS.pop(name, None) + + +def get_model(name: str) -> tuple[nn.Module, torch.device]: + """Return ``(model, device)`` for a deployed model, from inside a task. + + Raises + ------ + RuntimeError + If the model is not deployed where the task runs. + """ + try: + from distributed import get_worker + + worker = get_worker() + except ValueError: + worker = None + if worker is not None: + plugin = worker.plugins.get(name) + if plugin is not None and plugin.model is not None: + return plugin.model, plugin.device + elif name in _LOCAL_MODELS: + return _LOCAL_MODELS[name] + where = f"worker {worker.address}" if worker is not None else "this process" + raise RuntimeError( + f"Model '{name}' is not deployed on {where}. The prediction catalog must be computed " + "with the same Dask client that was active (or passed) when predict() was called, and " + "before the deployment is closed." + ) From f105a63e8d04b6710e207221d05f616b9343760d Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 12:27:27 -0700 Subject: [PATCH 7/8] Add predict() and ModelDeployment over catalog partitions Each partition is processed in order with task-local state only and predictions are attached by position to the same frame, so output is aligned with input even for duplicate _healpix_29 values. Tests cover a filtered and reordered partition layout, duplicate indices, and load-once-per-worker; a reversed-batch mutation fails them. Co-Authored-By: Claude Opus 5.5 --- src/lsdb_ml/__init__.py | 5 + src/lsdb_ml/inference/__init__.py | 7 + src/lsdb_ml/inference/predict.py | 214 ++++++++++++++++++++ tests/lsdb_ml/inference/inference_models.py | 70 +++++++ tests/lsdb_ml/inference/test_local.py | 26 +++ tests/lsdb_ml/inference/test_predict.py | 178 ++++++++++++++++ 6 files changed, 500 insertions(+) create mode 100644 src/lsdb_ml/inference/__init__.py create mode 100644 src/lsdb_ml/inference/predict.py create mode 100644 tests/lsdb_ml/inference/inference_models.py create mode 100644 tests/lsdb_ml/inference/test_local.py create mode 100644 tests/lsdb_ml/inference/test_predict.py diff --git a/src/lsdb_ml/__init__.py b/src/lsdb_ml/__init__.py index d514376..abb97bd 100644 --- a/src/lsdb_ml/__init__.py +++ b/src/lsdb_ml/__init__.py @@ -2,6 +2,7 @@ from .data import CatalogIterableDataset from .device import move_batch, resolve_device from .example_module import greetings, meaning +from .inference import ModelDeployment, ModelSpec, OutputSpec, predict from .spec import ColumnSpec, Field, Flat, Nested, ResolvedSpec, SpecError, TensorInfo from .training import SupervisedStep, TrainSession, TrainStep @@ -10,6 +11,10 @@ "SupervisedStep", "TrainSession", "TrainStep", + "ModelDeployment", + "ModelSpec", + "OutputSpec", + "predict", "ColumnSpec", "Field", "Flat", diff --git a/src/lsdb_ml/inference/__init__.py b/src/lsdb_ml/inference/__init__.py new file mode 100644 index 0000000..9b696a4 --- /dev/null +++ b/src/lsdb_ml/inference/__init__.py @@ -0,0 +1,7 @@ +"""Running models over LSDB catalogs on Dask workers.""" + +from .outputs import OutputSpec +from .plugin import ModelPlugin, ModelSpec +from .predict import ModelDeployment, predict + +__all__ = ["ModelDeployment", "ModelPlugin", "ModelSpec", "OutputSpec", "predict"] diff --git a/src/lsdb_ml/inference/predict.py b/src/lsdb_ml/inference/predict.py new file mode 100644 index 0000000..5becd9c --- /dev/null +++ b/src/lsdb_ml/inference/predict.py @@ -0,0 +1,214 @@ +"""Running a model over every row of an LSDB catalog.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping, Sequence + +import dask +import numpy as np +import pandas as pd +import torch +from torch import nn + +from lsdb_ml.data.dataset import iter_slices +from lsdb_ml.device import move_batch +from lsdb_ml.inference.outputs import OutputSpec +from lsdb_ml.inference.plugin import ModelPlugin, ModelSpec, deploy_locally, get_model, undeploy_locally +from lsdb_ml.inputs import InputsLike, model_args +from lsdb_ml.spec import ColumnSpec, ResolvedSpec + +logger = logging.getLogger(__name__) + +# (client id, plugin name) pairs registered from this process, so redeploying the same +# model to the same client does not rebuild it on every worker. +_REGISTERED: set[tuple[str, str]] = set() + + +def _active_client(): + try: + from distributed import get_client + + return get_client() + except ValueError: + return None + + +class ModelDeployment: + """A model installed on every Dask worker, ready to predict on catalogs. + + With a ``distributed`` client, a :class:`~lsdb_ml.inference.plugin.ModelPlugin` builds + the model once on each worker, including workers that join later. Without one (the + synchronous or threaded scheduler), the model is built once in this process. + + Use it as a context manager to remove the model from the workers afterwards:: + + with ModelDeployment(ModelSpec(make_model, weights="model.pt")) as deployed: + result = deployed.predict(catalog, spec, inputs="x", outputs={"z": "float32"}).compute() + + Parameters + ---------- + model : ModelSpec or nn.Module + The model to deploy. A module is wrapped with :meth:`ModelSpec.from_module`. + client : distributed.Client, optional + Client to deploy to. Defaults to the active client, if any. + """ + + def __init__(self, model: ModelSpec | nn.Module, client=None): + self.spec = ModelSpec.from_module(model) if isinstance(model, nn.Module) else model + if not isinstance(self.spec, ModelSpec): + raise TypeError(f"Expected a ModelSpec or torch.nn.Module, got {type(model).__name__}.") + self.name = f"lsdb-ml-model-{self.spec.token()}" + self.client = client if client is not None else _active_client() + self._closed = False + if self.client is None: + scheduler = dask.config.get("scheduler", None) + if scheduler in ("processes", "multiprocessing"): + raise RuntimeError( + "Inference without a distributed client needs the synchronous or threaded " + "scheduler; create a distributed Client to use multiple processes." + ) + deploy_locally(self.spec, self.name) + elif (self.client.id, self.name) not in _REGISTERED: + self.client.register_plugin(ModelPlugin(self.spec, self.name), name=self.name) + _REGISTERED.add((self.client.id, self.name)) + + def predict( + self, + catalog, + spec: ColumnSpec | ResolvedSpec, + outputs: OutputSpec | Mapping[str, object] | Sequence[str], + *, + inputs: InputsLike | None = None, + batch_size: int = 4096, + append: bool = True, + ): + """Return a lazy catalog with the model's predictions for every row. + + See :func:`predict` for the parameters. + """ + if self._closed: + raise RuntimeError("This deployment is closed.") + resolved = spec if isinstance(spec, ResolvedSpec) else spec.resolve(catalog) + if isinstance(batch_size, bool) or int(batch_size) < 1: + raise ValueError(f"batch_size must be a positive integer, got {batch_size!r}.") + outputs = outputs if isinstance(outputs, OutputSpec) else OutputSpec(outputs) + if inputs is None: + if len(resolved.fields) != 1: + raise ValueError( + f"The spec has several fields {list(resolved.fields)}; pass inputs= to say which " + "batch keys the model takes." + ) + inputs = next(iter(resolved.fields)) + if append: + clashes = sorted(set(outputs.columns) & set(catalog.columns)) + if clashes: + raise ValueError(f"Output columns {clashes} already exist in the catalog; rename them.") + + meta = outputs.meta(catalog.meta, append) + return catalog.map_partitions( + _predict_partition, + model_name=self.name, + resolved=resolved, + inputs=inputs, + outputs=outputs, + batch_size=int(batch_size), + append=append, + meta=meta, + ) + + def close(self) -> None: + """Remove the model from the workers (or this process). Safe to call twice.""" + if self._closed: + return + self._closed = True + if self.client is None: + undeploy_locally(self.name) + else: + self.client.unregister_worker_plugin(self.name) + _REGISTERED.discard((self.client.id, self.name)) + + def __enter__(self) -> ModelDeployment: + return self + + def __exit__(self, *exc_info) -> None: + self.close() + + +def predict( + catalog, + model: ModelSpec | nn.Module, + spec: ColumnSpec | ResolvedSpec, + outputs: OutputSpec | Mapping[str, object] | Sequence[str], + *, + inputs: InputsLike | None = None, + batch_size: int = 4096, + append: bool = True, + client=None, +): + """Run a model over every row of a catalog, returning a lazy catalog of predictions. + + The model is installed once per Dask worker (see :class:`ModelDeployment`) and stays + installed so the returned catalog can be computed later; use + :class:`ModelDeployment` directly to remove it when done. + + Each partition is processed in order, in batches of up to ``batch_size`` rows, under + ``torch.inference_mode()``. Predictions are attached to the rows they were computed + from by position within the partition, never by joining on the index, so the output + is aligned with the input even when ``_healpix_29`` values repeat. + + Parameters + ---------- + catalog : lsdb.Catalog + The input catalog. + model : ModelSpec or nn.Module + The model, or how to build it on each worker. + spec : ColumnSpec or ResolvedSpec + Which columns become which tensors. Validated against ``catalog.meta`` now, + before any data is read. + outputs : OutputSpec, mapping or sequence + Output column names and dtypes, e.g. ``{"z_phot": "float32"}``. + inputs : str, sequence of str or callable, optional + Batch key(s) passed positionally to the model, or ``callable(batch)``. Defaults to + the spec's only field. + batch_size : int, default 4096 + Maximum rows per forward pass. + append : bool, default True + Keep the input columns. With False, the result holds only the predictions (on the + same index). + client : distributed.Client, optional + Client to deploy to. Defaults to the active client; without one, the model is + built in this process and the catalog must be computed with the synchronous or + threaded scheduler. + + Returns + ------- + lsdb.Catalog + Lazy catalog whose ``meta`` includes the output columns. + """ + deployment = ModelDeployment(model, client=client) + return deployment.predict(catalog, spec, outputs, inputs=inputs, batch_size=batch_size, append=append) + + +def _predict_partition( + partition: pd.DataFrame, + *, + model_name: str, + resolved: ResolvedSpec, + inputs: InputsLike, + outputs: OutputSpec, + batch_size: int, + append: bool, +) -> pd.DataFrame: + """Predict on one partition. All state is local to this call except the read-only model.""" + if len(partition) == 0: + empty = {name: np.empty(0) for name in outputs.columns} + return outputs.build_frame(partition, empty, append) + model, device = get_model(model_name) + pieces: list[dict[str, np.ndarray]] = [] + with torch.inference_mode(): + for rows in iter_slices(partition, batch_size): + batch = move_batch(resolved.collate(rows), device) + pieces.append(outputs.to_columns(model(*model_args(batch, inputs)), len(rows))) + columns = {name: np.concatenate([piece[name] for piece in pieces]) for name in outputs.columns} + return outputs.build_frame(partition, columns, append) diff --git a/tests/lsdb_ml/inference/inference_models.py b/tests/lsdb_ml/inference/inference_models.py new file mode 100644 index 0000000..3364717 --- /dev/null +++ b/tests/lsdb_ml/inference/inference_models.py @@ -0,0 +1,70 @@ +"""Tiny models and checks shared by the inference tests.""" + +import numpy as np +import torch +from torch import nn + +from lsdb_ml import ColumnSpec, Flat + +# Counts factory calls; with processes=False every worker shares this module. +BUILDS = {"count": 0} + + +class AffineMag(nn.Module): + """pred = 2 * mag + 1, exactly, so every row can be checked against its own input.""" + + def __init__(self): + super().__init__() + self.linear = nn.Linear(1, 1, dtype=torch.float64) + with torch.no_grad(): + self.linear.weight.fill_(2.0) + self.linear.bias.fill_(1.0) + + def forward(self, mag): + """Map (B,) magnitudes to (B,) predictions.""" + return self.linear(mag.unsqueeze(-1)).squeeze(-1) + + +def counting_factory(): + """Build AffineMag and count how often a model is built.""" + BUILDS["count"] += 1 + return AffineMag() + + +class MaskedMean(nn.Module): + """Mean of the real (unpadded) light-curve values; 0 for empty light curves.""" + + def forward(self, flux, mask): + """Average padded values using the mask.""" + return (flux * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1) + + +class Stacked(nn.Module): + """Two outputs per row, as a (B, 2) tensor.""" + + def forward(self, mag): + """Return [mag, -mag].""" + return torch.stack([mag, -mag], dim=1) + + +class DropsRow(nn.Module): + """A broken model that returns one prediction too few.""" + + def forward(self, mag): + """Drop the last row.""" + return mag[:-1] + + +MAG_SPEC = ColumnSpec({"mag": Flat("mag", dtype=torch.float64)}) + + +def check_alignment(result, source_frame): + """Every prediction matches its own row's input, and rows match the source exactly once.""" + np.testing.assert_allclose( + result["pred"].to_numpy(), 2 * result["mag"].to_numpy(np.float64) + 1, rtol=1e-6 + ) + # Also join on the object id against an eager computation on the raw frame. + expected = 2 * source_frame.set_index("id")["mag"].astype(np.float64) + 1 + assert result["id"].is_unique + joined = result.set_index("id")["pred"] + np.testing.assert_allclose(joined.to_numpy(), expected.loc[joined.index].to_numpy(), rtol=1e-6) diff --git a/tests/lsdb_ml/inference/test_local.py b/tests/lsdb_ml/inference/test_local.py new file mode 100644 index 0000000..a32ffeb --- /dev/null +++ b/tests/lsdb_ml/inference/test_local.py @@ -0,0 +1,26 @@ +"""Inference without a distributed client: run in a module with no client fixture.""" + +import dask +import pytest +from inference_models import BUILDS, MAG_SPEC, AffineMag, check_alignment, counting_factory + +from lsdb_ml import ModelDeployment, ModelSpec +from lsdb_ml.inference.plugin import get_model + + +def test_without_distributed_client(nested_catalog, nested_frame): + """Without a client the model is built once in-process and works with local schedulers.""" + BUILDS["count"] = 0 + with dask.config.set(scheduler="threads"), ModelDeployment(ModelSpec(counting_factory)) as deployed: + assert deployed.client is None + result = deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False) + assert BUILDS["count"] == 1 + check_alignment(result, nested_frame) + with pytest.raises(RuntimeError, match="is not deployed on this process"): + get_model(deployed.name) + + +def test_processes_scheduler_rejected(): + """The multiprocessing scheduler cannot share an in-process model.""" + with dask.config.set(scheduler="processes"), pytest.raises(RuntimeError, match="synchronous or threaded"): + ModelDeployment(AffineMag()) diff --git a/tests/lsdb_ml/inference/test_predict.py b/tests/lsdb_ml/inference/test_predict.py new file mode 100644 index 0000000..1eb1652 --- /dev/null +++ b/tests/lsdb_ml/inference/test_predict.py @@ -0,0 +1,178 @@ +import numpy as np +import pandas as pd +import pytest +import torch +from distributed import Client, LocalCluster +from inference_models import ( + BUILDS, + MAG_SPEC, + AffineMag, + DropsRow, + MaskedMean, + Stacked, + check_alignment, + counting_factory, +) + +from lsdb_ml import ColumnSpec, Flat, ModelDeployment, ModelSpec, Nested, OutputSpec, predict + + +@pytest.fixture(scope="module") +def client(): + """Two in-process workers with two threads each, so tasks run concurrently.""" + with ( + LocalCluster(n_workers=2, threads_per_worker=2, processes=False, dashboard_address=None) as cluster, + Client(cluster) as dask_client, + ): + yield dask_client + + +@pytest.fixture +def scrambled_catalog(nested_catalog): + """A non-trivial layout: a filter that empties some partitions, then a reordered subset.""" + filtered = nested_catalog.query("mag > 21") + n = filtered.npartitions + return filtered.partitions[[n - 1, 3, 0, 7, 5, 2, 9, 11, 1, 4, 6, 8, 10] + list(range(12, n - 1))] + + +def test_index_aligned_after_nontrivial_layout(client, scrambled_catalog, nested_frame): + """Predictions line up with their input rows across empty, reordered partitions.""" + computed_input = scrambled_catalog.compute(progress_bar=False) + assert 0 in [len(p.compute(progress_bar=False)) for p in scrambled_catalog.partitions] + + with ModelDeployment(ModelSpec(AffineMag)) as deployed: + lazy = deployed.predict(scrambled_catalog, MAG_SPEC, {"pred": "float64"}, batch_size=3) + result = lazy.compute(progress_bar=False) + assert len(result) == len(computed_input) + assert result.index.equals(computed_input.index) + check_alignment(result, nested_frame) + + +def test_alignment_with_duplicate_index(client, nested_frame): + """Rows sharing a _healpix_29 value still get their own predictions.""" + import lsdb + + frame = nested_frame.iloc[:60].copy() + frame["ra"] = 10.0 # every object at the same position: one index value repeated + frame["dec"] = 10.0 + catalog = lsdb.from_dataframe(frame, partition_rows=100) + result = predict(catalog, AffineMag(), MAG_SPEC, {"pred": "float64"}, batch_size=7).compute( + progress_bar=False + ) + assert not result.index.is_unique + check_alignment(result, nested_frame) + + +def test_model_loaded_once_per_worker(client, nested_catalog): + """The worker plugin builds the model once per worker, not once per partition.""" + BUILDS["count"] = 0 + spec = ModelSpec(counting_factory) + with ModelDeployment(spec) as deployed: + deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False) + deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False) + # Deploying the same spec again reuses the installed plugin. + predict(nested_catalog, spec, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False) + assert nested_catalog.npartitions > len(client.scheduler_info()["workers"]) + assert BUILDS["count"] == len(client.scheduler_info()["workers"]) == 2 + + +def test_meta_matches_computed(client, nested_catalog): + """The lazy result's meta has the same columns and dtypes as the computed data.""" + for append in (True, False): + lazy = predict(nested_catalog, AffineMag(), MAG_SPEC, {"pred": "float32"}, append=append) + computed = lazy.compute(progress_bar=False) + assert list(lazy.columns) == list(computed.columns) + assert lazy.dtypes.to_dict() == computed.dtypes.to_dict() + assert computed["pred"].dtype == np.float32 + assert list(lazy.columns) == ["pred"] + assert len(computed) == 500 + + +def test_nested_inputs(client, nested_catalog): + """Nested columns are padded and masked, so a masked mean matches a per-row mean.""" + spec = ColumnSpec({"flux": Nested("lc", "flux", dtype=torch.float64)}) + result = predict( + nested_catalog, + MaskedMean(), + spec, + {"mean_flux": "float64"}, + inputs=["flux", "flux_mask"], + batch_size=9, + ).compute(progress_bar=False) + expected = [ + float(np.mean(v["flux"])) if isinstance(v, pd.DataFrame) and len(v) else 0.0 + for v in result["lc"].array + ] + np.testing.assert_allclose(result["mean_flux"].to_numpy(), expected, rtol=1e-6) + + +def test_multi_column_output(client, nested_catalog): + """A (B, k) output fills k columns in order.""" + result = predict(nested_catalog, Stacked(), MAG_SPEC, ["plus", "minus"], append=False).compute( + progress_bar=False + ) + np.testing.assert_allclose(result["plus"], -result["minus"]) + + +def test_wrong_row_count_fails(client, nested_catalog): + """A model that drops rows fails loudly instead of misaligning predictions.""" + lazy = predict(nested_catalog, DropsRow(), MAG_SPEC, {"pred": "float32"}) + with pytest.raises(Exception, match="exactly one prediction per input row"): + lazy.compute(progress_bar=False) + + +def test_validation(client, unloadable_catalog): + """Bad specs, outputs and arguments fail before any data is read.""" + with pytest.raises(ValueError, match="already exist"): + predict(unloadable_catalog, AffineMag(), MAG_SPEC, {"mag": "float32"}) + with pytest.raises(ValueError, match="pass inputs="): + predict(unloadable_catalog, AffineMag(), ColumnSpec({"a": Flat("ra"), "b": Flat("dec")}), ["p"]) + with pytest.raises(ValueError, match="not in the catalog"): + predict(unloadable_catalog, AffineMag(), ColumnSpec({"a": Flat("nope")}), ["p"]) + with pytest.raises(ValueError, match="batch_size"): + predict(unloadable_catalog, AffineMag(), MAG_SPEC, ["p"], batch_size=0) + with pytest.raises(TypeError, match="ModelSpec or torch.nn.Module"): + ModelDeployment("model") + + +def test_closed_deployment(client, nested_catalog): + """After close, the model is gone from the workers and computing fails clearly.""" + deployed = ModelDeployment(ModelSpec(AffineMag, device="cpu", torch_threads=None)) + lazy = deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float32"}) + deployed.close() + deployed.close() + with pytest.raises(Exception, match="is not deployed"): + lazy.compute(progress_bar=False) + with pytest.raises(RuntimeError, match="closed"): + deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float32"}) + + +def test_weights_from_training_checkpoint(client, nested_catalog, tmp_path): + """ModelSpec loads a state dict or a TrainSession checkpoint on each worker.""" + from lsdb_ml.training import save_checkpoint + + trained = AffineMag() + with torch.no_grad(): + trained.linear.bias.fill_(5.0) + save_checkpoint(tmp_path / "ckpt.pt", trained) + torch.save(trained.state_dict(), tmp_path / "state.pt") + for weights in ("ckpt.pt", "state.pt"): + spec = ModelSpec(AffineMag, weights=tmp_path / weights) + result = predict(nested_catalog, spec, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False) + np.testing.assert_allclose(result["pred"], 2 * result["mag"].astype(np.float64) + 5, rtol=1e-6) + + +def test_output_spec_errors(): + """Outputs that do not match the declared columns are rejected.""" + spec = OutputSpec({"a": "float32", "b": "float32"}) + with pytest.raises(ValueError, match="expected \\(3, 2\\)"): + spec.to_columns(torch.zeros(3), 3) + with pytest.raises(ValueError, match="returned keys"): + spec.to_columns({"a": torch.zeros(3)}, 3) + with pytest.raises(ValueError, match="expected \\(3,\\)"): + spec.to_columns({"a": torch.zeros(3), "b": torch.zeros(3, 2)}, 3) + with pytest.raises(TypeError, match="torch.Tensor"): + spec.to_columns([1, 2, 3], 3) + with pytest.raises(ValueError, match="at least one column"): + OutputSpec([]) + assert set(spec.to_columns({"a": torch.zeros(3, 1), "b": torch.zeros(3)}, 3)) == {"a", "b"} From fe19bcf0ad57757fe8afbfebb836d57ed04614d2 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 12:27:27 -0700 Subject: [PATCH 8/8] Add inference example script Co-Authored-By: Claude Opus 5.5 --- docs/examples/predict_on_catalog.py | 80 +++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 docs/examples/predict_on_catalog.py diff --git a/docs/examples/predict_on_catalog.py b/docs/examples/predict_on_catalog.py new file mode 100644 index 0000000..b68b051 --- /dev/null +++ b/docs/examples/predict_on_catalog.py @@ -0,0 +1,80 @@ +"""Run a PyTorch model over every row of an LSDB catalog on Dask workers. + +Run with ``python docs/examples/predict_on_catalog.py``. It uses a synthetic catalog; +for a real one, replace ``make_catalog()`` with ``lsdb.open_catalog(...)``. + +The model is built once per Dask worker (by a worker plugin) and the result is a lazy +catalog with prediction columns, aligned row by row with the input. +""" + +from pathlib import Path + +import lsdb +import nested_pandas as npd +import numpy as np +import pandas as pd +import torch +from distributed import Client, LocalCluster +from torch import nn + +from lsdb_ml import ColumnSpec, Flat, ModelDeployment, ModelSpec, Nested + + +def make_catalog(n_objects=5000, seed=0): + """Objects with a magnitude and a ragged light curve.""" + rng = np.random.default_rng(seed) + lengths = rng.integers(0, 40, n_objects) + objects = pd.DataFrame( + { + "object_id": np.arange(n_objects), + "ra": rng.uniform(0, 360, n_objects), + "dec": rng.uniform(-60, 60, n_objects), + "mag": rng.normal(20, 1, n_objects).astype(np.float32), + } + ) + light_curves = pd.DataFrame( + {"mjd": rng.uniform(0, 100, lengths.sum()), "flux": rng.normal(1, 0.1, lengths.sum())}, + index=np.repeat(np.arange(n_objects), lengths), + ) + return lsdb.from_dataframe(npd.NestedFrame(objects).join_nested(light_curves, "lc"), partition_rows=500) + + +class Scorer(nn.Module): + """Toy model combining a flat column with a masked summary of a nested column.""" + + def __init__(self): + super().__init__() + self.head = nn.Linear(2, 2) + + def forward(self, mag, flux, mask): + """Return two scores per object from the magnitude and the mean light-curve flux.""" + mean_flux = (flux * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1) + return self.head(torch.stack([mag, mean_flux], dim=1)) + + +def main(weights_path): + """Save some weights, deploy the model and predict on the catalog.""" + torch.manual_seed(0) + torch.save(Scorer().state_dict(), weights_path) # stands in for a trained model + + catalog = make_catalog().query("mag < 21") # any lazy lsdb pipeline works here + spec = ColumnSpec({"mag": Flat("mag"), "flux": Nested("lc", "flux", max_length=32, truncate="tail")}) + + # Workers build Scorer() and load the weights themselves; only the spec is shipped. + with ModelDeployment(ModelSpec(Scorer, weights=weights_path, device="cpu")) as deployed: + scored = deployed.predict( + catalog, + spec, + outputs={"score_a": "float32", "score_b": "float32"}, + inputs=["mag", "flux", "flux_mask"], + batch_size=1024, + ) + print(scored.dtypes[["object_id", "mag", "score_a", "score_b"]]) # meta, nothing computed yet + result = scored.compute() + print(result[["object_id", "mag", "score_a", "score_b"]].head()) + print(f"{len(result)} rows scored") + + +if __name__ == "__main__": + with LocalCluster(n_workers=2, threads_per_worker=2, dashboard_address=None) as cluster, Client(cluster): + main(Path("scorer_weights.pt").resolve())