←back to thread

Using uv with PyTorch

(docs.astral.sh)
167 points charliermarsh | 4 comments | 19 Nov 24 21:59 UTC | HN request time: 0.001s | source
Show context
gdiamos ◴[19 Nov 24 23:06 UTC] No.42189027[source]▶
uv significantly speeds up my pytorch in docker builds

  # Setup virtual env
  ENV VIRTUAL_ENV=/app/.venv
  ENV PATH="$VIRTUAL_ENV/bin:$PATH"
  RUN python3 -m venv $VIRTUAL_ENV
  RUN . $VIRTUAL_ENV/bin/activate

  # install using uv
  RUN pip install uv
  RUN uv pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cpu
The index-url makes it really convenient.
replies(3): >>42189326 #>>42189506 #>>42190557 #
1. ttyprintk ◴[20 Nov 24 03:20 UTC] No.42190557[source]▶
Apparently, uv respects a key in project.toml:

    [[tool.uv.index]]
    name = "pytorch-cu124"
    url = "https://download.pytorch.org/whl/cu124"
    explicit = true
replies(1): >>42190578 #
2. gdiamos ◴[20 Nov 24 03:24 UTC] No.42190578[source]▶
Nice, I wonder if there is a way to make it conditional, e.g. to pick a different key for a cpu vs cuda build.
replies(2): >>42190631 #>>42190644 #
3. ttyprintk ◴[20 Nov 24 03:39 UTC] No.42190631[source]▶
Eh, in that case the environment variable is a bit more evident in a Dockerfile.
4. kristjansson ◴[20 Nov 24 03:42 UTC] No.42190644[source]▶
That’s basically what TFA is about? Set up multiple index urls, and use a bog-standard platform marker on the PyTorch dep to pick between them.