diff --git a/docs/design/mvp-design.md b/docs/design/mvp-design.md new file mode 100644 index 0000000..4280582 --- /dev/null +++ b/docs/design/mvp-design.md @@ -0,0 +1,167 @@ +# lsdb-ml MVP: catalog inference + +Status: implemented on this branch. This document covers only what the MVP ships. The fuller +design, which also covers training, is on the `design/initial-design` branch; the parts of it +that this MVP leaves out are listed in §7. + +## 1. Scope + +This MVP runs a PyTorch model over every row of an LSDB catalog on Dask workers. It returns +a lazy catalog of predictions that has correct `meta` and is aligned row by row with the +input. Two pieces make that up: + +- **Inputs by column name:** the model's inputs are plain catalog column names. Whether + each is flat or nested, and its dtype, are read from catalog `meta`, and everything is + checked before any data is read. +- **Inference:** a Dask `WorkerPlugin` installs the model once per worker, and + `map_partitions` runs it partition by partition. + +Training, the Hyrax backend, and running without a `distributed` client are out of scope. +`hyrax` remains only an optional extra. + +Data preparation (filtering, derived columns, unit conversion, encoding strings) is left to +ordinary LSDB operations before calling `predict`. lsdb-ml has no column-spec layer of its +own; one can be added later if a clear need appears. + +## 2. API + +```python +import torch +from distributed import Client +from lsdb_ml import ModelDeployment, ModelSpec + +client = Client(...) # a dask.distributed client is required + +class MyModel(torch.nn.Module): + def forward(self, batch): # batch: dict[str, Tensor] keyed by column name + mag = batch["mag"] # (B,) + flux, mask = batch["lc.flux"], batch["lc._mask"] # (B, L), (B, L) + ... + +with ModelDeployment(ModelSpec(MyModel, weights="model.pt")) as deployed: + scored = deployed.predict( + catalog, + inputs=["mag", "lc.flux"], # flat column, nested sub-column (LSDB dotted name) + outputs={"score": "float32"}, + max_length=32, truncate="tail", + batch_size=4096, + ) + result = scored.compute() +``` + +- **One-call form.** `predict(catalog, model, inputs, outputs, ...)` does the same thing. + It leaves the model installed, because the returned catalog is lazy. + `ModelDeployment.close()` removes the model from the workers. +- **Models with positional arguments.** These use an adapter, + `forward=lambda m, b: m(b["mag"], b["lc.flux"], b["lc._mask"])`. + +## 3. Inputs and batches + +- **Names.** A flat column is named as itself (`"mag"`). A nested sub-column uses LSDB's + dotted name (`"lc.flux"`), the same form `catalog["lc.flux"]` accepts. Batch keys are + those same names. +- **Validation before compute.** `plan_inputs(catalog, inputs)` reads only `catalog.meta`. + Meta is an empty `NestedFrame` with full dtypes, and `NestedDtype.column_dtypes` gives the + sub-column types. That replaces the hand-built column preview in hyrax#1062. It checks: + - each name exists; + - a whole nested column isn't named where a sub-column is needed; + - dtypes are numeric or bool. + + Every problem is reported in one `InputError`. `predict` then dry-runs the conversion on + `catalog.meta` itself, which yields zero-row tensors and still reads no data. +- **Dtypes are inferred.** Float columns become `float_dtype` (float32 by default), integers + int64, and booleans bool. +- **Nested representation: zero padding with a shared mask.** Each nested sub-column + becomes a `(B, L)` tensor, padded to the longest row in the batch. Sub-columns of one + nested column share offsets, so they share one `"._mask"` (bool, `(B, L)`) and + one `"._lengths"` (int64, `(B,)`). This works directly with attention masks, + masked pooling and `pack_padded_sequence`. `torch.nested` is not mature enough to be the + default. +- **Sub-columns stay separate tensors.** They aren't stacked into `(B, L, F)`, to keep + things simple. Models can stack them if they want. +- **Conversion is vectorized.** It is built from `series.nest.len()` (null rows count as 0) + and `series.nest.to_flat()`, with no per-row Python. Missing values raise; drop or fill + them in the catalog first. +- **Batch-level padding is always correct, but has two costs.** + - *Width follows the longest row.* One very long light curve makes its whole batch that + wide, so memory and compute scale with `B × L_max` per batch. `max_length` with + `truncate` caps this. Both take one value for every nested input, or a mapping per nested + column for multi-modal catalogs, e.g. `max_length={"lc_g": 200, "spectrum": 4000}` and + `truncate={"lc_g": "tail", "spectrum": "head"}`. Keys are nested column names, not + dotted sub-column names, because a column's sub-columns share one mask. + - *Width varies from batch to batch.* Models must accept a variable `L`. That's natural + for sequence models and pooling, but a model that needs a fixed `L` (e.g. an MLP over a + flattened sequence) isn't supported yet. A fixed-width padding option (`pad_to`, or a + catalog-level `pad_nested`) is the natural follow-up if that's needed. + + Because padding is masked, predictions don't depend on which other rows share a batch, + as long as the model uses the mask. + +## 4. Inference lifecycle + +- **Model spec.** `ModelSpec(factory, weights=None, device="cpu", torch_threads=None)`. Only + the factory and the weights path go to the scheduler. Each worker builds the model, loads + the `state_dict`, calls `eval()` and moves it to the device. The device is resolved *on the + worker*: `"auto"` picks CUDA, then MPS, then CPU. A bare `nn.Module` is also accepted; it is + pickled into the plugin once per worker. +- **Plugin.** `ModelPlugin(WorkerPlugin)` is named `lsdb-ml-model-`. Concurrent + deployments therefore don't collide, and redeploying the same spec to the same client + doesn't rebuild the model. Workers that join later get the plugin automatically. +- **Tasks.** Each task looks the model up with `get_worker().plugins[name]`. If the model is + missing, the task raises a clear error, e.g. when the catalog is computed on another client + or after `close()`. +- **Concurrency.** Tasks on one worker share only the read-only model, under + `torch.inference_mode()`. GPU workers should use `threads_per_worker=1`. + +## 5. Alignment (the hyrax#1062 "shuffled results" issue) + +The prototype kept one mutable dataset per worker, overwrote it in each task, and wrote +predictions back by position into a buffer. Concurrent tasks on a worker could race on that +shared state. Here, each task: + +1. holds only task-local state; +2. walks its partition in order, in `iloc` slices of `batch_size` rows; +3. checks every batch's output for exactly one row per input row; +4. assigns the concatenated predictions **by position** to the same partition frame they came + from. + +Predictions are never joined on `_healpix_29`, which is not guaranteed unique. The output +`meta` comes from `catalog.meta` plus the `OutputSpec` dtypes. Columns are built with +`pd.Series(values, dtype=d).array`, so computed dtypes match the meta exactly. + +## 6. Tests + +All tests run on CPU, take about 16 seconds, and use a synthetic nested catalog and an +in-process 2-worker × 2-thread cluster. + +- **Input validation:** every error case, including a fixture catalog that raises if any + partition is computed, and dtype inference from meta. +- **Conversion:** checked against a per-row reference, covering the shared mask, truncation, + empty and null nested rows, missing values, and row order. +- **Alignment:** a filter that empties partitions, then a reordered partition subset; and a + catalog where every row shares one `_healpix_29` value. Both are checked per row and joined + on object id. Injecting a reversed batch order makes these tests fail. +- **Load once per worker:** the model is built exactly once per worker across many + partitions and repeated predictions. +- **Meta:** dtypes of the lazy result match the computed result. +- **`forward=` adapter:** a model taking positional tensors works through it. +- **Failure modes:** wrong output row count, closed deployment, no client, bad arguments. + +## 7. Deferred (in the full design) + +| Item | Why deferred | +|---|---| +| Training: `CatalogIterableDataset` (partition sharding across DataLoader workers and DDP) and `TrainSession` | Larger surface; depends on upstream `CatalogStream` changes. Planned as two follow-up PRs. | +| A column-spec layer (`ColumnSpec`, transforms, column grouping, missing-value filling) | Data preparation belongs to LSDB operations; add a middle layer only if a clear need appears. | +| Fixed-width padding (`pad_to` or catalog-level `pad_nested`), stacking sub-columns into `(B, L, F)` | Needed only for fixed-`L` models or as an optimization; batch-level padding covers correctness. | +| Dict-valued and list-column model outputs | Not needed for the first use cases. | +| Running without a `distributed` client | A second code path; requiring a `Client` is a reasonable MVP constraint. | +| Hyrax backend | The seams already exist: the model factory in `ModelSpec` and the `forward=` adapter. | +| Automatic output meta from a dummy forward pass | Fragile for models that can't take an empty or synthetic batch. | + +## 8. CI note + +The workflows only trigger on `main`, so PRs into `dev` get no CI. Add `dev` to the +`push`/`pull_request` branch lists in `testing-and-coverage.yml`, `pre-commit-ci.yml` and +`build-documentation.yml`. With torch as a dependency, CI should also install the CPU build +of torch, because Linux defaults to the multi-GB CUDA wheels. diff --git a/docs/examples/predict_on_catalog.py b/docs/examples/predict_on_catalog.py new file mode 100644 index 0000000..b29f952 --- /dev/null +++ b/docs/examples/predict_on_catalog.py @@ -0,0 +1,81 @@ +"""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 ModelDeployment, ModelSpec + + +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, batch): + """Return two scores per object from the magnitude and the mean light-curve flux.""" + # Batches are keyed by column name; nested sub-columns share a "._mask". + mag, flux, mask = batch["mag"], batch["lc.flux"], batch["lc._mask"] + 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 + # Workers build Scorer() and load the weights themselves; only the factory and path are shipped. + with ModelDeployment(ModelSpec(Scorer, weights=weights_path, device="cpu")) as deployed: + scored = deployed.predict( + catalog, + inputs=["mag", "lc.flux"], # flat vs nested is read from the catalog's meta + outputs={"score_a": "float32", "score_b": "float32"}, + max_length=32, # keep at most the last 32 observations per light curve + truncate="tail", + 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/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..2232ac5 100644 --- a/src/lsdb_ml/__init__.py +++ b/src/lsdb_ml/__init__.py @@ -1,4 +1,20 @@ from ._version import __version__ +from .batch import InputError, InputPlan, plan_inputs +from .device import move_batch, resolve_device from .example_module import greetings, meaning +from .inference import ModelDeployment, ModelSpec, OutputSpec, predict -__all__ = ["greetings", "meaning", "__version__"] +__all__ = [ + "predict", + "ModelDeployment", + "ModelSpec", + "OutputSpec", + "InputError", + "InputPlan", + "plan_inputs", + "move_batch", + "resolve_device", + "greetings", + "meaning", + "__version__", +] diff --git a/src/lsdb_ml/batch.py b/src/lsdb_ml/batch.py new file mode 100644 index 0000000..5031b99 --- /dev/null +++ b/src/lsdb_ml/batch.py @@ -0,0 +1,313 @@ +"""Turning catalog columns into batches of tensors. + +Inputs are named with plain column names. Flat columns use their name (``"mag"``), and +nested sub-columns use LSDB's dotted names (``"lc.flux"``). Whether a column is flat or +nested, and its dtype, are read from the catalog's ``meta``, so inputs can be checked +before any data is computed. + +A batch is a ``dict[str, torch.Tensor]`` keyed by those same names: + +- a flat column becomes a ``(B,)`` tensor; +- a nested sub-column becomes a ``(B, L)`` tensor, zero-padded to the longest row in the + batch (after optional truncation); +- each nested column with selected sub-columns also gets ``"._mask"`` (bool, + ``(B, L)``, True where a value is real) and ``"._lengths"`` (int64, ``(B,)``). + Sub-columns of one nested column share offsets, so they share one mask. + +Floating-point columns are cast to ``float_dtype`` (float32 by default). Integer columns +become int64 and boolean columns stay bool. +""" + +from __future__ import annotations + +from collections.abc import 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 TruncatePolicy, frame_to_array, pad_ragged + +_TRUNCATE_POLICIES = ("raise", "head", "tail") + +MASK_SUFFIX = "._mask" +LENGTHS_SUFFIX = "._lengths" + + +class InputError(ValueError): + """Raised when requested input columns do not match a catalog's schema.""" + + +def _kind(dtype) -> str | None: + """Classify a pandas/numpy/pyarrow dtype as "float", "int", "bool", or None (unsupported).""" + if isinstance(dtype, NestedDtype): + return None + if isinstance(dtype, pd.ArrowDtype): + dtype = dtype.pyarrow_dtype + if isinstance(dtype, pa.DataType): + if pa.types.is_floating(dtype): + return "float" + if pa.types.is_integer(dtype): + return "int" + if pa.types.is_boolean(dtype): + return "bool" + return None + if pd.api.types.is_bool_dtype(dtype): + return "bool" + if pd.api.types.is_float_dtype(dtype): + return "float" + if pd.api.types.is_integer_dtype(dtype): + return "int" + return None + + +def _torch_dtype(kind: str, float_dtype: torch.dtype) -> torch.dtype: + return {"float": float_dtype, "int": torch.int64, "bool": torch.bool}[kind] + + +@dataclass(frozen=True) +class _Column: + name: str + kind: str + + +class InputPlan: + """Validated inputs for a catalog: which columns become which tensors. + + Create it with :func:`plan_inputs`. It is picklable, so it can be shipped to Dask + tasks. + """ + + def __init__( + self, + flat: list[_Column], + nested: dict[str, list[_Column]], + float_dtype: torch.dtype, + max_length: dict[str, int | None], + truncate: dict[str, TruncatePolicy], + ): + self.flat = flat + self.nested = nested + self.float_dtype = float_dtype + #: Per nested column: longest sequence to keep (None for no limit), and the policy. + self.max_length = max_length + self.truncate = truncate + + def describe(self) -> dict[str, torch.dtype]: + """Return the dtype of every tensor in a batch, keyed by batch key.""" + info = {c.name: _torch_dtype(c.kind, self.float_dtype) for c in self.flat} + for parent, subs in self.nested.items(): + for sub in subs: + info[f"{parent}.{sub.name}"] = _torch_dtype(sub.kind, self.float_dtype) + info[parent + MASK_SUFFIX] = torch.bool + info[parent + LENGTHS_SUFFIX] = torch.int64 + return info + + def to_tensors(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]: + """Convert the rows of ``frame`` to a batch of tensors. + + Row ``i`` of every tensor comes from row ``i`` of ``frame``. + + Parameters + ---------- + frame : pd.DataFrame + Rows with the planned columns, e.g. a slice of a catalog partition. An empty + frame, such as ``catalog.meta``, gives zero-row tensors. + + Returns + ------- + dict[str, torch.Tensor] + One entry per key in :meth:`describe`. + + Raises + ------ + ValueError + If a value is missing, or a nested row is longer than its column's + ``max_length`` with ``truncate="raise"``. + """ + batch: dict[str, torch.Tensor] = {} + for column in self.flat: + dtype = _numpy_dtype(column.kind, self.float_dtype) + batch[column.name] = torch.from_numpy( + frame_to_array(frame[[column.name]], dtype, repr(column.name))[:, 0].copy() + ) + for parent, subs in self.nested.items(): + series = frame[parent] + # nest.len() counts null rows as 0; to_flat() returns every row's values + # concatenated in row order. Both stay consistent across iloc slices. + lengths = np.asarray(series.nest.len(), dtype=np.int64) + flat_values = series.nest.to_flat([sub.name for sub in subs]) + for sub in subs: + key = f"{parent}.{sub.name}" + values = frame_to_array( + flat_values[[sub.name]], _numpy_dtype(sub.kind, self.float_dtype), repr(key) + ) + padded, mask, kept = pad_ragged( + values, + lengths, + max_length=self.max_length[parent], + truncate=self.truncate[parent], + what=repr(key), + ) + batch[key] = torch.from_numpy(np.ascontiguousarray(padded[..., 0])) + batch[parent + MASK_SUFFIX] = torch.from_numpy(mask) + batch[parent + LENGTHS_SUFFIX] = torch.from_numpy(kept) + return batch + + +def _numpy_dtype(kind: str, float_dtype: torch.dtype) -> np.dtype: + return torch.empty(0, dtype=_torch_dtype(kind, float_dtype)).numpy().dtype + + +def _valid_max_length(value) -> bool: + return value is None or (isinstance(value, int) and not isinstance(value, bool) and value >= 1) + + +def _per_column(option, value, default, nested_columns, valid, describe, problems) -> dict: + """Expand a single value or a ``{nested column: value}`` mapping to every nested input.""" + if not isinstance(value, Mapping): + if not valid(value): + problems.append(f"{option} must be {describe}, got {value!r}.") + return dict.fromkeys(nested_columns, value) + resolved = dict.fromkeys(nested_columns, default) + for key, item in value.items(): + parent = str(key).partition(".")[0] + if key in nested_columns: + if valid(item): + resolved[key] = item + else: + problems.append(f"{option}['{key}'] must be {describe}, got {item!r}.") + elif parent in nested_columns: + problems.append( + f"{option} key '{key}': use the nested column name '{parent}'; its sub-columns share " + "one mask, so they are truncated together." + ) + else: + problems.append( + f"{option} key '{key}' is not a nested column among the inputs " + f"(nested inputs: {sorted(nested_columns)})." + ) + return resolved + + +def plan_inputs( + schema, + inputs: str | Sequence[str], + *, + float_dtype: torch.dtype = torch.float32, + max_length: int | Mapping[str, int] | None = None, + truncate: TruncatePolicy | Mapping[str, TruncatePolicy] = "raise", +) -> InputPlan: + """Check input column names against a catalog's schema, without computing data. + + Parameters + ---------- + schema : lsdb.Catalog or pd.DataFrame + A catalog (its ``meta`` is used), or a frame with the same columns and dtypes. + inputs : str or sequence of str + Column names: ``"mag"`` for a flat column, ``"lc.flux"`` for a nested sub-column. + float_dtype : torch.dtype, default torch.float32 + Dtype for floating-point columns. + max_length : int or Mapping[str, int], optional + Longest nested sequence to keep; see ``truncate``. An int applies to every nested + input. A mapping sets it per nested column, keyed by the column name (``"lc"``, + not ``"lc.flux"``, since sub-columns share one mask); unlisted columns have no + limit. + truncate : {"raise", "head", "tail"} or Mapping[str, str], default "raise" + For nested rows longer than ``max_length``: raise, keep the first values, or keep + the last values. A mapping sets it per nested column, like ``max_length``; + unlisted columns use "raise". + + Returns + ------- + InputPlan + The validated plan. + + Raises + ------ + InputError + Listing every problem found. + + 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") + >>> batch = plan_inputs(nf, ["a", "lc.f"]).to_tensors(nf) + >>> batch["lc.f"] + tensor([[1., 2.], + [3., 0.]]) + >>> batch["lc._mask"] + tensor([[ True, True], + [ True, False]]) + """ + meta = schema if isinstance(schema, pd.DataFrame) else getattr(schema, "meta", None) + if not isinstance(meta, pd.DataFrame): + raise TypeError(f"Expected an lsdb catalog or a DataFrame, got {type(schema).__name__}.") + names = [inputs] if isinstance(inputs, str) else list(inputs) + + problems: list[str] = [] + if not names: + problems.append("No input columns were given.") + if float_dtype not in (torch.float16, torch.float32, torch.float64): + problems.append(f"float_dtype must be torch.float16, float32 or float64, got {float_dtype}.") + + flat: list[_Column] = [] + nested: dict[str, list[_Column]] = {} + for name in dict.fromkeys(names): # drop duplicates, keep order + if name in meta.columns: + dtype = meta.dtypes[name] + if isinstance(dtype, NestedDtype): + problems.append( + f"'{name}' is a nested column; name its sub-columns instead, e.g. " + f"{[f'{name}.{sub}' for sub in list(dtype.column_dtypes)[:2]]}." + ) + elif _kind(dtype) is None: + problems.append(f"'{name}' has non-numeric dtype {dtype}.") + else: + flat.append(_Column(name, _kind(dtype))) + continue + parent, _, sub = name.partition(".") + if not sub or parent not in meta.columns: + problems.append(f"'{name}' is not a column of the catalog.") + elif not isinstance(meta.dtypes[parent], NestedDtype): + problems.append(f"'{name}': '{parent}' is not a nested column.") + elif sub not in meta.dtypes[parent].column_dtypes: + available = sorted(meta.dtypes[parent].column_dtypes) + problems.append( + f"'{name}': nested column '{parent}' has no sub-column '{sub}' (available: {available})." + ) + elif _kind(meta.dtypes[parent].column_dtypes[sub]) is None: + problems.append(f"'{name}' has non-numeric type {meta.dtypes[parent].column_dtypes[sub]}.") + else: + nested.setdefault(parent, []).append(_Column(sub, _kind(meta.dtypes[parent].column_dtypes[sub]))) + + lengths_limits = _per_column( + "max_length", max_length, None, nested, _valid_max_length, "a positive integer or None", problems + ) + policies = _per_column( + "truncate", + truncate, + "raise", + nested, + lambda v: v in _TRUNCATE_POLICIES, + "'raise', 'head' or 'tail'", + problems, + ) + + for parent, subs in nested.items(): + for suffix in (MASK_SUFFIX, LENGTHS_SUFFIX): + if any(f"{parent}.{sub.name}" == parent + suffix for sub in subs): + problems.append( + f"'{parent + suffix}' is reserved for the generated mask/lengths of '{parent}'." + ) + + if problems: + hint = "" + if any("is not a column" in p for p in problems) and getattr(schema, "all_columns", None) is not None: + hint = f"\nThe catalog was loaded with columns {list(meta.columns)}; others may exist on disk." + raise InputError("Input columns do not match the catalog:\n - " + "\n - ".join(problems) + hint) + return InputPlan(flat, nested, float_dtype, lengths_limits, policies) diff --git a/src/lsdb_ml/collate.py b/src/lsdb_ml/collate.py new file mode 100644 index 0000000..1efa96f --- /dev/null +++ b/src/lsdb_ml/collate.py @@ -0,0 +1,127 @@ +"""Vectorized conversion of (nested) dataframe columns into padded numpy arrays. + +These helpers do the numeric work behind :func:`lsdb_ml.batch.plan_inputs`. Nothing here +loops over rows in Python. +""" + +from __future__ import annotations + +from typing import Literal + +import numpy as np +import pandas as pd + +TruncatePolicy = Literal["raise", "head", "tail"] + + +def frame_to_array(frame: pd.DataFrame, dtype: np.dtype, 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. + 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))``. + + Raises + ------ + ValueError + If any value is missing. + """ + missing = frame.isna().to_numpy() + if missing.any(): + n_missing = int(missing.any(axis=1).sum()) + raise ValueError( + f"{what} contains {n_missing} row(s) with missing values. Drop or fill them before " + "collation, e.g. with catalog.dropna() or catalog.map_partitions()." + ) + # 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", + what: str = "nested column", +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Scatter flat ragged values into a zero-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 ``max_length``: raise, keep the first values, or keep the last values. + what : str, default "nested column" + Description of the data used in error messages. + + Returns + ------- + padded : np.ndarray + Shape ``(B, L, F)``, where ``L`` is the longest kept row. + 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 + kept = lengths + source_starts = starts + if max_length is not None and n_rows and int(lengths.max()) > max_length: + if truncate == "raise": + raise ValueError( + f"{what} has rows of length up to {int(lengths.max())}, over max_length={max_length}. " + "Pass truncate='head' or truncate='tail' to cut them." + ) + kept = np.minimum(lengths, max_length) + 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 = int(kept.max()) if n_rows else 0 + padded = np.zeros((n_rows, width, n_features), 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/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..43c80a0 --- /dev/null +++ b/src/lsdb_ml/inference/outputs.py @@ -0,0 +1,145 @@ +"""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 must be a tensor of shape ``(B,)`` or + ``(B, 1)`` for a single column, or ``(B, k)`` for ``k`` columns, in order. + + 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 + 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 not isinstance(output, torch.Tensor): + raise TypeError(f"Expected the model output to be a torch.Tensor, got {type(output).__name__}.") + array = output.detach().to("cpu").numpy() + 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 diff --git a/src/lsdb_ml/inference/plugin.py b/src/lsdb_ml/inference/plugin.py new file mode 100644 index 0000000..fd0f131 --- /dev/null +++ b/src/lsdb_ml/inference/plugin.py @@ -0,0 +1,141 @@ +"""Installing a model once per Dask worker.""" + +from __future__ import annotations + +import copy +import hashlib +import logging +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 + A saved ``state_dict`` (``torch.save(model.state_dict(), path)``), readable from + every worker. + 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: + model.load_state_dict(torch.load(self.weights, map_location=device, weights_only=True)) + 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() + + +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. + """ + from distributed import get_worker + + try: + worker = get_worker() + except ValueError as err: + raise RuntimeError( + f"Model '{name}' can only be used inside a Dask worker. Compute the prediction " + "catalog with the distributed Client it was deployed to." + ) from err + plugin = worker.plugins.get(name) + if plugin is None or plugin.model is None: + raise RuntimeError( + f"Model '{name}' is not deployed on worker {worker.address}. The prediction catalog " + "must be computed with the Dask client that was active (or passed) when predict() " + "was called, and before the deployment is closed." + ) + return plugin.model, plugin.device diff --git a/src/lsdb_ml/inference/predict.py b/src/lsdb_ml/inference/predict.py new file mode 100644 index 0000000..d0300c9 --- /dev/null +++ b/src/lsdb_ml/inference/predict.py @@ -0,0 +1,241 @@ +"""Running a model over every row of an LSDB catalog.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Iterator, Mapping, Sequence + +import numpy as np +import pandas as pd +import torch +from torch import nn + +from lsdb_ml.batch import InputPlan, plan_inputs +from lsdb_ml.collate import TruncatePolicy +from lsdb_ml.device import move_batch +from lsdb_ml.inference.outputs import OutputSpec +from lsdb_ml.inference.plugin import ModelPlugin, ModelSpec, get_model + +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() + +#: ``forward(model, batch) -> Tensor``: how to call the model on a batch dict. +Forward = Callable[[nn.Module, dict[str, torch.Tensor]], torch.Tensor] + + +def _active_client(): + from distributed import get_client + + try: + return get_client() + except ValueError as err: + raise RuntimeError( + "lsdb-ml inference needs a dask.distributed Client, e.g. Client(LocalCluster(...)) " + "or Client(processes=False); none is active." + ) from err + + +def _iter_row_slices(frame: pd.DataFrame, batch_size: int) -> Iterator[pd.DataFrame]: + """Consecutive slices of at most ``batch_size`` rows, in order, covering ``frame`` once.""" + for start in range(0, len(frame), batch_size): + yield frame.iloc[start : start + batch_size] + + +class ModelDeployment: + """A model installed on every Dask worker, ready to predict on catalogs. + + A :class:`~lsdb_ml.inference.plugin.ModelPlugin` builds the model once on each worker + of a ``dask.distributed`` client, including workers that join later. + + 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, ["mag", "lc.flux"], {"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. + + Raises + ------ + RuntimeError + If no client is given and none is active. + """ + + 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.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, + inputs: str | Sequence[str], + outputs: OutputSpec | Mapping[str, object] | Sequence[str], + *, + forward: Forward | None = None, + batch_size: int = 4096, + append: bool = True, + float_dtype: torch.dtype = torch.float32, + max_length: int | Mapping[str, int] | None = None, + truncate: TruncatePolicy | Mapping[str, TruncatePolicy] = "raise", + ): + """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.") + plan = plan_inputs(catalog, inputs, float_dtype=float_dtype, max_length=max_length, truncate=truncate) + plan.to_tensors(catalog.meta) # dry run on the empty meta frame: no data is read + if isinstance(batch_size, bool) or int(batch_size) < 1: + raise ValueError(f"batch_size must be a positive integer, got {batch_size!r}.") + if forward is not None and not callable(forward): + raise TypeError(f"forward must be callable(model, batch), got {type(forward).__name__}.") + outputs = outputs if isinstance(outputs, OutputSpec) else OutputSpec(outputs) + 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.") + + return catalog.map_partitions( + _predict_partition, + model_name=self.name, + plan=plan, + forward=forward, + outputs=outputs, + batch_size=int(batch_size), + append=append, + meta=outputs.meta(catalog.meta, append), + ) + + def close(self) -> None: + """Remove the model from the workers. Safe to call twice.""" + if self._closed: + return + self._closed = True + 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, + inputs: str | Sequence[str], + outputs: OutputSpec | Mapping[str, object] | Sequence[str], + *, + forward: Forward | None = None, + batch_size: int = 4096, + append: bool = True, + float_dtype: torch.dtype = torch.float32, + max_length: int | Mapping[str, int] | None = None, + truncate: TruncatePolicy | Mapping[str, TruncatePolicy] = "raise", + 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. + inputs : str or sequence of str + Input columns: ``"mag"`` for a flat column, ``"lc.flux"`` for a nested sub-column. + Checked against ``catalog.meta`` now, before any data is read. See + :mod:`lsdb_ml.batch` for the tensors each produces. + outputs : OutputSpec, mapping or sequence + Output column names and dtypes, e.g. ``{"z_phot": "float32"}``. + forward : callable, optional + ``forward(model, batch) -> Tensor``, for models that do not take the batch dict + directly, e.g. ``lambda m, b: m(b["mag"], b["lc.flux"], b["lc._mask"])``. By + default the model is called as ``model(batch)``. + 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). + float_dtype : torch.dtype, default torch.float32 + Dtype for floating-point input columns. + max_length : int or Mapping[str, int], optional + Longest nested sequence to keep; see ``truncate``. Nested inputs are padded to the + longest row in each batch, so one very long row makes the whole batch that wide. + An int applies to every nested input; a mapping sets it per nested column, keyed + by column name, e.g. ``{"lc_g": 200, "spectrum": 4000}`` (unlisted columns have no + limit). + truncate : {"raise", "head", "tail"} or Mapping[str, str], default "raise" + For nested rows longer than ``max_length``: raise, keep the first values, or keep + the last values. A mapping sets it per nested column; unlisted columns use + "raise". + client : distributed.Client, optional + Client to deploy to. Defaults to the active client; one is required. + + Returns + ------- + lsdb.Catalog + Lazy catalog whose ``meta`` includes the output columns. + """ + deployment = ModelDeployment(model, client=client) + return deployment.predict( + catalog, + inputs, + outputs, + forward=forward, + batch_size=batch_size, + append=append, + float_dtype=float_dtype, + max_length=max_length, + truncate=truncate, + ) + + +def _predict_partition( + partition: pd.DataFrame, + *, + model_name: str, + plan: InputPlan, + forward: Forward | None, + 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_row_slices(partition, batch_size): + batch = move_batch(plan.to_tensors(rows), device) + output = model(batch) if forward is None else forward(model, batch) + pieces.append(outputs.to_columns(output, len(rows))) + columns = {name: np.concatenate([piece[name] for piece in pieces]) for name in outputs.columns} + return outputs.build_frame(partition, columns, append) diff --git a/tests/lsdb_ml/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/inference/inference_models.py b/tests/lsdb_ml/inference/inference_models.py new file mode 100644 index 0000000..9d2fb0b --- /dev/null +++ b/tests/lsdb_ml/inference/inference_models.py @@ -0,0 +1,69 @@ +"""Tiny models and checks shared by the inference tests.""" + +import numpy as np +import torch +from torch import nn + +# 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) + with torch.no_grad(): + self.linear.weight.fill_(2.0) + self.linear.bias.fill_(1.0) + + def forward(self, batch): + """Map the (B,) "mag" tensor to (B,) predictions.""" + return self.linear(batch["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, batch): + """Average padded values using the mask.""" + flux, mask = batch["lc.flux"], batch["lc._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, batch): + """Return [mag, -mag].""" + return torch.stack([batch["mag"], -batch["mag"]], dim=1) + + +class DropsRow(nn.Module): + """A broken model that returns one prediction too few.""" + + def forward(self, batch): + """Drop the last row.""" + return batch["mag"][:-1] + + +MAG_INPUTS = "mag" + + +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_no_client.py b/tests/lsdb_ml/inference/test_no_client.py new file mode 100644 index 0000000..bb8960f --- /dev/null +++ b/tests/lsdb_ml/inference/test_no_client.py @@ -0,0 +1,12 @@ +"""Behavior without a distributed client; kept apart from the module-scoped client fixture.""" + +import pytest +from inference_models import AffineMag + +from lsdb_ml import ModelDeployment + + +def test_requires_distributed_client(): + """Without an active client, deploying fails with a clear message.""" + with pytest.raises(RuntimeError, match="needs a dask.distributed Client"): + 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..27a18d0 --- /dev/null +++ b/tests/lsdb_ml/inference/test_predict.py @@ -0,0 +1,217 @@ +import numpy as np +import pandas as pd +import pytest +import torch +from distributed import Client, LocalCluster +from inference_models import ( + BUILDS, + MAG_INPUTS, + AffineMag, + DropsRow, + MaskedMean, + Stacked, + check_alignment, + counting_factory, +) +from torch import nn + +from lsdb_ml import InputError, ModelDeployment, ModelSpec, 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_INPUTS, {"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_INPUTS, {"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_INPUTS, {"pred": "float64"}).compute(progress_bar=False) + deployed.predict(nested_catalog, MAG_INPUTS, {"pred": "float64"}).compute(progress_bar=False) + # Deploying the same spec again reuses the installed plugin. + predict(nested_catalog, spec, MAG_INPUTS, {"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_INPUTS, {"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 sub-columns arrive padded with a shared mask, so a masked mean matches a per-row mean.""" + result = predict( + nested_catalog, + MaskedMean(), + "lc.flux", + {"mean_flux": "float64"}, + float_dtype=torch.float64, + 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_forward_adapter(client, nested_catalog): + """forward= adapts models that take positional tensors instead of the batch dict.""" + result = predict( + nested_catalog, + nn.Identity(), + ["mag", "lc.flux"], + {"pred": "float64"}, + forward=lambda model, batch: 2 * model(batch["mag"]) + 1, + ).compute(progress_bar=False) + np.testing.assert_allclose(result["pred"], 2 * result["mag"].astype(np.float64) + 1, rtol=1e-6) + + +def test_per_column_max_length(client, nested_frame): + """predict passes per-column limits through: each nested input is truncated separately.""" + import lsdb + + rng = np.random.default_rng(2) + lengths = rng.integers(10, 30, len(nested_frame)) + spec = pd.DataFrame( + {"flux": rng.normal(0, 1, lengths.sum())}, index=np.repeat(nested_frame.index, lengths) + ) + catalog = lsdb.from_dataframe(nested_frame.join_nested(spec, "spec"), partition_rows=100) + + def widths(model, batch): + rows = batch["mag"].shape[0] + return torch.stack( + [ + torch.full((rows,), batch["lc.flux"].shape[1]), + torch.full((rows,), batch["spec.flux"].shape[1]), + ], + dim=1, + ) + + result = predict( + catalog, + nn.Identity(), + ["mag", "lc.flux", "spec.flux"], + {"lc_width": "int64", "spec_width": "int64"}, + forward=widths, + max_length={"lc": 2, "spec": 12}, + truncate="tail", + append=False, + ).compute(progress_bar=False) + assert result["lc_width"].max() <= 2 and result["spec_width"].max() == 12 + + +def test_multi_column_output(client, nested_catalog): + """A (B, k) output fills k columns in order.""" + result = predict(nested_catalog, Stacked(), MAG_INPUTS, ["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_INPUTS, {"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_INPUTS, {"mag": "float32"}) + with pytest.raises(InputError, match="'nope' is not a column"): + predict(unloadable_catalog, AffineMag(), ["mag", "nope"], ["p"]) + with pytest.raises(TypeError, match="forward must be callable"): + predict(unloadable_catalog, AffineMag(), MAG_INPUTS, ["p"], forward="mag") + with pytest.raises(ValueError, match="batch_size"): + predict(unloadable_catalog, AffineMag(), MAG_INPUTS, ["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_INPUTS, {"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_INPUTS, {"pred": "float32"}) + + +def test_weights_loaded_on_workers(client, nested_catalog, tmp_path): + """ModelSpec loads a saved state dict on each worker.""" + trained = AffineMag() + with torch.no_grad(): + trained.linear.bias.fill_(5.0) + torch.save(trained.state_dict(), tmp_path / "state.pt") + spec = ModelSpec(AffineMag, weights=tmp_path / "state.pt") + result = predict(nested_catalog, spec, MAG_INPUTS, {"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(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(torch.zeros(3, 2), 3)) == {"a", "b"} + assert OutputSpec(["z"]).to_columns(torch.zeros(3, 1), 3)["z"].shape == (3,) diff --git a/tests/lsdb_ml/test_batch.py b/tests/lsdb_ml/test_batch.py new file mode 100644 index 0000000..0c7c908 --- /dev/null +++ b/tests/lsdb_ml/test_batch.py @@ -0,0 +1,251 @@ +import nested_pandas as npd +import numpy as np +import pandas as pd +import pytest +import torch + +from lsdb_ml import InputError, plan_inputs +from lsdb_ml.collate import pad_ragged + + +def reference_rows(frame, column, sub): + """Per-row python reference: one array of sub-column values per row.""" + rows = [] + for value in frame[column].array: + rows.append(value[sub].to_numpy(np.float64) if isinstance(value, pd.DataFrame) else np.zeros(0)) + return rows + + +def reference_pad(rows, width, keep="head"): + """Pad reference rows to ``width`` the slow way.""" + padded = np.zeros((len(rows), width)) + 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.""" + 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]]) + + +# Validation against meta + + +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_plan_from_catalog_without_compute(unloadable_catalog): + """Flat vs nested and dtypes are inferred from meta; nothing is computed.""" + plan = plan_inputs(unloadable_catalog, ["mag", "id", "flag", "lc.t", "lc.flux"]) + assert plan.describe() == { + "mag": torch.float32, + "id": torch.int64, + "flag": torch.bool, + "lc.t": torch.float32, + "lc.flux": torch.float32, + "lc._mask": torch.bool, + "lc._lengths": torch.int64, + } + batch = plan.to_tensors(unloadable_catalog.meta) # the empty meta frame is a valid dry run + assert {k: v.shape[0] for k, v in batch.items()} == dict.fromkeys(plan.describe(), 0) + + +def test_all_problems_reported_together(unloadable_catalog): + """Every problem is reported in one error.""" + with pytest.raises(InputError) as err: + plan_inputs(unloadable_catalog, ["nope", "lc", "name", "lc.nope", "ra.x", "lc.band"], max_length=0) + message = str(err.value) + for expected in [ + "'nope' is not a column", + "'lc' is a nested column; name its sub-columns", + "'name' has non-numeric dtype", + "has no sub-column 'nope'", + "'ra' is not a nested column", + "'lc.band' has non-numeric type", + "max_length must be a positive integer", + "others may exist on disk", + ]: + assert expected in message + + +def test_other_argument_errors(unloadable_catalog): + """Bad truncate/float_dtype values, empty inputs and bad schemas are rejected.""" + with pytest.raises(InputError, match="truncate must be"): + plan_inputs(unloadable_catalog, "mag", truncate="middle") + with pytest.raises(InputError, match="float_dtype"): + plan_inputs(unloadable_catalog, "mag", float_dtype=torch.int32) + with pytest.raises(InputError, match="No input columns"): + plan_inputs(unloadable_catalog, []) + with pytest.raises(TypeError, match="Expected an lsdb catalog"): + plan_inputs(42, "mag") + + +def test_reserved_mask_name(): + """A real sub-column named like the generated mask is rejected.""" + nf = npd.NestedFrame({"a": [1.0]}, index=[0]).join_nested( + npd.NestedFrame({"_mask": [1.0]}, index=[0]), "lc" + ) + with pytest.raises(InputError, match="reserved"): + plan_inputs(nf, ["lc._mask"]) + + +# Conversion + + +def test_flat_columns(frame): + """Flat columns become (B,) tensors with inferred dtypes, in row order.""" + batch = plan_inputs(frame, ["ra", "mag", "id", "flag"]).to_tensors(frame) + np.testing.assert_allclose(batch["ra"].numpy(), frame["ra"].to_numpy(np.float32)) + assert batch["ra"].dtype == torch.float32 and batch["ra"].shape == (len(frame),) + np.testing.assert_array_equal(batch["id"].numpy(), frame["id"].to_numpy()) + assert batch["flag"].dtype == torch.bool + assert all(t.numpy().flags.writeable for t in batch.values()) + + +def test_float_dtype(frame): + """float_dtype controls floating-point columns only.""" + batch = plan_inputs(frame, ["ra", "id", "lc.t"], float_dtype=torch.float64).to_tensors(frame) + assert batch["ra"].dtype == torch.float64 and batch["lc.t"].dtype == torch.float64 + assert batch["id"].dtype == torch.int64 + + +def test_nested_matches_reference(frame): + """Padding matches a per-row reference, including empty and null nested rows.""" + batch = plan_inputs(frame, ["lc.t", "lc.flux"], float_dtype=torch.float64).to_tensors(frame) + lengths = np.array([len(r) for r in reference_rows(frame, "lc", "t")]) + assert (lengths == 0).any(), "fixture should include empty / null nested rows" + for sub in ("t", "flux"): + expected, expected_mask = reference_pad(reference_rows(frame, "lc", sub), lengths.max()) + np.testing.assert_allclose(batch[f"lc.{sub}"].numpy(), expected) + np.testing.assert_array_equal(batch["lc._mask"].numpy(), expected_mask) + np.testing.assert_array_equal(batch["lc._lengths"].numpy(), lengths) + + +@pytest.mark.parametrize("keep", ["head", "tail"]) +def test_truncation(frame, keep): + """Rows over max_length keep their first or last values.""" + batch = plan_inputs(frame, "lc.t", float_dtype=torch.float64, max_length=2, truncate=keep).to_tensors( + frame + ) + rows = reference_rows(frame, "lc", "t") + expected, expected_mask = reference_pad(rows, 2, keep=keep) + np.testing.assert_allclose(batch["lc.t"].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 max_length=2"): + plan_inputs(frame, "lc.t", max_length=2).to_tensors(frame) + + +def test_all_rows_empty_nested(frame): + """A batch where every nested row is empty pads to width 0.""" + empty_rows = frame[np.asarray(frame["lc"].nest.len()) == 0] + batch = plan_inputs(frame, "lc.flux").to_tensors(empty_rows) + assert batch["lc.flux"].shape == (len(empty_rows), 0) + assert not batch["lc._mask"].any() + + +def test_missing_values_raise(): + """Missing flat values and missing values inside nested lists raise.""" + 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"): + plan_inputs(nf, "a").to_tensors(nf) + with pytest.raises(ValueError, match="missing values"): + plan_inputs(nf, "lc.f").to_tensors(nf) + + +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 = plan_inputs(shuffled, ["id", "lc.t"], float_dtype=torch.float64).to_tensors(shuffled) + np.testing.assert_array_equal(batch["id"].numpy(), shuffled["id"].to_numpy()) + np.testing.assert_array_equal(batch["lc._lengths"].numpy(), np.asarray(shuffled["lc"].nest.len())) + + +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])) + + +# Several nested columns + + +@pytest.fixture(scope="module") +def two_nested(nested_frame): + """The test frame plus a second nested column ("spec") with longer rows (10-29 values).""" + rng = np.random.default_rng(1) + lengths = rng.integers(10, 30, len(nested_frame)) + spec = pd.DataFrame( + {"wave": rng.uniform(4000, 9000, lengths.sum()), "flux": rng.normal(0, 1, lengths.sum())}, + index=np.repeat(nested_frame.index, lengths), + ) + return nested_frame.join_nested(spec, "spec") + + +def test_per_column_max_length_and_truncate(two_nested): + """max_length and truncate can be set per nested column, keyed by column name.""" + plan = plan_inputs( + two_nested, + ["lc.t", "lc.flux", "spec.wave"], + float_dtype=torch.float64, + max_length={"lc": 2, "spec": 12}, + truncate={"lc": "head", "spec": "tail"}, + ) + batch = plan.to_tensors(two_nested) + for column, sub, width, keep in [("lc", "t", 2, "head"), ("spec", "wave", 12, "tail")]: + expected, expected_mask = reference_pad(reference_rows(two_nested, column, sub), width, keep=keep) + np.testing.assert_allclose(batch[f"{column}.{sub}"].numpy(), expected) + np.testing.assert_array_equal(batch[f"{column}._mask"].numpy(), expected_mask) + # Sub-columns of one nested column are truncated together and share the mask. + assert batch["lc.flux"].shape == batch["lc.t"].shape == batch["lc._mask"].shape + + +def test_unlisted_columns_are_unlimited(two_nested): + """Columns missing from the max_length mapping keep every value.""" + batch = plan_inputs( + two_nested, ["lc.t", "spec.wave"], max_length={"spec": 12}, truncate="tail" + ).to_tensors(two_nested) + assert batch["spec.wave"].shape[1] == 12 + assert batch["lc.t"].shape[1] == int(np.asarray(two_nested["lc"].nest.len()).max()) + + +def test_single_value_applies_to_every_nested_column(two_nested): + """An int max_length and a single policy apply to all nested inputs.""" + batch = plan_inputs(two_nested, ["lc.t", "spec.wave"], max_length=3, truncate="head").to_tensors( + two_nested + ) + assert batch["lc.t"].shape[1] == batch["spec.wave"].shape[1] == 3 + + +def test_per_column_key_errors(two_nested): + """Mapping keys must be nested input columns; dotted keys get a hint; values are checked.""" + with pytest.raises(InputError) as err: + plan_inputs( + two_nested, + ["lc.t", "spec.wave"], + max_length={"lc.t": 5, "lcc": 5, "spec": 0}, + truncate={"lc": "middle", "mag": "head"}, + ) + message = str(err.value) + for expected in [ + "max_length key 'lc.t': use the nested column name 'lc'", + "max_length key 'lcc' is not a nested column among the inputs", + "max_length['spec'] must be a positive integer", + "truncate['lc'] must be 'raise', 'head' or 'tail'", + "truncate key 'mag' is not a nested column among the inputs", + ]: + assert expected in message 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"