Build with uv instead of just pip #1324
Conversation
Signed-off-by: Jennifer Zhou <[email protected]>
There was a problem hiding this comment.
uv looks neat, but I'm uncomfortable hard-coding like this until it has wider adoption.
I've written up my thoughts on the root problems at #1323 (comment). I think the best approach for now is to adopt your changes to uninstall_te_wheel_packages and to build TE with build isolation disabled. uv should support this (astral-sh/uv#1715).
As an aside, I agree with @Taytay's assessment that build isolation is not a good fit for the ML ecosystem since it depends on custom, vendor-provided packages that are tricky to install. It's annoying the lengths that the Python developers went to to coerce package maintainers to use build isolation.
| def uninstall_te_wheel_packages(): | ||
| subprocess.check_call( | ||
| [ | ||
| sys.executable, | ||
| "-m", | ||
| "pip", | ||
| "uninstall", | ||
| "-y", | ||
| "transformer_engine_cu12", | ||
| "transformer_engine_torch", | ||
| "transformer_engine_paddle", | ||
| "transformer_engine_jax", | ||
| ] | ||
| ) | ||
| if find_spec("pip") is not None: | ||
| subprocess.check_call( | ||
| [ | ||
| sys.executable, | ||
| "-m", | ||
| "pip", | ||
| "uninstall", | ||
| "-y", | ||
| "transformer_engine_cu12", | ||
| "transformer_engine_torch", | ||
| "transformer_engine_paddle", | ||
| "transformer_engine_jax", | ||
| ] | ||
| ) |
There was a problem hiding this comment.
This is a reasonable change that I would accept on its own. It doesn't fully solve the root issue (the package structure is different when installing from PyPI or from source), so I could imagine users running into problems if they install TE multiple times. But it seems fair to expect users to handle this on their own if they are using a non-standard build environment.
build_tools/build_ext.py
Outdated
| ) | ||
|
|
||
| install_and_import("pybind11[global]") |
There was a problem hiding this comment.
This installation should be guarded depending on which framework the extensions are being built for
|
@timmoon10 Thanks for looking into this! I've updated the PR to be just the I think it's easier for me to just maintain a fork than Do you think we can (for the time being) support both build isolation and no build isolation? One way this might work is: if NVTE_FRAMEWORK contains frameworks that aren't in the current python environment, try installing them; otherwise use the existing package. |
Description
Fixes #1323
Type of change
Changes
Please list the changes introduced in this PR:
install_and_imports for uv's clean build environmentsChecklist: