From 29f5e23c773314bf38b6b2f148d6debfae633422 Mon Sep 17 00:00:00 2001
From: Ruksi Korpisara <me@ruk.si>
Date: Tue, 19 Nov 2024 10:00:38 +0200
Subject: [PATCH] Make the current dependencies and Docker image CPU-only

Also makes it clear that this is a CPU-only setup
---
 .gitignore                               |  2 +-
 DEVELOPMENT.md                           | 40 ++++++++++++++++++++++++
 Dockerfile                               |  8 -----
 Dockerfile.cpu                           | 15 +++++++++
 requirements.txt => requirements-cpu.txt |  2 +-
 valohai.yaml                             | 12 +++----
 6 files changed, 63 insertions(+), 16 deletions(-)
 create mode 100644 DEVELOPMENT.md
 delete mode 100644 Dockerfile
 create mode 100644 Dockerfile.cpu
 rename requirements.txt => requirements-cpu.txt (98%)

diff --git a/.gitignore b/.gitignore
index cf5a229..9edbfdc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,4 +2,4 @@
 .DS_Store
 .idea
 .valohai
-Dockerfile
+Dockerfile.cpu
diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md
new file mode 100644
index 0000000..d7b22c6
--- /dev/null
+++ b/DEVELOPMENT.md
@@ -0,0 +1,40 @@
+# CPU-Only Environments
+
+## Dependencies
+
+Resolve and lock CPU dependencies:
+
+```bash
+uv pip compile --extra-index-url https://download.pytorch.org/whl/cpu requirements.in -o requirements-cpu.txt
+```
+
+## Docker Image
+
+Build the CPU-only Docker image:
+
+```bash
+docker build -f Dockerfile.cpu -t llm-toolkit:dev-cpu .
+```
+
+Smoke test the Docker image:
+
+```bash
+docker run -it --rm -v $(pwd):/workspace llm-toolkit:dev-cpu /bin/bash
+python -c "import torch; print(torch.__version__)"
+python -c "from transformers import pipeline; print(pipeline('sentiment-analysis')('we love you'))"
+```
+
+Release a new version of the CPU-only Docker image:
+
+```bash
+export LLM_TOOLKIT_VERSION=0.2-cpu
+docker tag llm-toolkit:dev-cpu valohai/llm-toolkit:$LLM_TOOLKIT_VERSION
+docker push valohai/llm-toolkit:$LLM_TOOLKIT_VERSION
+```
+
+Cleanup:
+
+```bash
+docker rmi valohai/llm-toolkit:$LLM_TOOLKIT_VERSION
+docker rmi llm-toolkit:dev-cpu
+```
diff --git a/Dockerfile b/Dockerfile
deleted file mode 100644
index ace23b7..0000000
--- a/Dockerfile
+++ /dev/null
@@ -1,8 +0,0 @@
-FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime
-
-ENV PYTHONUNBUFFERED=1 \
-    PIP_DISABLE_PIP_VERSION_CHECK=1 \
-    PIP_ROOT_USER_ACTION=ignore
-
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt && rm requirements.txt
diff --git a/Dockerfile.cpu b/Dockerfile.cpu
new file mode 100644
index 0000000..a46ce2b
--- /dev/null
+++ b/Dockerfile.cpu
@@ -0,0 +1,15 @@
+FROM python:3.11.10-bookworm
+
+ENV PYTHONUNBUFFERED=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    PIP_ROOT_USER_ACTION=ignore
+
+WORKDIR /workspace
+
+COPY requirements-cpu.txt .
+
+RUN pip install \
+    --no-cache-dir \
+    --extra-index-url https://download.pytorch.org/whl/cpu \
+    -r requirements-cpu.txt \
+    && rm requirements-cpu.txt
diff --git a/requirements.txt b/requirements-cpu.txt
similarity index 98%
rename from requirements.txt
rename to requirements-cpu.txt
index a65fb0e..c462e30 100644
--- a/requirements.txt
+++ b/requirements-cpu.txt
@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements.in -o requirements.txt
+#    uv pip compile requirements.in -o requirements-cpu.txt
 accelerate==1.1.1
     # via
     #   -r requirements.in
diff --git a/valohai.yaml b/valohai.yaml
index e5f91f7..a2e63e8 100644
--- a/valohai.yaml
+++ b/valohai.yaml
@@ -1,9 +1,9 @@
 - step:
     name: data-preprocess
-    image: valohai/llm-toolkit:0.2
+    image: valohai/llm-toolkit:0.2-cpu
     environment: trial2023-g4dn-xlarge
     command:
-      - pip install -r requirements.txt
+      - pip install -r requirements-cpu.txt
       - python data-preprocess.py {parameters}
     parameters:
       - name: tokenizer
@@ -21,10 +21,10 @@
 
 - step:
     name: finetune
-    image: valohai/llm-toolkit:0.2
+    image: valohai/llm-toolkit:0.2-cpu
     environment: trial2023-g4dn-xlarge
     command:
-      - pip install -r requirements.txt
+      - pip install -r requirements-cpu.txt
       - python finetune-mistral.py {parameters}
     parameters:
       - name: base_mistral_model
@@ -57,10 +57,10 @@
 
 - step:
     name: inference
-    image: valohai/llm-toolkit:0.2
+    image: valohai/llm-toolkit:0.2-cpu
     environment: trial2023-g4dn-xlarge
     command:
-      - pip install -r requirements.txt
+      - pip install -r requirements-cpu.txt
       - python inference-mistral.py {parameters}
     parameters:
       - name: base_mistral_model