Switch to [uv](https://docs.astral.sh/uv/) (#71)

* Switch to [uv](https://docs.astral.sh/uv/) Signed-off-by: Fabrice Normandin <[email protected]> * Fix mkdocs and build.yml issue Signed-off-by: Fabrice Normandin <[email protected]> * Fix issue with check_docs workflow Signed-off-by: Fabrice Normandin <[email protected]> * Add missing mkdocstrings to default deps Signed-off-by: Fabrice Normandin <[email protected]> * Update changed values in regression file Signed-off-by: Fabrice Normandin <[email protected]> * Add xfail on rye-related tests Signed-off-by: Fabrice Normandin <[email protected]> --------- Signed-off-by: Fabrice Normandin <[email protected]>
mila-iqia · Oct 21, 2024 · f753a20 · f753a20
1 parent 6e2ff4c
commit f753a20
Show file tree

Hide file tree

Showing 12 changed files with 4,793 additions and 1,644 deletions.
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
@@ -6,11 +6,16 @@
 	"image": "mcr.microsoft.com/devcontainers/python:1-3.12-bullseye",
 	// Features to add to the dev container. More info: https://containers.dev/features.
 	"features": {
-		"ghcr.io/devcontainers-contrib/features/pre-commit:2": {},
+		"ghcr.io/devcontainers-contrib/features/pre-commit:2": {
+			"version": "3.8.0"
+		},
 		"ghcr.io/devcontainers-contrib/features/apt-get-packages": {
-			"packages": "vim"
+			"packages": [
+				"vim",
+				"sshfs"
+			]
 		},
-		"ghcr.io/e-gineering/devcontainer-features/rye:1": {}
+		"ghcr.io/va-h/devcontainers-features/uv:1": {}
 	},
 	// Use 'forwardPorts' to make a list of ports inside the container available locally.
 	// "forwardPorts": [],
@@ -74,7 +79,6 @@
 		"SCRATCH": "/home/vscode/scratch",
 		"SLURM_TMPDIR": "/tmp",
 		"NETWORK_DIR": "/network",
-		// used by rye (fixed issues with symlinked .cache dir to $SCRATCH)
 		"UV_LINK_MODE": "symlink"
 	},
 	"mounts": [
@@ -110,8 +114,7 @@
 		"pre-commit": "pre-commit install --install-hooks"
 	},
 	"updateContentCommand": {
-		// BUG: doesn't work atm. Perhaps some permission issues?
-		// "Sync dependencies": "rye sync --all-features"
+		"Sync dependencies": "uv sync --frozen"
 	},
 	// Use 'postCreateCommand' to run commands after the container is created.
 	"postCreateCommand": {

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -38,17 +38,18 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - name: Setup Rye with caching of venvs
-        uses: eifinger/setup-rye@v4
-        id: setup-rye
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v3
         with:
+          version: "latest"
           enable-cache: true
+          # https://github.com/astral-sh/setup-uv?tab=readme-ov-file#github-authentication-token
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          cache-prefix: "3.10"
+          cache-suffix: "3.10"
       - name: Install dependencies
-        run: rye sync --no-lock --features=docs
+        run: uv sync --frozen --extra docs
       - name: Build the documentation (strict mode)
-        run: rye run mkdocs build --strict
+        run: uv run mkdocs build --strict
 
   unit_tests:
     needs: [linting]
@@ -60,21 +61,22 @@ jobs:
         python-version: ["3.10"]
     steps:
       - uses: actions/checkout@v4
-      - name: Setup Rye with caching of venvs
-        uses: eifinger/setup-rye@v4
-        id: setup-rye
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v3
         with:
+          version: "latest"
           enable-cache: true
+          # https://github.com/astral-sh/setup-uv?tab=readme-ov-file#github-authentication-token
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          cache-prefix: ${{ matrix.python-version }}
+          cache-suffix: ${{ matrix.python-version }}
       - name: Pin python-version ${{ matrix.python-version }}
-        run: rye pin ${{ matrix.python-version }}
+        run: uv python pin ${{ matrix.python-version }}
       - name: Install dependencies
-        run: rye sync --no-lock
+        run: uv sync --frozen
       - name: Test with pytest
         env:
           JAX_PLATFORMS: cpu
-        run: rye run pytest -v --cov=project --cov-report=xml --cov-append --skip-if-files-missing
+        run: uv run pytest -v --cov=project --cov-report=xml --cov-append --skip-if-files-missing
 
       - name: Store coverage report as an artifact
         uses: actions/upload-artifact@v4
@@ -92,24 +94,25 @@ jobs:
         python-version: ["3.10"]
     steps:
       - uses: actions/checkout@v4
-      - name: Setup Rye with caching of venvs
-        uses: eifinger/setup-rye@v4
-        id: setup-rye
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v3
         with:
-          enable-cache: true
+          version: "latest"
+          enable-cache: false  # no need, uses the local uv cache.
+          # https://github.com/astral-sh/setup-uv?tab=readme-ov-file#github-authentication-token
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          cache-prefix: ${{ matrix.python-version }}
+
       - name: Pin python-version ${{ matrix.python-version }}
-        run: rye pin ${{ matrix.python-version }}
+        run: uv python pin ${{ matrix.python-version }}
 
       - name: Install dependencies
-        run: rye sync --no-lock
+        run: uv sync --all-extras --frozen
 
       - name: Test with pytest
-        run: rye run pytest -v --cov=project --cov-report=xml --cov-append
+        run: uv run pytest -v --cov=project --cov-report=xml --cov-append
 
       - name: Test with pytest (only slow tests)
-        run: rye run pytest -v -m slow --slow --cov=project --cov-report=xml --cov-append
+        run: uv run pytest -v -m slow --slow --cov=project --cov-report=xml --cov-append
 
       - name: Store coverage report as an artifact
         uses: actions/upload-artifact@v4
@@ -159,11 +162,11 @@ jobs:
       - name: Set up the repo using the setup script
         run: scripts/mila_setup.sh
       - name: Install dependencies
-        run: rye sync --all-features --no-lock
+        run: uv sync --all-extras --frozen
       - name: Show installed packages
-        run: rye list
+        run: uv pip list
       - name: Test with pytest
-        run: rye run pytest -v --cov=project --cov-report=xml --cov-append
+        run: uv run pytest -v --cov=project --cov-report=xml --cov-append
 
       # TODO: Disabling full regression tests on the cluster for now, because the worker is often
       # interrupted and we want to avoid using the unkillable partition to not interrupt other's

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
@@ -11,20 +11,23 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
-      - name: Setup Rye with caching of venvs
-        uses: eifinger/setup-rye@v4
-        id: setup-rye
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v3
         with:
-          enable-cache: true
+          version: "latest"
+          enable-cache: true  # no need, uses the local uv cache.
+          # https://github.com/astral-sh/setup-uv?tab=readme-ov-file#github-authentication-token
           github-token: ${{ secrets.GITHUB_TOKEN }}
+          cache-suffix: "3.10"
+
       - name: Pin python-version
-        run: rye pin 3.10
+        run: uv python pin 3.10
 
       - name: Install dependencies
-        run: rye sync --no-lock --features docs
+        run: uv sync --extra docs --frozen
 
       - name: Deploy docs
-        run: rye run mkdocs gh-deploy --force
+        run: uv run mkdocs gh-deploy --force
         # note: Checking if we really need the one below:
         # uses: mhausenblas/[email protected]
         # # Or use mhausenblas/mkdocs-deploy-gh-pages@nomaterial to build without the mkdocs-material theme

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -13,6 +13,7 @@ repos:
       # - id: check-docstring-first
       - id: check-added-large-files
         require_serial: true
+        exclude: uv.lock
       - id: check-ast
         require_serial: true
       - id: check-yaml
@@ -28,8 +29,6 @@ repos:
         require_serial: true
       - id: check-case-conflict
         require_serial: true
-      - id: check-added-large-files
-        require_serial: true
 
   - repo: https://github.com/charliermarsh/ruff-pre-commit
     # Ruff version.

diff --git a/...xample_test/test_backward_pass_is_reproducible/cpu/albert_base_v2_hf_text_hf_example.yaml b/...xample_test/test_backward_pass_is_reproducible/cpu/albert_base_v2_hf_text_hf_example.yaml
@@ -35,12 +35,12 @@ batch.token_type_ids:
   sum: 0
 grads.network.albert.embeddings.LayerNorm.bias:
   device: cpu
-  max: '9.494e-03'
-  mean: '-1.081e-05'
+  max: '9.495e-03'
+  mean: '-1.080e-05'
   min: '-1.796e-02'
   shape:
   - 128
-  sum: '-1.384e-03'
+  sum: '-1.383e-03'
 grads.network.albert.embeddings.LayerNorm.weight:
   device: cpu
   max: '1.186e-02'
@@ -51,72 +51,72 @@ grads.network.albert.embeddings.LayerNorm.weight:
   sum: '-3.360e-02'
 grads.network.albert.embeddings.position_embeddings.weight:
   device: cpu
-  max: '6.971e-01'
-  mean: '-9.095e-12'
+  max: '6.970e-01'
+  mean: '-3.638e-12'
   min: '-1.086e+00'
   shape:
   - 512
   - 128
-  sum: '-5.960e-07'
+  sum: '-2.384e-07'
 grads.network.albert.embeddings.token_type_embeddings.weight:
   device: cpu
   max: '6.053e-01'
-  mean: '-4.657e-09'
+  mean: '-1.863e-09'
   min: '-1.119e+00'
   shape:
   - 2
   - 128
-  sum: '-1.192e-06'
+  sum: '-4.768e-07'
 grads.network.albert.embeddings.word_embeddings.weight:
   device: cpu
   max: '1.541e+00'
-  mean: '3.492e-14'
-  min: '-6.232e-01'
+  mean: '-2.008e-13'
+  min: '-6.233e-01'
   shape:
   - 30000
   - 128
-  sum: '1.341e-07'
+  sum: '-7.711e-07'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm.bias:
   device: cpu
   max: '6.357e-02'
   mean: '-3.738e-04'
   min: '-6.593e-02'
   shape:
   - 768
-  sum: '-2.870e-01'
+  sum: '-2.871e-01'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm.weight:
   device: cpu
   max: '8.125e-02'
   mean: '1.121e-04'
-  min: '-5.810e-01'
+  min: '-5.811e-01'
   shape:
   - 768
-  sum: '8.613e-02'
+  sum: '8.612e-02'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense.bias:
   device: cpu
   max: '6.013e-02'
-  mean: '1.940e-11'
+  mean: '-1.940e-11'
   min: '-5.395e-02'
   shape:
   - 768
-  sum: '1.490e-08'
+  sum: '-1.490e-08'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense.weight:
   device: cpu
   max: '1.061e-01'
-  mean: '-2.324e-12'
+  mean: '4.042e-13'
   min: '-1.112e-01'
   shape:
   - 768
   - 768
-  sum: '-1.371e-06'
+  sum: '2.384e-07'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.bias:
   device: cpu
-  max: '7.916e-09'
-  mean: '-6.855e-12'
-  min: '-5.835e-09'
+  max: '1.275e-08'
+  mean: '-1.333e-11'
+  min: '-6.650e-09'
   shape:
   - 768
-  sum: '-5.265e-09'
+  sum: '-1.023e-08'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.weight:
   device: cpu
   max: '6.536e-01'
@@ -159,15 +159,15 @@ grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.attention.val
   shape:
   - 768
   - 768
-  sum: '-3.593e+00'
+  sum: '-3.592e+00'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.ffn.bias:
   device: cpu
-  max: '4.212e-02'
+  max: '4.213e-02'
   mean: '-3.888e-05'
   min: '-6.737e-02'
   shape:
   - 3072
-  sum: '-1.194e-01'
+  sum: '-1.195e-01'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.ffn.weight:
   device: cpu
   max: '2.953e-01'
@@ -180,20 +180,20 @@ grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.ffn.weight:
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.ffn_output.bias:
   device: cpu
   max: '5.003e-02'
-  mean: '8.731e-11'
+  mean: '-5.821e-11'
   min: '-5.843e-02'
   shape:
   - 768
-  sum: '6.706e-08'
+  sum: '-4.470e-08'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.ffn_output.weight:
   device: cpu
   max: '6.105e-01'
-  mean: '-5.053e-12'
+  mean: '-2.627e-12'
   min: '-5.125e-01'
   shape:
   - 768
   - 3072
-  sum: '-1.192e-05'
+  sum: '-6.199e-06'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm.bias:
   device: cpu
   max: '6.435e-02'
@@ -204,7 +204,7 @@ grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.full_layer_la
   sum: '-1.468e-01'
 grads.network.albert.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm.weight:
   device: cpu
-  max: '5.072e-02'
+  max: '5.071e-02'
   mean: '-6.398e-04'
   min: '-4.395e-01'
   shape:
@@ -214,7 +214,7 @@ grads.network.albert.encoder.embedding_hidden_mapping_in.bias:
   device: cpu
   max: '7.07e-03'
   mean: '-8.878e-05'
-  min: '-7.23e-03'
+  min: '-7.231e-03'
   shape:
   - 768
   sum: '-6.818e-02'
@@ -247,20 +247,20 @@ grads.network.albert.pooler.weight:
 grads.network.classifier.bias:
   device: cpu
   max: '2.129e-01'
-  mean: '0.e+00'
+  mean: '7.451e-09'
   min: '-2.129e-01'
   shape:
   - 2
-  sum: '0.e+00'
+  sum: '1.490e-08'
 grads.network.classifier.weight:
   device: cpu
   max: '2.222e-01'
-  mean: '-4.657e-10'
+  mean: '-3.444e-10'
   min: '-2.222e-01'
   shape:
   - 2
   - 768
-  sum: '-7.153e-07'
+  sum: '-5.29e-07'
 outputs.labels:
   device: cpu
   max: 1