[Akshay Sonawane (apsonawane)]

Summary

Extend the Olive MobiusBuilder pass so recipes can provide a custom Mobius export callback through Olive's standard user_script mechanism.

This enables model-specific, multi-component Mobius workflows—such as CLM and KEV—to run through olive run while retaining standard Olive configuration, system, accelerator, caching, and output handling.

Changes

• Add optional exporter_func configuration to MobiusBuilder.
  • Accepts a callable directly or a function name resolved from user_script.
  • Receives the input model, Olive output directory, precision, execution provider, and recipe-specific configuration.
• Add optional exporter_config for model-specific callback settings.
• Add standard user_script and script_dir configuration support to MobiusBuilder.
• Allow recipe callbacks to write complete multi-component packages directly into the Olive output directory.
• Require callbacks to return a non-empty component-name list.
• Preserve package-level sidecar files through the returned CompositeModelHandler.
• Preserve component subdirectory layouts through no_flatten.
• Correctly handle custom exports containing a single component without assuming a root-level model.onnx.
• Reject unsupported combinations of exporter_func with text_only or components_to_export.
• Preserve all existing MobiusBuilder behavior when no custom exporter is configured.

Example

{
  "input_model": {
    "type": "HfModel",
    "model_path": "Qwen/Qwen3.5-4B-Base",
    "load_kwargs": {
      "revision": "1001bb4d826a52d1f399e183466143f4da7b741b",
      "torch_dtype": "float16"
    }
  },
  "systems": {
    "local_system": {
      "type": "LocalSystem",
      "accelerators": [
        {
          "device": "gpu",
          "execution_providers": ["CUDAExecutionProvider"]
        }
      ]
    }
  },
  "passes": {
    "m": {
      "type": "MobiusBuilder",
      "precision": "fp16",
      "exporter_func": "export_kev_package",
      "exporter_config": {
        "artifact_path": "artifacts/kev",
        "recipe_precision": "mixed"
      },
      "user_script": "../user_script.py",
      "script_dir": "../.."
    }
  },
  "target": "local_system",
  "output_dir": "build",
  "cache_dir": "cache",
  "no_artifacts": true
}

[Copilot]

Copilot review overview

🟡 Changes recommended

Package directories and generated layout metadata can be lost during downstream processing.

Review effort: Balanced
Findings: 2

Open (2)

• Root-level directories are omitted from generated packages · New
• Input attributes override generated package metadata · New

What changed in this PR

Adds recipe-defined Mobius export callbacks through Olive’s user-script mechanism.

Changes:

• Adds callback and recipe configuration support.
• Builds composite handlers for custom package layouts.
• Tests direct and user-script callback resolution.

File	Description
olive/​passes/​onnx/​mobius_model_builder.py	Implements custom package export handling.
test/​passes/​onnx/​test_mobius_model_builder.py	Covers callback configuration and package creation.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot review overview

🔵 Needs a closer look

The implementation rejects valid non-dict mapping results despite documenting a general mapping contract.

Review effort: Balanced
Findings: None

Resolved since last review (2)

• Input attributes override generated package metadata
• Root-level directories are omitted from generated packages

Previously missed (1)

In code that hasn't changed since last review

Accept Mapping implementations instead of only dict instances

olive/​passes/​onnx/​mobius_model_builder.py:234

The documented callback contract says this result may be a “mapping,” but this check accepts only concrete dict instances. A valid collections.abc.Mapping implementation (for example, UserDict or a recipe-specific immutable mapping) is therefore rejected despite supporting get and indexing exactly as the following code requires. Check against Mapping (and import it alongside Callable) so the implementation matches the public contract.

[Copilot]

Copilot review overview

🔵 Needs a closer look

Child handlers can retain stale package metadata that overrides the corrected composite attributes.

Review effort: Balanced
Findings: None

[Xiaoyu Z (xiaoyu-work)]

can you specify a use case for this?

[Akshay Sonawane (apsonawane)]

can you specify a use case for this?

A concrete use case is olive-recipes#646, which exports CLM and KEV multi-component packages through olive run.

These models require recipe-specific orchestration beyond the standard mobius.build() path: loading pinned base and head artifacts, merging KEV LoRA weights, applying strict component-specific weights, and assembling provider-specific package manifests and sidecars.

We do not want to put that model-specific logic in Olive core. exporter_func, resolved through Olive’s existing user_script mechanism, keeps the implementation in the recipe while Olive still owns accelerator selection, pass execution, caching, and output handling. The callback contract itself remains generic and can support other specialized multi-component Mobius exports.