From 89a3410bd3f7338f484d33e6aefa0738302b94cb Mon Sep 17 00:00:00 2001
From: titaiwangms <titaiwang@microsoft.com>
Date: Fri, 8 May 2026 17:20:39 +0000
Subject: [PATCH 1/8] feat: add components_to_export filter to MobiusBuilder
 pass

Add an optional `components_to_export` parameter to MobiusBuilder that
allows callers to select a subset of components from a multi-component
model (e.g. only vision + embedding, skipping the decoder).

When set, pkg.save() is called with a components predicate that filters
out unneeded components before writing to disk, avoiding unnecessary I/O
for large model files.  When not set, all components are exported
(backward-compatible default).

Multi-component models filtered to one component still return a
CompositeModelHandler (sub-directory layout) rather than a plain
ONNXModelHandler (root layout), since the save layout is determined by
the model's architecture, not the filter result.

Add 5 new unit tests covering: subset filter, default exports-all,
single-component-via-filter, unknown-component error, and config defaults.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/passes/onnx/mobius_model_builder.py     |  42 ++++++-
 test/passes/onnx/test_mobius_model_builder.py | 103 ++++++++++++++++++
 2 files changed, 142 insertions(+), 3 deletions(-)

diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py
index 7b7e52bb3..e408ca6bf 100644
--- a/olive/passes/onnx/mobius_model_builder.py
+++ b/olive/passes/onnx/mobius_model_builder.py
@@ -107,6 +107,19 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
                     "configs alongside the ONNX models. 'none' to skip."
                 ),
             ),
+            "components_to_export": PassConfigParam(
+                type_=list,
+                required=False,
+                default_value=None,
+                description=(
+                    "Optional list of component names to export from a multi-component model "
+                    "(e.g. ['vision', 'embedding'] to skip the decoder). "
+                    "When set, only the named components are saved and returned; "
+                    "all others are discarded after the mobius build step. "
+                    "When not set, all components are exported (default, backward compatible). "
+                    "Has no effect on single-component models."
+                ),
+            ),
         }
 
     def _run_for_config(
@@ -163,10 +176,31 @@ def _run_for_config(
             trust_remote_code=trust_remote_code,
         )
 
+        # Determine which package components to export.
+        all_keys = list(pkg.keys())
+        if config.components_to_export:
+            requested = set(config.components_to_export)
+            unknown = requested - set(all_keys)
+            if unknown:
+                raise ValueError(
+                    f"MobiusBuilder: components_to_export contains unknown component(s): {sorted(unknown)}. "
+                    f"Available components from this model: {sorted(all_keys)}"
+                )
+            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],
+            )
+            components_filter = lambda name: name in requested  # noqa: E731
+        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))
+        pkg.save(str(output_dir), components=components_filter)
 
         # Generate ORT GenAI config artifacts (genai_config.json, tokenizer
         # files, processor configs) when runtime is set to ort-genai.
@@ -174,10 +208,12 @@ def _run_for_config(
         if config.runtime == self.MobiusRuntime.ORT_GENAI:
             genai_artifacts = self._write_genai_config(pkg, str(output_dir), model_id, ep_str)
 
-        package_keys = list(pkg.keys())
         logger.info("MobiusBuilder: saved components %s to '%s'", package_keys, output_dir)
 
-        if len(package_keys) == 1:
+        # 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:
             # Single-component model (most LLMs): return a plain ONNXModelHandler.
             onnx_path = output_dir / "model.onnx"
             if not onnx_path.exists():
diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py
index cd3d7338d..851019967 100644
--- a/test/passes/onnx/test_mobius_model_builder.py
+++ b/test/passes/onnx/test_mobius_model_builder.py
@@ -454,3 +454,106 @@ def test_no_warning_when_trust_remote_code_false(tmp_path):
 
     warning_messages = [call.args[0] for call in mock_logger.warning.call_args_list]
     assert not any("trust_remote_code" in msg for msg in warning_messages)
+
+
+# ---------------------------------------------------------------------------
+# components_to_export filter tests
+# ---------------------------------------------------------------------------
+
+
+def test_components_to_export_filters_subset(tmp_path):
+    """Only requested components are saved and returned when components_to_export is set."""
+    out = tmp_path / "out"
+    keys = ["decoder", "vision_encoder", "embedding"]
+    pkg = _fake_pkg(keys, out)
+
+    accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider)
+    p = create_pass_from_dict(
+        MobiusBuilder,
+        {"precision": "fp16", "components_to_export": ["vision_encoder", "embedding"]},
+        disable_search=True,
+        accelerator_spec=accelerator_spec,
+    )
+
+    with _patch_build(pkg):
+        result = p.run(_make_hf_model("org/vlm"), out)
+
+    assert isinstance(result, CompositeModelHandler)
+    assert result.model_component_names == ["vision_encoder", "embedding"]
+    # pkg.save must have been called with a components filter that excludes decoder
+    save_kwargs = pkg.save.call_args.kwargs
+    components_filter = save_kwargs.get("components")
+    assert components_filter is not None
+    assert components_filter("vision_encoder") is True
+    assert components_filter("embedding") is True
+    assert components_filter("decoder") is False
+
+
+def test_components_to_export_none_exports_all(tmp_path):
+    """All components are exported when components_to_export is None (default)."""
+    out = tmp_path / "out"
+    keys = ["decoder", "vision_encoder", "embedding"]
+    pkg = _fake_pkg(keys, out)
+
+    with _patch_build(pkg):
+        result = _make_pass().run(_make_hf_model("org/vlm"), out)
+
+    assert isinstance(result, CompositeModelHandler)
+    assert result.model_component_names == keys
+    # pkg.save must have been called without a filter (components=None)
+    save_kwargs = pkg.save.call_args.kwargs
+    assert save_kwargs.get("components") is None
+
+
+def test_components_to_export_single_component_via_filter(tmp_path):
+    """Filtering a multi-component model to one component returns CompositeModelHandler with one component.
+
+    Unlike an architecturally single-component model (which uses root layout), a
+    filtered multi-component model still uses the component sub-directory layout,
+    so we always return CompositeModelHandler for multi-component packages.
+    """
+    out = tmp_path / "out"
+    keys = ["decoder", "vision_encoder", "embedding"]
+    pkg = _fake_pkg(keys, out)
+
+    accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider)
+    p = create_pass_from_dict(
+        MobiusBuilder,
+        {"precision": "fp16", "components_to_export": ["decoder"]},
+        disable_search=True,
+        accelerator_spec=accelerator_spec,
+    )
+
+    with _patch_build(pkg):
+        result = p.run(_make_hf_model("org/vlm"), out)
+
+    # Multi-component model filtered to 1 → still CompositeModelHandler (component sub-dir layout)
+    assert isinstance(result, CompositeModelHandler)
+    assert result.model_component_names == ["decoder"]
+
+
+def test_components_to_export_unknown_component_raises(tmp_path):
+    """ValueError when components_to_export names a component not in the package."""
+    out = tmp_path / "out"
+    keys = ["decoder", "vision_encoder"]
+    pkg = _fake_pkg(keys, out)
+
+    accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider)
+    p = create_pass_from_dict(
+        MobiusBuilder,
+        {"precision": "fp16", "components_to_export": ["nonexistent"]},
+        disable_search=True,
+        accelerator_spec=accelerator_spec,
+    )
+
+    with _patch_build(pkg), pytest.raises(ValueError, match="unknown component"):
+        p.run(_make_hf_model("org/vlm"), out)
+
+
+def test_components_to_export_in_default_config():
+    """components_to_export parameter must appear in _default_config with None default."""
+    accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider)
+    config = MobiusBuilder._default_config(accelerator_spec)  # pylint: disable=protected-access
+    assert "components_to_export" in config
+    assert config["components_to_export"].default_value is None
+    assert config["components_to_export"].required is False

From 6bfdff40cbb8bcb9d55f3adb7f67135091479ce3 Mon Sep 17 00:00:00 2001
From: titaiwangms <titaiwang@microsoft.com>
Date: Fri, 8 May 2026 17:26:44 +0000
Subject: [PATCH 2/8] fix: use named function instead of lambda for
 components_filter

Replace lambda assigned to variable (E731 linting violation) with a
proper named inner function, matching Google Python style guide.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/passes/onnx/mobius_model_builder.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py
index e408ca6bf..f4021526d 100644
--- a/olive/passes/onnx/mobius_model_builder.py
+++ b/olive/passes/onnx/mobius_model_builder.py
@@ -192,7 +192,8 @@ def _run_for_config(
                 package_keys,
                 [k for k in all_keys if k not in requested],
             )
-            components_filter = lambda name: name in requested  # noqa: E731
+            def components_filter(name: str) -> bool:
+                return name in requested
         else:
             package_keys = all_keys
             components_filter = None

From d1207e1eaeaac0752b189f72ba78532d54b49513 Mon Sep 17 00:00:00 2001
From: titaiwangms <titaiwang@microsoft.com>
Date: Fri, 8 May 2026 18:22:18 +0000
Subject: [PATCH 3/8] fix: address PR review comments for components_to_export

- Use type_=list[str] instead of bare list for components_to_export
- Use 'is not None' check to handle empty-list edge case correctly
- Update docstring: remove inaccurate 'Has no effect on single-component
  models' and clarify that ValueError is raised for unknown names
- Update _fake_pkg._save to respect the components filter kwarg so
  skipped component directories are not written to disk
- Add disk assertions: verify skipped component dir is absent on disk
- Wrap long AcceleratorSpec construction lines (Black 120-char limit)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/passes/onnx/mobius_model_builder.py     |  8 ++---
 test/passes/onnx/test_mobius_model_builder.py | 35 ++++++++++++++-----
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py
index f4021526d..e74e440b0 100644
--- a/olive/passes/onnx/mobius_model_builder.py
+++ b/olive/passes/onnx/mobius_model_builder.py
@@ -108,7 +108,7 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
                 ),
             ),
             "components_to_export": PassConfigParam(
-                type_=list,
+                type_=list[str],
                 required=False,
                 default_value=None,
                 description=(
@@ -116,8 +116,8 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
                     "(e.g. ['vision', 'embedding'] to skip the decoder). "
                     "When set, only the named components are saved and returned; "
                     "all others are discarded after the mobius build step. "
-                    "When not set, all components are exported (default, backward compatible). "
-                    "Has no effect on single-component models."
+                    "When not set (None), all components are exported (default, backward compatible). "
+                    "Raises ValueError if any specified name is not found in the model's components."
                 ),
             ),
         }
@@ -178,7 +178,7 @@ def _run_for_config(
 
         # Determine which package components to export.
         all_keys = list(pkg.keys())
-        if config.components_to_export:
+        if config.components_to_export is not None:
             requested = set(config.components_to_export)
             unknown = requested - set(all_keys)
             if unknown:
diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py
index 851019967..127f233e8 100644
--- a/test/passes/onnx/test_mobius_model_builder.py
+++ b/test/passes/onnx/test_mobius_model_builder.py
@@ -74,9 +74,13 @@ def _make_pass(ep: str = ExecutionProvider.CPUExecutionProvider) -> MobiusBuilde
 
 
 def _fake_pkg(keys: list[str], _output_dir: Path) -> MagicMock:
-    """Create a fake ModelPackage that writes dummy .onnx files when .save() is called."""
+    """Create a fake ModelPackage that writes dummy .onnx files when .save() is called.
 
-    def _save(directory: str, **_kwargs):
+    Respects the optional ``components`` filter kwarg passed to ``save()``: only writes
+    files for components for which ``components(name)`` returns True (or all if None).
+    """
+
+    def _save(directory: str, components=None, **_kwargs):
         out = Path(directory)
         if len(keys) == 1:
             # Single-component: saved as <dir>/model.onnx
@@ -84,8 +88,9 @@ def _save(directory: str, **_kwargs):
         else:
             # Multi-component: saved as <dir>/<key>/model.onnx
             for k in keys:
-                (out / k).mkdir(parents=True, exist_ok=True)
-                (out / k / "model.onnx").write_text("dummy")
+                if components is None or components(k):
+                    (out / k).mkdir(parents=True, exist_ok=True)
+                    (out / k / "model.onnx").write_text("dummy")
 
     pkg = MagicMock()
     pkg.keys.return_value = keys
@@ -467,7 +472,9 @@ def test_components_to_export_filters_subset(tmp_path):
     keys = ["decoder", "vision_encoder", "embedding"]
     pkg = _fake_pkg(keys, out)
 
-    accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider)
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
+    )
     p = create_pass_from_dict(
         MobiusBuilder,
         {"precision": "fp16", "components_to_export": ["vision_encoder", "embedding"]},
@@ -480,6 +487,7 @@ def test_components_to_export_filters_subset(tmp_path):
 
     assert isinstance(result, CompositeModelHandler)
     assert result.model_component_names == ["vision_encoder", "embedding"]
+
     # pkg.save must have been called with a components filter that excludes decoder
     save_kwargs = pkg.save.call_args.kwargs
     components_filter = save_kwargs.get("components")
@@ -488,6 +496,11 @@ def test_components_to_export_filters_subset(tmp_path):
     assert components_filter("embedding") is True
     assert components_filter("decoder") is False
 
+    # Verify skipped component directory is absent from disk
+    assert (out / "vision_encoder" / "model.onnx").exists(), "vision_encoder should be on disk"
+    assert (out / "embedding" / "model.onnx").exists(), "embedding should be on disk"
+    assert not (out / "decoder").exists(), "decoder directory should not exist on disk (was skipped)"
+
 
 def test_components_to_export_none_exports_all(tmp_path):
     """All components are exported when components_to_export is None (default)."""
@@ -516,7 +529,9 @@ def test_components_to_export_single_component_via_filter(tmp_path):
     keys = ["decoder", "vision_encoder", "embedding"]
     pkg = _fake_pkg(keys, out)
 
-    accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider)
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
+    )
     p = create_pass_from_dict(
         MobiusBuilder,
         {"precision": "fp16", "components_to_export": ["decoder"]},
@@ -538,7 +553,9 @@ def test_components_to_export_unknown_component_raises(tmp_path):
     keys = ["decoder", "vision_encoder"]
     pkg = _fake_pkg(keys, out)
 
-    accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider)
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
+    )
     p = create_pass_from_dict(
         MobiusBuilder,
         {"precision": "fp16", "components_to_export": ["nonexistent"]},
@@ -552,7 +569,9 @@ def test_components_to_export_unknown_component_raises(tmp_path):
 
 def test_components_to_export_in_default_config():
     """components_to_export parameter must appear in _default_config with None default."""
-    accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider)
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
+    )
     config = MobiusBuilder._default_config(accelerator_spec)  # pylint: disable=protected-access
     assert "components_to_export" in config
     assert config["components_to_export"].default_value is None

From c6bb2c04d1519ca93296082e674efd91a9c8e771 Mon Sep 17 00:00:00 2001
From: titaiwangms <titaiwang@microsoft.com>
Date: Fri, 8 May 2026 18:42:49 +0000
Subject: [PATCH 4/8] Apply lintrunner fixes: add blank line before nested
 function definition

---
 olive/passes/onnx/mobius_model_builder.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py
index e74e440b0..ec88d90f6 100644
--- a/olive/passes/onnx/mobius_model_builder.py
+++ b/olive/passes/onnx/mobius_model_builder.py
@@ -192,6 +192,7 @@ def _run_for_config(
                 package_keys,
                 [k for k in all_keys if k not in requested],
             )
+
             def components_filter(name: str) -> bool:
                 return name in requested
         else:

From 87e41ff414c25e01fa6aa52747d85252df0666bd Mon Sep 17 00:00:00 2001
From: titaiwangms <titaiwang@microsoft.com>
Date: Fri, 8 May 2026 18:55:02 +0000
Subject: [PATCH 5/8] fix: raise ValueError when components_to_export is empty
 list
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

An empty list [] is always a programming mistake — it would silently
produce a model with zero components. Explicitly check before the
unknown-names validation so the error is clear and actionable.

Also adds test_components_to_export_empty_list_raises to cover this case.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/passes/onnx/mobius_model_builder.py     |  5 +++++
 test/passes/onnx/test_mobius_model_builder.py | 20 +++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py
index ec88d90f6..409d29966 100644
--- a/olive/passes/onnx/mobius_model_builder.py
+++ b/olive/passes/onnx/mobius_model_builder.py
@@ -179,6 +179,11 @@ def _run_for_config(
         # Determine which package components to export.
         all_keys = list(pkg.keys())
         if config.components_to_export is not None:
+            if len(config.components_to_export) == 0:
+                raise ValueError(
+                    "MobiusBuilder: components_to_export cannot be empty. "
+                    "Pass None to export all components, or specify at least one component name."
+                )
             requested = set(config.components_to_export)
             unknown = requested - set(all_keys)
             if unknown:
diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py
index 127f233e8..33e745fb4 100644
--- a/test/passes/onnx/test_mobius_model_builder.py
+++ b/test/passes/onnx/test_mobius_model_builder.py
@@ -567,6 +567,26 @@ def test_components_to_export_unknown_component_raises(tmp_path):
         p.run(_make_hf_model("org/vlm"), out)
 
 
+def test_components_to_export_empty_list_raises(tmp_path):
+    """components_to_export=[] must raise ValueError — empty list is always a mistake."""
+    out = tmp_path / "out"
+    keys = ["decoder", "vision_encoder"]
+    pkg = _fake_pkg(keys, out)
+
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
+    )
+    p = create_pass_from_dict(
+        MobiusBuilder,
+        {"precision": "fp16", "components_to_export": []},
+        disable_search=True,
+        accelerator_spec=accelerator_spec,
+    )
+
+    with _patch_build(pkg), pytest.raises(ValueError, match="cannot be empty"):
+        p.run(_make_hf_model("org/vlm"), out)
+
+
 def test_components_to_export_in_default_config():
     """components_to_export parameter must appear in _default_config with None default."""
     accelerator_spec = AcceleratorSpec(

From ec35cecf3488e3e423c9c8ead4176a3957c8e637 Mon Sep 17 00:00:00 2001
From: titaiwangms <titaiwang@microsoft.com>
Date: Fri, 8 May 2026 23:10:46 +0000
Subject: [PATCH 6/8] fix: add TypeError guard for pkg.save() components kwarg
 and fix _fake_pkg filter

- Wrap pkg.save(components=...) in try/except TypeError to gracefully fall back
  when an older mobius version does not support the 'components' kwarg
- Log a warning when fallback is triggered so users know to upgrade mobius
- Fix _fake_pkg._save in tests: single-component branch now respects the
  'components' filter kwarg consistently with the multi-component branch
- Update components_to_export docstring to explicitly mention empty list raises ValueError
- Add test_pkg_save_typeerror_falls_back_gracefully to cover the fallback path

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/passes/onnx/mobius_model_builder.py     | 14 ++++-
 test/passes/onnx/test_mobius_model_builder.py | 51 ++++++++++++++++++-
 2 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py
index 409d29966..cddee2f09 100644
--- a/olive/passes/onnx/mobius_model_builder.py
+++ b/olive/passes/onnx/mobius_model_builder.py
@@ -117,7 +117,8 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
                     "When set, only the named components are saved and returned; "
                     "all others are discarded after the mobius build step. "
                     "When not set (None), all components are exported (default, backward compatible). "
-                    "Raises ValueError if any specified name is not found in the model's components."
+                    "Raises ValueError if the list is empty or if any specified name is not found in "
+                    "the model's components."
                 ),
             ),
         }
@@ -207,7 +208,16 @@ def components_filter(name: str) -> bool:
         # 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)
+        # Older mobius releases may not support the `components` kwarg — fall back gracefully.
+        try:
+            pkg.save(str(output_dir), components=components_filter)
+        except TypeError:
+            if components_filter is not None:
+                logger.warning(
+                    "MobiusBuilder: installed mobius version does not support the 'components' filter kwarg; "
+                    "all components will be saved. Upgrade mobius to enable selective export."
+                )
+            pkg.save(str(output_dir))
 
         # Generate ORT GenAI config artifacts (genai_config.json, tokenizer
         # files, processor configs) when runtime is set to ort-genai.
diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py
index 33e745fb4..e5fa6cc91 100644
--- a/test/passes/onnx/test_mobius_model_builder.py
+++ b/test/passes/onnx/test_mobius_model_builder.py
@@ -83,8 +83,11 @@ def _fake_pkg(keys: list[str], _output_dir: Path) -> MagicMock:
     def _save(directory: str, components=None, **_kwargs):
         out = Path(directory)
         if len(keys) == 1:
-            # Single-component: saved as <dir>/model.onnx
-            (out / "model.onnx").write_text("dummy")
+            # Single-component: saved as <dir>/model.onnx.
+            # Apply the components filter consistently with multi-component behaviour.
+            key = keys[0]
+            if components is None or components(key):
+                (out / "model.onnx").write_text("dummy")
         else:
             # Multi-component: saved as <dir>/<key>/model.onnx
             for k in keys:
@@ -596,3 +599,47 @@ def test_components_to_export_in_default_config():
     assert "components_to_export" in config
     assert config["components_to_export"].default_value is None
     assert config["components_to_export"].required is False
+
+
+def test_pkg_save_typeerror_falls_back_gracefully(tmp_path):
+    """When pkg.save() raises TypeError for the components= kwarg (old mobius), fall back without filter."""
+    out = tmp_path / "out"
+    keys = ["decoder", "vision_encoder", "embedding"]
+
+    # Build a pkg whose save() raises TypeError only when components= is passed (old mobius API).
+    def _save_old_api(directory: str, **kwargs):
+        if "components" in kwargs:
+            raise TypeError("unexpected keyword argument 'components'")
+        # Old API: save all components unconditionally.
+        d = Path(directory)
+        for k in keys:
+            (d / k).mkdir(parents=True, exist_ok=True)
+            (d / k / "model.onnx").write_text("dummy")
+
+    pkg = MagicMock()
+    pkg.keys.return_value = keys
+    pkg.__iter__ = MagicMock(return_value=iter(keys))
+    pkg.items.return_value = [(k, MagicMock()) for k in keys]
+    pkg.save.side_effect = _save_old_api
+
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
+    )
+    p = create_pass_from_dict(
+        MobiusBuilder,
+        {"precision": "fp16", "components_to_export": ["vision_encoder", "embedding"]},
+        disable_search=True,
+        accelerator_spec=accelerator_spec,
+    )
+
+    with (
+        _patch_build(pkg),
+        patch("olive.passes.onnx.mobius_model_builder.logger") as mock_logger,
+    ):
+        result = p.run(_make_hf_model("org/vlm"), out)
+
+    # Old mobius saved all; pass returns CompositeModelHandler with all keys (filter not enforced).
+    assert isinstance(result, CompositeModelHandler)
+    # A warning must have been logged about the missing kwarg support.
+    warning_messages = [str(call) for call in mock_logger.warning.call_args_list]
+    assert any("components" in msg and "mobius" in msg for msg in warning_messages)

From 79eb44fec6836a7394e4fb560d964c0671afe081 Mon Sep 17 00:00:00 2001
From: titaiwangms <titaiwang@microsoft.com>
Date: Mon, 11 May 2026 21:16:48 +0000
Subject: [PATCH 7/8] fix: use inspect.signature to detect components kwarg,
 move empty-list check before build()

Addresses PR review comments:

1. Replace try/except TypeError with inspect.signature check:
   - inspect.signature(pkg.save).parameters is checked before calling save()
   - Avoids masking real TypeErrors and prevents orphan dirs that the
     try/except approach could produce if save() partially executed
   - Log a warning when falling back to old API (no filter applied)

2. Move empty-list validation before mobius.build():
   - Fail fast with a clear ValueError before the expensive build step
   - Unknown component check still runs after build (needs pkg.keys())

3. Fix test for old-API fallback path:
   - Simulate old mobius API by setting pkg.save.__signature__ without
     'components' parameter (inspect.signature detects this at call time)
   - Assert all 3 component dirs are on disk (old API writes all)
   - Assert result only has requested 2 components (package_keys controls handler)
   - Assert pkg.save was NOT called with 'components=' kwarg

4. Add regression test test_pkg_save_components_kwarg_detected_and_filter_applied:
   - Explicitly verifies that when pkg.save signature includes 'components',
     the filter IS passed and only requested components land on disk

Also update _fake_pkg to set pkg.save.__signature__ = inspect.signature(_save)
so production code's inspect.signature check resolves 'components' in parameters.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/passes/onnx/mobius_model_builder.py     | 20 +++--
 test/passes/onnx/test_mobius_model_builder.py | 75 +++++++++++++++++--
 2 files changed, 80 insertions(+), 15 deletions(-)

diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py
index cddee2f09..4763310d7 100644
--- a/olive/passes/onnx/mobius_model_builder.py
+++ b/olive/passes/onnx/mobius_model_builder.py
@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import inspect
 import logging
 from pathlib import Path
 from typing import TYPE_CHECKING, ClassVar
@@ -166,6 +167,13 @@ def _run_for_config(
         if trust_remote_code:
             logger.warning("MobiusBuilder: trust_remote_code=True — only use with trusted model sources.")
 
+        # Validate components_to_export early (before the expensive build step).
+        if config.components_to_export is not None and len(config.components_to_export) == 0:
+            raise ValueError(
+                "MobiusBuilder: components_to_export cannot be empty. "
+                "Pass None to export all components, or specify at least one component name."
+            )
+
         output_dir = Path(output_model_path)
         output_dir.mkdir(parents=True, exist_ok=True)
 
@@ -180,11 +188,6 @@ def _run_for_config(
         # Determine which package components to export.
         all_keys = list(pkg.keys())
         if config.components_to_export is not None:
-            if len(config.components_to_export) == 0:
-                raise ValueError(
-                    "MobiusBuilder: components_to_export cannot be empty. "
-                    "Pass None to export all components, or specify at least one component name."
-                )
             requested = set(config.components_to_export)
             unknown = requested - set(all_keys)
             if unknown:
@@ -208,10 +211,11 @@ def components_filter(name: str) -> bool:
         # 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
-        # Older mobius releases may not support the `components` kwarg — fall back gracefully.
-        try:
+        # Check the installed mobius version supports the 'components' kwarg before passing it,
+        # rather than using a try/except (which could mask real errors or leave orphan dirs).
+        if "components" in inspect.signature(pkg.save).parameters:
             pkg.save(str(output_dir), components=components_filter)
-        except TypeError:
+        else:
             if components_filter is not None:
                 logger.warning(
                     "MobiusBuilder: installed mobius version does not support the 'components' filter kwarg; "
diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py
index e5fa6cc91..db16bdc02 100644
--- a/test/passes/onnx/test_mobius_model_builder.py
+++ b/test/passes/onnx/test_mobius_model_builder.py
@@ -78,7 +78,12 @@ def _fake_pkg(keys: list[str], _output_dir: Path) -> MagicMock:
 
     Respects the optional ``components`` filter kwarg passed to ``save()``: only writes
     files for components for which ``components(name)`` returns True (or all if None).
+
+    ``pkg.save.__signature__`` is set explicitly so that ``inspect.signature(pkg.save)``
+    reports ``components`` in its parameters.  This lets production code use signature
+    introspection (instead of a try/except TypeError) to detect kwarg support.
     """
+    import inspect as _inspect
 
     def _save(directory: str, components=None, **_kwargs):
         out = Path(directory)
@@ -100,6 +105,9 @@ def _save(directory: str, components=None, **_kwargs):
     pkg.__iter__ = MagicMock(return_value=iter(keys))
     pkg.items.return_value = [(k, MagicMock()) for k in keys]
     pkg.save.side_effect = _save
+    # Set __signature__ so inspect.signature(pkg.save) sees 'components' in parameters,
+    # matching a real modern-API mobius ModelPackage.save().
+    pkg.save.__signature__ = _inspect.signature(_save)
     return pkg
 
 
@@ -601,16 +609,20 @@ def test_components_to_export_in_default_config():
     assert config["components_to_export"].required is False
 
 
-def test_pkg_save_typeerror_falls_back_gracefully(tmp_path):
-    """When pkg.save() raises TypeError for the components= kwarg (old mobius), fall back without filter."""
+def test_pkg_save_old_api_no_components_kwarg_falls_back_gracefully(tmp_path):
+    """When pkg.save() has no 'components' kwarg (old mobius), fall back to saving all and log a warning.
+
+    The production code uses inspect.signature to detect kwarg support rather than a
+    try/except TypeError.  This test exercises that detection path when the installed
+    mobius version exposes a save() without the 'components' parameter.
+    """
+    import inspect as _inspect
+
     out = tmp_path / "out"
     keys = ["decoder", "vision_encoder", "embedding"]
 
-    # Build a pkg whose save() raises TypeError only when components= is passed (old mobius API).
+    # Old API: no 'components' kwarg — saves all components unconditionally.
     def _save_old_api(directory: str, **kwargs):
-        if "components" in kwargs:
-            raise TypeError("unexpected keyword argument 'components'")
-        # Old API: save all components unconditionally.
         d = Path(directory)
         for k in keys:
             (d / k).mkdir(parents=True, exist_ok=True)
@@ -621,6 +633,8 @@ def _save_old_api(directory: str, **kwargs):
     pkg.__iter__ = MagicMock(return_value=iter(keys))
     pkg.items.return_value = [(k, MagicMock()) for k in keys]
     pkg.save.side_effect = _save_old_api
+    # Set __signature__ WITHOUT 'components' so inspect.signature detects the old API.
+    pkg.save.__signature__ = _inspect.signature(_save_old_api)
 
     accelerator_spec = AcceleratorSpec(
         accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
@@ -638,8 +652,55 @@ def _save_old_api(directory: str, **kwargs):
     ):
         result = p.run(_make_hf_model("org/vlm"), out)
 
-    # Old mobius saved all; pass returns CompositeModelHandler with all keys (filter not enforced).
+    # Old API saved all 3 components to disk (no filter applied), but the returned
+    # CompositeModelHandler only includes the requested components (built from package_keys).
     assert isinstance(result, CompositeModelHandler)
+    assert result.model_component_names == ["vision_encoder", "embedding"]
+
+    # All 3 component directories are on disk — old API couldn't filter, so decoder is an orphan.
+    for k in keys:
+        assert (out / k / "model.onnx").exists(), f"{k}/model.onnx should be on disk"
+
+    # pkg.save must NOT have been called with a 'components=' kwarg (fallback path).
+    assert "components" not in pkg.save.call_args.kwargs
+
     # A warning must have been logged about the missing kwarg support.
     warning_messages = [str(call) for call in mock_logger.warning.call_args_list]
     assert any("components" in msg and "mobius" in msg for msg in warning_messages)
+
+
+def test_pkg_save_components_kwarg_detected_and_filter_applied(tmp_path):
+    """When pkg.save() signature includes 'components', the filter is passed and applied.
+
+    Regression test for the inspect.signature detection path: verifies that the
+    production code correctly calls pkg.save(components=filter) when the kwarg
+    is present in the signature, and that only the requested components land on disk.
+    """
+    out = tmp_path / "out"
+    keys = ["decoder", "vision_encoder", "embedding"]
+    pkg = _fake_pkg(keys, out)  # _fake_pkg uses spec=_save which has 'components' in signature
+
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
+    )
+    p = create_pass_from_dict(
+        MobiusBuilder,
+        {"precision": "fp16", "components_to_export": ["vision_encoder", "embedding"]},
+        disable_search=True,
+        accelerator_spec=accelerator_spec,
+    )
+
+    with _patch_build(pkg):
+        result = p.run(_make_hf_model("org/vlm"), out)
+
+    # Only the requested components should be returned.
+    assert isinstance(result, CompositeModelHandler)
+    assert result.model_component_names == ["vision_encoder", "embedding"]
+
+    # pkg.save must have been called WITH the 'components=' kwarg (modern API path).
+    assert "components" in pkg.save.call_args.kwargs
+
+    # Requested components must be on disk; decoder must not be.
+    assert (out / "vision_encoder" / "model.onnx").exists()
+    assert (out / "embedding" / "model.onnx").exists()
+    assert not (out / "decoder").exists(), "decoder must not be written when filtered out"

From dcd23b4dc909161798046176554b8fe3c4aee0ec Mon Sep 17 00:00:00 2001
From: titaiwangms <titaiwang@microsoft.com>
Date: Mon, 11 May 2026 22:41:26 +0000
Subject: [PATCH 8/8] docs: add MobiusBuilder to pass reference docs and update
 class docstring

- Add MobiusBuilder entry to docs/source/reference/pass.rst (autoconfigclass)
- Add MobiusBuilder row to pass table in docs/source/reference/options.md
- Update MobiusBuilder class docstring to document components_to_export
  parameter with example JSON config and ValueError behavior

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 docs/source/reference/options.md          |  1 +
 docs/source/reference/pass.rst            |  6 ++++++
 olive/passes/onnx/mobius_model_builder.py | 13 +++++++++++++
 3 files changed, 20 insertions(+)

diff --git a/docs/source/reference/options.md b/docs/source/reference/options.md
index a55b9cc67..acc07514b 100644
--- a/docs/source/reference/options.md
+++ b/docs/source/reference/options.md
@@ -289,6 +289,7 @@ Please also find the detailed options from following table for each pass:
 | [OnnxOpVersionConversion](pass.rst#onnxopversionconversion) | Convert a Onnx model to target op version |
 | [OnnxScriptFusion](pass.rst#onnxscriptfusion) | Fuse Ops using onnxscript |
 | [ModelBuilder](pass.rst#modelbuilder) | Convert a generative PyTorch model to ONNX model using [ONNX Runtime Generative AI](https://github.com/microsoft/onnxruntime-genai) module |
+| [MobiusBuilder](pass.rst#mobius_builder) | Build ONNX models from a HuggingFace model ID using the mobius package; supports optional component filtering via `components_to_export` |
 | [OnnxPeepholeOptimizer](pass.rst#onnxpeepholeoptimizer) | Optimize ONNX model by fusing nodes. |
 | [OnnxTransformersOptimization](pass.rst#onnxtransformersoptimization) | Optimize transformer based models in scenarios where ONNX Runtime does not apply the optimization at load time. It is based on onnxruntime.transformers.optimizer. |
 | [OrtSessionParamsTuning](pass.rst#ortsessionparamstuning) | Optimize ONNX Runtime inference settings. |
diff --git a/docs/source/reference/pass.rst b/docs/source/reference/pass.rst
index 7a96ede68..b5b702736 100644
--- a/docs/source/reference/pass.rst
+++ b/docs/source/reference/pass.rst
@@ -206,6 +206,12 @@ ModelBuilder
 ------------
 .. autoconfigclass:: olive.passes.ModelBuilder
 
+.. _mobius_builder:
+
+MobiusBuilder
+-------------
+.. autoconfigclass:: olive.passes.MobiusBuilder
+
 .. _aimet_quantization:
 
 AimetQuantization
diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py
index 4763310d7..22e2f43ff 100644
--- a/olive/passes/onnx/mobius_model_builder.py
+++ b/olive/passes/onnx/mobius_model_builder.py
@@ -48,6 +48,19 @@ class MobiusBuilder(Pass):
     whose components are individual :class:`~olive.model.ONNXModelHandler` objects.
     Single-component models return a plain :class:`~olive.model.ONNXModelHandler`.
 
+    Use ``components_to_export`` to export only a subset of components.  This is
+    useful when some components (e.g. a text decoder) are already exported and
+    you only need the remaining ones (e.g. vision encoder and embedding)::
+
+        {
+            "type": "MobiusBuilder",
+            "model_path": "mistralai/Ministral-3B-Instruct-2512",
+            "components_to_export": ["vision_encoder", "embedding"]
+        }
+
+    Raises :class:`ValueError` if ``components_to_export`` is an empty list or
+    contains names not present in the built package.
+
     Requires ``mobius-ai`` to be installed::
 
         pip install mobius-ai