feat: add a two-stage inference runner - #241
Open
codename1995 wants to merge 1 commit into
Open
Conversation
Some policies cannot be exported as one graph. A vision-language-action model may have a backbone no exporter can capture and an action head that exports cleanly, and it is the head that is evaluated repeatedly per action chunk, so running it on an accelerated backend is where the time goes. TwoStage runs the primary adapter, matches its output names against the second model's declared input names and feeds them across, reporting a mismatch up front rather than leaving it to surface as a shape error inside the backend. The second model is loaded on first use, so building the runner while the manifest is read never compiles a graph. The runner needs its own artifact, which a manifest names relative to the export directory, so get_runner() takes an optional export_dir and applies resolve_artifact() to the spec - the same path preprocessor artifacts already take. Instantiation and the runner type check moved into one helper shared by both branches. Note that a runner spec carrying manifest metadata must use class_path mode: in type mode every extra field is forwarded to the constructor, so a chunk_size annotation reaches the adapter as a backend option. Signed-off-by: Weijie Wei <codename1995@gmail.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This pull request introduces a new InferenceRunner implementation, physicalai.inference.runners.TwoStage, to support hybrid exports where a policy must be executed as two separate artifacts (potentially on different backends) while keeping the export directory layout and InferenceModel loading flow unchanged.
Changes:
- Added
TwoStagerunner that runs a primary adapter and feeds selected intermediate outputs into a lazily-loaded second-stage adapter. - Updated runner factory +
InferenceModelintegration so runner artifact paths can be resolved relative to the export directory (with traversal protection viaresolve_artifact). - Added a dedicated unit test suite for runner behavior, factory selection, lazy loading, and artifact resolution.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/inference/test_runners.py | Adds unit tests for SinglePass, TwoStage, and get_runner (including lazy-load and artifact resolution behavior). |
| src/physicalai/inference/runners/two_stage.py | Implements the new TwoStage runner with lazy second-stage loading and input-name-based wiring/validation. |
| src/physicalai/inference/runners/factory.py | Extends get_runner to accept export_dir and resolves runner artifacts via resolve_artifact before instantiation. |
| src/physicalai/inference/runners/init.py | Exports TwoStage from the runners package. |
| src/physicalai/inference/model.py | Passes self.export_dir into get_runner so runner artifacts resolve like other manifest components. |
| src/physicalai/inference/component_factory.py | Registers the "two_stage" runner type in the component registry. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
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
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.
Summary
physicalai.inference.runners.TwoStage, a runner that executes two artifacts in sequence: the first stage's output feeds the second's input, and each stage can sit on a different backend.class_pathlike any other runner.tests/unit/inference/test_runners.py.Why
A vision-language-action policy is not always exportable as one graph. Concretely, a Qwen3-VL backbone under
transformers5.5.4 is capturable by neithertorch.exportnortorch.jit.trace, while the action expert that runs several times per chunk exports cleanly and benefits most from an optimized backend.Splitting execution recovers the value: the backbone runs on the Torch adapter, the action expert on the OpenVINO one, and
InferenceModelloads the directory unchanged. Measured end to end on such an export, the hybrid path matches the pure-Torch path to 7.9e-08 absolute, with the action expert 2.9x faster on CPU (2.6 ms vs 7.6 ms, AMD Threadripper 3960X).The runner is deliberately generic. It takes any two adapters, has no policy name in its logic, and knows nothing about the model that motivated it. An earlier draft added a
bind_export_dirhook toInferenceRunner; it was dropped once it became clear the existingresolve_artifactalready solved the problem. One new concept, in the place runners already live.Consumer side: open-edge-platform/physical-ai-studio PR adding the XR-1 policy, whose hybrid export writes a manifest naming this runner.
Validation
pytest tests/unit/inference/test_runners.py— new coverage for two-stage execution, including artifact resolution and stage wiring.pytest tests/unit— full unit suite passes.prek run --all-filesandpyrefly checkclean.InferenceModel(...)loading a two-artifact export, backbone on Torch and action expert on OpenVINO, 7.9e-08 against the single-backend export.Breaking changes
Related issues