From d3f8ccc8e6e5078dea20ddfc26d52667b4044a50 Mon Sep 17 00:00:00 2001
From: Liana <liana@lianas.org>
Date: Thu, 5 Dec 2024 13:49:46 -0600
Subject: [PATCH] Test caching improvements, test bug fix

---
 apps/chai/Dockerfile    | 40 +++++++++++++++++++++++-----------------
 apps/chai/metadata.yaml |  3 ++-
 2 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/apps/chai/Dockerfile b/apps/chai/Dockerfile
index 22faf09..6131659 100644
--- a/apps/chai/Dockerfile
+++ b/apps/chai/Dockerfile
@@ -11,8 +11,7 @@ ARG DEBIAN_FRONTEND=noninteractive
 
 ENV \
   NVIDIA_DRIVER_CAPABILITIES="compute,video,utility,graphics" \
-  VIRTUAL_ENV=/opt/venv \
-  PATH="${VIRTUAL_ENV}:$PATH" \
+  PATH="/opt/venv/bin:$PATH" \
   UMASK="0002" \
   LANG=C.UTF-8 \
   TZ="Etc/UTC" \
@@ -37,18 +36,20 @@ RUN \
   && apt-get update && apt-get install -y --no-install-recommends \
   curl unzip build-essential catatonit jq \
   gnupg ca-certificates lsb-release \
-  nano vim tree \
-  # -----------------------------------
-  # TODO: Build images with extras
-  # -----------------------------------
-  # htop tmux psmisc \
-  # socat rsync aria2 openssh-server \
-  # -----------------------------------
-  # TODO: Build images with RDMA & InfiniBand
-  # -----------------------------------
-  # libibverbs1 librdmacm1 \
-  # -----------------------------------
-  && curl -fsSL -o /tmp/app.zip "https://github.com/chaidiscovery/chai-lab/archive/refs/tags/v${VERSION}.zip" \
+  nano vim tree
+# -----------------------------------
+# TODO: Build images with extras
+# -----------------------------------
+# htop tmux psmisc \
+# socat rsync aria2 openssh-server \
+# -----------------------------------
+# TODO: Build images with RDMA & InfiniBand
+# -----------------------------------
+# libibverbs1 librdmacm1 \
+# -----------------------------------
+
+RUN \
+  curl -fsSL -o /tmp/app.zip "https://github.com/chaidiscovery/chai-lab/archive/refs/tags/v${VERSION}.zip" \
   && unzip -q /tmp/app.zip -d /tmp/app \
   && cp -R /tmp/app/chai-lab-${VERSION}/LICENSE /app \
   && cp -R /tmp/app/chai-lab-${VERSION}/assets /app \
@@ -58,10 +59,15 @@ RUN \
   && chown -R ${UID}:${GID} /app && chmod -R 755 /app \
   && curl -LsSf https://astral.sh/uv/0.5.6/install.sh | sh \
   && . $HOME/.local/bin/env \
-  && uv venv --no-python-downloads ${VIRTUAL_ENV} \
-  && . ${VIRTUAL_ENV}/bin/activate \
+  && uv venv --no-python-downloads /opt/venv \
+  && . /opt/venv/bin/activate \
+  && cd /tmp/app/chai-lab-${VERSION} \
+  && uv pip install torch torchvision torchaudio \
   && uv pip install -r /tmp/app/chai-lab-${VERSION}/requirements.in \
-  && chown -R ${UID}:${GID} ${VIRTUAL_ENV}  && chmod -R 755 ${VIRTUAL_ENV} \
+  && uv pip install .
+
+RUN \
+  chown -R ${UID}:${GID} /opt/venv && chmod -R 755 /opt/venv \
   && apt-get purge -y build-essential \
   && apt-get autoremove -y \
   && apt-get clean \
diff --git a/apps/chai/metadata.yaml b/apps/chai/metadata.yaml
index 608541f..6a558c7 100644
--- a/apps/chai/metadata.yaml
+++ b/apps/chai/metadata.yaml
@@ -4,7 +4,8 @@ app: chai
 semver: true
 channels:
   - name: stable
-    platforms: ["linux/amd64", "linux/arm64"]
+    #platforms: ["linux/amd64", "linux/arm64"]
+    platforms: ["linux/amd64"]
     stable: true
     tests:
       enabled: false