Add component inspection API - #378
Open
xiaoyu-work wants to merge 16 commits into
Open
Conversation
Expose inspect_components so external tools can query the ModelPackage components mobius would build without constructing graphs or loading weights. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Xiaoyu <xiaoyuzhang@microsoft.com>
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a lightweight public inspection API to determine which ModelPackage components a given HuggingFace model_id would produce (and their optimization “roles”) without constructing ONNX graphs or loading weights. This supports external tooling (e.g., Olive) that wants to plan per-component processing ahead of mobius.build().
Changes:
- Introduces
mobius.inspect_components()andComponentInfoto report component names/kinds (and optional HFsource_pathhints for Qwen VL families). - Implements model-type/task resolution logic sufficient for component inspection (including special-cases like Qwen3.5-MoE-VL detection and wav2vec/CTC →
mmsmapping). - Adds unit tests validating single-component, VLM, encoder-decoder, explicit-task behavior, and error handling; exports API from
mobius.__init__.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| src/mobius/_inspect.py | New inspect_components API + ComponentInfo, plus Qwen VL source_path hints and task/model_type resolution. |
| src/mobius/_inspect_test.py | Unit tests covering expected component layouts and resolution edge cases. |
| src/mobius/init.py | Exposes inspect_components and ComponentInfo from the package root. |
Comment on lines
+84
to
+87
| raise ValueError( | ||
| f"Could not load a HuggingFace config for {model_id!r}. inspect_components supports " | ||
| "transformers/registry models; diffusers pipelines are not supported." | ||
| ) from None |
|
The author of this PR, xiaoyu-work, is not an activated member of this organization on Codecov. |
justinchuby
reviewed
Jun 26, 2026
Move component source metadata to registered model classes and resolve layout-specific named_modules paths without instantiating models. Expand static inspection coverage across multimodal, seq2seq, and speech models.
Use the active CUDA 12 package feed as the sole source for ONNX Runtime GPU wheels. Install the GenAI CUDA package without dependency resolution so it cannot replace ORT with the CUDA 13 PyPI wheel. Copilot-Session: 8f087f94-098a-4cca-8653-436c45150324
Exercise the wav2vec2 ForCTC remap to the MMS registration. This validates the task-resolution branch and restores patch coverage above the CI threshold. Copilot-Session: 8f087f94-098a-4cca-8653-436c45150324
Bring the PR onto the latest base before addressing the cross-platform test regression introduced on main. Copilot-Session: 8f087f94-098a-4cca-8653-436c45150324
Save parity-test models into temporary directories instead of reopening live NamedTemporaryFile handles. Windows denies reopening those handles, causing every Windows test matrix job to fail. Copilot-Session: 8f087f94-098a-4cca-8653-436c45150324
Comment on lines
+145
to
+150
| resolved_task, model_type, hf_config = _resolve_task_model_type_and_config( | ||
| model_id, task, trust_remote_code | ||
| ) | ||
| task_obj = get_task(resolved_task) | ||
| roles = task_obj.model_roles or {} | ||
|
|
Comment on lines
+235
to
+242
| def test_qwen3_5_moe_vl_detected_from_vision_config(monkeypatch): | ||
| _patch_autoconfig( | ||
| monkeypatch, SimpleNamespace(model_type="qwen3_5_moe", vision_config=SimpleNamespace()) | ||
| ) | ||
| names = {c.name for c in inspect_components("fake/qwen-vl")} | ||
| assert "vision_encoder" in names and "decoder" in names | ||
|
|
||
|
|
justinchuby
reviewed
Jul 21, 2026
| - name: Install onnxruntime (GPU) | ||
| if: inputs.device == 'cuda' | ||
| run: | | ||
| # Install after project dependencies so the CPU ORT package cannot overwrite it. |
Member
There was a problem hiding this comment.
Do we need to uninstall the onnxruntime cpu first?
Contributor
Author
There was a problem hiding this comment.
onnxruntime cpi is skipped in line 77 for cuda case
justinchuby
reviewed
Jul 21, 2026
| pip install -e '.[testing,transformers]' | ||
| pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda | ||
| # PyPI's ORT 1.27+ wheel requires CUDA 13, while these runners use CUDA 12. | ||
| pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu |
justinchuby
reviewed
Jul 21, 2026
Integrate the latest main branch while preserving the branch's CUDA 12 CI setup and Windows-safe ONNX session handling.\n\nAdd HuggingFace source metadata for the new Qwen3-TTS step and prefill embedder roles so inspect_components stays aligned with TTSTask. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>\nCopilot-Session: ad41ef64-d67a-49f7-b340-200b666497f2 Signed-off-by: Xiaoyu <xiaoyuzhang@microsoft.com>
Document the two embedder models added on main so the architecture summary matches TTSTask and inspection metadata. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ad41ef64-d67a-49f7-b340-200b666497f2 Signed-off-by: Xiaoyu <xiaoyuzhang@microsoft.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 29 out of 29 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
src/mobius/_inspect.py:150
- inspect_components currently derives the component list directly from task_obj.model_roles. For Gemma4Task, build() only includes "audio_encoder" when config.audio is set (driven by HF config.audio_config), but model_roles declares audio_encoder unconditionally. This makes inspect_components report an audio_encoder component even for text+vision-only Gemma4 checkpoints, which contradicts the docstring/PR intent (“plus audio when config.audio is set”) and can mislead downstream tooling.
task_obj = get_task(resolved_task)
roles = task_obj.model_roles or {}
|
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Contributor
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 29 out of 29 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
src/mobius/tasks/_gemma4.py:434
- File still contains unresolved Git merge conflict markers (<<<<<<<, =======, >>>>>>>), which will break imports and any task resolution relying on Gemma4Task. Please resolve the conflict by keeping both the new
model_rolesmapping and the existing__init__implementation, and removing the conflict markers.
#: decoder + vision + embedding, plus audio when ``config.audio`` is set.
#: ``audio_encoder`` is declared statically (it is config-gated at build time).
model_roles: ClassVar[dict[str, str]] = {
"decoder": "decoder",
"vision_encoder": "encoder",
"audio_encoder": "encoder",
"embedding": "embedding",
}
def __init__(
self,
*,
static_cache: bool = False,
max_seq_len: int | None = None,
):
self._static_cache = static_cache
self._max_seq_len = max_seq_len
def build(
src/mobius/_inspect.py:55
- The return type annotation for
_resolve_task_model_type_and_configis too narrow: whentaskis passed (especially as aModelTaskinstance), the function returns that value directly (seereturn task, None, None), so the first tuple element is not guaranteed to bestr. This can cause type-checking failures and is inconsistent withget_task(task: str | ModelTask).
def _resolve_task_model_type_and_config(
model_id: str, task, trust_remote_code: bool
) -> tuple[str, str | None, object | None]:
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 29 out of 29 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
src/mobius/_inspect.py:166
- inspect_components currently returns one ComponentInfo per entry in task_obj.model_roles, but some tasks build components conditionally based on config. For example, Gemma4Task only adds "audio_encoder" to the ModelPackage when config.audio is not None, yet Gemma4Task.model_roles now always includes "audio_encoder". As a result, inspect_components can over-report components that mobius.build will not actually produce for a given checkpoint/config.
task_obj = get_task(resolved_task)
roles = task_obj.model_roles or {}
# Runtime HF paths are owned by the registered model class. Most classes
# declare a fixed ``HF_COMPONENT_SOURCES`` mapping; classes shared by
# several HF layouts can resolve paths from the already-loaded config.
component_sources: dict[str, tuple[str, ...]] = {}
if model_type is not None and hf_config is not None and model_type in registry:
module_class = registry.get(model_type)
component_sources = _get_hf_component_sources(module_class, model_type, hf_config)
components = [
ComponentInfo(
name=name,
role=role,
source_paths=tuple(component_sources.get(name, ())),
)
for name, role in roles.items()
]
src/mobius/_inspect.py:56
- _resolve_task_model_type_and_config is annotated as returning tuple[str, ...], but it returns
taskverbatim when an explicit task is provided (which can be a ModelTask instance per inspect_components docstring). This is a type mismatch that will fail strict type checking and makes the helper’s contract inaccurate.
def _resolve_task_model_type_and_config(
model_id: str, task, trust_remote_code: bool
) -> tuple[str, str | None, object | None]:
"""Resolve the mobius task name for a model id without building it.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 32 out of 32 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (3)
src/mobius/_inspect_test.py:436
with torch.device("meta"):is not a valid / reliable way to instantiate modules on the meta device;torch.deviceis not a general context manager (and this will fail on common torch versions). Since these configs are already tiny, instantiate the HF model normally to make the test robust.
for model_type, hf_model_class, hf_config in cases:
with torch.device("meta"):
hf_model = hf_model_class(hf_config)
runtime_paths = {name for name, _ in hf_model.named_modules()}
src/mobius/_testing/ort_inference.py:264
OnnxModelSession.close()unconditionally callsself._tmpdir.cleanup(). If__init__raises before_tmpdiris created (e.g. when a multi-modelModelPackageis passed and the ValueError is raised early),__del__will still run andclose()will raiseAttributeErrorduring GC. Guard_tmpdirthe same way_sessionis guarded so__del__is exception-safe.
def close(self) -> None:
# Release ORT resources eagerly (not only at Python GC time) so
# long-running GPU test suites do not accumulate session memory.
if hasattr(self, "_session"):
del self._session
self._tmpdir.cleanup()
src/mobius/models/unet_parity_test.py:51
import onnx_irinsidetest_cross_attention_block_matches_diffusers()is now unused after switching to_run_onnx()(which imports/saves internally). This will trip unused-import linting in the test suite.
def test_cross_attention_block_matches_diffusers():
pytest.importorskip("diffusers")
import onnx_ir
import torch
from diffusers.models.transformers.transformer_2d import Transformer2DModel
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 32 out of 32 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (7)
src/mobius/_inspect.py:166
- inspect_components() builds its return list by iterating over task_obj.model_roles, which is a static map and can include config-gated or optional components (e.g., Gemma4Task only builds audio_encoder when config.audio is set). This means inspect_components can report components that mobius.build will not actually produce for a particular checkpoint, contradicting the API docstring that these are the ModelPackage keys build() produces.
resolved_task, model_type, hf_config = _resolve_task_model_type_and_config(
model_id, task, trust_remote_code
)
task_obj = get_task(resolved_task)
roles = task_obj.model_roles or {}
# Runtime HF paths are owned by the registered model class. Most classes
# declare a fixed ``HF_COMPONENT_SOURCES`` mapping; classes shared by
# several HF layouts can resolve paths from the already-loaded config.
component_sources: dict[str, tuple[str, ...]] = {}
if model_type is not None and hf_config is not None and model_type in registry:
module_class = registry.get(model_type)
component_sources = _get_hf_component_sources(module_class, model_type, hf_config)
components = [
ComponentInfo(
name=name,
role=role,
source_paths=tuple(component_sources.get(name, ())),
)
for name, role in roles.items()
]
src/mobius/tasks/_gemma4.py:423
- Gemma4Task.model_roles declares an
audio_encoderrole unconditionally, but build() only adds models["audio_encoder"] when config.audio is not None. Because inspect_components currently uses model_roles keys directly, Gemma4 vision-only checkpoints will be reported as having an audio_encoder component even though mobius.build() won't emit it.
src/mobius/_inspect.py:56 - _resolve_task_model_type_and_config() is annotated as returning tuple[str, ...], but it can return the caller-provided
taskwhich may be a ModelTask instance. This makes the type annotation inaccurate and can break strict type checking.
def _resolve_task_model_type_and_config(
model_id: str, task, trust_remote_code: bool
) -> tuple[str, str | None, object | None]:
"""Resolve the mobius task name for a model id without building it.
.github/workflows/main.yml:288
- This installs onnxruntime-genai-cuda from the default index (PyPI), while onnxruntime-gpu is installed from the CUDA-12 internal feed. Mixing sources can lead to ABI/CUDA-version mismatches (the comment above already notes PyPI CUDA requirements differ). Consider installing onnxruntime-genai-cuda from the same CUDA-12 feed (or pin both to a known-compatible pair).
# PyPI's ORT 1.27+ wheel requires CUDA 13, while these runners use CUDA 12.
pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu
pip install --no-deps onnxruntime-genai-cuda
.github/workflows/validation_examples_gpu.yml:88
- This installs onnxruntime-genai-cuda from the default index (PyPI), while onnxruntime-gpu is installed from the CUDA-12 internal feed. Mixing sources can lead to ABI/CUDA-version mismatches; consider installing onnxruntime-genai-cuda from the same CUDA-12 feed (or pin both to a known-compatible pair).
# PyPI's ORT 1.27+ wheel requires CUDA 13, while these runners use CUDA 12.
pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu
pip install --no-deps onnxruntime-genai-cuda
.github/workflows/gpu_l4_golden_parity.yml:64
- This installs onnxruntime-genai-cuda from the default index (PyPI), while onnxruntime-gpu is installed from the CUDA-12 internal feed. Mixing sources can lead to ABI/CUDA-version mismatches; consider installing onnxruntime-genai-cuda from the same CUDA-12 feed (or pin both to a known-compatible pair).
# PyPI's ORT 1.27+ wheel requires CUDA 13, while these runners use CUDA 12.
pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu
pip install --no-deps onnxruntime-genai-cuda
.github/workflows/gpu_l5_generation_e2e.yml:64
- This installs onnxruntime-genai-cuda from the default index (PyPI), while onnxruntime-gpu is installed from the CUDA-12 internal feed. Mixing sources can lead to ABI/CUDA-version mismatches; consider installing onnxruntime-genai-cuda from the same CUDA-12 feed (or pin both to a known-compatible pair).
# PyPI's ORT 1.27+ wheel requires CUDA 13, while these runners use CUDA 12.
pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu
pip install --no-deps onnxruntime-genai-cuda
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 32 out of 32 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (5)
src/mobius/_inspect.py:22
- _inspect.py uses ModelTask instances as valid
taskinputs (per docstring and tasks.get_task signature), but this module doesn't import ModelTask for type checking. Adding a TYPE_CHECKING-only import avoids NameError in annotations and keeps runtime imports minimal.
import dataclasses
import logging
src/mobius/_inspect.py:122
inspect_componentsaccepts either a task name or a ModelTask instance, but the signature leavestaskuntyped. Typing it asstr | ModelTask | Nonemakes the new public API easier to consume and keeps it consistent withmobius.tasks.get_task.
def inspect_components(
model_id: str,
task=None,
trust_remote_code: bool = False,
) -> list[ComponentInfo]:
src/mobius/_inspect.py:150
inspect_componentsderives its component list directly fromtask_obj.model_roles, but several tasks build components conditionally (e.g. Gemma4Task only addsaudio_encoderwhenconfig.audio is not None, and TTSTask only addsspeaker_encoderwhen the module exposes it). As a result,inspect_componentscan over-report components thatbuild()would not emit for the given model_id.
task_obj = get_task(resolved_task)
roles = task_obj.model_roles or {}
src/mobius/tasks/_gemma4.py:422
- Gemma4Task declares
audio_encoderinmodel_roles, butbuild()only includes theaudio_encodermodel whenconfig.audio is not None. Any tooling that treatsmodel_rolesas “the produced package keys” (including the new inspect_components API) will report an audio component for vision-only Gemma4 checkpoints.
src/mobius/_inspect.py:56 _resolve_task_model_type_and_configis annotated as returningtuple[str, ...], but it can return aModelTaskinstance when an explicit task object is provided (it forwardstaskthrough). This will fail mypy/pyright and is inconsistent with the public API contract.
def _resolve_task_model_type_and_config(
model_id: str, task, trust_remote_code: bool
) -> tuple[str, str | None, object | None]:
"""Resolve the mobius task name for a model id without building it.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Validation