Merge branch 'NVIDIA:main' into build-test-publish-wheel

.github/workflows/gpuci.yml

@@ -1,6 +1,9 @@
 name: "GPU CI/CD"
 
 on:
+  push:
+    branches:
+      - main
   pull_request:
     branches:
       # We can run gpuCI on any PR targeting these branches
@@ -11,12 +14,25 @@ on:
     # If new commits are added, the "gpuCI" label has to be removed and re-added to rerun gpuCI
     types: [ labeled ]
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
 jobs:
   # First, we build and push a NeMo-Curator container
   build-container:
     # "build-container" job is run if the "gpuci" label is added to the PR
-    if: ${{ github.event.label.name == 'gpuci' }}
-    uses: ./.github/workflows/_build_container.yml
+    if: ${{ github.event.label.name == 'gpuci' || github.ref == 'refs/heads/main' }}
+    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/[email protected]
+    with:
+      image-name: nemo_curator_container
+      dockerfile: Dockerfile
+      image-label: nemo-curator
+      build-args: |
+        IMAGE_LABEL=nemo-curator
+        FORKED_REPO_URL=https://github.com/${{ github.event.pull_request.head.repo.full_name }}.git
+        CURATOR_COMMIT=${{ github.event.pull_request.head.sha || github.sha }}
+      prune-filter-timerange: 24h
 
   # Then, we run our PyTests in the container we just built
   run-gpu-tests:
@@ -25,7 +41,7 @@ jobs:
     # It has 2 A100 GPUs
     runs-on: self-hosted-azure
     # "run-gpu-tests" job is run if the "gpuci" label is added to the PR
-    if: ${{ github.event.label.name == 'gpuci' }}
+    if: ${{ github.event.label.name == 'gpuci' || github.ref == 'refs/heads/main' }}
 
     steps:
       # If something went wrong during the last cleanup, this step ensures any existing container is removed

Dockerfile

@@ -3,8 +3,28 @@
 ARG CUDA_VER=12.5.1
 ARG LINUX_VER=ubuntu22.04
 ARG PYTHON_VER=3.10
-FROM rapidsai/ci-conda:cuda${CUDA_VER}-${LINUX_VER}-py${PYTHON_VER}
+ARG IMAGE_LABEL
+ARG FORKED_REPO_URL
+ARG CURATOR_COMMIT
 
+FROM rapidsai/ci-conda:cuda${CUDA_VER}-${LINUX_VER}-py${PYTHON_VER} as curator-update
+# Needed to navigate to and pull the forked repository's changes
+ARG FORKED_REPO_URL
+ARG CURATOR_COMMIT
+
+# Clone the user's repository, find the relevant commit, and install everything we need
+RUN bash -exu <<EOF
+  mkdir -p /opt/NeMo-Curator
+  cd /opt/NeMo-Curator
+  git init
+  git remote add origin $FORKED_REPO_URL
+  git fetch origin '+refs/pull/*/merge:refs/remotes/pull/*/merge'
+  git checkout $CURATOR_COMMIT
+EOF
+
+
+FROM rapidsai/ci-conda:cuda${CUDA_VER}-${LINUX_VER}-py${PYTHON_VER}
+LABEL "nemo.library"=${IMAGE_LABEL}
 WORKDIR /opt
 
 # Install the minimal libcu* libraries needed by NeMo Curator
@@ -15,19 +35,21 @@ RUN conda create -y --name curator -c conda-forge -c nvidia \
   libcublas \
   libcurand \
   libcusparse \
-  libcusolver
+  libcusolver && \
+  source activate curator && \
+  pip install --upgrade cython pytest pip
 
-# Needed to navigate to and pull the forked repository's changes
-ARG FORKED_REPO_URL
-ARG CURATOR_COMMIT
+RUN \
+  --mount=type=bind,source=/opt/NeMo-Curator/nemo_curator/__init__.py,target=nemo_curator/__init__.py,from=curator-update \
+  --mount=type=bind,source=/opt/NeMo-Curator/pyproject.toml,target=pyproject.toml,from=curator-update \
+  source activate curator && \
+  pip install ".[all]"
+
+COPY --from=curator-update /opt/NeMo-Curator/ /opt/NeMo-Curator/
 
 # Clone the user's repository, find the relevant commit, and install everything we need
 RUN bash -exu <<EOF
-  git clone $FORKED_REPO_URL
-  cd NeMo-Curator
-  git fetch origin $CURATOR_COMMIT --depth=1
-  git checkout $CURATOR_COMMIT
   source activate curator
-  pip install --upgrade cython pytest pip
+  cd /opt/NeMo-Curator/
   pip install --extra-index-url https://pypi.nvidia.com ".[all]"
 EOF