Split one SpatialData store across GPUs in run_inference_flash.py

[rushin682]

What

run_inference_flash.py can now split a single store across several GPUs with Lightning DDP.

• --devices N / --num-nodes M, launched with one SLURM task per GPU (srun --ntasks-per-node=N).
• Each rank predicts every N-th cell. Rank 0 gathers the rows, restores table order and writes the pred_table with the existing code, so the table format is the same for 1 and N GPUs.
• Each rank seeds with seed + rank, so ranks do not share noise.
• Progress is a tqdm bar (Lightning's default Rich bar prints nothing to a log file until the store is done).
• FlowPipeline.__call__ is split into sample_batch and denormalize so the Lightning path reuses the sampling code. __call__ is unchanged; a new test pins it to those two.

Why

A store with ~800k cells takes more than 24 h on one GPU, past the QoS wall clock, and the per-store queue cannot help with a single slide.

Heads-up

• With N > 1 the values are not bitwise equal to a 1-GPU run: each cell gets different noise and batches hold different cells. The difference is the size of a seed change (numbers below). With one device the output is bitwise identical to the previous script.
• With several devices a failed store aborts the job instead of being skipped (the other ranks would hang). Stores that already have a pred_table are skipped, so a rerun resumes.
• The SLURM script is not in this PR (scripts/*.sbatch is gitignored). It needs --gpu-bind=none so each task sees all four GPUs.

Tested

• pytest tests: 47 passed, 2 failed (test_models.py: llama3_forward, mixer_autoencoder_roundtrip; CUDA errors on the GPU-less login node, not re-run on a GPU).

• prek hooks on the changed files: all pass.

• H100, 243 cells, batch size 16, mean |Δ|:

  Comparison	mean |Δ|
  --devices 1 vs old script	0 (bitwise, default and --fast)
  4 GPUs vs 1 GPU, seed 0 / seed 1	0.183 / 0.182
  seed 1 vs seed 0, 1 GPU	0.188

  Table format (dtype, shape, obs, var, obsm["spatial"], uns["spatialdata_attrs"]) is identical.

• 56,510 cells, 4 × H100: 1838 s vs 7191 s on 1 GPU (3.9×). Format identical, mean |Δ| 0.153.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 48.25%. Comparing base (e96e260) to head (1d93b38).

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #28      +/-   ##
==========================================
+ Coverage   48.20%   48.25%   +0.05%     
==========================================
  Files          10       10              
  Lines        1029     1030       +1     
==========================================
+ Hits          496      497       +1     
  Misses        533      533              

Files with missing lines	Coverage Δ
src/phoenix/helpers/inference.py	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.