From c5ae19b223ebf3b7155f4de67dd26c0f7781eed7 Mon Sep 17 00:00:00 2001 From: Christophe Favergeon Date: Wed, 9 Sep 2026 10:14:23 +0200 Subject: [PATCH] Add support for uv in setup scripts. --- .vscode.d/tasks.json | 27 ++++++++++++ README.md | 51 +++++++++++++++++++++-- setup_venv.bat | 30 +++++++++++++- setup_venv.py | 99 ++++++++++++++++++++++++++++++++++++++++---- setup_venv.sh | 27 +++++++++++- 5 files changed, 219 insertions(+), 15 deletions(-) diff --git a/.vscode.d/tasks.json b/.vscode.d/tasks.json index cee1153e..93d6281f 100644 --- a/.vscode.d/tasks.json +++ b/.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" } ] } diff --git a/README.md b/README.md index ddbac7df..22d9f3c2 100644 --- a/README.md +++ b/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. diff --git a/setup_venv.bat b/setup_venv.bat index a5c4e9c9..82cbe038 100644 --- a/setup_venv.bat +++ b/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% diff --git a/setup_venv.py b/setup_venv.py index e5551f51..19527d7e 100755 --- a/setup_venv.py +++ b/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, ) diff --git a/setup_venv.sh b/setup_venv.sh index de5895eb..fd874fa7 100755 --- a/setup_venv.sh +++ b/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" "$@"