pyproject.toml

[ksivaman]

Description

Move to a pyproject.toml setup for build (pep517).

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Use pypyproject.toml setup for wheel and source builds. Related fixes to wheels.
• Remove nvidia-* libs as dependencies. These are currently not being used and will be added to build dependencies in a future PR supporting full e2e compilation and usage without CUDA toolkit.
• The manual installation of pybind11 for JAX extensions is removed, for installations without build isolation, the user must have it preinstalled in the environment.
• Remove redundant custom Pybind11CPPExtension implementation.
• Update install instructions to include --no-build-isolation missed previously.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[ksivaman]

/te-ci

[timmoon10]

Does this mean PyTorch is a dependency for JAX-only builds, and vice versa? That doesn't seem right.

One problem with the pyproject.toml approach is that build-time dependencies are static, but we would really like to have dynamic dependencies. Should we migrate away from supporting a monolithic TE package and commit to using framework-specific subpackages? The TE core package only has runtime dependencies on the framework subpackages, and runtime dependencies can be dynamic.

[ksivaman]

A couple of important points to note:

1. These static build dependencies are only relevant when using build isolation, which is to be viewed purely as an extra feature on top of the current recommendation of using --no-build-isolation.
2. This is also only the case for source installations. As installations from PyPI circumvents this by having a wheel for the core build and specific dependencies from the framework packages.

Feature 1 will help users who want to be able to do a simple build without installing any dependencies beforehand and in that case they will get both frameworks. The main issue that I see with having the framework packages completely separate even for a source install is that we won't be able to support the workflow of doing a simple pip install from the top level directory.

[timmoon10]

I think requiring --no-build-isolation is fine for now (I wonder if there's a way to error out if it's not set), but doesn't this mean you also need --no-deps for the source build? Or else you might install PyTorch/JAX unnecessarily. Maybe better to assume that the user has PyTorch/JAX in their environment.

[ksivaman]

doesn't this mean you also need --no-deps for the source build?

No, because of the distinction between build vs runtime dependencies. --no-build-isolation is strictly for build time dependencies, whereas --no-deps is for runtime dependencies. With --no-build-isolation (the current recommended install method), the user must preinstall all needed build deps for the necessary framework extension they want to compile (include the fw itself). Few cases:

1. If they only have torch installed then the Jax extensions won't be built and the install deps for TE/Jax will not be installed, and vice-versa.
2. If the user has both frameworks but doesn't set NVTE_FRAMEWORK then everything will be built and all runtime deps will be installed.
3. If the user has both frameworks and sets NVTE_FRAMEWORK then only that framework's extension will be built and thus only it's runtime deps will be installed.

I wonder if there's a way to error out if it's not set

I think this is not needed because:

1. We updated all installation instructions and docs in v2.3 to ensure that --no-build-isolation is the default everywhere.
2. This PR does add support for build isolation, purely as an new feature or fallback. Worst case the user env has a few extra libs/dependencies installed that they might not need, but this is a good trade-off for the basic user that doesn't want to manage deps and worry about the env. This will be powerful especially once we remove the toolkit restriction because then we will not longer have any system package dependency and user can simply pip install ..

[ksivaman]

/te-ci

[ksivaman]

/te-ci

[timmoon10]

LGTM