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()) 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/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/__init__.py b/src/lsdb_ml/__init__.py index e42fde6..abb97bd 100644 --- a/src/lsdb_ml/__init__.py +++ b/src/lsdb_ml/__init__.py @@ -1,4 +1,30 @@ from ._version import __version__ +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 -__all__ = ["greetings", "meaning", "__version__"] +__all__ = [ + "CatalogIterableDataset", + "SupervisedStep", + "TrainSession", + "TrainStep", + "ModelDeployment", + "ModelSpec", + "OutputSpec", + "predict", + "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/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/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() + } 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/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." + ) 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/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/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/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/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/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 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"} 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) 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)