Add train_stream verb for streaming training on live data sources by drewoldag · Pull Request #987 · lincc-frameworks/hyrax

drewoldag · 2026-07-16T00:56:03Z

Change Description

Adds a new train_stream verb that enables training models on open-ended streaming data sources (e.g., Kafka topics) using a context-manager session pattern. This mirrors the existing infer_stream verb but for training workflows.

Solution Description

New Components

TrainStream verb (src/hyrax/verbs/train_stream.py):

Loads a model once and returns a TrainStreamSession for on-demand batch training
Supports two modes:
1. Data-source driven: Configure [data_request.train_stream] with a streaming dataset; the session is iterable and yields (batch, metrics) pairs
2. Manual: Pass a representative sample_batch and feed batches via session.train_batch(batch)
Reuses existing infrastructure: create_process_func("train_batch", ...) for the per-batch training step, setup_model / setup_model_from_sample for model initialization, and MLflow + TensorBoard logging

TrainStreamSession context manager:

Iterates a streaming data loader or accepts manual batches
Skips empty batches and optionally enforces min_batch_size (guards against size-1 batches breaking BatchNorm in user models)
Logs per-batch metrics to TensorBoard and MLflow
Periodically checkpoints weights (every N batches if save_weights_every is set)
Persists final weights and ends MLflow run on close()
Returns the trained model

Configuration

Added [train_stream] section to hyrax_default_config.toml:

model_weights_file: warm-start weights path (default: false = train from scratch)
weights_filename: checkpoint filename (default: "example_model.pth")
save_weights_every: checkpoint interval in batches (default: false = only on close)
experiment_name: MLflow experiment (default: "notebook")
run_name: MLflow run name (default: false = use results-dir name)
min_batch_size: skip batches smaller than this (default: false = train all)

Design Rationale

No epoch loop: Streams are open-ended, so we drive a per-batch loop directly instead of using Ignite's create_trainer / trainer.run(..., max_epochs=...) path
Partial batches: Streaming sources yield ragged batches (flushed on timeout); we train every non-empty batch as-is, with optional min_batch_size guard
Gradient weighting: Each batch steps the optimizer independently (no accumulation); accepted for streaming workflows
Code sharing: Mirrors infer_stream pattern; no shared-base refactor to avoid over-engineering

Testing

Added comprehensive test suite (tests/hyrax/test_train_stream.py):

Data-source-driven iteration with KafkaStreamDataset (using FakeConsumer stand-in)
Manual session with spy process_func to verify empty/small-batch skipping
Idempotent close() and error handling
CLI raises NotImplementedError (programmatic API only)

Code Quality

I have read the Contribution Guide and agree to the Code of Conduct
My code follows the code style of this project
My code builds cleanly without errors or warnings
My code contains relevant comments and necessary documentation

https://claude.ai/code/session_01Vhimu4YPAyMBPj3dPzEmp4

Introduce a `train_stream` verb, the training counterpart to the recently merged `infer_stream`. It loads a model once and trains on batches from an open-ended IterableDataset (e.g. KafkaStreamDataset) through a context-manager session object, mirroring the InferStream pattern. Because a live stream has no epochs, length, or train/validate/test splits, the verb bypasses the Ignite create_trainer / max_epochs path and drives a per-batch loop directly, reusing create_process_func("train_batch", ...) — the model's self-contained train_batch already runs the full optimizer step. Highlights: - TrainStream verb + TrainStreamSession (iterable and manual train_batch/process modes); run_cli raises NotImplementedError (programmatic API only). - Weights persisted periodically (save_weights_every) and on close; close() returns the trained model. - Full MLflow + TensorBoard logging: a session-spanning MLflow run started in run() and ended in close(), per-batch loss logged to both back-ends. - Ragged/partial batches (streaming timeout flushes) are trained as-is, with a defensive empty-batch skip and an optional min_batch_size guard. - New [train_stream] config section; verb auto-exposed via @hyrax_verb. Adds tests/hyrax/test_train_stream.py and specs/train_stream.spec. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Vhimu4YPAyMBPj3dPzEmp4

codecov · 2026-07-17T21:35:06Z

Codecov Report

❌ Patch coverage is 87.10526% with 49 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.47%. Comparing base (cdcb697) to head (1c15930).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/hyrax/datasets/lsdb_stream_dataset.py	86.19%	29 Missing ⚠️
src/hyrax/verbs/train_stream.py	88.09%	20 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #987      +/-   ##
==========================================
+ Coverage   67.52%   69.47%   +1.94%     
==========================================
  Files          87       89       +2     
  Lines        8309     8689     +380     
==========================================
+ Hits         5611     6037     +426     
+ Misses       2698     2652      -46

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

…reaming data.

drewoldag · 2026-07-21T23:12:33Z

Note for Drew for later to run the hyrax dataset producer to emit HSC images:
python hyrax_dataset_producer.py --dataset hsc --data-dir <base_path>/hsc_8asec_1000 --burst-min=75 --burst-max=100 --num-bursts=100 --delay-min=1 --delay-max=5

review-notebook-app · 2026-08-04T23:03:22Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

…ple notebook.

…basic lsdb dataset to dynamically generate getters for simple and nested columns.

…ing TESS data, dim. red. and scatter plotted 33k points.

…o claude/train-stream-verb-planning-8qjyzm

…thub.com/lincc-frameworks/hyrax into claude/train-stream-verb-planning-8qjyzm

drewoldag · 2026-09-08T22:15:18Z

+            if num_samples == 0:
+                logger.debug("Skipping empty batch.")
+                return None
+            min_batch_size = self._config["train_stream"]["min_batch_size"]


Make sure this is build up the batch and not drop it on the floor.

drewoldag · 2026-09-08T22:16:31Z

+        return self._model
+
+    def checkpoint(self, model_metrics=None) -> None:
+        """Save the current model weights if loss is lower than the previous best.


Double check on this to see if it's still used.

drewoldag · 2026-09-08T22:45:08Z

+
+        return lsdb.open_catalog(self.data_location, **open_catalog_kwargs)
+
+    def _requested_columns_from_config(self, config: dict):


Can we just get rid of this entirely? We should just load all the columns.

This functionality is in HATSDataset, so we could just call it directly.

drewoldag · 2026-09-08T22:49:45Z

+        open_catalog_kwargs = dict(ds_config["open_catalog_kwargs"] or {})
+        if requested_columns and "columns" not in open_catalog_kwargs:
+            open_catalog_kwargs["columns"] = requested_columns
+
+        return lsdb.open_catalog(self.data_location, **open_catalog_kwargs)


I would argue for removing this functionality. We should just expect users to hand Hyrax the pandas future. We should then be able to remove _requested_columns_from_config

drewoldag · 2026-09-08T22:51:48Z

+        self._warn_on_small_chunks(client)
+        return client
+
+    def _warn_on_small_chunks(self, client) -> None:


Maybe don't do this warning thing. This is copied from the LSDB documentation. We don't need to warn users about it. This is not our fight, this is LSDB and Dask.

drewoldag · 2026-09-08T22:55:23Z

+        rows: list[dict] = []
+        taken = 0
+
+        try:


Make sure that left over rows are held until the next chunk, and then combined with the first batch.

drewoldag · 2026-09-08T22:56:17Z

+                self._iterator = None
+            self._timings.log()
+
+    def collate_lightcurve(self, samples):


This should be tossed out. It was for experimentation.

drewoldag · 2026-09-08T23:12:55Z

@maxwest-uw - @mtauraso and I came to the conclusion that the best path forward for this PR is likely to leave it as a draft, create a fresh branch and pull over only the important pieces. There has been a lot of experimentation along the way and as a result also a lot of superfluous code, notebooks, etc.

Feel free to take a look over it and comment if you have time though :)

mtauraso assigned drewoldag Jul 17, 2026

Adding stream producer for cifar and hsc datasets.

ac4ab58

Investigating using the pytorch ignite Engine to run training on st…

77a5bc1

…reaming data.

LSDB CatalogStream example notebook.

197c2f7

drewoldag and others added 10 commits August 5, 2026 15:37

Worked with Claude to produce initial lsdb_streaming_dataset and exma…

2aab34e

…ple notebook.

Merge branch 'main' into claude/train-stream-verb-planning-8qjyzm

6c576e1

point in time commit - may not be worth keeping.

1c15930

Merge branch 'main' into claude/train-stream-verb-planning-8qjyzm

4426e58

Added basic tess streaming dataset with collation functions. Updated …

9ea5699

…basic lsdb dataset to dynamically generate getters for simple and nested columns.

Initial implementation of ts2vec and demo notebook training and infer…

a9dadd2

…ing TESS data, dim. red. and scatter plotted 33k points.

:Merge branch 'main' of https://github.com/lincc-frameworks/hyrax int…

14f8ccb

…o claude/train-stream-verb-planning-8qjyzm

Small change to train_stream, updates to the notebook.

fd69468

Adding buffer for LSDB CatalogStream.

7a95d7f

Merge branch 'claude/train-stream-verb-planning-8qjyzm' of https://gi…

6fa1546

…thub.com/lincc-frameworks/hyrax into claude/train-stream-verb-planning-8qjyzm

drewoldag commented Sep 8, 2026

View reviewed changes

drewoldag requested review from maxwest-uw and mtauraso September 8, 2026 23:07

drewoldag mentioned this pull request Sep 9, 2026

Clean implementation of train_stream #1053

Merged

Navigation Menu

Add train_stream verb for streaming training on live data sources - #987

Add train_stream verb for streaming training on live data sources#987
drewoldag wants to merge 14 commits into
mainfrom
claude/train-stream-verb-planning-8qjyzm

drewoldag commented Jul 16, 2026

Uh oh!

codecov Bot commented Jul 17, 2026 •

edited

Loading

Uh oh!

drewoldag commented Jul 21, 2026

Uh oh!

review-notebook-app Bot commented Aug 4, 2026

Uh oh!

drewoldag Sep 8, 2026

Uh oh!

drewoldag Sep 8, 2026

Uh oh!

drewoldag Sep 8, 2026

Uh oh!

drewoldag Sep 8, 2026

Uh oh!

drewoldag Sep 8, 2026 •

edited

Loading

Uh oh!

drewoldag Sep 8, 2026 •

edited

Loading

Uh oh!

drewoldag Sep 8, 2026

Uh oh!

drewoldag Sep 8, 2026

Uh oh!

drewoldag commented Sep 8, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants


		return lsdb.open_catalog(self.data_location, **open_catalog_kwargs)

		def _requested_columns_from_config(self, config: dict):

Conversation

drewoldag commented Jul 16, 2026

Change Description

Solution Description

New Components

Configuration

Design Rationale

Testing

Code Quality

Uh oh!

codecov Bot commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

drewoldag commented Jul 21, 2026

Uh oh!

review-notebook-app Bot commented Aug 4, 2026

Uh oh!

drewoldag Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag Sep 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

drewoldag Sep 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

drewoldag Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag commented Sep 8, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Jul 17, 2026 •

edited

Loading

drewoldag Sep 8, 2026 •

edited

Loading

drewoldag Sep 8, 2026 •

edited

Loading