Skip to content

feat(vit): true vit-only mode — skip the iterative engine entirely#43

Merged
Garrett Bischof (gwbischof) merged 1 commit into
mainfrom
h4b/vit-only-mode
Jun 9, 2026
Merged

feat(vit): true vit-only mode — skip the iterative engine entirely#43
Garrett Bischof (gwbischof) merged 1 commit into
mainfrom
h4b/vit-only-mode

Conversation

@gwbischof

Copy link
Copy Markdown
Collaborator

H4b of the #37 decomposition — make recon_mode='vit' actually skip the iterative engine.

Problem

recon_mode='vit' previously still created the PtychoRecon / StreamingPtychoRecon engine (it only skipped wiring its flows), which opens a CuPy context on the ViT GPU. PyCUDA (ViT) + CuPy (recon) on the same device from different threads can crash with SIGABRT.

Changes (ptycho_holo.py, preprocess.py)

  • Skip the enginePtychoRecon, SaveResult, SaveLiveResult are created only for recon_mode in (iterative, both). The reset_for_scan block is likewise guarded. (The iterative add-flows were already guarded on main.)
  • Config-derived geometry when there's no engine:
    • config_ops gains a has_recon branch; in vit-only it sizes ops from param.nx/ny and computes the position-conversion pixel size via ptychoml.compute_sample_pixel_size (the same far-field formula the engine uses). x/y_direction already come from config at point_proc construction.
    • compose computes _x_pixel_m/_y_pixel_m/_lambda_nm once (engine when present, else config) for the run metadata and SaveViTResult. Ranges use abs(param.x/y_range) — the same values the engine supplies, no axis swap (that swap is a separate orientation change, H5).
  • preprocess.pyImageSendOp/PointProcessorOp skip their GPU device copies when diff_d_target/point_info_target are None (the vit-only case; both default to None).

Verified every self.pty reference is inside a has_recon/recon_mode guard, and the vit-only range/pixel-size fallbacks reproduce what StreamingPtychoRecon computes (x_range_um = abs(param.x_range), x_pixel_m = compute_sample_pixel_size(...)) — so vit-only geometry matches the iterative path on main.

Testing

This is compose-time Holoscan behavior — not unit-testable in CI (the ptycho_holo smoke import already fails on a missing TILED_BASE_URL, before compose). It's exercised by real vit-only runs, faithfully reproduces #37's working vit-only path, and the geometry fallback delegates to the ptychoml-tested compute_sample_pixel_size. Safety net here is the audited guards + no regression:

pixi run test   # 85 passed; the 2 vit_inference/ptycho_holo smoke failures are pre-existing.

recon_mode='vit' previously still created the PtychoRecon /
StreamingPtychoRecon engine (it only skipped wiring its flows), opening a
CuPy context on the ViT GPU — PyCUDA + CuPy on the same device from
different threads can crash (SIGABRT). Now the engine, SaveResult, and
SaveLiveResult are created only for recon_mode in (iterative, both).

With no engine, the geometry it used to supply is derived from config:
- config_ops gains a has_recon branch; in vit-only it sizes the ops from
  param.nx/ny and computes the position-conversion pixel size via
  ptychoml.compute_sample_pixel_size (same far-field formula the engine
  uses), leaving x/y_direction at the values point_proc already got at
  construction.
- compose computes _x_pixel_m/_y_pixel_m/_lambda_nm once (engine when
  present, else config) for run metadata and SaveViTResult; SaveViTResult's
  ranges come from abs(param.x/y_range) — the same values recon supplies, no
  axis swap (that's a separate orientation change).
- the reset_for_scan block is guarded.

preprocess.py: ImageSendOp and PointProcessorOp skip their device copies
when diff_d_target / point_info_target are None (the vit-only case).

All self.pty references are now inside has_recon / recon_mode guards. This is
compose-time Holoscan behavior exercised by real vit-only runs; the geometry
fallback delegates to the ptychoml-tested helper, and the existing suite
passes (85, modulo the 2 pre-existing TILED smoke imports).

H4b of the #37 decomposition.

Co-authored-by: Himanshu Goel <4122621+himanshugoel2797@users.noreply.github.com>
@gwbischof Garrett Bischof (gwbischof) merged commit 012dd9c into main Jun 9, 2026
5 checks passed
@gwbischof Garrett Bischof (gwbischof) deleted the h4b/vit-only-mode branch June 9, 2026 16:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant