Skip to content

Migrate Python packaging to maturin - #4730

Draft
dulinriley wants to merge 2 commits into
meta-pytorch:mainfrom
dulinriley:use-maturin-build-system
Draft

Migrate Python packaging to maturin#4730
dulinriley wants to merge 2 commits into
meta-pytorch:mainfrom
dulinriley:use-maturin-build-system

Conversation

@dulinriley

@dulinriley dulinriley commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR replaces the setuptools/setuptools-rust packaging path with a project-owned PEP 517 backend layered on maturin. It preserves the torchmonarch distribution name, existing Python import paths, the native monarch-tui command, build feature environment variables, and MONARCH_VERSION overrides.

The two Torch C++ modules are now compiled into the main PyO3 library and registered as monarch.common._C and monarch.gradient._gradient_generator during initialization. Cargo builds and stages the TUI executable as part of the same dependency graph as the extension. The wrapper builds frontend assets and rewrites wheel/sdist metadata and RECORD entries where necessary.

Benefits

  • Reduces the installed Python native-module count from 3 shared libraries to 1: _rust_bindings, _C, and _gradient_generator become one _rust_bindings library. monarch-tui remains a separate native executable.
  • Removes the 551-line setup.py and all setuptools/setuptools-rust build integration.
  • Reduces declared isolated build dependencies from 5 to 4: setuptools, setuptools-rust, wheel, numpy, and torch become maturin, packaging, numpy, and torch.
  • Keeps the existing user-facing commands: uv build, python -m build, and editable installs continue using standard PEP 517 hooks.
  • Keeps one shared implementation for Torch discovery, CUDA/ROCm selection, Cargo features, compiler flags, RPATH flags, and frontend behavior instead of duplicating those decisions between build systems.
  • Builds monarch-tui as a Cargo artifact dependency of the PyO3 extension, so the extension and executable share one Cargo invocation, dependency graph, and target cache.
  • Makes release frontend packaging strict through MONARCH_REQUIRE_FRONTEND=1, while local builds retain the previous warning-and-continue behavior when npm is unavailable.
  • Adds 11 unit tests covering slim, CPU tensor-engine, CUDA, ROCm, invalid-platform, dual-toolchain, version normalization, wheel rewriting, executable modes, RECORD, and sdist rewriting.

Costs and tradeoffs

  • The project now owns more packaging code. The replacement consists of 970 production lines (build_backend.py: 533, build_support.py: 437), versus the removed 551-line setup.py: a net increase of 419 production lines. Including tests and related changes, the branch is approximately +1,501/-672 lines (+829 net).
  • Wheel post-processing is now our responsibility. The wrapper must maintain filename/dist-info consistency, metadata versions, executable permissions, and hashes for every RECORD entry.
  • Source distributions require explicit rewriting because Cargo uses SemVer while Python exposes PEP 440 versions. The wrapper patches PKG-INFO, monarch_extension/Cargo.toml, and the monarch_extension entry in Cargo.lock.
  • Enabling the TUI target adds its UI dependencies to the package build graph. This avoids the previous duplicate target tree, but changes to the shared feature set can still invalidate a substantial portion of the Rust graph.
  • Cargo artifact dependencies currently require the repository's existing nightly Cargo bindeps support; this is not yet a stable-Cargo-only packaging path.
  • Embedding the C++ modules means they no longer have independent .so/.dylib files or meaningful module __file__ paths. Their public import paths and APIs remain unchanged, but low-level debugging and tooling will see symbols inside _rust_bindings.
  • Maturin emits a warning that build-backend is not literally maturin; this is expected because the project intentionally uses a wrapper backend.
  • Local end-to-end validation covered macOS CPU tensor-engine builds. CUDA, ROCm, Linux, and all supported Python versions rely on the existing CI matrices and are the primary reason to run this PR through CI before merging.

Validation completed locally

  • All 11 new packaging/configuration unit tests pass.
  • cargo check passes for the CPU tensor-engine configuration with embedded-cpp enabled.
  • Prepared metadata correctly changed 0.0.0 to 1.2.3rc1.
  • A generated sdist was approximately 3.26 MB and contained the backend files, Cargo nightly configuration, and patched PKG-INFO, Cargo manifest, and lockfile versions. A wheel rebuilt successfully from that sdist with --locked.
  • A real custom-version wheel was built with version 9.8.7.dev3; it contained 431 entries and was approximately 74 MB compressed on macOS.
  • The wheel contained exactly 1 native Python extension, retained both .pyi files, had a valid regenerated RECORD, and installed monarch-tui with executable permissions.
  • From an isolated extracted wheel, import monarch, monarch.common._C, gradient generation, Torch build-version checking, and monarch-tui --help all succeeded.
  • Editable-wheel generation also retained the native monarch-tui executable with executable permissions.
  • Frontend compilation could not be exercised locally because npm is unavailable. Frontend injection is covered by the artifact rewrite test, and release CI now installs Node and treats missing assets as an error.

Local build timings

These are indicative measurements from one macOS arm64 workspace, not controlled CI benchmarks:

  • Previous separate optimized builds: approximately 15m22s for the extension plus 16m43s for the TUI, or 32m05s total.
  • First unified optimized wheel build after this change: 7m45s, approximately 4.1x faster or 76% less wall time than the two-build sequence.
  • No-change unified wheel rebuild: 15.2s.
  • Rebuild after touching one TUI source file: 26.4s.
  • Rebuild after touching the extension Rust source: 22.8s.

The large improvement comes from sharing one Cargo graph and cache, not from wheel assembly itself. CI should repeat cold and warm measurements before treating these figures as portable across Linux, CUDA, and ROCm builders.

CI coverage requested

The existing build matrices exercise slim Linux, macOS CPU tensor-engine, Linux CUDA, and opt-in ROCm builds across the supported Python versions (3.10-3.13). The build jobs now run the packaging unit tests, and release-oriented wheel jobs require frontend assets.

@dulinriley dulinriley self-assigned this Aug 21, 2026
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 21, 2026
@dulinriley dulinriley added the ciflow/macos Run additional CI on a PR for MacOS label Aug 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/macos Run additional CI on a PR for MacOS CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant