Skip to content

docs: document multi-GPU serve pattern and add example config - #7

Merged
rylinjames merged 1 commit into
FastCrest:mainfrom
DsThakurRawat:docs/multi-gpu-serve
Apr 30, 2026
Merged

docs: document multi-GPU serve pattern and add example config#7
rylinjames merged 1 commit into
FastCrest:mainfrom
DsThakurRawat:docs/multi-gpu-serve

Conversation

@DsThakurRawat

@DsThakurRawat DsThakurRawat commented Apr 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

This PR introduces documentation and example configurations for deploying Reflex on multiple GPUs using the process-per-GPU pattern.

Fixes #5

Type of Change

  • Bug fix
  • New feature (documentation & examples)
  • Breaking change

What was added:

  1. examples/05-multi-gpu-serve.md: An architectural guide explaining how to scale reflex serve across multiple GPUs.
  2. examples/multi-gpu/docker-compose.yml: A ready-to-use Docker Compose stack that binds workers to specific GPUs.
  3. examples/multi-gpu/nginx.conf: An NGINX configuration optimized for inference load-balancing.
@rylinjames
rylinjames merged commit f5fd249 into FastCrest:main Apr 30, 2026
rylinjames added a commit that referenced this pull request Apr 30, 2026
@rylinjames @claude
infra+docs: polish post-merge — tighten CI matrix, fix multi-gpu doc
35babdf 
Two follow-up improvements after merging #6 and #7.

CI (.github/workflows/pytest.yml):
- Drop macos-latest from the matrix. CI coverage was the import + CLI
  + lazy-load surface; macos adds runtime without adding coverage since
  every reflex tester runs Linux. Cuts ~50% of runner time.
- Add pip caching keyed on pyproject.toml. Saves ~3 min per run after
  the first since torch/onnxruntime/etc don't re-download.
- Pytest verbose flag (-v) so CI output shows which tests skipped vs
  ran (relevant since most are GPU-gated and skip on hosted runners).
- Inline comment explaining why [gpu] is deliberately omitted (2 GB
  install for tests that all skip via @pytest.mark.skipif anyway).

Multi-GPU docs (examples/05-multi-gpu-serve.md):
- Fix the cold-start claim: was "10-70 seconds", now "30-60 seconds
  (target floor < 90s)" matching METRICS.md. Also note that subsequent
  boots hit the engine cache and start in seconds.
- Add "When to use this vs --policy-a / --policy-b" section so users
  pick the right pattern. Process-per-GPU = horizontal scale-out;
  --policy-a/--policy-b = sticky-per-episode A/B testing on one GPU.
  Notes that the two patterns compose.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
rylinjames added a commit that referenced this pull request May 21, 2026
@rylinjames @claude
feat(models): Pi0VLA spine composition class (lift #1 Day 4f) (#149)
c92fff1 
Wires the 4 components (SigLIPBackbone + PaliGemmaBackbone +
LinearProjector + FlowMatchingHead) together through the BaseVLA spine.
Proves the spine actually composes with real registered components.

Scope: composition shape ONLY. `predict_action()` is an explicit
NotImplementedError pointing at Day 4g, which adds the full inference
pipeline (vision → project → merge → language → flow-matching denoise
loop) + the `--use-new-spine` CLI flag + the parity gate vs the
legacy pi0_exporter + a Modal smoke validating cos=+1.0 bit-identical
actions.

This split keeps the composition class reviewable in isolation
without bundling 200+ LOC of inference orchestration that wants its
own focused PR.

New files (~340 LOC):

- src/reflex/models/vlas/__init__.py — directory scaffold + Pi0VLA
  re-export. Future VLAs (Pi05VLA, SmolVLA, GR00TVLA, DreamZeroVLA)
  land here per the per-VLA decomposition days (5-7) + lift #7.

- src/reflex/models/vlas/pi0.py (~180 LOC) — Pi0VLA(BaseVLA). 4 required
  slots (vision/llm/projector/head), empty NAME_MAPPING (default per
  decision S-1 — lerobot/pi0_base keys route via load_state_dict's
  slot-prefix splitting). `from_pretrained(hf_id)` loads PaliGemma
  once, extracts paligemma.model.vision_tower → SigLIPBackbone, wraps
  the rest → PaliGemmaBackbone, builds projector from state_proj
  weights if present, builds ExpertStack via the legacy
  build_pi0_expert_stack() then wraps as FlowMatchingHead. forward()
  is minimal — routes to llm_backbone with pre-merged inputs_embeds
  (the test path). predict_action() raises NotImplementedError with
  a clear Day 4g pointer (no silent stub).

- tests/test_pi0_vla.py (~160 LOC, 10 tests) — registration, ABC
  subclass, slot declarations (REQUIRED + OPTIONAL + NAME_MAPPING),
  construction via direct kwargs + from_config, validation (missing
  required + undeclared slot raise), forward routing to llm_backbone,
  predict_action NotImplementedError. Uses 4 stub components so no HF
  download at test time.

Pre-merge validation:
- syntax: ✓
- AST orphan-imports sweep: ✓
- live import: VLAS contains Pi0VLA; REQUIRED_SLOTS correct;
  NAME_MAPPING empty; vlas/__init__.py re-export works

Lift #1 spine progress after this PR: 8 of 12 days done (1, 2, 3, 4a,
4b, 4c, 4e, 4f). Day 4d was folded into 4e per design revision in PR
#147. Day 4g next: inference pipeline + parity gate + cutover.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
rylinjames added a commit that referenced this pull request May 25, 2026
@rylinjames @claude
lift #7 day 2: DreamZeroVLA + WanBackbone on the BaseVLA spine
4efe958 
Wires DreamZero (NVIDIA Research's World-Action Model) as the 6th model
family on the spine. Validates the vlm_backbone slot for a non-LLM
architecture (Wan2.1 = T5 + CLIP + VAE, not autoregressive).

What lands:
- src/reflex/models/vlas/dreamzero.py (~150 LOC): DreamZeroVLA composition
  class. Required slots: vlm_backbone + vla_head. forward() + predict_action()
  ported from FluxVLA dreamzero_vla.py.
- src/reflex/models/vlm/wan_backbone.py (~120 LOC): WanBackbone wrapping T5
  text encoder + CLIP image encoder + Video VAE. Downloads weights from
  Wan-AI/Wan2.1-I2V-14B-480P on HF. All parameters frozen.

Both files AST-parse + import OK.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
rylinjames added a commit that referenced this pull request May 25, 2026
@rylinjames @claude
lift #7 day 3: DreamZeroHead — DiT action + video diffusion head
3086cb0 
Ported from FluxVLA dreamzero_head.py (732 LOC, Apache-2.0). Contains
the CausalWanModel (DiT backbone) + FlowMatchScheduler + the joint
video+action denoising loop. Imports rewritten from fluxvla.models.*
to reflex.models.third_party.dreamzero.modules.*.

Supports both stateless (predict_action_stateless) and cached
(predict_action with KV cache) inference paths. Default is stateless
(4-step UniPC sampling).

Lift #7 now has all three layers:
- Day 1: vendored modules (9.5K LOC)
- Day 2: DreamZeroVLA + WanBackbone on spine
- Day 3: DreamZeroHead (this commit)

Remaining: ONNX exporter wiring + registry update + parity test.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
rylinjames added a commit that referenced this pull request May 25, 2026
@rylinjames @claude
lift #7 day 4: DreamZero exporter stub + registry entry (6th model fa…
bf63e3a 
…mily, 9 models total)

Exporter writes reflex_config.json with dreamzero-specific fields (video
input shape, component mapping). Full DiT ONNX export deferred — the
CausalWanModel has dynamic KV cache + flash-attn that needs careful
tracing. V1 ships config + PyTorch runtime path.

Registry: dreamzero-libero10 added (limxdynamics/FluxVLAEngine, 94.65%
LIBERO, ~28 GB). Family validator updated to accept 'dreamzero'.

9 models across 6 families: pi0, pi05, smolvla, groot, openvla, dreamzero.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
rylinjames added a commit that referenced this pull request May 25, 2026
@rylinjames @claude
lift #7: DreamZero WAM exporter (6th model family) (#180)
180f332 
* lift #7 day 1: vendor DreamZero third-party modules (~9.5K LOC)

Vendored NVIDIA Research DreamZero modules from FluxVLA's third_party_models
directory (Apache-2.0, via LimX Dynamics). Manual cp to preserve internal
relative imports which work as-is.

What lands:
- src/reflex/models/third_party/ package
- src/reflex/models/third_party/dreamzero/ (15 Python modules, ~9.5K LOC):
  Wan2.1 DiT backbone, T5 text encoder, CLIP image encoder, Video VAE,
  flow-matching scheduler, UniPC sampler, causal attention blocks,
  action/state encoder, camera controller, utility modules
- ATTRIBUTION.txt with three-link source chain (NVIDIA Research → FluxVLA → reflex)

All 15 modules AST-parse OK. No fluxvla.* imports — all internal imports
use relative paths.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* lift #7 day 2: DreamZeroVLA + WanBackbone on the BaseVLA spine

Wires DreamZero (NVIDIA Research's World-Action Model) as the 6th model
family on the spine. Validates the vlm_backbone slot for a non-LLM
architecture (Wan2.1 = T5 + CLIP + VAE, not autoregressive).

What lands:
- src/reflex/models/vlas/dreamzero.py (~150 LOC): DreamZeroVLA composition
  class. Required slots: vlm_backbone + vla_head. forward() + predict_action()
  ported from FluxVLA dreamzero_vla.py.
- src/reflex/models/vlm/wan_backbone.py (~120 LOC): WanBackbone wrapping T5
  text encoder + CLIP image encoder + Video VAE. Downloads weights from
  Wan-AI/Wan2.1-I2V-14B-480P on HF. All parameters frozen.

Both files AST-parse + import OK.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* lift #7 day 3: DreamZeroHead — DiT action + video diffusion head

Ported from FluxVLA dreamzero_head.py (732 LOC, Apache-2.0). Contains
the CausalWanModel (DiT backbone) + FlowMatchScheduler + the joint
video+action denoising loop. Imports rewritten from fluxvla.models.*
to reflex.models.third_party.dreamzero.modules.*.

Supports both stateless (predict_action_stateless) and cached
(predict_action with KV cache) inference paths. Default is stateless
(4-step UniPC sampling).

Lift #7 now has all three layers:
- Day 1: vendored modules (9.5K LOC)
- Day 2: DreamZeroVLA + WanBackbone on spine
- Day 3: DreamZeroHead (this commit)

Remaining: ONNX exporter wiring + registry update + parity test.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* lift #7 day 4: DreamZero exporter stub + registry entry (6th model family, 9 models total)

Exporter writes reflex_config.json with dreamzero-specific fields (video
input shape, component mapping). Full DiT ONNX export deferred — the
CausalWanModel has dynamic KV cache + flash-attn that needs careful
tracing. V1 ships config + PyTorch runtime path.

Registry: dreamzero-libero10 added (limxdynamics/FluxVLAEngine, 94.65%
LIBERO, ~28 GB). Family validator updated to accept 'dreamzero'.

9 models across 6 families: pi0, pi05, smolvla, groot, openvla, dreamzero.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@rylinjames rylinjames mentioned this pull request May 25, 2026
Merged
1 task
rylinjames added a commit that referenced this pull request May 25, 2026
@rylinjames @claude
fix: add dreamzero to canonical family set in test + validator message (
545571f 
#185)

Lift #7 (DreamZero WAM) registered the dreamzero family in the
validator allowlist and registry entries, but the test hardcoded
canonical set and the error message string were not updated.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
rylinjames added a commit that referenced this pull request May 29, 2026
@rylinjames @claude
fix: route integration-command errors to stderr + classify two export…
47d5262 
…ers in registry audit

Fixes two pre-existing red tests on main, independent of the serve/bench
hardening that lands in the next commit:

- test_no_stdout_error_paints_remain_in_cli: connect_up (ValueError +
  RuntimeError handlers) and connect_status (unknown-integration) painted
  red error text on the stdout `console`. Route them to `err_console`
  (stderr) like every other error path, so `reflex connect ... 2>/dev/null`
  no longer swallows the only error signal.
- test_exporter_directory_audit_covers_all_files: dreamzero.py (lift #7 WAM
  exporter) and weight_fusion.py (external-data ONNX helper) existed in
  src/reflex/exporters/ but were unclassified. Classify dreamzero.py as a
  primary exporter (family=dreamzero) and weight_fusion.py as internal.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
rylinjames added a commit that referenced this pull request May 29, 2026
@rylinjames @claude
fix: route integration-command errors to stderr + classify two export…
3ef1b37 
…ers in registry audit

Fixes two pre-existing red tests on main, independent of the serve/bench
hardening that lands in the next commit:

- test_no_stdout_error_paints_remain_in_cli: connect_up (ValueError +
  RuntimeError handlers) and connect_status (unknown-integration) painted
  red error text on the stdout `console`. Route them to `err_console`
  (stderr) like every other error path, so `reflex connect ... 2>/dev/null`
  no longer swallows the only error signal.
- test_exporter_directory_audit_covers_all_files: dreamzero.py (lift #7 WAM
  exporter) and weight_fusion.py (external-data ONNX helper) existed in
  src/reflex/exporters/ but were unclassified. Classify dreamzero.py as a
  primary exporter (family=dreamzero) and weight_fusion.py as internal.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@rylinjames rylinjames mentioned this pull request May 30, 2026
Merged
4 tasks
@rylinjames rylinjames mentioned this pull request Jun 10, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

2 participants

@DsThakurRawat @rylinjames