feat(apps): compile a multi-model bundle to a program.json for a Python-free runtime - #66
Merged
Conversation
A bundle with several checkpoints exports one <fold>.onnx per checkpoint plus a program.json -- the multi-model dataflow the konfai-rs runtime runs. The ensemble reduction is read from the config (MergeLabels -> merge_labels, InferenceStack -> mean), never a per-app flag; in_channels falls back to the loaded model. export_to_onnx returns the manifest so folds embed it in the program; export_onnx_into_bundle stays a single-checkpoint wrapper.
A prediction config with an auxiliary mask group (a nested KonfAIInference model plus a Mask on the output) compiles to the full dataflow program the konfai-rs runtime runs: the nested model is resolved and exported, its buffer is resampled onto the primary grid and dilated, the primary folds run over flip test-time-augmentation passes, the copies are averaged, masked, inverse resampled, and cast. Everything is read from the config -- the is_input group, the program-level transforms (Mask / KonfAIInference / Dilate / Save), the ensemble reduction, and the TTA. The flip passes are drawn (identity + nb draws, reproducible from manual_seed): randomness is embraced, so the result is close, not byte-identical, to a torch-seeded run; a non-Flip augmentation is skipped rather than blocking the export.
InferenceStack's reduction mode is read from the config (mean vs median) instead of always assuming mean, and both are runtime buffer ops -- so a median ensemble is reproduced rather than silently averaged. The masked/TTA program takes the derived reduction for its copy reduction too.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughPortable bundle export now supports single and multi-checkpoint artifacts, ensemble reduction, masked TTA synthesis, nested models, and configurable ONNX output contracts. Transform compilation and CLI wiring share the portable export path. ChangesPortable runtime bundle export
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant PortableExporter
participant ONNXExporter
participant BundleProgram
CLI->>PortableExporter: provide checkpoints and prediction configuration
PortableExporter->>ONNXExporter: export checkpoint ONNX and manifest
ONNXExporter-->>PortableExporter: return artifact path and manifest
PortableExporter->>BundleProgram: assemble ensemble or masked-TTA program
BundleProgram-->>CLI: return portable bundle artifact
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
konfai/export.py (1)
134-169: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winConstrain
model_filenameto a bare.onnxfilename.
out_dir / model_filenameaccepts absolute and../paths, soexport_to_onnx()can write outsideoutput_dir. Rejectmanifest.jsontoo; withwrite_manifest=True, the manifest write will overwrite the exported ONNX file.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@konfai/export.py` around lines 134 - 169, Validate model_filename in export_to_onnx before constructing onnx_path: require a bare filename with an .onnx suffix, reject absolute paths, path components such as directories or .., and manifest.json, and raise an appropriate error for invalid values. Keep writing the validated filename beneath output_dir.
🧹 Nitpick comments (1)
konfai/export.py (1)
134-136: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the new export options.
The changed tests exercise only the defaults. Add a focused case covering a custom
.onnxfilename, the returned manifest, andwrite_manifest=Falseensuring nomanifest.jsonis written.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@konfai/export.py` around lines 134 - 136, Add a focused test for the export function covering a custom .onnx filename, validation of the returned manifest, and write_manifest=False. Assert the custom model file is created, the manifest is returned as expected, and manifest.json is not written.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@konfai-apps/konfai_apps/bundle.py`:
- Around line 484-486: Update the docstring describing non-Flip TTA handling to
state that such augmentations are skipped rather than refused or rejected. Keep
the existing behavior and the inline comment near the draw loop unchanged, and
ensure the documentation accurately reflects the `f_prob`/`range(0)` path.
- Around line 637-639: Update _export_nested_model around
export_portable_into_bundle so nested mask models with multiple checkpoints do
not unconditionally read manifest.json or copy model.onnx. Validate that nested
exports use exactly one checkpoint, or explicitly handle the program.json
multi-checkpoint output before copying the model and loading metadata.
---
Outside diff comments:
In `@konfai/export.py`:
- Around line 134-169: Validate model_filename in export_to_onnx before
constructing onnx_path: require a bare filename with an .onnx suffix, reject
absolute paths, path components such as directories or .., and manifest.json,
and raise an appropriate error for invalid values. Keep writing the validated
filename beneath output_dir.
---
Nitpick comments:
In `@konfai/export.py`:
- Around line 134-136: Add a focused test for the export function covering a
custom .onnx filename, validation of the returned manifest, and
write_manifest=False. Assert the custom model file is created, the manifest is
returned as expected, and manifest.json is not written.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2083cc40-78d5-4753-8828-517f9cadbce3
📒 Files selected for processing (5)
konfai-apps/konfai_apps/bundle.pykonfai-apps/tests/unit/test_bundle.pykonfai-apps/tests/unit/test_onnx_fold.pykonfai/export.pytests/unit/test_onnx_export.py
mypy failed on three call sites in the program export: the TTA draw read an optional f_prob, the checkpoint list mixed None with str, and the ensemble reduction stayed optional at assemble_program. Narrow each at its source: skip a non-Flip augmentation explicitly instead of looping zero times, name the resolved checkpoint list, and resolve the ensemble reduction once (the earlier guard already rejects a missing one). A nested mask model with several checkpoints exported a program.json while the caller copied model.onnx, failing on a missing file; refuse it with an actionable error. Align the _tta_passes docstring with the skip behaviour.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Follow-up to #64 (single-model onnx + manifest). Adds multi-model export so a whole
bundle — an ensemble's folds or a nested pipeline — compiles to one
program.jsonaPython-free runtime can execute.
assemble_program(models, *, reduce, classes)emits an ordered dataflow program(
{steps, output}) over named buffers: oneModelstep per checkpoint, then areduction step.
mean,median(new), andmerge_labels(disjoint-label ensembleslike TotalSegmentator's 5 tasks — cumulative per-model offset, last-model-wins).
model), matching what the KonfAI predictor does at runtime.
InferenceStackmode →
mean/median), so the emitted program mirrors the app's own config.Why
manifest.json(from #64) covers one model; real apps are ensembles or nested pipelines.program.jsonis the portable counterpart that lets a non-Python runtime run the wholebundle.
Validation
pixi run checkgreen (lint + format-check + core tests + apps tests).test_program.py,test_bundle.py,test_onnx_fold.py,test_onnx_export.py(45 pass on the export/program surface).disjoint heads → 118 labels via
merge_labels): 98.78% foreground voxel agreementvs the KonfAI-native app output on a real CT.
Scope
konfai-apps + konfai export only. No Rust runtime in this PR.
Summary by CodeRabbit