Repository navigation
Conversation
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ColumnSpec maps flat and nested catalog columns to tensors. resolve() validates against catalog meta (nested sub-column dtypes come from NestedDtype.column_dtypes) without computing data, and reports every problem at once. Nested columns are padded to the batch maximum with a mask and lengths, fully vectorized via nest.len() and nest.to_flat(). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Partitions are sharded across DDP ranks and DataLoader workers from a shared per-epoch permutation, so every row is read exactly once per epoch for any worker count and world size. Each shard streams its partitions through lsdb's CatalogStream; batches are cut within a chunk and never span chunks. Supports fixed epochs and open-ended streaming. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Context-manager session modeled on Hyrax's train_stream: iterate (batch, metrics), stop(), close() returns the model. The optimization step is a user callable (SupervisedStep covers the common case). Periodic and final atomic checkpoints, rank-0 only; DDP runs under model.join() so uneven batch counts across ranks do not hang. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ModelSpec ships only a factory and weights path; each worker builds the model once in its WorkerPlugin setup, loads weights (plain state dicts or TrainSession checkpoints), switches to eval and moves to the device resolved on the worker. OutputSpec builds output meta and validates every batch's output shape against it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Each partition is processed in order with task-local state only and predictions are attached by position to the same frame, so output is aligned with input even for duplicate _healpix_29 values. Tests cover a filtered and reordered partition layout, duplicate indices, and load-once-per-worker; a reversed-batch mutation fails them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## dev #8 +/- ##
===========================================
- Coverage 100.00% 96.06% -3.94%
===========================================
Files 2 18 +16
Lines 7 864 +857
===========================================
+ Hits 7 830 +823
- Misses 0 34 +34 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Change Description
Inference path from the initial design (§7). Stacked on the training PR; merge that first.
Solution Description
predict(catalog, model, spec, outputs, ...)returns a lazylsdb.Catalogwith prediction columns and correctmeta.ModelDeployment/ModelPlugin: a Dask WorkerPlugin builds the model once per worker from aModelSpec(factory + weights path, or an in-memory module). It loads plain state dicts orTrainSessioncheckpoints, switches toeval(), and moves to the device resolved on the worker. Without a distributed client, the model is built once in-process for the synchronous/threaded schedulers._healpix_29, and row counts are checked per batch.docs/examples/predict_on_catalog.py.Tests: alignment after a filter that empties partitions plus reordering, duplicate
_healpix_29values, one model build per worker across many partitions, meta dtypes matching computed dtypes, nested inputs, and clear failures for wrong row counts and closed deployments.🤖 Generated with Claude Code