←back to thread

Using uv with PyTorch

(docs.astral.sh)
167 points charliermarsh | 3 comments | | HN request time: 0s | source
Show context
gdiamos ◴[] 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 #
ttyprintk ◴[] 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 #
1. gdiamos ◴[] 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 #
2. ttyprintk ◴[] No.42190631[source]
Eh, in that case the environment variable is a bit more evident in a Dockerfile.
3. kristjansson ◴[] 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.