From 7ab99a8ca41db160fd8293148fe947bdcdd62157 Mon Sep 17 00:00:00 2001
From: Liana <liana@lianas.org>
Date: Thu, 5 Dec 2024 11:37:00 -0600
Subject: [PATCH] Features: Add Chai-1, optimize builds, fix bugs

Add Chai-1
Add Astral UV for fast python package management
Fix bug with Boltz env/flag translation
Fix bug with Boltz default env variables
Enable python fault handling by default (`PYTHONFAULTHANDLER`)
Update Boltz readme
Add VIRTUAL_ENV variable for builds
Add 'lsb-release tree vim nano' to defaults
---
 apps/boltz/Dockerfile    | 48 +++++++++++++------------
 apps/boltz/entrypoint.sh |  7 +++-
 apps/chai/Dockerfile     | 77 ++++++++++++++++++++++++++++++++++++++++
 apps/chai/README.md      | 18 ++++++++++
 apps/chai/ci/goss.yaml   |  7 ++++
 apps/chai/ci/latest.sh   |  5 +++
 apps/chai/entrypoint.sh  | 23 ++++++++++++
 apps/chai/metadata.yaml  | 11 ++++++
 8 files changed, 172 insertions(+), 24 deletions(-)
 create mode 100644 apps/chai/Dockerfile
 create mode 100644 apps/chai/README.md
 create mode 100644 apps/chai/ci/goss.yaml
 create mode 100755 apps/chai/ci/latest.sh
 create mode 100644 apps/chai/entrypoint.sh
 create mode 100644 apps/chai/metadata.yaml

diff --git a/apps/boltz/Dockerfile b/apps/boltz/Dockerfile
index 97366be..a7296d7 100644
--- a/apps/boltz/Dockerfile
+++ b/apps/boltz/Dockerfile
@@ -11,7 +11,8 @@ ARG DEBIAN_FRONTEND=noninteractive
 
 ENV \
   NVIDIA_DRIVER_CAPABILITIES="compute,video,utility,graphics" \
-  PATH="/opt/venv/bin:$PATH" \
+  VIRTUAL_ENV=/opt/venv \
+  PATH="${VIRTUAL_ENV}:$PATH" \
   UMASK="0002" \
   LANG=C.UTF-8 \
   TZ="Etc/UTC" \
@@ -20,10 +21,12 @@ ENV \
   GID=900 \
   PYTHONDONTWRITEBYTECODE=1 \
   PYTHONUNBUFFERED=1 \
+  PYTHONFAULTHANDLER=1 \
   PIP_ROOT_USER_ACTION=ignore \
   PIP_NO_CACHE_DIR=1 \
   PIP_DISABLE_PIP_VERSION_CHECK=1 \
-  PIP_BREAK_SYSTEM_PACKAGES=1
+  PIP_BREAK_SYSTEM_PACKAGES=1 \
+  UV_HTTP_TIMEOUT=1000
 
 USER root
 WORKDIR /app
@@ -32,29 +35,24 @@ RUN \
   groupadd --gid ${GID} ${USERNAME} \
   && useradd --uid ${UID} --gid ${GID} --create-home --shell /bin/bash ${USERNAME} \
   && apt-get update && apt-get install -y --no-install-recommends \
-  curl \
-  jq \
-  unzip \
-  build-essential \
-  catatonit \
-  python3-venv \
-  && python3 -m venv /opt/venv \
-  && . /opt/venv/bin/activate \
+  curl unzip build-essential catatonit jq lsb-release \
+  nano vim tree \
   && curl -fsSL -o /tmp/app.zip "https://github.com/jwohlwend/boltz/archive/refs/tags/v${VERSION}.zip" \
-  && unzip -q /tmp/app.zip -d /app \
-  && mv /app/boltz-${VERSION}/LICENSE /app \
-  && mv /app/boltz-${VERSION}/README.md /app \
-  && mv /app/boltz-${VERSION}/pyproject.toml /app \
-  && mv /app/boltz-${VERSION}/docs /app \
-  && mv /app/boltz-${VERSION}/examples /app \
-  && mv /app/boltz-${VERSION}/scripts /app \
-  && mv /app/boltz-${VERSION}/src /app \
+  && unzip -q /tmp/app.zip -d /tmp/app \
+  && cp -R /tmp/app/boltz-${VERSION}/LICENSE /app \
+  && cp -R /tmp/app/boltz-${VERSION}/README.md /app \
+  && cp -R /tmp/app/boltz-${VERSION}/pyproject.toml /app \
+  && cp -R /tmp/app/boltz-${VERSION}/docs /app \
+  && cp -R /tmp/app/boltz-${VERSION}/examples /app \
+  && cp -R /tmp/app/boltz-${VERSION}/scripts /app \
+  && cp -R /tmp/app/boltz-${VERSION}/src /app \
   && printf "UpdateMethod=docker\nBranch=master\nPackageVersion=%s\nPackageAuthor=[RareCompute](https://github.com/RareCompute)\n" "${VERSION}" > /app/package_info \
-  && pip install --no-cache-dir \
-  torch \
-  torchvision \
-  torchaudio \
-  && pip install --no-cache-dir --upgrade . \
+  && 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 pip install torch torchvision torchaudio \
+  && uv pip install -r /tmp/app/boltz-${VERSION}/pyproject.toml \
   && chown -R ${UID}:${GID} /app && chmod -R 755 /app \
   && chown -R ${UID}:${GID} /opt/venv  && chmod -R 755 /opt/venv \
   && apt-get purge -y build-essential \
@@ -72,6 +70,10 @@ VOLUME ["/cache"]
 WORKDIR /data
 VOLUME ["/data"]
 
+ENV \
+  CACHE_PATH="/cache" \
+  DATA_PATH="/data"
+
 WORKDIR /app
 
 ENTRYPOINT ["/usr/bin/catatonit", "--", "/entrypoint.sh"]
diff --git a/apps/boltz/entrypoint.sh b/apps/boltz/entrypoint.sh
index eaaca3c..908bdd6 100644
--- a/apps/boltz/entrypoint.sh
+++ b/apps/boltz/entrypoint.sh
@@ -24,7 +24,12 @@ declare -A FLAG_MAP=(
 
 for VAR in "${!FLAG_MAP[@]}"; do
     if [ -n "${!VAR}" ]; then
-        EXTRA_FLAGS+=" ${FLAG_MAP[$VAR]} ${!VAR}"
+        VALUE_LOWER=$(echo "${!VAR}" | tr '[:upper:]' '[:lower:]')
+        if [[ "$VALUE_LOWER" == "true" ]]; then
+            EXTRA_FLAGS+=" ${FLAG_MAP[$VAR]}"
+        elif [[ "$VALUE_LOWER" != "false" ]]; then
+            EXTRA_FLAGS+=" ${FLAG_MAP[$VAR]} ${!VAR}"
+        fi
     fi
 done
 
diff --git a/apps/chai/Dockerfile b/apps/chai/Dockerfile
new file mode 100644
index 0000000..22faf09
--- /dev/null
+++ b/apps/chai/Dockerfile
@@ -0,0 +1,77 @@
+FROM docker.io/library/python:3.12-slim-bookworm
+
+LABEL \
+  maintainer="Liana64" \
+  org.opencontainers.image.source="https://github.com/chaidiscovery/chai-lab"
+
+ARG TARGETPLATFORM
+ARG VERSION
+ARG CHANNEL
+ARG DEBIAN_FRONTEND=noninteractive
+
+ENV \
+  NVIDIA_DRIVER_CAPABILITIES="compute,video,utility,graphics" \
+  VIRTUAL_ENV=/opt/venv \
+  PATH="${VIRTUAL_ENV}:$PATH" \
+  UMASK="0002" \
+  LANG=C.UTF-8 \
+  TZ="Etc/UTC" \
+  USERNAME=rare \
+  UID=900 \
+  GID=900 \
+  PYTHONDONTWRITEBYTECODE=1 \
+  PYTHONUNBUFFERED=1 \
+  PYTHONFAULTHANDLER=1 \
+  PIP_ROOT_USER_ACTION=ignore \
+  PIP_NO_CACHE_DIR=1 \
+  PIP_DISABLE_PIP_VERSION_CHECK=1 \
+  PIP_BREAK_SYSTEM_PACKAGES=1 \
+  UV_HTTP_TIMEOUT=1000
+
+USER root
+WORKDIR /app
+
+RUN \
+  groupadd --gid ${GID} ${USERNAME} \
+  && useradd --uid ${UID} --gid ${GID} --create-home --shell /bin/bash ${USERNAME} \
+  && 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" \
+  && 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 \
+  && cp -R /tmp/app/chai-lab-${VERSION}/examples /app \
+  && cp -R /tmp/app/chai-lab-${VERSION}/tests /app \
+  && printf "UpdateMethod=docker\nBranch=master\nPackageVersion=%s\nPackageAuthor=[RareCompute](https://github.com/RareCompute)\n" "${VERSION}" > /app/package_info \
+  && 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 pip install -r /tmp/app/chai-lab-${VERSION}/requirements.in \
+  && chown -R ${UID}:${GID} ${VIRTUAL_ENV}  && chmod -R 755 ${VIRTUAL_ENV} \
+  && apt-get purge -y build-essential \
+  && apt-get autoremove -y \
+  && apt-get clean \
+  && rm -rf /root/.cache /var/lib/apt/lists/* /tmp/* /var/tmp/* \
+  && chsh -s /bin/bash
+
+COPY --chown=${UID}:${GID} ./apps/chai/entrypoint.sh /entrypoint.sh
+RUN chmod -R 755 /entrypoint.sh
+
+USER ${USERNAME}
+WORKDIR /app
+
+ENTRYPOINT ["/usr/bin/catatonit", "--", "/entrypoint.sh"]
diff --git a/apps/chai/README.md b/apps/chai/README.md
new file mode 100644
index 0000000..347701c
--- /dev/null
+++ b/apps/chai/README.md
@@ -0,0 +1,18 @@
+# Chai
+
+From the [upstream repository](https://github.com/chaidiscovery/chai-lab)
+
+> Chai-1 is a multi-modal foundation model for molecular structure prediction that performs at the state-of-the-art across a variety of benchmarks. Chai-1 enables unified prediction of proteins, small molecules, DNA, RNA, glycosylations, and more.
+
+## Environment variables
+
+You can configure the docker image using the below environment variables
+
+| Environment Variable | CLI Flag           | Type      | Default Value    | Description                                     |
+| -------------------- | ------------------ | --------- | ---------------- | ----------------------------------------------- |
+| `USERNAME`           | N/A                | `STR`     | `rare`           | Username for the container                      |
+| `UID`                | N/A                | `INTEGER` | `900`            | UID for the container                           |
+| `GID`                | N/A                | `INTEGER` | `900`            | GID for the container                           |
+| `CHAI_DOWNLOADS_DIR` | N/A                | `PATH`    | `/tmp/downloads` | The path where downloads will be cached.        |
+| `USE_MSA_SERVER`     | `--use-msa-server` | `BOOLEAN` | `False`          | Whether to use the MSA server to generate MSAs. |
+| `MSA_SERVER_URL`     | `--msa-server-url` | `STR`     | N/A              | Use a custom MSA server                         |
diff --git a/apps/chai/ci/goss.yaml b/apps/chai/ci/goss.yaml
new file mode 100644
index 0000000..b635bdb
--- /dev/null
+++ b/apps/chai/ci/goss.yaml
@@ -0,0 +1,7 @@
+---
+# yaml-language-server: $schema=https://raw.githubusercontent.com/goss-org/goss/master/docs/schema.yaml
+file:
+  /opt/venv/bin/chai:
+    exists: true
+  /app/LICENSE:
+    exists: true
diff --git a/apps/chai/ci/latest.sh b/apps/chai/ci/latest.sh
new file mode 100755
index 0000000..83c2d0d
--- /dev/null
+++ b/apps/chai/ci/latest.sh
@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+version=$(curl -sX GET "https://api.github.com/repos/chaidiscovery/chai-lab/releases/latest" | jq --raw-output '.tag_name' 2>/dev/null)
+version="${version#*v}"
+version="${version#*release-}"
+printf "%s" "${version}"
diff --git a/apps/chai/entrypoint.sh b/apps/chai/entrypoint.sh
new file mode 100644
index 0000000..077799b
--- /dev/null
+++ b/apps/chai/entrypoint.sh
@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+
+EXTRA_FLAGS=""
+declare -A FLAG_MAP=(
+    [USE_MSA_SERVER]="--use-msa-server"
+    [MSA_SERVER_URL]="--msa-server-url"
+)
+
+for VAR in "${!FLAG_MAP[@]}"; do
+    if [ -n "${!VAR}" ]; then
+        VALUE_LOWER=$(echo "${!VAR}" | tr '[:upper:]' '[:lower:]')
+        if [[ "$VALUE_LOWER" == "true" ]]; then
+            EXTRA_FLAGS+=" ${FLAG_MAP[$VAR]}"
+        elif [[ "$VALUE_LOWER" != "false" ]]; then
+            EXTRA_FLAGS+=" ${FLAG_MAP[$VAR]} ${!VAR}"
+        fi
+    fi
+done
+
+exec \
+  chai \
+  "$@" \
+  $EXTRA_FLAGS
diff --git a/apps/chai/metadata.yaml b/apps/chai/metadata.yaml
new file mode 100644
index 0000000..608541f
--- /dev/null
+++ b/apps/chai/metadata.yaml
@@ -0,0 +1,11 @@
+---
+#yamllint disable
+app: chai
+semver: true
+channels:
+  - name: stable
+    platforms: ["linux/amd64", "linux/arm64"]
+    stable: true
+    tests:
+      enabled: false
+      type: cli