Migrate onnx-genai metadata emitters to the current pipeline/speculative contract - #554
Merged
justinchuby merged 3 commits intoAug 23, 2026
Merged
Conversation
|
|
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR migrates Mobius’s onnx-genai metadata emitters and conformance tests to match the runtime’s current published JSON schema, particularly the typed SSA pipeline.workflow (PipelineSpec) and the new SpeculativeContract, and makes schema validation deterministic in CI by vendoring the schema.
Changes:
- Replace legacy diffusion/VLM/MTP metadata shapes with the current
pipeline.workflowcontract and updated speculative metadata contract. - Vendor the onnx-genai
inference_metadata.schema.jsonundersrc/mobius/integrations/onnx_genai/_schema/and update tests to always validate against it (withONNX_GENAI_SCHEMAoverride support). - Update ComfyUI conversion to emit the new workflow-based diffusion metadata and ship required policy component artifacts.
Reviewed changes
Copilot reviewed 13 out of 14 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py | Stops skipping schema validation; uses vendored/override schema path helper. |
| src/mobius/integrations/onnx_genai/inference_metadata.py | Major contract migration: diffusion emits typed SSA workflow + policy components; MTP speculative metadata moved to SpeculativeContract; VLM writer publishes workflow document. |
| src/mobius/integrations/onnx_genai/inference_metadata_test.py | Updates schema lookup to vendored default; rewrites diffusion/MTP/VLM tests for new contracts and adds new regression assertions. |
| src/mobius/integrations/onnx_genai/duplex_workflow_metadata_test.py | Stops skipping schema validation; uses vendored/override schema path helper. |
| src/mobius/integrations/onnx_genai/decoder_metadata_test.py | Stops skipping schema validation; uses vendored/override schema path helper. |
| src/mobius/integrations/onnx_genai/convert.py | Ensures diffusion workflow metadata includes VAE scaling factor and saves generated policy component artifacts alongside metadata. |
| src/mobius/integrations/onnx_genai/convert_test.py | Updates assertions to the workflow-based diffusion schema and component model. |
| src/mobius/integrations/onnx_genai/comfyui.py | Tightens ComfyUI parsing (reject latent-only) and translates to workflow-based diffusion metadata. |
| src/mobius/integrations/onnx_genai/comfyui_test.py | Updates ComfyUI translation tests to the workflow contract and new failure-closed behaviors. |
| src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py | Stops skipping schema validation; uses vendored/override schema path helper. |
| src/mobius/integrations/onnx_genai/_schema/README.md | Documents why the schema is vendored and how to update/override it. |
| src/mobius/integrations/onnx_genai/init.py | Updates diffusion metadata docstring to reflect workflow-based emission and shipped policy artifacts. |
| pyproject.toml | Packages the vendored schema JSON in wheel/sdist via package-data. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
…ive contract
mobius's legacy onnx-genai emitters had drifted from two upstream contract
redesigns, and the drift was invisible because every schema conformance test
skipped in CI.
Why it stayed hidden: the conformance tests searched a couple of hard-coded
local onnx-genai checkout paths and `pytest.skip`ped when none existed. CI has
no such checkout, so six tests never ran there, and on the machines that did
have one the result depended on whatever revision that clone sat on.
The two contract redesigns:
1. `pipeline` is now `PipelineSpec`, whose only property is `workflow` (a typed
SSA graph) with `additionalProperties: false`. The legacy emitters produced
`{models, dataflow, strategy, phases}`, which the schema now rejects
outright.
2. `speculative` is now `SpeculativeContract`, requiring
`{proposer, target, vocabulary, max_proposal_width}` and forbidding every
field of the old flat MTP block. The flat `SpeculatorConfig` that block was
modelled on still exists upstream, but it describes a HuggingFace
`config.json` section, not `InferenceMetadata.speculative`.
What changed:
* Diffusion (`build_diffusion_pipeline_metadata`). Now emits the denoise loop
as an explicit workflow: the solver, schedule, timestep table, guidance
combine and output clamp are real ONNX components built from mobius's policy
library and shipped with the document. It cannot delegate to
`build_diffusion_workflow_metadata`, because the ComfyUI conversion path has
no component graphs at all -- it deliberately does not build or export them
-- so it builds the workflow directly while reusing that module's `_invoke`
and `_publish_workflow_v1`. The schedule is derived from the scheduler's own
betas via the same diffusers-compatible helpers the package exporter uses, so
a ComfyUI conversion and a package export of one checkpoint describe the same
dynamics. img2img's `start_step` lowers to a sliced schedule.
Three things now fail closed rather than being silently mis-described: an
ancestral sampler (no deterministic solver exists), Karras/exponential sigma
spacing (the workflow ships the sigma table as a constant, so a hint field is
not enough), and a latent-only graph with no VAE decode.
* MTP (`write_mtp_speculator_metadata`). Now emits a `SpeculativeContract`
anchored to the backbone's workflow: it registers the head as a workflow
component, names the target by its declared `logits` port role, states the
hidden handoff as `port_bindings.target_hidden_context` plus a
`hidden_states` role on the target output, and completes the rollback
capacity its own claim requires. It is `block` rather than `chained` because
a chained proposer must expose a `logits_output` and this sidecar emits only
`mtp_hidden` -- the runtime decodes it through the shared LM head, which is
why that initializer is in `shared_weights`.
* VLM. `build_native_vlm_package_metadata` produces mobius's internal
structural descriptor, not a publishable document -- `build_vlm_workflow_metadata`
already consumes it and republishes only `preprocessing` under a real
`pipeline.workflow`. The bug was that `write_native_vlm_package_metadata`
wrote the descriptor to `inference_metadata.yaml`; it now writes the workflow
document, and the tests validate that published document instead.
* CI visibility. The schema is vendored under `_schema/` and is the default, so
conformance never skips and drift is a test failure. A local checkout is no
longer consulted implicitly, since one that is ahead of or behind `main`
reintroduces the same machine-dependent result; set `ONNX_GENAI_SCHEMA` to
validate against a specific revision. Three further conformance tests in the
codec, speech-to-text and duplex suites were skipping for the same reason and
now run.
Tests updated in step: the MTP tests assert the new contract (with the
banned-name test repointed at the now-legacy field names) and are anchored to a
workflow mobius actually emits, and the diffusion/ComfyUI tests assert workflow
structure rather than the removed `strategy` block.
Full suite: 8 failed / 4613 passed -> 1 failed / 4640 passed. The remaining
failure (`qwen_image_test.py::test_deterministic_l4_l5_image_edit_golden`) is
pre-existing and handled separately.
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Drop the before/after narration comparing against the superseded pipeline and speculative shapes; the docstrings describe the contract the code emits today. The banned-name test keeps its list of now-legacy fields, since rejecting them is the point of that test. Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
justinchuby
force-pushed
the
justinchuby-fix-onnx-genai-metadata-schema-drift
branch
from
August 23, 2026 01:46
cd1f69e to
e5a00bf
Compare
The translated document declares the sampler as executable components under ``policies/*.onnx``, but ``parse_comfyui_workflow`` built those graphs and threw them away, so a caller holding only the metadata could not produce a loadable package. ``ComfyUIWorkflow`` now carries them and exposes ``save_policy_components``; the dict-only wrappers point at it. ``convert_comfyui_workflow`` keeps building its own package rather than reusing the parse-time one, because the checkpoint's scheduler config can reconcile to a different solver than the ComfyUI sampler implied -- reusing it would leave that run's components behind for a document that never references them. That is now stated where someone would otherwise "simplify" it, and covered by a test asserting the written set equals the referenced set exactly. Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
justinchuby
deleted the
justinchuby-fix-onnx-genai-metadata-schema-drift
branch
August 23, 2026 02:02
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.
Fixes six long-standing schema-conformance failures caused by mobius's legacy onnx-genai emitters drifting from the runtime's published JSON schema.
Why the drift was invisible
The conformance tests searched a couple of hard-coded local onnx-genai checkout paths and
pytest.skipped when none was found. CI has no such checkout, so six tests never ran there; on the machines that did have one, the result depended on whatever revision that clone happened to sit on. Two upstream contract redesigns accumulated unnoticed.The two contract redesigns
pipelineis nowPipelineSpec— a single property,workflow(a typed SSA graph), withadditionalProperties: false. The legacy emitters produced{models, dataflow, strategy, phases}, all of which the schema now rejects.speculativeis nowSpeculativeContract— requires{proposer, target, vocabulary, max_proposal_width}and forbids every field of the old flat MTP block. The flatSpeculatorConfigthat block was modelled on still exists upstream, but it describes a HuggingFaceconfig.jsonsection, notInferenceMetadata.speculative.What was migrated
Diffusion —
build_diffusion_pipeline_metadataNow emits the denoise loop explicitly: the solver, sigma/alpha schedule, timestep table, guidance combine and output clamp are real ONNX components built from mobius's policy library and shipped with the document.
It cannot delegate to
build_diffusion_workflow_metadata: the ComfyUI conversion path has no component graphs at all (it deliberately does not build or export them), and that builder derives everything from liveir.Models. So the workflow is built directly, reusingworkflow_metadata's_invoke/_publish_workflow_v1andmobius.generation's solver builders. The schedule is derived from the scheduler's own betas through the same diffusers-compatible helpers the package exporter uses, so a ComfyUI conversion and a package export of the same checkpoint describe the same dynamics. img2img'sstart_steplowers to a sliced schedule; the VAEscaling_factorand the sigma-space initial-state scale are emitted as explicit components.Three cases now fail closed instead of being silently mis-described:
use_karras_sigmashint is no longer sufficient),MTP —
write_mtp_speculator_metadataEmits a
SpeculativeContractanchored to the backbone's workflow: it registers the head as a workflow component, names the target by its declaredlogitsport role, states the hidden handoff asport_bindings.target_hidden_contextplus ahidden_statesrole on the target output, and completes the rollback capacity its own claim requires.It is declared
block, notchained: a chained proposer must expose alogits_outputcarrying the next-token distribution, and this sidecar emits onlymtp_hidden— the runtime decodes it through the target's shared LM head, which is why that initializer is listed inshared_weights.VLM
build_native_vlm_package_metadataproduces mobius's internal structural descriptor, not a publishable document —build_vlm_workflow_metadataalready consumes it and republishes onlypreprocessingunder a realpipeline.workflow. The actual bug was thatwrite_native_vlm_package_metadatawrote that descriptor toinference_metadata.yaml. It now writes the workflow document, and the tests validate the published document instead of the descriptor.CI visibility
The upstream schema is vendored under
src/mobius/integrations/onnx_genai/_schema/and is the default, so conformance never skips and drift becomes a test failure. A local checkout is no longer consulted implicitly — one that is ahead of or behindmainreintroduces exactly the machine-dependent result that hid this bug. SetONNX_GENAI_SCHEMAto validate against a specific revision.Three further conformance tests (codec, speech-to-text, duplex workflows) were skipping for the same reason and now run.
Tests
Migrated in the same change rather than loosened:
TestMtpSpeculatorMetadata::test_exact_schema_keys_and_valuesasserts the new contract;test_no_legacy_field_namesis repointed at the now-legacy field names.write_decoder_workflow_metadataactually emits (which ships ~11 generated policy components), so target selection can't regress to "the only ONNX component".strategyblock; scheduler facts that were duplicated into the document are asserted on the parsedComfyUIWorkflowwhere they actually live.Verification
Rebased onto current
main(includes #552 and #553).mainat branch point)src/mobius/integrations/onnx_genai/withONNX_GENAI_SCHEMAat upstreammainSix of those eight failures are the schema-drift ones fixed here. The other two —
qwen_image_test.py::test_deterministic_l4_l5_image_edit_goldenandTestNativeVlmPackageMetadata::test_cached_gemma_processor_matches_emitted_patch_budget— were fixed by #553, which is now merged; this branch preserves both of its changes through the rebase.lintrunneris clean.Sequencing note
workflow_metadata.pyandauto_export.pyare not modified (both are read from only), so this does not conflict with #551.