Add support for uv in setup scripts.

.vscode.d/tasks.json

@@ -18,6 +18,33 @@
                 "cwd": "${workspaceFolder}"
             },
             "problemMatcher": []
+        },
+        {
+            "label": "Setup Python virtual environment (uv)",
+            "type": "shell",
+            "command": "./setup_venv.sh",
+            "windows": {
+                "command": ".\\setup_venv.bat"
+            },
+            "linux": {
+                "command": "./setup_venv.sh"
+            },
+            "osx": {
+                "command": "./setup_venv.sh"
+            },
+            "args": ["--uv", "--python", "${input:setupPythonVersion}"],
+            "options": {
+                "cwd": "${workspaceFolder}"
+            },
+            "problemMatcher": []
+        }
+    ],
+    "inputs": [
+        {
+            "id": "setupPythonVersion",
+            "type": "promptString",
+            "description": "Python version for uv (>=3.10,<3.15), e.g. 3.12 or 3.12.10",
+            "default": "3.12"
         }
     ]
 }

README.md

@@ -22,7 +22,8 @@ application.
 
 ## Prerequisites
 
-- Python `>=3.10,<3.15`.
+- Python `>=3.10,<3.15`, or [uv](https://docs.astral.sh/uv/getting-started/installation/)
+  to obtain Python automatically with the setup wrappers' `--uv` option.
 - [Keil Studio for VS Code](https://marketplace.visualstudio.com/items?itemName=Arm.keil-studio-pack) from the VS Code marketplace.
 - Tools listed in [`vcpkg-configuration.json`](./vcpkg-configuration.json).
 - Keil Studio manages the required license; the free Keil MDK Community edition can be used for evaluation.
@@ -45,9 +46,13 @@ command-line commands are required.
 
 1. Install [Keil Studio for VS Code](https://marketplace.visualstudio.com/items?itemName=Arm.keil-studio-pack) and [Python extension](https://marketplace.visualstudio.com/items?itemName=ms-python.python) from the VS Code marketplace.
 2. Clone or download this repository, then open its folder in VS Code.
-3. Before using the example for the first time, select **Terminal > Run Task >
-   Setup Python virtual environment**. Wait for the task to create the `.venv`
-   environment and install the packages required to export the model.
+3. Before using the example for the first time, select **Terminal > Run Task**,
+   then **Setup Python virtual environment** to use the installed Python and pip,
+   or **Setup Python virtual environment (uv)** to use uv. The uv task prompts for
+   a Python version (default `3.12`) and requires uv on `PATH`; it can download
+   Python automatically. Wait for the task to create `.venv` and install the
+   packages required to export the model. To change an existing environment's
+   Python version, use the command-line setup with `--recreate` as described below.
 4. Use the CMSIS action buttons to build the application, then select **Run** or
    **Debug**. Keil Studio starts the Corstone-320 FVP automatically.
 
@@ -86,6 +91,44 @@ The setup script creates `.venv/` and installs the packages required to
 quantize and export the model. It is safe to run again; use `--recreate` when
 you want a completely new environment.
 
+By default, setup uses Python's built-in `venv` and installs packages with pip.
+To use [uv](https://docs.astral.sh/uv/getting-started/installation/), install it
+on your `PATH`, then pass `--uv`. Add `--python VERSION` to select the Python
+version for the environment:
+
+```bash
+# Linux/macOS
+./setup_venv.sh --uv --python 3.12
+```
+
+```powershell
+# Windows
+.\setup_venv.bat --uv --python 3.12
+```
+
+If Python is already available, you can also invoke the Python script directly:
+
+```bash
+python setup_venv.py --uv --python 3.12
+```
+
+`--uv` creates the environment with `uv venv` and installs all packages with
+`uv pip`. `--python` requires `--uv` and accepts a major/minor version such as
+`3.12`, or an exact patch version such as `3.12.10`, within `>=3.10,<3.15`.
+With `--uv`, both wrappers use `uv run` to launch the setup script, so no
+preinstalled Python or working `python` command is needed. uv can download the
+requested interpreter if needed. Without `--python`, uv selects a supported
+Python version for the launcher, and a new environment uses that version.
+The launcher runs in an isolated environment so `--recreate` can safely replace
+`.venv`. Without `--uv`, the `PYTHON` environment variable overrides the wrapper's
+Python launcher as before.
+
+An existing usable `.venv` is reused. To change its Python version, add
+`--recreate`, for example `./setup_venv.sh --uv --python 3.12 --recreate`
+(or `.\setup_venv.bat --uv --python 3.12 --recreate` on Windows).
+These options can also be combined with `--executorch-ref REF` to install
+ExecuTorch from a Git ref. Run either wrapper with `--help` for all options.
+
 > [!Note]
 > On Windows, enable long-path support or keep the repository close to the drive
 > root. PyTorch packages can otherwise exceed the legacy 260-character path limit.

setup_venv.bat

@@ -2,14 +2,42 @@
 REM Copyright 2026 Arm Limited and/or its affiliates.
 REM SPDX-License-Identifier: Apache-2.0
 REM
-REM Windows wrapper. All the logic lives in setup_venv.py.
+REM Windows wrapper. Environment setup lives in setup_venv.py.
+REM All options are forwarded, e.g. setup_venv.bat --uv --python 3.12
+REM Add --recreate to change the Python version of an existing environment.
+REM With --uv, uv supplies the launcher too; no system Python is required.
+REM Otherwise PYTHON selects the launcher.
 REM "python" rather than "python3": python3.exe is not reliably present on
 REM Windows, while python3 is the reliable name on Linux/macOS -- which is why
 REM this wrapper and setup_venv.sh differ.
 setlocal
+set "SETUP_USE_UV="
+set "SETUP_UV_PYTHON=>=3.10,<3.15"
+REM SHIFT only changes numbered arguments; %%* still forwards the original list.
+:scan_args
+if "%~1"=="" goto launch
+if "%~1"=="--uv" set "SETUP_USE_UV=1"
+if "%~1"=="--python" if not "%~2"=="" set "SETUP_UV_PYTHON=%~2"
+set "SETUP_ARG=%~1"
+if "%SETUP_ARG:~0,9%"=="--python=" set "SETUP_UV_PYTHON=%SETUP_ARG:~9%"
+shift /1
+goto scan_args
+
+:launch
+if defined SETUP_USE_UV goto launch_uv
 if defined PYTHON (
     "%PYTHON%" "%~dp0setup_venv.py" %*
 ) else (
     python "%~dp0setup_venv.py" %*
 )
 exit /b %ERRORLEVEL%
+
+:launch_uv
+where uv >nul 2>nul
+if errorlevel 1 (
+    echo error: --uv requires uv on PATH; install it from https://docs.astral.sh/uv/getting-started/installation/ 1>&2
+    exit /b 2
+)
+REM Isolation lets --recreate remove .venv without removing the running launcher.
+uv run --no-project --isolated --python "%SETUP_UV_PYTHON%" "%~dp0setup_venv.py" %*
+exit /b %ERRORLEVEL%

setup_venv.py

@@ -12,6 +12,7 @@
 
 import argparse
 import os
+import re
 import shutil
 import subprocess
 import sys
@@ -111,8 +112,50 @@ def venv_is_usable(venv_dir: Path) -> bool:
     return True
 
 
-def pip(python: Path, *args: str, env: dict[str, str] | None = None) -> None:
-    cmd = [str(python), "-m", "pip", *args]
+def python_version(value: str) -> tuple[int, ...]:
+    """Accept a supported major.minor version, optionally with a patch version."""
+    if not re.fullmatch(r"[0-9]+\.[0-9]+(?:\.[0-9]+)?", value):
+        raise argparse.ArgumentTypeError("expected a Python version such as 3.12 or 3.12.10")
+    version = tuple(int(n) for n in value.split("."))
+    if not (MIN_PYTHON <= version[:2] < MAX_PYTHON_EXCLUSIVE):
+        raise argparse.ArgumentTypeError("ExecuTorch needs Python >=3.10,<3.15")
+    return version
+
+
+def check_venv_python(python: Path, requested: tuple[int, ...] | None) -> None:
+    result = subprocess.run(
+        [str(python), "-c", "import sys; print('.'.join(map(str, sys.version_info[:3])))"],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    have = result.stdout.strip()
+    version = tuple(int(n) for n in have.split("."))
+    if not (MIN_PYTHON <= version[:2] < MAX_PYTHON_EXCLUSIVE):
+        sys.exit(
+            f"error: {python} uses Python {have}; ExecuTorch needs >=3.10,<3.15. "
+            "Re-run with --recreate and a supported Python version."
+        )
+    if requested and version[:len(requested)] != requested:
+        want = ".".join(map(str, requested))
+        sys.exit(
+            f"error: .venv uses Python {have}, but --python {want} was requested. "
+            "Re-run with --recreate to change the environment's Python version."
+        )
+
+
+def pip(
+    python: Path, *args: str, uv: str | None = None, env: dict[str, str] | None = None
+) -> None:
+    if uv:
+        # Match pip's selection across PyPI and the PyTorch nightly index:
+        # uv's default first-index strategy can hide the pinned nightly wheels.
+        cmd = [
+            uv, "pip", *args, "--python", str(python),
+            "--index-strategy", "unsafe-best-match",
+        ]
+    else:
+        cmd = [str(python), "-m", "pip", *args]
     print(f"+ {' '.join(cmd)}", flush=True)
     subprocess.run(cmd, check=True, env=env)
 
@@ -155,6 +198,23 @@ def smoke_test(python: Path) -> None:
 
 def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--uv",
+        action="store_true",
+        help=(
+            "create the environment with uv venv and install packages with uv pip "
+            "(requires uv on PATH)"
+        ),
+    )
+    parser.add_argument(
+        "--python",
+        metavar="VERSION",
+        type=python_version,
+        help=(
+            "Python version for uv, e.g. 3.12 or 3.12.10 "
+            "(requires --uv; defaults to the setup interpreter)"
+        ),
+    )
     parser.add_argument(
         "--executorch-ref",
         metavar="REF",
@@ -171,7 +231,16 @@ def main() -> int:
     )
     args = parser.parse_args()
 
-    check_host_python()
+    if args.python and not args.uv:
+        parser.error("--python requires --uv")
+    uv = shutil.which("uv") if args.uv else None
+    if args.uv and not uv:
+        parser.error(
+            "--uv requires uv on PATH; install it from "
+            "https://docs.astral.sh/uv/getting-started/installation/"
+        )
+    if not args.python:
+        check_host_python()
     warn_windows_long_paths()
 
     if args.recreate and VENV_DIR.exists():
@@ -184,14 +253,24 @@ def main() -> int:
 
     if not VENV_DIR.exists():
         print(f"Creating venv at {VENV_DIR}")
-        venv.EnvBuilder(with_pip=True, symlinks=os.name != "nt").create(VENV_DIR)
+        if uv:
+            # The wrapper may run us in uv's temporary isolated environment.
+            # Request its Python version, not a path inside that environment.
+            requested = ".".join(map(str, args.python or sys.version_info[:3]))
+            cmd = [uv, "venv", "--python", requested, str(VENV_DIR)]
+            print(f"+ {' '.join(cmd)}", flush=True)
+            subprocess.run(cmd, check=True)
+        else:
+            venv.EnvBuilder(with_pip=True, symlinks=os.name != "nt").create(VENV_DIR)
 
     python = venv_python(VENV_DIR)
-    pip(python, "install", "--upgrade", "pip")
+    check_venv_python(python, args.python)
+    if not uv:
+        pip(python, "install", "--upgrade", "pip")
 
     # Pass 1: everything that resolves from PyPI. Kept free of any index
     # directive so pip cannot prefer a nightly torch over the pinned release.
-    pip(python, "install", "-r", str(HERE / "requirements.txt"))
+    pip(python, "install", "-r", str(HERE / "requirements.txt"), uv=uv)
 
     # Pass 2: executorch + torchao. From the PyTorch nightly index (see the
     # file header), or from a git ref when the caller asked for one.
@@ -202,7 +281,7 @@ def main() -> int:
             "takes tens of minutes. Ctrl-C now to use the pinned wheel instead.\n",
             file=sys.stderr,
         )
-        pip(python, "install", f"git+{EXECUTORCH_REPO}@{args.executorch_ref}")
+        pip(python, "install", f"git+{EXECUTORCH_REPO}@{args.executorch_ref}", uv=uv)
         # The git install brings no torchao pin; take the one 1.4 expects.
         pip(
             python,
@@ -212,9 +291,10 @@ def main() -> int:
             "--extra-index-url",
             "https://pypi.org/simple",
             "torchao==0.18.0.dev20260715",
+            uv=uv,
         )
     else:
-        pip(python, "install", "-r", str(HERE / "requirements-executorch.txt"))
+        pip(python, "install", "-r", str(HERE / "requirements-executorch.txt"), uv=uv)
 
     # Pass 3: the TOSA serializer, without dependencies. See the header of
     # requirements-arm-tosa.txt for why --no-dependencies is load-bearing.
@@ -223,9 +303,10 @@ def main() -> int:
     pip(
         python,
         "install",
-        "--no-dependencies",
+        "--no-deps",
         "-r",
         str(HERE / "requirements-arm-tosa.txt"),
+        uv=uv,
         env=env,
     )
 

setup_venv.sh

@@ -2,9 +2,34 @@
 # Copyright 2026 Arm Limited and/or its affiliates.
 # SPDX-License-Identifier: Apache-2.0
 #
-# Linux/macOS wrapper. All the logic lives in setup_venv.py so the same setup
+# Linux/macOS wrapper. Environment setup lives in setup_venv.py so the same setup
 # runs on Windows too; this only picks an interpreter. Override with e.g.
 #   PYTHON=python3.12 ./setup_venv.sh
+# All options are forwarded, e.g. ./setup_venv.sh --uv --python 3.12
+# Add --recreate to change the Python version of an existing environment.
+# With --uv, uv supplies the launcher too; no system Python is required.
+# Otherwise PYTHON selects the launcher.
 set -euo pipefail
 HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+use_uv=false
+uv_python='>=3.10,<3.15'
+previous=''
+for arg in "$@"; do
+    if [[ "$previous" == --python ]]; then
+        uv_python="$arg"
+    fi
+    case "$arg" in
+        --uv) use_uv=true ;;
+        --python=*) uv_python="${arg#--python=}" ;;
+    esac
+    previous="$arg"
+done
+if "$use_uv"; then
+    if ! command -v uv >/dev/null 2>&1; then
+        echo 'error: --uv requires uv on PATH; install it from https://docs.astral.sh/uv/getting-started/installation/' >&2
+        exit 2
+    fi
+    # Isolation lets --recreate remove .venv without removing the running launcher.
+    exec uv run --no-project --isolated --python "$uv_python" "${HERE}/setup_venv.py" "$@"
+fi
 exec "${PYTHON:-python3}" "${HERE}/setup_venv.py" "$@"