iTranslated by AI

The content below is an AI-generated translation. This is an experimental feature, and may contain errors. View original article
🐍

Simplifying PyTorch Environment Setup with uv: A Unified Approach for Cross-Platform and CPU/GPU Support

に公開
Python
機械学習
PyTorch
uv
tech

3-line Summary

  • In PyTorch environment setup, procedures vary by OS (Mac/Linux) and GPU presence, making it difficult to unify across team development.
  • By combining uv's optional-dependencies and Makefile, we built a system that automatically determines environmental differences and switches the installation source.
  • With just one command, make install, the optimal PyTorch (CPU/GPU version) is automatically installed on macOS/Linux and any CPU/GPU environment, simplifying environment setup.

The sample code for this article is available in the following repository.

https://github.com/haru-256/blog-install-pytorch-cross-platform-20250501

1. Introduction: Common Issues in PyTorch Environment Setup

In machine learning projects, especially development using PyTorch, it's not uncommon to run into trouble with environment setup.

  • Development members' OSs are scattered (Mac, Linux), making it impossible to unify procedures.
  • The PyTorch version to install varies by environment—for example, the local PC is the CPU version, while the experimental server is the GPU version.
  • As a result, environment-specific instruction manuals are required, leading to longer setup times and a higher likelihood of errors.

In this article, I will introduce a method to solve these issues using the Python package management tool uv, which is becoming the modern de facto standard.

Target Audience

  • Those developing machine learning projects using PyTorch.
  • Those who want to unify environment setup across a team.
  • Those who know the basic usage of uv.

*Note: Please refer to the official documentation for instructions on how to install uv itself.

Prerequisites

To execute the steps in this article, you must have the following installed:

  • uv (≥ 0.5.3)
  • make command (pre-installed on Linux/macOS)
  • (For GPU environments) NVIDIA CUDA Toolkit 12.6

Goal of this Article

The goal of this article is to create a state where simply executing a single command, make install, automatically installs the optimal PyTorch (CPU or GPU version) regardless of the environment.

Time required: Approx. 15 minutes (excluding uv installation)

For example, in team development, "Person A's macOS/CPU machine" and "Person B's Linux/GPU machine" can use the exact same repository and command to build a PyTorch environment optimized for their respective machines.

2. Overview of Automating PyTorch Installation with uv

The overall architecture of the system we are building is shown in the following flow diagram.

By combining uv and Makefile, the following processes are automated:

  1. The user executes the make install command.
  2. The Makefile determines if an NVIDIA GPU (CUDA) exists in the execution environment.
  3. Based on the result, it instructs uv to use either the "CPU version (--extra=cpu)" or the "GPU version (--extra=gpu)".
  4. Based on the current OS (Mac/Linux) and the specified type, uv automatically selects the optimal PyTorch download source (Index) configured in pyproject.toml.
  5. The appropriate PyTorch is installed from the selected download source.

The core of this mechanism lies in the definition of dependencies within pyproject.toml. The next section will explain the specific configuration method.

3. Why use uv?

uv is a Python package management tool, similar to pip, poetry, or pipenv. Its main feature is its processing speed due to being written in Rust, and it can improve the development experience while maintaining compatibility with existing pip and pip-tools workflows.

The reason for using uv in this article is that while traditional tools like pipenv required splitting complex dependencies into multiple files, uv allows managing them in a single file.

Challenges with package management in pipenv

For packages like PyTorch, where the download source (index URL) differs between CPU and GPU versions, pipenv faced significant challenges. Since pipenv cannot describe settings to switch download sources based on the OS within a Pipfile, workarounds like managing separate Pipfile.lock files for Linux and macOS were necessary. This leads to increased repository complexity and management costs.

Solution with uv

On the other hand, uv allows you to write settings in pyproject.toml to dynamically switch download sources for each OS or environment using markers (conditional specifications based on environment variables, etc.).

This makes it possible to include dependency information for all environments—including macOS/Linux and CPU/GPU—within a single uv.lock file.

This capability is why uv was chosen for this article.

4. Cross-platform PyTorch Installation Procedure

From here, I will explain the specific configuration method in three steps.

Step 1: Define Dependencies in pyproject.toml

You will describe the settings in pyproject.toml to instruct uv on "which package to download from where for which environment." This is a crucial setting for achieving cross-platform installation.

The overall configuration is shown below, followed by a detailed explanation of each section.

[project]
name = "blog-install-pytorch-cpu-gpu"
version = "0.1.0"
description = "Sample code for a blog explaining how to install CPU/GPU PyTorch"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "numpy~=2.2.5",
    "loguru~=0.7.3"
]

# --- Main part of this article starts here ---

[project.optional-dependencies]
cpu = ["torch==2.7.0"]
gpu = ["torch==2.7.0"]

[tool.uv]
# Prevents both "cpu" and "gpu" from being specified at the same time
conflicts = [[{ extra = "cpu" }, { extra = "gpu" }]]

[tool.uv.sources]
torch = [
  # macOS (CPU version) is obtained from PyPI
  { index = "pytorch-cpu-mac", extra = "cpu", marker = "platform_system == 'Darwin'" },
  # Linux (CPU version) is obtained from the PyTorch-specific index
  { index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
  # GPU version is obtained from the dedicated index for all OSs
  { index = "pytorch-gpu", extra = "gpu" },
]

[[tool.uv.index]]
name = "pytorch-cpu-mac"
url = "https://pypi.python.org/simple"
explicit = true # This index is used only when explicitly specified

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "pytorch-gpu"
url = "https://download.pytorch.org/whl/cu126" # Modification is needed if using a different CUDA version
explicit = true

Explanation: The Role of Each Section

  1. [project.optional-dependencies]
    Here, we define two additional dependency groups ("extras"): cpu and gpu.
    While torch itself is a required library, by deliberately placing it in optional-dependencies, we allow the choice between cpu and gpu at the time of installation.

  2. [tool.uv.sources]
    Specifies where to download the torch package based on conditions.

    • If extra = "cpu" and platform_system == 'Darwin' (specified by marker), the index named pytorch-cpu-mac is used.
    • If extra = "cpu" and the OS is not macOS, the index named pytorch-cpu is used.
    • If extra = "gpu", the index named pytorch-gpu is used.
      In this way, the combination of extra and marker enables switching the download source according to the environment.

  3. [[tool.uv.index]]
    Defines the specific URLs for each index specified in [tool.uv.sources].

    • pytorch-cpu-mac: Standard PyPI (https://pypi.python.org/simple)
    • pytorch-cpu: PyTorch's dedicated repository for the CPU version
    • pytorch-gpu: PyTorch's dedicated repository for the GPU (CUDA 12.6) version
      explicit = true ensures that uv does not automatically search this index; it is only used when explicitly specified in [tool.uv.sources].

Why put torch in optional-dependencies?

Although torch is a required library, the reason for writing it in optional-dependencies is to dynamically select "which version to obtain from where" during installation. If placed in standard dependencies, this flexibility is lost.

With this configuration, uv will automatically select the CPU index if --extra=cpu is specified in the installation command, or the GPU index if --extra=gpu is specified, taking the OS environment into account.

Step 2: Create Lock Files for All Environments with uv lock

Once pyproject.toml is ready, run the following command in your terminal:

# Execute this command in the project root directory
uv lock

When you run this command, uv analyzes the settings in pyproject.toml and writes dependency information corresponding to all combinations of CPU/GPU and macOS/Linux into a file called uv.lock.

This uv.lock file acts as a "blueprint" to guarantee the same installation results in any environment. Project members can share this file via a Git repository so that everyone uses the same versions of the packages.

Step 3: Automate Installation with a Makefile

Finally, write the command executed by the user and the logic to determine the presence of a GPU in a Makefile. This allows users to use a unified command, make install, without having to worry about their environment.

Create a file named Makefile in the root directory of your project with the following content:

# Note: Use tab characters for indentation in Makefiles
# Set HAS_CUDA variable to 1 if the `nvcc` command exists, otherwise 0
HAS_CUDA := $(shell command -v nvcc 2>&1 /dev/null && echo 1 || echo 0)

.PHONY: install
install: ## Project setup
	@if [ $(HAS_CUDA) -eq 1 ]; then \
		echo "✅ GPU (CUDA) environment detected. Installing GPU version of PyTorch."; \
		uv sync --all-groups --extra=gpu; \
	else \
		echo "✅ CPU environment detected. Installing CPU version of PyTorch."; \
		uv sync --all-groups --extra=cpu; \
	fi

Explanation: The Role of the Makefile

This Makefile performs the following processes when the make install command is executed:

  1. HAS_CUDA := ...

    • First, it checks if the nvcc command is executable (i.e., if it's in the PATH). nvcc is the NVIDIA CUDA compiler command; if it exists, it can be determined that the environment is capable of using a GPU (CUDA). The result is stored in the HAS_CUDA variable as 1 (exists) or 0 (does not exist).

  2. install:

    • The main body of make install. It uses a conditional branch if [ $(HAS_CUDA) -eq 1 ] to switch the command to be executed based on the value of the HAS_CUDA variable.
      • If a GPU exists: Runs uv sync --all-groups --extra=gpu. Because --extra=gpu is specified, uv includes the gpu group (GPU version of torch) defined in [project.optional-dependencies] of pyproject.toml in the installation targets.
      • If no GPU exists: Runs uv sync --all-groups --extra=cpu. Similarly, --extra=cpu is specified, and the CPU version of torch is installed.

uv sync is a command that installs (synchronizes) packages into the environment according to the contents of the uv.lock file.

Now, simply running make install will automatically call the appropriate uv sync command for the environment.

5. How to Use: Just Run make install

Everything is now ready.

Project setup can be completed in the following two steps.

  1. Generate the lock file (first time only)
    If you edit pyproject.toml or are working with a newly acquired project, first run uv lock to generate or update the uv.lock file.
uv lock
  1. Build the environment
    After that, simply run the following command in any environment (macOS/CPU, Linux/GPU, etc.).
    When executed, the Makefile will automatically determine the presence of a GPU and install the PyTorch version best suited for your environment.
    When a new member joins the team, they can start developing immediately just by cloning the repository and running make install.
make install

Verifying the Installation

You can check if the environment was built correctly with the following command:

uv run python -c "import torch; print(f'PyTorch {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}')"

Expected output:

  • CPU environment: CUDA available: False
  • GPU environment: CUDA available: True

6. Troubleshooting

Here are some anticipated issues and their solutions.

Error: No matching distribution found for torch

If this error occurs during uv sync or uv lock, there are several possible causes.

  • Cause: Incorrect index settings in pyproject.toml or network issues.
  • Solutions:
    1. Try running uv lock --refresh to regenerate the lock file.
    2. Verify that the URL configured in [[tool.uv.index]] of pyproject.toml is correct (especially the CUDA version, etc.).
    3. Check for network connection issues, such as firewalls.

CPU version is installed on a GPU machine

  • Cause: The Makefile cannot correctly identify the GPU environment. This happens when the nvcc command is not in your PATH.
  • Solutions:
    1. Run which nvcc (macOS/Linux) in your terminal to check the location of the nvcc command.
    2. If the command is not found, please check if the CUDA Toolkit is correctly installed.
    3. If it is installed but not found, review your PATH environment variable settings in files like .bashrc or .zshrc.
# Example: Add to .bashrc
export PATH=/usr/local/cuda/bin:$PATH

7. Appendix: Comparison with the --torch-backend=auto Option

uv also features a preview capability: uv pip install torch --torch-backend=auto.

This feature is extremely convenient because it automatically detects the local environment's CUDA version and installs the optimal PyTorch. Users do not need to worry about specific CUDA versions (e.g., 12.1, 12.6).

However, this method is an approach that "discovers" the optimal package at installation time, rather than locking the dependencies for all environments in advance via uv lock. Consequently, using this feature means you cannot achieve the goal of this article: "supporting all OS, CPU, and GPU environments with a single lock file."

If you need to guarantee strictly identical dependencies across all members and environments in team development, the method of defining detailed settings in pyproject.toml as introduced in this article is more suitable.

8. Summary

In this article, we explained how to use uv to abstract environmental differences such as OS and the presence of a GPU, and automate PyTorch installation with a single command.

The benefits of this approach are as follows:

  • Simplified environment setup: Anyone can set up the environment with a single make install command.
  • Centralized configuration management: Dependencies for all environments can be managed via pyproject.toml and uv.lock.
  • CI/CD compatibility: The workflow becomes simpler, improving maintainability.

9. References

Discussion