torch-gpu-demo / Dockerfile
ejschwartz's picture
bla
fbd4fbf
raw
history blame
1.48 kB
# Base image with CUDA for GPU support
#FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
FROM ubuntu:latest
ENV PIP_BREAK_SYSTEM_PACKAGES=1
# Install dependencies
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
--mount=type=cache,target=/var/lib/apt,sharing=locked \
apt-get update && apt-get install -y software-properties-common && \
apt-add-repository ppa:deadsnakes/ppa && \
apt-get update && apt-get install -y \
git \
build-essential \
cmake \
python3.11-dev python3.11-pip python3.11-distutils \
wget
RUN sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1
# Install llvm
RUN wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
--mount=type=cache,target=/var/lib/apt,sharing=locked \
apt-get -y update && apt-get install -y mlir-18-tools llvm-18
# Install pytorch-mlir
RUN pip install --pre torch-mlir torch \
--extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
# Add PyTorch-MLIR tools to PATH
ENV PATH="/root/.local/bin:${PATH}"
# Verify installations
RUN python -c "import torch; print(torch.__version__)" && \
mlir-opt --version && \
llvm-as --version
# Working directory for MLIR operations
WORKDIR /workspace
# Command to keep the container running
CMD ["bash"]