uv support in CI

[santacodes]

Description

Added support for uv backend in nox-session and changed to default uv pip install in CI to fasten the installation times.

Note: This is a work in progress and uv fails to install iree-compiler from the source link specified using the --find-links flag (logs). It is an upstream problem, and we would have to wait until it gets resolved.

Reference issues - astral-sh/uv/issues/6082, y21/tl/issues/68

Fixes #3825

Type of change

Please add a line in the relevant section of CHANGELOG.md to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

• New feature (non-breaking change which adds functionality)
• Optimization (back-end change that speeds up the code)
• Bug fix (non-breaking change which fixes an issue)

Key checklist:

• No style issues: $ pre-commit run (or $ nox -s pre-commit) (see CONTRIBUTING.md for how to set this up to run automatically when committing locally, in just two lines of code)
• All tests pass: $ python run-tests.py --all (or $ nox -s tests)
• The documentation builds: $ python run-tests.py --doctest (or $ nox -s doctests)

You can run integration tests, unit tests, and doctests together at once, using $ python run-tests.py --quick (or $ nox -s quick).

Further checks:

• Code is commented, particularly in hard-to-understand areas
• Tests added that prove fix is effective or that feature works

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.42%. Comparing base (ad41dbe) to head (4b57947).
Report is 207 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff            @@
##           develop    #4353   +/-   ##
========================================
  Coverage    99.42%   99.42%           
========================================
  Files          292      292           
  Lines        22215    22215           
========================================
  Hits         22087    22087           
  Misses         128      128           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[santacodes]

The tests pass now as the upstream bug got resolved.

[kratman]

@santacodes, @agriyakhetarpal If we are being impacted by bugs in UV, maybe we should keep using pip until it is more stable

[agriyakhetarpal]

I don't think there should be any more bugs that are relevant to our current use case. I would suggest that we be open to trying this out and reporting more bugs upstream – uv has worked for pybamm-cookie quite well, which pulls in PyBaMM as a dependency when testing the generated projects

[kratman]

I think it is fine if people want to use UV locally, but it is still new and it is more important to keep our CI stable than to experiment with new tools. We can easily revisit UV in the future and we don't have to risk more bugs impacting our deliveries

[agriyakhetarpal]

How about keeping it for PR builds? It's more than just a new tool at this point. The improvement is going to be substantial even without a cache because we can avoid two to three minutes of pip installing packages in CI, and more so when we get to separate out the IDAKLU solver. The existence of warm caches when pursuing development work outside of CI means that the act of installing dependencies locally is instantaneous: we can even turn off the nox.options.reuse_existing_virtualenvs boolean to always keep virtual environments fresh and let uv spin up a virtual environment + install the dependencies in a matter of milliseconds. WIth scikits.odes now not used, all our dependencies have wheels and install quickly once downloaded.

Also, I didn't get what you meant by our deliveries, since uv is currently not integrated with our release procedure – our userbase of scientists/researchers won't notice anything different; this is an improvement specific to just developers and maintainers. The upstream bug mentioned by @santacodes in the PR description, in this case, was with uv being unable to find iree as a dependency for the IREE-MLIR code in an alternative PyPI-like index, but once fixed and up in v0.3.4 (and now 0.4.0 has been released), there were no problems with the dependency itself or with installing it. Our binaries built with the IREE-JIT compiler enabled for the code are only tested in CI, and are not distributed with the wheels we publish on PyPI¹

Other projects such as NumPy and SciPy as mentioned in #3825 (comment) are unlikely to use uv in the near future, the reason being they have various BLAS vendors² to deal with for linkage and actions such as setup-miniconda and setup-micromamba are used instead (plus, conda as of 24.5 and later uses libmambapy as its resolver, so it's almost as fast). However, there are some plans to use pixi in the coming times: https://github.com/rgommers/pixi-dev-scipystack, and pixi itself ended up scrapping away rip (its implementation of pip in rust) and later adopted uv (see astral-sh/uv#1572).

Another reason I advocate for this PR, besides the reasons posted by @Saransh-cpp in #3825 (comment) is that we didn't take as much time to implement ruff as a linter earlier when it came out, and we haven't had much of a problem with it ever since we started using it. Everything seems to work in this PR with uv without issues, and it most likely will continue to do so since we are not advanced users of pip. It's always beneficial to try out new tools, and in the case of any future bugs, help make them better in an altruistic sense of the world, and us being receptive to the implementation would be quite helpful :)

Footnotes

1. This is because IREE is available on macOS >13.0 + Python==3.11 (and not on Ubuntu/Windows) – hence, we have an asymmetric distribution where we wouldn't want that a wheel for a specific platform and specific Python version(s) gets the binaries compiled with IREE solver code while others don't. This is also something to keep in mind for the IDAKLU solver separation, since we would like to distribute the binaries with equivalent features and yet have a way to test the PyBaMM code. ↩

2. BLAS distributions such as OpenBLAS are nowadays available on PyPI as well (https://pypi.org/project/scipy-openblas64/), but they aren't as stable as the system-level counterparts right now (and yet continue to be tested). ↩

[santacodes]

I think uv is pretty stable and reliable because even if we run into some errors or bugs, they get addressed quickly. One example of that would be the issue I linked in this comment. The issue was resolved within a week, when clearly the issue was in one of their dependencies and not their own and they still bothered to resolve it upstream. Technically the upstream dependency fix wasn't merged so they fixed it by including their own fork of that dependency which makes me think they take their users seriously making uv reliable to depend on. Even if uv is not as mature as pip, I think it offers better support in case of bugs even though I don't think we would run into any bugs that would break the testing infrastructure because it is just a package manager and not some kind of build system.

[Saransh-cpp]

Looks quite nice, thanks @santacodes! Have there been any speedups in the CI with this PR?

[santacodes]

Looks quite nice, thanks @santacodes! Have there been any speedups in the CI with this PR?

Looking at the test actions, specifically in the run unit tests job (this is where uv resolves all the packages) it looks like we would be saving around ~5-6 minutes on average. In the future, if we get caching support, that would potentially save more time.

[arjxn-py]

Looks pretty good, thanks @santacodes