Public API design proposal: promote curated surface, demote internals to `_internal`

[justinchuby]

Summary

While auditing the project structure I noticed that almost the entire codebase lives under _-prefixed modules (_builder.py, _build_context.py, _registry.py, _execution_providers.py, _model_package.py, _optimizations.py, _weight_loading.py, _configs/, _passes/, _constants.py, …), yet mobius/__init__.py re-exports ~35 symbols from those private modules as the de-facto public API. This pattern, combined with a few other issues observed across the repo, will hurt long-term maintainability and user experience. I'd like to propose a concrete public-API design and raise a few smaller items for discussion.

---

Issues observed

1. Public API surface is hidden behind private modules

Every concrete implementation module starts with _, but __init__.py exports them as the public API. This creates two problems:

• For users: IDE "Go to definition" lands in _builder.py / _registry.py, signalling "do not depend on this", contradicting the re-export. Type stubs, docs, and downstream tooling can't reliably tell what's public.
• For maintainers: there is no enforced boundary — internal refactors silently break user code because everything reachable from mobius.* is fair game. Already in Unreleased there is a breaking removal of EpCapabilities.supports_shape at 0.1.x.

2. SemVer vs. 0.1.0 + breaking changes

The CHANGELOG advertises SemVer, but pre-1.0 the project is already shipping breaking changes (supports_shape removal). Either the policy should be relaxed explicitly for 0.x, or breaking changes should be gated by deprecation warnings for ≥1 minor.

3. models/ is becoming flat

97 files in a single directory. Future scaling (text / vision / audio / diffusion / multimodal) will be easier if this is split now, before user code starts importing mobius.models.<name> directly.

4. examples/ discoverability

~30 example scripts, only 4 referenced from README. No examples/README.md index — newcomers miss most of them.

5. requirements/ci/requirements.txt is monolithic

One file installs everything (torch + transformers + diffusers + accelerate + onnxruntime-genai + librosa + Pillow). CI jobs that only need, say, the text path pay the full install cost. Split by job (text / vision / audio / diffusion / lint) would shorten CI and clarify minimal install footprints.

6. Branding / ownership clarity

Copyright says Microsoft Corporation; repo lives at onnxruntime/mobius; package is mobius-ai. Worth one paragraph in README clarifying the relationship to ONNX Runtime so downstream users know what they're depending on.

---

Proposed public API design

The goal: promote a small, intentional surface to mobius.*, demote everything else to clearly-private mobius._internal.*, and group related symbols into sub-namespaces so the top-level __init__ stops growing linearly with features.

Proposed top-level layout

mobius/
├── __init__.py                # ~10 symbols, curated
├── api/                       # public, stable
│   ├── __init__.py
│   ├── build.py               # build(), build_from_module(), build_diffusers_pipeline()
│   ├── context.py             # BuildContext, get_build_dtype()
│   ├── weights.py             # apply_weights(), WeightSource
│   ├── optimize.py            # optimize_model(), OptimizationLevel
│   └── package.py             # ModelPackage, save(), load()
├── config/                    # public configs (was _configs/)
│   ├── __init__.py            # ArchitectureConfig, BaseModelConfig, ...
│   └── per_model/             # already exists
├── tasks/                     # already public; keep
├── components/                # already public; keep
├── models/                    # split into subpackages (see §3)
├── execution_providers/       # public (was _execution_providers.py)
│   ├── __init__.py            # EpCapabilities, get_ep, register_ep, ep_registry
│   └── builtin/               # cuda, dml, webgpu, cpu
├── registry/                  # public (was _registry.py)
│   └── __init__.py            # ModelRegistry, ModelRegistration, registry
└── _internal/                 # EVERYTHING ELSE — explicitly private
    ├── builder_impl.py
    ├── diffusers_builder_impl.py
    ├── weight_loading_impl.py
    ├── weight_utils.py
    ├── graph_diff.py
    ├── flags.py
    ├── passes/                # was _passes/
    └── rewrite_rules/         # already mostly internal, move here

Curated mobius.__init__.py

# mobius/__init__.py
"""Mobius: declarative ONNX graph construction for generative AI models."""

from mobius._version import __version__

# Sub-namespaces (lazy-importable; users do `from mobius import build, config, ...`)
from mobius import api, components, config, execution_providers, models, registry, tasks

# Top-level convenience aliases — the 80% path
from mobius.api.build import build, build_from_module, build_diffusers_pipeline
from mobius.api.context import BuildContext
from mobius.api.package import ModelPackage

__all__ = [
    "__version__",
    # sub-namespaces
    "api", "components", "config", "execution_providers",
    "models", "registry", "tasks",
    # convenience
    "build", "build_from_module", "build_diffusers_pipeline",
    "BuildContext", "ModelPackage",
]

Everything else (the 30+ specific *Config classes, OPSET_VERSION, optimize_model, apply_weights, EpCapabilities, …) moves behind a sub-namespace. Users write:

import mobius
from mobius.config import Gemma4Config, VisionLanguageConfig
from mobius.execution_providers import EpCapabilities, register_ep
from mobius.api.weights import apply_weights
from mobius.api.optimize import optimize_model, OptimizationLevel

pkg = mobius.build(config, weights=...)        # 80% path
pkg.save("model.onnx")

Why this shape

• Discoverable: dir(mobius) returns ~10 things, not 35. Sub-namespaces make autocomplete useful again.

• Stable: any name reachable without traversing _internal is part of the public contract. CI can enforce this with a lint rule (tests/test_public_api_surface.py snapshotting __all__).

• Refactor-friendly: implementation modules can be renamed / split freely as long as mobius.api.* re-exports stay stable.

• No mass break: keep shims in mobius/__init__.py for one minor version with DeprecationWarning:

  def __getattr__(name: str):
      if name in _DEPRECATED_TOPLEVEL:
          warnings.warn(
              f"mobius.{name} moved to {_DEPRECATED_TOPLEVEL[name]}; "
              "the top-level alias will be removed in 0.3.",
              DeprecationWarning, stacklevel=2,
          )
          return _resolve(_DEPRECATED_TOPLEVEL[name])
      raise AttributeError(name)

Public-API contract (proposed addition to CONTRIBUTING.md)

1. Anything reachable from mobius, mobius.api, mobius.config, mobius.components, mobius.tasks, mobius.models, mobius.registry, mobius.execution_providers is public and SemVer-governed once 1.0 ships.
2. Anything under mobius._internal or any module starting with _ is private — no SemVer guarantee, may change in any release.
3. Adding a public symbol requires updating tests/test_public_api_surface.py (snapshot test of __all__).
4. Removing or renaming a public symbol requires ≥1 minor of DeprecationWarning.

Migration plan (suggested, non-breaking until 0.3)

Step	Action	Release
1	Create mobius.api, mobius.config, mobius.execution_providers, mobius.registry re-export packages	0.2.0
2	Move private modules under mobius._internal (keep import shims)	0.2.0
3	Add __getattr__ deprecation warnings for old top-level names	0.2.0
4	Add tests/test_public_api_surface.py snapshot test	0.2.0
5	Update docs + examples to use new paths	0.2.x
6	Remove top-level shims	0.3.0

---

Additional smaller suggestions

• models/ subpackages: models/text/, models/vision/, models/audio/, models/diffusion/, models/multimodal/. Provide an aggregator models/__init__.py so existing from mobius.models import LlamaModel keeps working.
• examples/README.md: one-line description per script, grouped by task (causal-lm, diffusion, multimodal, asr, EP comparison).
• Split CI requirements: requirements/ci/text.txt, vision.txt, audio.txt, diffusion.txt, lint.txt, docs.txt. Compose via -r files.
• README.md "What is Mobius vs. ONNX Runtime": one paragraph clarifying it's a Microsoft / ORT-team project that produces ONNX models consumed by ORT.

---

Questions for maintainers before any of this lands

1. Is there an existing internal stance on what counts as public? (I couldn't find one in CONTRIBUTING.md or docs/design.)
2. Is 1.0 on the roadmap, or is 0.x explicitly "no stability promise"? This decides whether the migration plan above is worth the deprecation churn.
3. Are there downstream consumers (ORT-GenAI, Olive, internal MS teams) whose import paths must be preserved? If so, please list them so the shim layer covers them explicitly.

Happy to send a PR for steps 1–4 of the migration plan if there is directional agreement.