From 75a647538de3ab07af7418c8d5dbc0fa8c31a417 Mon Sep 17 00:00:00 2001
From: Lorenzo Balzani <balzanilo@gmail.com>
Date: Mon, 6 May 2024 16:46:03 +0200
Subject: [PATCH] SERVICE: update Dockerfile and add requirements.txt

---
 docker/Dockerfile       | 44 +++++++++++++----------------------------
 docker/requirements.txt | 14 ++++++-------
 2 files changed, 20 insertions(+), 38 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 5e96898..325c3e7 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -8,10 +8,6 @@ ENV DEBIAN_FRONTEND=noninteractive
 # Set work directory
 WORKDIR /workspace
 
-# Define torch version
-ARG TORCH_VERSION="2.2.0+cu118"
-ARG TORCHVISION_VERSION="0.17.0+cu118"
-
 # Install general-purpose packages
 RUN apt-get update -y && \
     apt-get install -y curl \
@@ -27,32 +23,20 @@ RUN apt-get update -y && \
     apt-get clean -y && \
     rm -rf /var/lib/apt/lists/*
 
-# Upgrade pip
-RUN python3.9 -m pip install --upgrade pip
-
-# Install project dependencies
-RUN python3.9 -m pip install \
-    torch==$TORCH_VERSION torchvision==$TORCHVISION_VERSION torchaudio==$TORCH_VERSION -f https://download.pytorch.org/whl/cu118/torch_stable.html \
-    git+https://github.com/huggingface/transformers \
-    pykeen==1.10.1 \
-    scikit-learn==1.3.1 \
-    networkx \
-    huggingface_hub \
-    datasets \
-    evaluate \
-    nvidia-ml-py3 \
-    accelerate \
-    sentence-transformers \
-    nltk \
-    tqdm \
-    wget \
-    gdown \
-    wrapt -U --ignore-installed
-
-# Install RAPIDS dependencies
-RUN python3.9 -m pip install \
-    cudf-cu11 --extra-index-url https://pypi.nvidia.com \
-    cugraph-cu11 --extra-index-url=https://pypi.nvidia.com
+COPY requirements.txt .
+
+# Install project dependens
+RUN python3.9 -m pip install -r requirements.txt
+
+# Install CUDA packages
+## Install correct PyTorch packages
+ARG TORCH_VERSION="2.2.0+cu118"
+ARG TORCHVISION_VERSION="0.17.0+cu118"
+RUN python3.9 -m pip install torch==$TORCH_VERSION torchvision==$TORCHVISION_VERSION  \
+    torchaudio==$TORCH_VERSION -f https://download.pytorch.org/whl/cu118/torch_stable.html
+
+## Install RAPIDS dependencies
+RUN python3.9 -m pip install cudf-cu11 cugraph-cu11 --extra-index-url https://pypi.nvidia.com
 
 # Back to default frontend
 ENV DEBIAN_FRONTEND=dialog
diff --git a/docker/requirements.txt b/docker/requirements.txt
index cddc0c9..38f6eda 100644
--- a/docker/requirements.txt
+++ b/docker/requirements.txt
@@ -1,9 +1,11 @@
-torch==2.2.0+cu118
-torchvision==0.17.0+cu118
-torchaudio==2.2.0+cu118
-git+https://github.com/huggingface/transformers
+pip
+torch
+torchvision
+torchaudio
+wrapt==1.12.1
 pykeen==1.10.1
 scikit-learn==1.3.1
+git+https://github.com/huggingface/transformers
 networkx
 huggingface_hub
 datasets
@@ -15,7 +17,3 @@ nltk
 tqdm
 wget
 gdown
-wrapt==1.12.1
--i https://pypi.nvidia.com
-cudf-cu11
-cugraph-cu11