From 17b19999de2e4e33dcc6cd2e86cab2bdceebfb16 Mon Sep 17 00:00:00 2001 From: Steve Cliff Date: Mon, 6 Apr 2026 16:13:41 +0100 Subject: [PATCH] Switch nvidia and rocm Dockerfiles from onnxruntime to torch Nvidia: install torch+torchvision from PyTorch cu130 index, drop onnxruntime-gpu. ROCm: use local torch wheel with rocm6.4 index for torchvision, clean up nvidia remnants from the venv. Co-Authored-By: Claude Opus 4.6 (1M context) --- engine/Dockerfile.nvidia | 4 ++-- engine/Dockerfile.rocm | 7 ++++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/engine/Dockerfile.nvidia b/engine/Dockerfile.nvidia index 1064f07..e9f602c 100644 --- a/engine/Dockerfile.nvidia +++ b/engine/Dockerfile.nvidia @@ -20,8 +20,8 @@ COPY VERSION ./ RUN uv venv .venv && \ . .venv/bin/activate && \ - uv pip install -e . && \ - uv pip install --no-deps onnxruntime-gpu + UV_HTTP_TIMEOUT=600 uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu130 && \ + uv pip install -e . ENV PATH="/app/.venv/bin:$PATH" ENV VIRTUAL_ENV="/app/.venv" diff --git a/engine/Dockerfile.rocm b/engine/Dockerfile.rocm index adc4008..4a45274 100644 --- a/engine/Dockerfile.rocm +++ b/engine/Dockerfile.rocm @@ -18,10 +18,15 @@ COPY kb/ kb/ COPY main.py ./ COPY VERSION ./ +COPY torch-*.whl /tmp/ + RUN uv venv .venv && \ . .venv/bin/activate && \ + uv pip install /tmp/torch-*.whl --extra-index-url https://download.pytorch.org/whl/rocm6.4 && \ + uv pip install torchvision --index-url https://download.pytorch.org/whl/rocm6.4 && \ uv pip install -e . && \ - uv pip install --no-deps onnxruntime-rocm + rm /tmp/torch-*.whl && \ + rm -rf .venv/lib/python3.12/site-packages/nvidia # Stage 2: Runtime — minimal ROCm runtime libs only FROM ubuntu:24.04