Villar-PSO light curve fitting on GPU in ZTF enrichment workers

[frenbox]

Summary

Adds Villar light curve fitting on GPU (via frenbox/villar-pso) to the ZTF enrichment pipeline, on top of gpu-setup-fix. Each enrichment worker runs batch_pso_multi_seed on its batch's fittable alerts after ONNX classification and writes villar_fit.* fields to the ZTF alerts collection.

• New gpu feature gates dep:villar-pso (its CUDA build needs nvcc at build time, unlike ORT which is load-dynamic). ORT GPU is unchanged.
  Two-gate runtime: villar runs only when both the build and the config allow it.

• Build flag (--features gpu): controls whether villar code is included in the binary.

• Config (config.gpu.enabled): controls whether that code runs when the binary starts.

Villar is GPU-only — there's no CPU fallback. When the config disables it, workers skip the villar step entirely while the rest of enrichment runs normally.

• Per-worker GpuContext, round-robin over config.gpu.device_ids (mirrors upstream SharedModelPool).
• Graceful failure: preprocess errors and GPU batch failures both write NaN-valued villar_fit.* docs so downstream consumers always see the expected keys.

Files

File	Change
Cargo.toml / Cargo.lock	gpu feature + optional villar-pso dep (Linux)
Dockerfile.gpu	Add --features gpu to the final cargo build
src/enrichment/ztf.rs	Villar integration in process_alerts

• villar-pso submits to the CUDA NULL stream, so N workers on one GPU serialize at the driver level (~1× throughput, not N×). Multi-GPU scales linearly. Per-stream support in villar-pso or a boom-side dispatcher would unlock single-GPU concurrency — follow-up work.
• Only G/R bands are mapped; other bands fall through to NaN.

Additional Change 1

• Share a single CUDA stream between ONNX Runtime sessions (ACAI/btsbot) and
  villar-pso PSO for each enrichment device, replacing the previous mix of
  the default (NULL) stream + 6 ORT-internal streams. Eliminates the
  default-stream's implicit cross-stream fences and makes H2D/D2H transfers
  actually async.

• SharedModels now owns the per-device Stream and the villar-pso
  GpuContext. Field declaration order ensures the stream is destroyed
  after every session and context that uses it. SharedModelPool already
  fans this out per device.

• Drops the parallel round-robin (GPU_DEVICE_COUNTER) in
  ZtfEnrichmentWorker, which was assigning Villar to a different device
  than ACAI/btsbot for the same worker. Both now come from the same
  Arc, so device alignment is guaranteed.

• Removes the dead boom/src/gpu/ module (GpuPool was never wired up).

Additional Change 2

• Configurable batch size (config.yaml, src/conf.rs): workers.ztf.enrichment.batch_size is now a single config variable (default 750) controlling both the RPOP cap and the fixed ONNX inference shape. Zero-padding keeps the GPU memory arena stable across partial batches. I tested multiple benchmark for this and the GPU memory usage remained stable.

[github-actions]

Throughput results (8db738a4d491cef918646f947f3a2509344a1803):

New wall time	Baseline wall time	Difference
246.3	252.4	-2.00%

[github-actions]

Throughput results (44566880c2c2db6f7b47f88918c5899682b7ea94):

New wall time	Baseline wall time	Difference
236.5	252.4	-6.00%

[frenbox]

I curated the following dataset of longest classified source on fritz for the LC fitting test

• AGN  : 662
• Stellar variable   :  177
• Cataclysmic : 57
• QSO : 30
• SN Ia : 21
• RR Lyrae : 10
• And other few sources counting upto 1000

The memory alone used for computation is quite small

Metric	Single-seed	Multi-seed (×3)
Wall time	1.05 s	3.13 s
ms/source	1.10	3.28
Peak GPU (process)	30.0 MB	30.0 MB
Obs upload	4.0 MB	4.0 MB
χ² mean / median	0.101 / 0.046	0.092 / 0.041

however nvidia-smi reports ~312 MB for the process likely due to CUDA context over heads

[Theodlz]

My immediate thought, looking at this, was: why are we having one GPU algorithm behind a rust feature flag (villar fits) and not the other (ml with ort). Obviously the difference is that there is no villar fits happening on CPU (unlike the ml work), but still it feels cumbersome to have one env var to enable GPU for ML but one rust feature flag to enable GPU for lc fitting. My suggestion for this PR would be to undertake the following work: put both villar fits and GPU-enabled ML behind a feature flag(s). The "s" is in parenthesis because if we think as soon as we have GPUs all work should be done there then we should couple ml and villar fits, but if not we could have a feature flag for each. I'd suggest coupling them for now. BUT, if what you are looking for is:

• feature flag to specify if something is included in the compiled binary
• and separately, config to tell the code if it should run the GPU stuff
  then I'm fine keeping both! But even then, I wonder if we should apply the same logic to the ML stuff? (feature flag to use ort with dynamic loading but then at runtime we can go back to CPU inference is the config says its disabled, and that way the CPU only version - when no feature flag is set - doesn't need to bother with the ort dynamic loading).

Another comment, related to this: I see that in your implementation "Workers are round-robin-assigned a device from config.gpu.device_ids at init time.". May I ask why this choice? Also, how does that play w/ the GPU session pool we use for ML? Instinctively I'd think a shared pool + context would make more sense for throughput (because here your enrichment worker might want to use GPU that was assigned to it at init rather than whichever is actually available now at runtime)? If you think there is no advantage to this and round robin makes more sense, then may I ask why we would not apply the same treatment to the ML stuff? I guess I just want to make sure we don't add complexity and different GPU-handling systems unless we have a specific reason for it.

Last but not least (related to the round-robin villar vs gpu ml pool stuff), let's say one GPU is running an inference task: can that GPU also run the villar fits in parallel? My understanding is that GPUs do have the ability to parallelize over different compute (which varies between cards of course, some more capable than others, and it seems server cards are best at this). If so, then maybe that answer the round-robin question. Also what happens when 2 workers try to use the same GPU context (which seems to be analogous to a device in this implementation) at the same time? Can CUDA decide to parallelize which means higher mem usage? If so, that could be an issue.

I think my understanding of parallel compute on GPUs is a little poor here, so these are all genuine questions I am asking myself as well!

[frenbox]

Hi Theo,

Thanks for the comments! I am also new to figuring out this GPU stuff so we are on similar boat. Here are my thoughts on your comments

1. Cargo feature for villar fit but not ORT

I feel like it is forced by the dependencies:

• ORT uses load-dynamic which requires no CUDA toolkit needed at build time. The CUDA EP is used during the runtime. This makes one compiled binary works on CPU and GPU hosts.
• villar-pso compiles a .cu file via nvcc at build time. This is not equivalent to load-dynamic . A CPU-only build host literally can't compile a binary with villar in it.

So the Cargo feature exists because villar-pso requires it, not because we wanted asymmetry. The runtime semantics are unified such that both ORT and villar honor config.gpu.enabled.

2. Round-robin device assignment at init

The current existing pattern introduced by Michael's PR is: SharedModelPool::next_model_set kind of does the same round-robin at worker init.

However, this has its implication to your third concern

3. How is ML inference + villar handled on the same GPU

This is how the current CUDA streams work:

• ORT inference uses ORT's own cudaStream_t
• villar-pso uses the CUDA NULL stream (default) since i wrote the villar-pso without thinking of dedicated stream (which was later introduced by Michael's PR)

From NVIDIA CUDA docs, it seems NULL stream has legacy "sync everything" semantics: any work on it implicitly waits for all other streams, and other streams wait for it. So ORT and villar serialize in practice, even though they're on different streams.

So in practice:

• When villar fires a kernel on the NULL stream, it implicitly waits for ORT's stream to drain.
• When ORT fires the next kernel on its own stream, it implicitly waits for villar's NULL-stream work.

To actually unlock parallelism on a single GPU, villar-pso would need to use a non-default stream per GpuContext.
This change has to done by me on the villar-pso end. I could try to do it and then change this PR or try to make another PR for this. If we leave the PR as is, we might be able to do fix Inference on one GPU and fits on another (given say we have 2 GPU system). I think we need to think of design choices as well moving forward.

[github-actions]

Throughput results (779c609c82cc0596c61c51ff632f02300f14b1a4):

New wall time	Baseline wall time	Difference
252.5	242.4	4.00%

[github-actions]

Throughput results (2d3a42bbd34be924c2bf1d12bffee7aeeceefb14):

New wall time	Baseline wall time	Difference
251.9	253.2	0%

[frenbox]

Additional Changes Summary

• Share a single CUDA stream between ONNX Runtime sessions (ACAI/btsbot) and
  villar-pso PSO for each enrichment device, replacing the previous mix of
  the default (NULL) stream + 6 ORT-internal streams. Eliminates the
  default-stream's implicit cross-stream fences and makes H2D/D2H transfers
  actually async.

• SharedModels now owns the per-device Stream and the villar-pso
  GpuContext. Field declaration order ensures the stream is destroyed
  after every session and context that uses it. SharedModelPool already
  fans this out per device.

• Drops the parallel round-robin (GPU_DEVICE_COUNTER) in
  ZtfEnrichmentWorker, which was assigning Villar to a different device
  than ACAI/btsbot for the same worker. Both now come from the same
  Arc<SharedModels>, so device alignment is guaranteed.

• Removes the dead boom/src/gpu/ module (GpuPool was never wired up).

[frenbox]

@Theodlz this is now ready for you to test Theo.

[github-actions]

Throughput results (5194ff9333a3ee6a9f24a958053ed24566405177):

New wall time	Baseline wall time	Difference
252.3	253.2	0%

[github-actions]

Throughput results (0eb4af7d4f58d09d6615591743d395100da91622):

New wall time	Baseline wall time	Difference
246.7	253.2	-2.00%

[github-actions]

Throughput results (47717347a0df117e6c9c0822e06272e743e2a7c1):

New wall time	Baseline wall time	Difference
243.3	253.2	-3.00%

[frenbox]

It seems after a combined cuda stream, in local testing my GPU is hitting OOM.
I am looking more into it.

[github-actions]

Throughput results (6df8a2d0a299329774cb61790ee52ddd85470e84):

New wall time	Baseline wall time	Difference
248.5	245.1	1.00%

[github-actions]

Throughput results (99f181e72427ad04668a8e2f51952077d3c7b5c0):

Storage	New wall time	Baseline wall time	Difference
mongo	246.6	248.7	0%
s3	272.7	281.0	-2.00%

[github-actions]

Throughput results (b660b6e2c8ef6dc664192e100059e6cfb6c24dda):

Storage	New wall time	Baseline wall time	Difference
mongo	256.7	248.7	3.00%
s3	276.2	281.0	-1.00%

[frenbox]

Final round of changes:

Configurable batch size (config.yaml, src/conf.rs): workers.ztf.enrichment.batch_size is now a single config variable (default 750) controlling both the RPOP cap and the fixed ONNX inference shape. Zero-padding keeps the GPU memory arena stable across partial batches. I tested multiple benchmark for this and the GPU memory usage remained stable.

@Theodlz Finally this is done!

[github-actions]

Throughput results (9898831c7774c73d0978ab2b48afc41986078393):

Storage	New wall time	Baseline wall time	Difference
mongo	249.5	252.5	-1.00%
s3	262.4	236.5	10.00%

[github-actions]

Throughput results (9898831c7774c73d0978ab2b48afc41986078393):

Storage	New wall time	Baseline wall time	Difference
mongo	249.5	252.5	-1.00%
s3	262.4	236.5	10.00%

[github-actions]

Throughput results (9898831c7774c73d0978ab2b48afc41986078393):

Storage	New wall time	Baseline wall time	Difference
mongo	249.5	252.5	-1.00%
s3	294.2	236.5	24.00%

[github-actions]

Throughput results (aa9592677cdec02183aa10f4ffed096d5c8976a2):

Storage	New wall time	Baseline wall time	Difference
mongo	257.9	246.0	4.00%
s3	262.0	294.0	-10.00%

Baseline run: 28239525849