refactor docker image

Dockerfile

@@ -1,10 +1,12 @@
-# build with: docker build . --tag sctg/roco-idefics3:0.0.2 --tag sctg/roco-idefics3:latest --push
-FROM ovhcom/ai-training-pytorch:latest
-RUN source /workspace/.miniconda3/bin/activate \
-    && pip install -U "safetensors>=0.4.5" bitsandbytes\
-    && pip install -U git+https://github.com/huggingface/transformers.git\
-    && pip install huggingface_hub accelerate datasets peft\
-    && pip install -U Pillow
+# build with: docker build . --tag sctg/roco-idefics3:0.0.4 --tag sctg/roco-idefics3:latest --push
+# run with 
+# docker run --gpus all --user=42420:42420 -e HF_TOKEN=hf_TOKEN -it sctg/roco-idefics3:0.0.4 bash -i /start.sh hf_TOKEN
+FROM nvidia/cuda:11.6.1-devel-ubuntu20.04
+# FROM nvidia/cuda:11.0.3-devel-ubuntu20.04
+# RUN mkdir -p /workspace
+RUN /usr/sbin/addgroup --gid 42420 ovh
+RUN /usr/sbin/useradd -u 42420 --gid 42420 -m -d /workspace -s /bin/bash ovh
+RUN apt update -y && apt-get install -y curl git git-lfs screen
 COPY --chmod=777 start.sh /start.sh
 COPY learn.py /learn.py
 COPY preload.py /preload.py
@@ -12,6 +14,21 @@ COPY preload.py /preload.py
 USER root
 RUN chown -R 42420:42420 /workspace
 USER 42420
-RUN source /workspace/.miniconda3/bin/activate \
-    && mkdir -p /workspace/data \
-    && python /preload.py
+RUN curl -L https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh > /workspace/miniconda.sh
+RUN /bin/bash /workspace/miniconda.sh -b -p /workspace/.miniconda3
+RUN . /workspace/.miniconda3/bin/activate && conda init --all
+RUN . /workspace/.miniconda3/bin/activate \
+    && pip install -U "safetensors>=0.4.5" \
+    && pip install -U "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-0.44.2.dev0-py3-none-manylinux_2_24_x86_64.whl" \
+    && pip install -U git+https://github.com/huggingface/transformers.git\
+    && pip install huggingface_hub[cli] accelerate datasets peft\
+    && pip install -U Pillow
+RUN . /workspace/.miniconda3/bin/activate && conda install -y jupyter
+# Mandatory to run the jobs in rootless mode
+# USER root
+# RUN chown -R 42420:42420 /workspace
+USER 42420
+WORKDIR /workspace
+# RUN export HOME=/workspace && cd /workspace && . /workspace/.miniconda3/bin/activate \
+#     && mkdir -p /workspace/data \
+#     && python /preload.py

learn.py

@@ -148,4 +148,5 @@ trainer = Trainer(
     train_dataset = train_dataset,
 )
 
-trainer.train()
+trainer.train()
+

preload.py

@@ -3,4 +3,5 @@
 # License: Apache License 2.0
 from datasets import load_dataset
 dataset_id = "eltorio/ROCO-radiology"
+cache_dir = "/workspace/data"
 train_dataset = load_dataset(dataset_id, split="train", cache_dir=cache_dir)

start.sh

@@ -1,10 +1,26 @@
 #!/bin/bash
+JOB_URL_SCHEME=${JOB_URL_SCHEME:-"http://"}
+JOB_ID=${JOB_ID:-'localhost'}
+JOB_HOST=${JOB_HOST:-'local'}
+export HOME=/workspace
 cd /workspace
-git config --global credential.helper store
 git lfs install
 export HF_TOKEN=$1
+unset $1
 echo "HF_TOKEN: $HF_TOKEN"
-huggingface-cli login --add-to-git-credential --token $HF_TOKEN
-git clone https://huggingface.co/eltorio/IDEFICS3_ROCO
+. /workspace/.bashrc
 . /workspace/.miniconda3/bin/activate
+git clone https://huggingface.co/eltorio/IDEFICS3_ROCO
+git config --global credential.helper store
+
+huggingface-cli login --add-to-git-credential --token $HF_TOKEN
+
+
+jupyter lab --ip=0.0.0.0 --port=8080 --no-browser --allow-root \
+  --notebook-dir=/workspace \
+  --LabApp.token='' \
+  --LabApp.custom_display_url=${JOB_URL_SCHEME}${JOB_ID}-8080.${JOB_HOST} \
+  --LabApp.allow_remote_access=True \
+  --LabApp.allow_origin='*' \
+  --LabApp.disable_check_xsrf=True &
 python /learn.py