From a4f4a56d77fc1e6a9211737343d2b5c119502bfc Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Mon, 2 Dec 2024 18:27:46 -0500
Subject: [PATCH] build causal_conv1d and mamba-ssm into the base image (#2113)

* build causal_conv1d and mamba-ssm into the base image

* also build base images on changes to Dockerfile-base and base workflow yaml
---
 .github/workflows/base.yml | 10 ++++++++++
 docker/Dockerfile-base     |  4 +++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
index e086a5f961..640d2cd7ae 100644
--- a/.github/workflows/base.yml
+++ b/.github/workflows/base.yml
@@ -1,6 +1,16 @@
 name: ci-cd-base
 
 on:
+  push:
+    branches:
+      - "main"
+    paths:
+      - 'Dockerfile-base'
+      - '.github/workflows/base.yml'
+  pull_request:
+    paths:
+      - 'Dockerfile-base'
+      - '.github/workflows/base.yml'
   workflow_dispatch:
 
 jobs:
diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
index 3f13bba30a..7eab3b3e43 100644
--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -29,7 +29,9 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
 WORKDIR /workspace
 
 RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
-    python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA
+    python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA && \
+    python3 -m pip install --no-cache-dir "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" && \
+    python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main"
 
 RUN git lfs install --skip-repo && \
     pip3 install awscli && \