Skip to content

feat(export): support exporting sub-models from composite models#1031

Merged
xieofxie merged 6 commits into
mainfrom
hualxie/export_submodule
Jul 10, 2026
Merged

feat(export): support exporting sub-models from composite models#1031
xieofxie merged 6 commits into
mainfrom
hualxie/export_submodule

Conversation

@xieofxie

@xieofxie xieofxie commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Add --submodel option to winml export for selecting a specific sub-model from composite models (e.g., encoder, decoder).

When a composite model is detected:

  • --submodel : exports only that component as {stem}_{name}.onnx
  • No --submodel: exports all sub-models automatically
  • Non-composite models: unchanged single-model export

Extract _export_single_model() helper shared by both paths and add _resolve_composite_model_components() for registry/auto-detection.

Resolves #1014

hualxie added 5 commits July 2, 2026 17:34
Add --submodel option to `winml export` for selecting a specific
sub-model from composite models (e.g., encoder, decoder).

When a composite model is detected:
- --submodel <name>: exports only that component as {stem}_{name}.onnx
- No --submodel: exports all sub-models automatically
- Non-composite models: unchanged single-model export

Extract _export_single_model() helper shared by both paths and add
_resolve_composite_model_components() for registry/auto-detection.

Resolves #1014
@xieofxie xieofxie marked this pull request as ready for review July 10, 2026 02:32
@xieofxie xieofxie requested a review from a team as a code owner July 10, 2026 02:32

@DingmaomaoBJTU DingmaomaoBJTU left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Took a look — small, focused feature and the tests cover the happy path plus the two error cases. A few things worth addressing before merge: a test-hermeticity concern (the new --submodel docstring example is auto-run and uniquely depends on live Hub resolution), an exception-type consistency nit, and a --submodel + --input-specs dead end. Details inline.

Also, the PR description mentions extracting _export_single_model() and adding _resolve_composite_model_components(), but neither is in the diff — the code uses the existing _run_component_export() / resolve_composite_components(). Worth updating the description so it matches what actually changed.

Comment thread src/winml/modelkit/commands/export.py
Comment thread src/winml/modelkit/commands/export.py Outdated
Comment thread src/winml/modelkit/commands/export.py
Comment thread src/winml/modelkit/commands/export.py
- --submodel help now notes output is written to {stem}_{name}.onnx
- --submodel validation raises BadParameter (param_hint) for consistent
  usage errors/exit codes instead of ClickException
- allow --input-specs when --submodel narrows to a single sub-model;
  reword the multi-sub-model rejection to point at --submodel
- patch resolve_composite_components in test_export_help_examples_run and
  assert it is called, so the --submodel example no longer hits the Hub
- add test_submodel_allows_input_specs

@DingmaomaoBJTU DingmaomaoBJTU left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-reviewed the fix(export): address PR #1031 review comments commit (3bf05092). All four earlier points are resolved correctly:

  1. Test hermeticityresolve_composite_components is now patched in test_export_help_examples_run (with a .called assertion), so the --submodel example no longer depends on the live Hub. Nice bonus: this makes all the docstring examples hermetic, not just the composite one.
  2. Exception type — both --submodel validations now raise click.BadParameter(param_hint="--submodel"), consistent with the rest of the command (usage error + exit code 2).
  3. --submodel + --input-specs dead end — the specs are now allowed for a single selected sub-model (input_specs and submodel is None guard), and the multi-component message points at --submodel instead of the --task dead end. I checked test_composite_rejects_input_specs still holds (single-component mock, no --submodel → still rejected, message still contains "composite").
  4. Help text — now documents the {stem}_{name}.onnx output.

One follow-up inline (test-strength, non-blocking). Also very minor: the reworded rejection says "exporting multiple sub-models", but the guard fires for any composite without --submodel; it reads slightly off against the single-component mock in test_composite_rejects_input_specs (real composites are all ≥2, so it's only cosmetic). LGTM once the test point is considered.

Comment thread tests/unit/commands/test_export.py

@DingmaomaoBJTU DingmaomaoBJTU left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving. All four points from the earlier review are addressed correctly and I confirmed there is no regression (test_composite_rejects_input_specs still holds; changes are within the 100-char line limit). The one remaining inline note (strengthening test_submodel_allows_input_specs to assert the specs actually apply) is non-blocking and can be picked up as a follow-up.

Note: I could not execute the suite locally (win_arm64 has no torch wheel), so this is based on code-path analysis; please make sure CI is green before merging.

@xieofxie xieofxie merged commit defa3f6 into main Jul 10, 2026
9 checks passed
@xieofxie xieofxie deleted the hualxie/export_submodule branch July 10, 2026 06:23
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.

feat: Composite model: support export, build, perf etc. one or any of it sub models

2 participants