switch from poetry to uv

[mivanit]

NOTE: re-opening #986 but with docs in the name so that the proper tests run

TODOs

• update pyproject.toml
• update makefile
• test everything
• update github actions
  • checks.yml
  • release.yml (changed this, but quite unsure whether what I did is a good idea. needs review for sure)
• update README.md
• update contributing.md
• update .devcontainer/devcontainer.json and .devcontainer/Dockerfile to use uv instead of poetry. not sure how to do this
  • I think in the dockerfile we can probably set rid of installing/updating system python as uv should handle this automatically?
• using uv breaks pip install inside notebooks
  • added a new "notebooks" group for those deps
  • shouldn't we be avoiding using torchtyping, and use jaxtyping instead?
  • shouldn't linters be catching unused imports?

Description

Poetry can be quite slow for dependency resolution and environment setup. uv is much faster and has rapidly become the new standard.

Type of change

• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Screenshots

Tool	Lockfile	Real time (s)
Poetry	present	33.265
Poetry	absent	43.167
uv	present	0.232
uv	absent	0.578

raw timing data

All installs are tested with a warm cache (everything downloaded) and the virtual environment absent.

poetry

lockfile present:

$ poetry env remove ...
$ time poetry install
real    0m33.265s
user    0m49.167s
sys     0m9.027s

lockfile absent:

$ poetry env remove ...
$ rm poetry.lock
$ time poetry install
real    0m43.167s
user    0m59.352s
sys     0m9.196s

uv

lockfile present:

$ rm -rf .venv
$ time uv sync
real    0m0.232s
user    0m0.203s
sys     0m0.721s

lockfile absent:

$ rm -rf .venv
$ rm poetry.lock
$ time uv sync
real    0m0.578s
user    0m0.586s
sys     0m0.908s

Checklist:

• 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
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[mivanit]

ok, looks like nearly everything is passing! The few tests still failing:

• Checks / Compatibility Checks (3.10) (pull_request) is failing due to huggingface_hub.errors.HfHubHTTPError: 429 Client Error: Too Many Requests for url: https://huggingface.co/api/models/NeelNanda/Attn_Only_2L512W_C4_Code/tree/main?recursive=True&expand=False so hopefully will work once we re-run the tests
• Checks / Code Checks (pull_request) is still failing, but it's also still failing on dev-3.x

remaining TODOs:

• assistance needed in updating .devcontainer/devcontainer.json and .devcontainer/Dockerfile
• careful review of changes to release.yml would be appreciated
• wait for tests to be passing on dev-3.x, merge into add-uv, and then all tests should pass?

[mivanit]

looks like building docs is failing due to a lack of a huggingface token secret?

[mivanit]

errors when building docs look like

You are trying to access a gated repo.
Make sure to have access to it at https://huggingface.co/google/gemma-2-27b.
401 Client Error. (Request ID: Root=1-68bab32d-2ede034e3d22c9473ae54888;3ce59a07-e52d-4daa-b743-a0c63c4f0567)

Cannot access gated repo for url https://huggingface.co/google/gemma-2-27b/resolve/main/config.json.
Access to model google/gemma-2-27b is restricted. You must have access to it and be authenticated to access it. Please log in.

which are then followed by 429 too many requests errors. It's not immediately clear to me if resolving gated repo access will resolve the 429 errors, or why we are getting repo access errors in the first place.