Use astral uv on CI testing. Add uv lock and lock check

goanpeca · goanpeca · commit abd2dda0a4c2 · 2025-06-16T18:49:43.000-05:00
diff --git a/.github/workflows/mps-tests.yml b/.github/workflows/mps-tests.yml
@@ -76,15 +76,16 @@ jobs:
         run: |
           conda shell.bash hook
           conda activate $CONDA_ENV
-          pip install -U torch torchvision
+          pip install uv
+          uv pip install torch torchvision
 
       - name: Install PyTorch (nightly)
         if: ${{ matrix.pytorch-channel == 'pytorch-nightly' }}
         shell: bash -l {0}
         run: |
           conda shell.bash hook
           conda activate $CONDA_ENV
-          pip install --pre -U torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
+          uv pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
 
       - name: Install dependencies
         shell: bash -l {0}
@@ -94,11 +95,11 @@ jobs:
           # TODO: We add set -xe to explicitly fail the CI if one of the commands is failing.
           # Somehow the step is passing even if a subcommand failed
           set -xe
-          pip install -r requirements-dev.txt
+          uv pip install -r requirements-dev.txt
           echo "1 returned code: $?"
-          pip install -e .
+          uv pip install -e .
           echo "2 returned code: $?"
-          pip list
+          uv pip list
           echo "3 returned code: $?"
 
       # Download MNIST: https://github.com/pytorch/ignite/issues/1737
@@ -132,4 +133,4 @@ jobs:
         working-directory: ${{ github.repository }}
         run: |
           conda activate $CONDA_ENV
-          python examples/mnist/mnist.py --epochs=1
+          uv run python examples/mnist/mnist.py --epochs=1
diff --git a/.github/workflows/tpu-tests.yml b/.github/workflows/tpu-tests.yml
@@ -6,6 +6,7 @@ on:
       - "*.*.*"
     paths:
       - "ignite/**"
+      - "pyproject.toml"
       - "tests/ignite/**"
       - "tests/run_tpu_tests.sh"
       - "tests/run_code_style.sh"
@@ -14,6 +15,7 @@ on:
   pull_request:
     paths:
       - "ignite/**"
+      - "pyproject.toml"
       - "tests/ignite/**"
       - "tests/run_tpu_tests.sh"
       - "tests/run_code_style.sh"
@@ -32,25 +34,33 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        xla-version: [nightly]
+        xla-version: ["nightly"]
+        PYTHON_VERSION: ["3.10"]
 
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python 3.10
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-          architecture: "x64"
 
       - name: Get year & week number
         id: get-date
         run: echo "date=$(/bin/date "+%Y-%U")" >> $GITHUB_OUTPUT
         shell: bash -l {0}
 
+      - name: Set up Python 3.10
+        uses: astral-sh/setup-uv@v6
+        with:
+          # architecture: "x64"
+          version: "latest"
+          python-version: "${{ matrix.PYTHON_VERSION }}"
+          activate-environment: true
+          enable-cache: true
+          cache-suffix: "${{ steps.get-date.outputs.date }}-pytorch-${{ runner.os }}"
+          cache-dependency-glob: |
+            **/requirements-dev.txt
+            **/pyproject.toml
+
       - name: Get pip cache dir
         id: pip-cache
         run: |
-          pip3 install -U "pip<24"
           echo "pip_cache=$(pip cache dir)" >> $GITHUB_OUTPUT
         shell: bash -l {0}
 
@@ -67,18 +77,19 @@ jobs:
 
           ## Install mkl (alternative approach to https://github.com/pytorch/xla/blob/b0ba29f98a695671972d4a4cc07441014dba2892/.kokoro/common.sh#L31-L32)
           sudo apt-get update && sudo apt-get install -y libopenblas-dev libomp5
-          pip install mkl==2021.4.0
+          uv pip install mkl==2021.4.0
 
           ## Install torch & xla and torchvision
-          pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
-          pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl
+          uv pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
+          # https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl
+          uv run python scripts/install_torch_xla.py "${{ matrix.PYTHON_VERSION }}" "${{ matrix.xla-version }}"
 
           # Check installation
-          python -c "import torch"
+          uv run python -c "import torch"
 
           ## Install test deps and Ignite
-          pip install -r requirements-dev.txt
-          pip install .
+          uv pip install -r requirements-dev.txt
+          uv pip install .
 
       # Download MNIST: https://github.com/pytorch/ignite/issues/1737
       # to "/tmp" for tpu tests
@@ -94,7 +105,7 @@ jobs:
           timeout_minutes: 25
           shell: bash
           command: |
-            python -c "import torch_xla; print('torch xla version:', torch_xla.__version__)"
+            uv run python -c "import torch_xla; print('torch xla version:', torch_xla.__version__)"
             bash tests/run_tpu_tests.sh
           new_command_on_retry: USE_LAST_FAILED=1 bash tests/run_tpu_tests.sh
         env:
diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml
@@ -59,44 +59,37 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-
       - name: Get year & week number
         id: get-date
         run: |
           echo "date=$(/bin/date "+%Y-%U")" >> $GITHUB_OUTPUT
-
-      - name: Get pip cache dir
-        id: pip-cache
-        run: |
-          pip install -U pip || python -m pip install -U pip
-          echo "pip_cache=$(pip cache dir)" >> $GITHUB_OUTPUT
-
-      - uses: actions/cache@v4
+      - uses: astral-sh/setup-uv@v6
         with:
-          path: |
-            ${{ steps.pip-cache.outputs.pip_cache }}
-          key: ${{ steps.get-date.outputs.date }}-pytorch-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.pytorch-channel }}-${{ hashFiles('requirements-dev.txt') }}
-          restore-keys: |
-            ${{ steps.get-date.outputs.date }}-pytorch-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.pytorch-channel }}-
+          version: "latest"
+          python-version: ${{ matrix.python-version }}
+          activate-environment: true
+          enable-cache: true
+          cache-suffix: "${{ steps.get-date.outputs.date }}-pytorch-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.pytorch-channel }}"
+          cache-dependency-glob: |
+            **/requirements-dev.txt
+            **/pyproject.toml
 
-      - run: pip install pip wheel setuptools -Uqq
+      - name: Install dev tools
+        run: uv pip install wheel setuptools -Uqq
 
       - name: Install PyTorch
         if: ${{ matrix.pytorch-channel == 'pytorch' }}
-        run: pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+        run: uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
 
       - name: Install PyTorch (nightly)
         if: ${{ matrix.pytorch-channel == 'pytorch-nightly' }}
-        run: pip install torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu --pre
+        run: uv pip install torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu --pre
 
       - name: Install dependencies
         run: |
-          pip install -r requirements-dev.txt
-          pip install .
-          pip list
+          uv pip install -r requirements-dev.txt
+          uv pip install .
+          uv pip list
 
       - name: Check code formatting
         run: |
@@ -142,63 +135,63 @@ jobs:
         run: |
           # MNIST
           # 1) mnist.py
-          python examples/mnist/mnist.py --epochs=1
+          uv run python examples/mnist/mnist.py --epochs=1
 
       - name: Run MNIST with loggers Examples
         if: ${{ matrix.os == 'ubuntu-latest' }}
         run: |
           # 2) mnist_with_visdom.py
-          python -c "from visdom.server.build import download_scripts; download_scripts()" # download scripts : https://github.com/facebookresearch/visdom/blob/master/py/server.py#L929
-          python -m visdom.server &
+          uv run python -c "from visdom.server.build import download_scripts; download_scripts()" # download scripts : https://github.com/facebookresearch/visdom/blob/master/py/server.py#L929
+          uv run python -m visdom.server &
           sleep 10
-          python examples/mnist/mnist_with_visdom.py --epochs=1
+          uv run python examples/mnist/mnist_with_visdom.py --epochs=1
           kill %1
           # 3.1) mnist_with_tensorboard.py with tbX
-          python examples/mnist/mnist_with_tensorboard.py --epochs=1
+          uv run python examples/mnist/mnist_with_tensorboard.py --epochs=1
           # 3.2) mnist_with_tensorboard.py with native torch tb
-          pip uninstall -y tensorboardX
-          python examples/mnist/mnist_with_tensorboard.py --epochs=1
+          uv pip uninstall tensorboardX
+          uv run python examples/mnist/mnist_with_tensorboard.py --epochs=1
 
       - name: Run MNIST Example With Crash
         if: ${{ matrix.os == 'ubuntu-latest' }}
         continue-on-error: true
         run: |
           # 4) mnist_save_resume_engine.py
-          python examples/mnist/mnist_save_resume_engine.py --epochs=2 --crash_iteration 1100
+          uv run python examples/mnist/mnist_save_resume_engine.py --epochs=2 --crash_iteration 1100
 
       - name: Resume MNIST from previous crash
         if: ${{ matrix.os == 'ubuntu-latest' }}
         run: |
-          python examples/mnist/mnist_save_resume_engine.py --epochs=2 --resume_from=/tmp/mnist_save_resume/checkpoint_1.pt
+          uv run python examples/mnist/mnist_save_resume_engine.py --epochs=2 --resume_from=/tmp/mnist_save_resume/checkpoint_1.pt
 
       - name: Run GAN example
         if: ${{ matrix.os == 'ubuntu-latest' }}
         run: |
           # DCGAN
-          python examples/gan/dcgan.py --dataset fake --dataroot /tmp/fakedata --output-dir /tmp/outputs-dcgan --batch-size 2 --epochs 2  --workers 0
+          uv run python examples/gan/dcgan.py --dataset fake --dataroot /tmp/fakedata --output-dir /tmp/outputs-dcgan --batch-size 2 --epochs 2  --workers 0
 
       - name: Run RL Examples
         if: ${{ matrix.os == 'ubuntu-latest' }}
         run: |
           # RL
           # 1) Actor-Critic
-          python examples/reinforcement_learning/actor_critic.py --max-episodes=2
+          uv run python examples/reinforcement_learning/actor_critic.py --max-episodes=2
           # 2) Reinforce
-          python examples/reinforcement_learning/reinforce.py --max-episodes=2
+          uv run python examples/reinforcement_learning/reinforce.py --max-episodes=2
 
       - name: Run Neural Style Example
         if: ${{ matrix.os == 'ubuntu-latest' }}
         run: |
           #fast-neural-style
           #train
           mkdir -p ~/.cache/torch/checkpoints/ && wget "https://download.pytorch.org/models/vgg16-397923af.pth" -O ~/.cache/torch/checkpoints/vgg16-397923af.pth
-          python examples/fast_neural_style/neural_style.py train --epochs 1 --cuda 0 --dataset test --dataroot . --image_size 32 --style_image examples/fast_neural_style/images/style_images/mosaic.jpg --style_size 32
+          uv run python examples/fast_neural_style/neural_style.py train --epochs 1 --cuda 0 --dataset test --dataroot . --image_size 32 --style_image examples/fast_neural_style/images/style_images/mosaic.jpg --style_size 32
       - name: Run SR Example
         if: ${{ matrix.os == 'ubuntu-latest' }}
         run: |
           set -e
           echo "Testing Caltech101 dataset availability..."
-          if python -c "import torchvision; torchvision.datasets.Caltech101(root='./data', download=True)"; then
+          if uv run python -c "import torchvision; torchvision.datasets.Caltech101(root='./data', download=True)"; then
             echo "Caltech101 dataset downloaded successfully. Please remove this workaround and restore dataset check."
             exit 1
             # python examples/super_resolution/main.py --upscale_factor 3 --crop_size 180 --batch_size 4 --test_batch_size 100 --n_epochs 1 --lr 0.001 --threads 2 --debug
diff --git a/pyproject.toml b/pyproject.toml
@@ -14,6 +14,7 @@ license-files = ["LICENSE"]
 classifiers = [
     "Programming Language :: Python :: 3",
 ]
+requires-python = ">=3.9,<3.13"
 dependencies = [
     "torch>=1.3,<3",
     "packaging"
diff --git a/scripts/install_torch_xla.py b/scripts/install_torch_xla.py
@@ -0,0 +1,37 @@
+import os
+import sys
+import tempfile
+import zipfile
+from pathlib import Path
+from urllib.request import urlopen
+
+
+def main(python_version: str, xla_version: str) -> None:
+    temp_dir = Path(tempfile.mkdtemp()) / "torch_xla"
+    temp_dir.mkdir(parents=True, exist_ok=True)
+    pyver = python_version.replace(".", "")
+    url = f"https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-{xla_version}-cp{pyver}-cp{pyver}-linux_x86_64.whl"
+    suffix = f"-cp{pyver}-cp{pyver}-linux_x86_64.whl"
+    r = urlopen(url)
+    wheel_path = temp_dir / "torch_xla.zip"
+
+    with open(wheel_path, "b+w") as f:
+        f.write(r.read())
+
+    with zipfile.ZipFile(wheel_path, "r") as zip_ref:
+        zip_ref.extractall(wheel_path.parent)
+
+    dist = [p for p in os.listdir(temp_dir) if "dist-info" in p]
+    if dist:
+        pkg_name = dist[0].split(".dist-info")[0]
+        wheel_path = wheel_path.rename(f"{pkg_name}{suffix}")
+        cmd = ["uv", "pip", "install", str(wheel_path)]
+        os.system(" ".join(cmd))
+
+
+if __name__ == "__main__":
+    if len(sys.argv) == 3:
+        python_version = sys.argv[1]
+        xla_version = sys.argv[2]
+
+    main(python_version, xla_version)
diff --git a/tests/common_test_functionality.sh b/tests/common_test_functionality.sh
@@ -102,8 +102,8 @@ run_tests() {
 
     # Run the command
     if [ "$trap_deselected_exit_code" -eq "1" ]; then
-        eval "pytest ${pytest_args}" || { exit_code=$?; if [ "$exit_code" -eq ${last_failed_no_failures_code} ]; then echo "All tests deselected"; else exit $exit_code; fi; }
+        eval "uv run pytest ${pytest_args}" || { exit_code=$?; if [ "$exit_code" -eq ${last_failed_no_failures_code} ]; then echo "All tests deselected"; else exit $exit_code; fi; }
     else
-        eval "pytest ${pytest_args}"
+        eval "uv run pytest ${pytest_args}"
     fi
 }
diff --git a/tests/run_code_style.sh b/tests/run_code_style.sh
@@ -3,12 +3,12 @@
 set -xeu
 
 if [ $1 = "lint" ]; then
-    flake8 ignite tests examples --config setup.cfg
+    uv run flake8 ignite tests examples --config setup.cfg
     ufmt diff .
 elif [ $1 = "fmt" ]; then
-    ufmt format .
+    uv run ufmt format .
 elif [ $1 = "mypy" ]; then
-    mypy --config-file mypy.ini
+    uv run mypy --config-file mypy.ini
 elif [ $1 = "install" ]; then
-    pip install --upgrade flake8 "black==24.10.0" "usort==1.0.8.post1" "ufmt==2.7.3" "mypy"
+    uv pip install flake8 "black==24.10.0" "usort==1.0.8.post1" "ufmt==2.7.3" "mypy"
 fi

Original file line number	Diff line number	Diff line change
`@@ -14,6 +14,7 @@ license-files = ["LICENSE"]`
`14`	`14`	`classifiers = [`
`15`	`15`	`"Programming Language :: Python :: 3",`
`16`	`16`	`]`
	`17`	`+requires-python = ">=3.9,<3.13"`
`17`	`18`	`dependencies = [`
`18`	`19`	`"torch>=1.3,<3",`
`19`	`20`	`"packaging"`