Support recipe-defined Mobius exports

olive/passes/onnx/mobius_model_builder.py

@@ -8,6 +8,7 @@
 
 import json
 import logging
+from collections.abc import Callable
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -18,6 +19,7 @@
 from olive.model.handler.composite import CompositeModelHandler
 from olive.passes import Pass
 from olive.passes.olive_pass import PassConfigParam
+from olive.passes.pass_config import ParamCategory, get_user_script_data_config
 
 if TYPE_CHECKING:
     from olive.hardware.accelerator import AcceleratorSpec
@@ -107,6 +109,7 @@ def is_accelerator_agnostic(cls, accelerator_spec: AcceleratorSpec) -> bool:
     @classmethod
     def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]:
         return {
+            **get_user_script_data_config(),
             "precision": PassConfigParam(
                 type_=Precision,
                 required=False,
@@ -152,6 +155,24 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
                     "the model's components."
                 ),
             ),
+            "exporter_func": PassConfigParam(
+                type_=str | Callable,
+                required=False,
+                default_value=None,
+                category=ParamCategory.OBJECT,
+                description=(
+                    "Optional recipe-local function that exports a Mobius package "
+                    "directly to output_model_path and returns a mapping containing "
+                    "a non-empty 'components' list. The function receives model, "
+                    "output_dir, precision, execution_provider, and exporter_config."
+                ),
+            ),
+            "exporter_config": PassConfigParam(
+                type_=dict,
+                required=False,
+                default_value=None,
+                description="Configuration forwarded to exporter_func.",
+            ),
         }
 
     def _run_for_config(
@@ -160,20 +181,20 @@ def _run_for_config(
         config: type[BasePassConfig],
         output_model_path: str,
     ) -> ONNXModelHandler | CompositeModelHandler:
+        try:
+            import mobius
+        except ImportError as exc:
+            raise ImportError(
+                "mobius-onnx is required to run MobiusBuilder. Install with: pip install mobius-onnx"
+            ) from exc
+
         if not isinstance(model, HfModelHandler):
             raise ValueError(f"MobiusBuilder requires an HfModelHandler input, got {type(model).__name__}.")
 
         from olive.common.quant.hf_utils import validate_olive_onnx_export
 
         validate_olive_onnx_export(model.get_hf_model_config().to_dict().get("quantization_config") or {})
 
-        try:
-            from mobius import build
-        except ImportError as exc:
-            raise ImportError(
-                "mobius-onnx is required to run MobiusBuilder. Install with: pip install mobius-onnx"
-            ) from exc
-
         # Map Olive EP to mobius EP. If unsupported/unknown, fall back to mobius default EP.
         requested_ep = self.accelerator_spec.execution_provider
         ep_str: str = EXECUTION_PROVIDER_TO_MOBIUS_EP.get(requested_ep, self.MobiusEP.DEFAULT)
@@ -187,7 +208,6 @@ def _run_for_config(
 
         dtype_str: str = _PRECISION_TO_DTYPE.get(config.precision, "f32")
         model_id: str = model.model_name_or_path
-
         load_kwargs = model.get_load_kwargs()
         revision: str | None = load_kwargs.get("revision")
         trust_remote_code: bool = load_kwargs.get("trust_remote_code", False)
@@ -213,82 +233,90 @@ def _run_for_config(
 
         output_dir = Path(output_model_path)
         output_dir.mkdir(parents=True, exist_ok=True)
-
-        text_only_kwargs = {"text_only": True} if config.text_only else {}
-        revision_kwargs = {"revision": revision} if revision is not None else {}
-        pkg = build(
-            model_id,
-            dtype=dtype_str,
-            execution_provider=ep_str,
-            load_weights=True,
-            trust_remote_code=trust_remote_code,
-            **text_only_kwargs,
-            **revision_kwargs,
-        )
-
-        # Determine which package components to export.
-        all_keys = list(pkg.keys())
-        if config.components_to_export is not None:
-            requested = set(config.components_to_export)
-            unknown = requested - set(all_keys)
-            if unknown:
+        custom_exporter = config.exporter_func is not None
+        if custom_exporter:
+            if config.components_to_export is not None or config.text_only:
                 raise ValueError(
-                    f"MobiusBuilder: components_to_export contains unknown component(s): {sorted(unknown)}. "
-                    f"Available components from this model: {sorted(all_keys)}"
+                    "MobiusBuilder: exporter_func cannot be combined with components_to_export or text_only."
                 )
-            package_keys = [k for k in all_keys if k in requested]
-            logger.info(
-                "MobiusBuilder: exporting subset of components %s (skipping %s)",
-                package_keys,
-                [k for k in all_keys if k not in requested],
+            exporter = self._user_module_loader.load_object(config.exporter_func)
+            result = exporter(
+                model=model,
+                output_dir=output_dir,
+                precision=config.precision,
+                execution_provider=ep_str,
+                exporter_config=config.exporter_config or {},
             )
-
-            def components_filter(name: str) -> bool:
-                return name in requested
+            if not isinstance(result, dict) or not isinstance(result.get("components"), list):
+                raise ValueError("MobiusBuilder: exporter_func must return a mapping with a 'components' list.")
+            package_keys = result["components"]
+            if not package_keys or not all(isinstance(name, str) and name for name in package_keys):
+                raise ValueError("MobiusBuilder: exporter_func returned invalid components.")
+            all_keys = package_keys
+            component_paths = {output_dir / name for name in package_keys}
+            genai_artifacts = {path.name: str(path) for path in output_dir.iterdir() if path not in component_paths}
         else:
-            package_keys = all_keys
-            components_filter = None
-
-        # ModelPackage.save() handles both single and multi-component layouts:
-        #   single component  → <output_dir>/model.onnx
-        #   multi-component   → <output_dir>/<name>/model.onnx  for each key
-        pkg.save(str(output_dir), components=components_filter)
-
-        # ORT GenAI config generation assumes every component in `pkg` was actually saved to
-        # disk (e.g. it unconditionally writes a "decoder/model.onnx" filename reference for
-        # multimodal packages). For a partial export (components_to_export set), that produces
-        # a genai_config.json/tokenizer set that references components we deliberately omitted
-        # — invalid artifacts, since mobius has no API to generate GenAI config for a subset of
-        # a package. Skip config generation entirely in that case and let the caller (e.g. a
-        # recipe combining this partial export with another tool's output) assemble the final
-        # genai_config.json itself.
-        if components_filter is not None:
-            logger.info(
-                "MobiusBuilder: components_to_export is set; skipping ORT GenAI config generation "
-                "since it cannot be scoped to a subset of components. The caller is responsible for "
-                "producing a genai_config.json that covers the full pipeline."
-            )
-            genai_artifacts = {}
-        else:
-            # Generate ORT GenAI config artifacts (genai_config.json, tokenizer
-            # files, processor configs) alongside the ONNX models.
-            genai_artifacts = self._write_genai_config(
-                pkg,
-                str(output_dir),
+            text_only_kwargs = {"text_only": True} if config.text_only else {}
+            revision_kwargs = {"revision": revision} if revision is not None else {}
+            pkg = mobius.build(
                 model_id,
-                ep_str,
-                revision=revision,
+                dtype=dtype_str,
+                execution_provider=ep_str,
+                load_weights=True,
                 trust_remote_code=trust_remote_code,
+                **text_only_kwargs,
+                **revision_kwargs,
             )
-            if config.genai_config:
-                self._update_genai_config(output_dir, config.genai_config)
+
+            all_keys = list(pkg.keys())
+            if config.components_to_export is not None:
+                requested = set(config.components_to_export)
+                unknown = requested - set(all_keys)
+                if unknown:
+                    raise ValueError(
+                        "MobiusBuilder: components_to_export contains unknown "
+                        f"component(s): {sorted(unknown)}. Available components "
+                        f"from this model: {sorted(all_keys)}"
+                    )
+                package_keys = [key for key in all_keys if key in requested]
+                logger.info(
+                    "MobiusBuilder: exporting subset of components %s (skipping %s)",
+                    package_keys,
+                    [key for key in all_keys if key not in requested],
+                )
+
+                def components_filter(name: str) -> bool:
+                    return name in requested
+            else:
+                package_keys = all_keys
+                components_filter = None
+
+            pkg.save(str(output_dir), components=components_filter)
+            if components_filter is not None:
+                logger.info(
+                    "MobiusBuilder: components_to_export is set; skipping ORT "
+                    "GenAI config generation since it cannot be scoped to a "
+                    "subset of components."
+                )
+                genai_artifacts = {}
+            else:
+                genai_artifacts = self._write_genai_config(
+                    pkg,
+                    str(output_dir),
+                    model_id,
+                    ep_str,
+                    revision=revision,
+                    trust_remote_code=trust_remote_code,
+                )
+                if config.genai_config:
+                    self._update_genai_config(output_dir, config.genai_config)
 
         logger.info("MobiusBuilder: saved components %s to '%s'", package_keys, output_dir)
 
         # Use the single-component (root layout) path only when the model is
         # architecturally single-component.  A multi-component model filtered
         # down to one component still uses component sub-directories on disk.
-        if len(all_keys) == 1:
+        if len(all_keys) == 1 and not custom_exporter:
             # Single-component model (most LLMs): return a plain ONNXModelHandler.
             onnx_path = output_dir / "model.onnx"
             if not onnx_path.exists():
@@ -305,9 +333,9 @@ def components_filter(name: str) -> bool:
                 model_path=str(output_dir),
                 onnx_file_name="model.onnx",
                 model_attributes={
+                    **(model.model_attributes or {}),
                     "mobius_package_keys": package_keys,
                     "additional_files": additional_files,
-                    **(model.model_attributes or {}),
                 },
             )
 
@@ -336,9 +364,9 @@ def components_filter(name: str) -> bool:
                     model_path=str(component_dir),
                     onnx_file_name="model.onnx",
                     model_attributes={
+                        **(model.model_attributes or {}),
                         "mobius_component": key,
                         "additional_files": component_additional_files,
-                        **(model.model_attributes or {}),
                     },
                 )
             )
@@ -348,6 +376,7 @@ def components_filter(name: str) -> bool:
             model_component_names=package_keys,
             model_path=str(output_dir),
             model_attributes={
+                **(model.model_attributes or {}),
                 "mobius_package_keys": package_keys,
                 # Preserve the <component>/model.onnx subdirectory layout so
                 # ORT GenAI can resolve each component by its "filename" key.
@@ -358,7 +387,6 @@ def components_filter(name: str) -> bool:
                 # so they end up at the package root (alongside genai_config.json),
                 # not duplicated into each <component>/ subdirectory.
                 "additional_files": sorted(set(genai_artifacts.values())),
-                **(model.model_attributes or {}),
             },
         )
 

test/passes/onnx/test_mobius_model_builder.py

@@ -78,11 +78,16 @@ def _make_hf_model(model_path: str, load_kwargs: dict | None = None, task: str |
     return model
 
 
-def _make_pass(ep: str = ExecutionProvider.CPUExecutionProvider, text_only: bool | None = None) -> MobiusBuilder:
+def _make_pass(
+    ep: str = ExecutionProvider.CPUExecutionProvider,
+    text_only: bool | None = None,
+    **extra_config,
+) -> MobiusBuilder:
     accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ep)
     pass_config = {"precision": "fp32"}
     if text_only is not None:
         pass_config["text_only"] = text_only
+    pass_config.update(extra_config)
     return create_pass_from_dict(
         MobiusBuilder,
         pass_config,
@@ -293,10 +298,92 @@ def test_default_config_params():
     assert "precision" in config
     assert config["text_only"].default_value is False
     assert config["text_only"].required is False
+    assert config["exporter_func"].default_value is None
+    assert config["exporter_config"].default_value is None
+    assert "user_script" in config
+    assert "script_dir" in config
     assert "execution_provider" not in config
     assert "trust_remote_code" not in config
 
 
+def test_recipe_exporter_function_creates_composite_package(tmp_path):
+    """A recipe callback can own multi-component package export."""
+
+    def exporter(**kwargs):
+        output = Path(kwargs["output_dir"])
+        for name in ("backbone", "pointer_head"):
+            (output / name).mkdir(parents=True)
+            (output / name / "model.onnx").write_text("dummy")
+        (output / "assets").mkdir()
+        (output / "assets" / "metadata.json").write_text("{}")
+        (output / "component_manifest.json").write_text("{}")
+        assert kwargs["exporter_config"] == {"artifact_path": "artifact"}
+        return {"components": ["backbone", "pointer_head"]}
+
+    p = _make_pass(
+        exporter_func=exporter,
+        exporter_config={"artifact_path": "artifact"},
+    )
+    stale_file = tmp_path / "stale.txt"
+    stale_file.write_text("stale")
+    model = _make_hf_model("Qwen/Qwen3.5-4B-Base")
+    model.model_attributes = {
+        "additional_files": [str(stale_file)],
+        "mobius_package_keys": ["stale"],
+        "no_flatten": False,
+        "preserved": True,
+    }
+    result = p.run(model, tmp_path / "out")
+
+    assert isinstance(result, CompositeModelHandler)
+    assert result.model_component_names == ["backbone", "pointer_head"]
+    assert result.model_attributes == {
+        "additional_files": [
+            str(tmp_path / "out" / "assets"),
+            str(tmp_path / "out" / "component_manifest.json"),
+            str(tmp_path / "out" / "stale.txt"),
+        ],
+        "mobius_package_keys": ["backbone", "pointer_head"],
+        "no_flatten": True,
+        "preserved": True,
+    }
+    for name, component in result.get_model_components():
+        assert component.model_attributes["additional_files"] == []
+        assert component.model_attributes["mobius_component"] == name
+        assert component.model_attributes["preserved"] is True
+
+
+def test_recipe_exporter_function_requires_component_list(tmp_path):
+    """Invalid recipe callback results fail before handler construction."""
+    p = _make_pass(exporter_func=lambda **_: {})
+    with pytest.raises(ValueError, match="'components' list"):
+        p.run(_make_hf_model("Qwen/Qwen3-8B"), tmp_path / "out")
+
+
+def test_recipe_exporter_function_loads_from_user_script(tmp_path):
+    """Recipe configs can resolve their exporter callback by name."""
+    script = tmp_path / "user_script.py"
+    script.write_text(
+        """
+from pathlib import Path
+
+def export_package(**kwargs):
+    output = Path(kwargs["output_dir"])
+    (output / "encoder").mkdir(parents=True)
+    (output / "encoder" / "model.onnx").write_text("dummy")
+    return {"components": ["encoder"]}
+"""
+    )
+    p = _make_pass(
+        exporter_func="export_package",
+        user_script=str(script),
+    )
+    result = p.run(_make_hf_model("Qwen/Qwen3-8B"), tmp_path / "out")
+
+    assert isinstance(result, CompositeModelHandler)
+    assert result.model_component_names == ["encoder"]
+
+
 def test_is_not_accelerator_agnostic():
     """Pass must be EP-specific because it chooses fused ops based on the EP."""
     accelerator_spec = AcceleratorSpec(