Add train_stream verb for streaming training on live data sources

[drewoldag]

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

[codecov]

Codecov Report

❌ Patch coverage is 88.16568% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 66.23%. Comparing base (8f32125) to head (77a5bc1).
⚠️ Report is 47 commits behind head on main.

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

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #987      +/-   ##
==========================================
+ Coverage   65.69%   66.23%   +0.54%     
==========================================
  Files          85       86       +1     
  Lines        8170     8373     +203     
==========================================
+ Hits         5367     5546     +179     
- Misses       2803     2827      +24     

☔ 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.

[drewoldag]

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