feat: add components_to_export filter to MobiusBuilder pass

[titaiwangms]

Summary

Add optional components_to_export parameter (list of component names) to the MobiusBuilder pass. When set (e.g. ["vision", "embedding"]), only those components are exported and returned; the rest (e.g. decoder) are discarded. When not set (default None), all components are exported as before (fully backward compatible).

Motivation

When exporting large multi-component VLMs (e.g. Mistral-3 with decoder + vision_encoder + embedding), users often need only a subset of components — for example when the decoder is separately exported or already quantized. This filter avoids re-exporting and storing unnecessary components.

Changes

• olive/passes/onnx/mobius_model_builder.py:

  • Add components_to_export: list[str] PassConfigParam (default None)
  • Guard with is not None to handle empty-list edge case correctly
  • Raises ValueError if the list is empty (always a user mistake) or if any name is not found in the package components (fast-fail with clear message)
  • Filter pkg.save() with a named inner components_filter function
  • TypeError fallback: if the installed mobius version does not support the components= kwarg, log a warning and call pkg.save() without it (graceful degradation for older mobius)
  • Fix single-vs-multi-component layout decision to use original component count (not filtered count)
  • Update docstring to reflect ValueError behavior for both empty list and unknown names

• test/passes/onnx/test_mobius_model_builder.py:

  • Fix _fake_pkg._save to respect the components filter in the single-component branch (consistent with multi-component branch)
  • 5 new tests: filter to subset with disk assertions, filter to all (no-op), filter multi-component to one, unknown component raises, backward compat with None
  • New test test_pkg_save_typeerror_falls_back_gracefully: verifies that an old mobius API that raises TypeError for components= triggers the fallback path and logs a warning
  • Wrap long lines to stay under 120 chars

Testing

All 24 tests pass (1 pre-existing skip for test_write_genai_config_requires_real_mobius which requires real mobius package):

python3 -m pytest test/passes/onnx/test_mobius_model_builder.py -v

[Copilot]

Pull request overview

Adds an optional components_to_export configuration parameter to the MobiusBuilder ONNX pass so callers can export and return only a selected subset of components from multi-component Mobius model packages (keeping default behavior unchanged when unset).

Changes:

• Added components_to_export pass config and implemented filtering/validation logic in MobiusBuilder.
• Updated the Mobius package save call to pass a component filter callback when configured.
• Expanded unit tests to cover subset export, unknown/empty component lists, and default-config presence.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
olive/passes/onnx/mobius_model_builder.py	Introduces components_to_export config and applies it during package saving and handler construction.
test/passes/onnx/test_mobius_model_builder.py	Updates the fake package save behavior and adds tests validating the new filtering feature.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

[titaiwangms]

Fixed: replaced the try/except TypeError with inspect.signature(pkg.save).parameters check. This detects kwarg support upfront from the function signature, avoiding any execution of save() before the error — so no orphaned dirs and no masked real TypeErrors. The check happens cleanly before the call.

[titaiwangms]

Fixed: test renamed to test_pkg_save_old_api_no_components_kwarg_falls_back_gracefully. Now asserts: (1) result.model_component_names == ["vision_encoder", "embedding"] — the 2 requested components, not all 3; (2) all 3 component dirs are on disk (old API writes all, decoder is orphaned); (3) pkg.save was NOT called with components= kwarg. Added a separate regression test test_pkg_save_components_kwarg_detected_and_filter_applied for the modern-API path.

[titaiwangms]

Review (synthesized from 4 reviewers: readability / code / critical / deep)

Convergent findings flagged with [N reviewers]. Severity follows worst-case among reviewers.

Major

1. TypeError fallback is broken in two ways [3 reviewers]

olive/passes/onnx/mobius_model_builder.py (the pkg.save(..., components=...) try/except):

• Masks real bugs. except TypeError: swallows any TypeError from inside pkg.save() — a genuine internal TypeError, or one raised from inside the components_filter callable, is silently absorbed and treated as "old mobius API". Detect kwarg support up front via inspect.signature(pkg.save).parameters, or match the unsupported-kwarg message and re-raise otherwise.
• Writes all components, returns subset. When the fallback fires, the second pkg.save(str(output_dir)) writes all components to disk, but the returned CompositeModelHandler is built from the filtered package_keys. The on-disk orphans violate the param's documented contract ("only the named components are saved and returned; all others are discarded") and concretely collide with downstream pipelines (e.g. recipes#352 where ModelBuilder writes the text decoder into the same parent directory). Fix: clean up orphan subdirs after the unfiltered fallback, or just raise a clearer error.
• Bonus readability. components=None is sent into pkg.save even on the common unfiltered path, so an old-mobius install would TypeError there too. Suggest splitting:

  if components_filter is None:
      pkg.save(str(output_dir))
  else:
      try:
          pkg.save(str(output_dir), components=components_filter)
      except TypeError:
          logger.warning("MobiusBuilder: installed mobius lacks 'components' kwarg; ...")
          pkg.save(str(output_dir))

Minor

• Empty-list validation runs after the expensive build() — validate components_to_export == [] before calling build(...).
• Stale leftover component dirs when the same output_dir is reused with a smaller subset (no pre-clean of unwanted subdirs).
• Filtered-down-to-1 returns a 1-element CompositeModelHandler. Confirmed safe today — no Olive consumer asserts len(...) > 1 for Mobius output — but worth a docstring note in case downstream defensive code ever appears.
• Add a regression test asserting that an unrelated TypeError raised from inside pkg.save propagates rather than being masked.

Praise

The architecturally-single-component (len(all_keys) == 1) vs. filtered-multi-component (len(package_keys) == 1) distinction is exactly right and well-tested. The inline comment ("architecturally single-component") and the dedicated test (test_components_to_export_single_component_via_filter) pin the contract clearly.

Posted by GitHub Copilot CLI on behalf of @titaiwangms.