Adding steps to cache qutip wheel

Signed-off-by: Sachin Pisal <spisal@nvidia.com>
NVIDIA · Dec 20, 2024 · 8046fe9 · 8046fe9
1 parent 1a4f99e
commit 8046fe9
Show file tree

Hide file tree

Showing 4 changed files with 78 additions and 2 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -211,9 +211,38 @@ jobs:
       devdeps_archive: ${{ fromJson(needs.config.outputs.json).tar_archive[format('{0}-{1}', matrix.platform, matrix.toolchain)] }}
       export_environment: ${{ github.event_name == 'workflow_dispatch' && inputs.export_environment }}
 
+  cache_qutip_wheel:
+    name: Build and Cache qutip Wheel
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        platform: [amd64, arm64]
+    steps:
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+
+      - name: Install Dependencies
+        run: |
+          python -m pip install --upgrade pip setuptools wheel
+
+      - name: Build qutip Wheel
+        run: |
+          mkdir -p wheelhouse
+          python -m pip wheel qutip --wheel-dir=./wheelhouse
+          ls -lh ./wheelhouse
+
+      - name: Upload Wheel Artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: qutip-wheel-${{ matrix.platform }}
+          path: ./wheelhouse
+          retention-days: 7
+
   docker_image:
     name: Create Docker images
-    needs: config
+    needs: [config, cache_qutip_wheel]
     strategy:
       matrix:
         platform: [amd64, arm64]
@@ -227,6 +256,19 @@ jobs:
       devdeps_image: ${{ fromJson(needs.config.outputs.json).image_hash[format('{0}-gcc11', matrix.platform)] }}
       devdeps_cache: ${{ fromJson(needs.config.outputs.json).cache_key[format('{0}-gcc11', matrix.platform)] }}
       devdeps_archive: ${{ fromJson(needs.config.outputs.json).tar_archive[format('{0}-gcc11', matrix.platform)] }}
+    steps:
+      - name: Download Cached qutip Wheel
+        uses: actions/download-artifact@v4
+        with:
+          name: qutip-wheel-${{ matrix.platform }}
+          path: ./wheelhouse
+
+      - name: Build Docker Image
+        run: |
+          docker build \
+            --build-arg QUTIP_WHEEL=./wheelhouse/qutip*.whl \
+            --build-arg DEVDEPS_IMAGE=${{ fromJson(needs.config.outputs.json).image_hash[format('{0}-gcc11', matrix.platform)] }} \
+            --tag myrepo/cudaq:${{ matrix.platform }} .
 
   python_wheels:
     name: Create Python wheels

diff --git a/.github/workflows/docker_images.yml b/.github/workflows/docker_images.yml
@@ -321,6 +321,12 @@ jobs:
             org.opencontainers.image.title=cuda-quantum-dev
             org.opencontainers.image.description=Dev environment for CUDA Quantum (debug build)
 
+      - name: Download Cached qutip Wheel
+        uses: actions/download-artifact@v4
+        with:
+          name: qutip-wheel-${{ inputs.platforms }}
+          path: ./wheelhouse
+
       - name: Build cuda-quantum-dev image (debug)
         id: docker_build
         uses: docker/build-push-action@v5
@@ -329,6 +335,7 @@ jobs:
           file: ./docker/build/cudaq.dev.Dockerfile
           build-args: |
             base_image=${{ steps.prereqs.outputs.base_image }}
+            QUTIP_WHEEL=./wheelhouse/qutip*.whl
             install="CMAKE_BUILD_TYPE=Debug"
             git_source_sha=${{ github.sha }}
           tags: ${{ steps.metadata.outputs.tags }}
@@ -582,6 +589,12 @@ jobs:
         run: |
           docker run --privileged multiarch/qemu-user-static:latest --reset -p yes --credential yes
 
+      - name: Download Cached qutip Wheel
+        uses: actions/download-artifact@v4
+        with:
+          name: qutip-wheel-${{ inputs.platforms }}
+          path: ./wheelhouse
+
       - name: Build cuda-quantum image
         id: cudaq_build
         if: success() && !cancelled()
@@ -592,6 +605,7 @@ jobs:
           build-args: |
             cudaqdev_image=${{ steps.prereqs.outputs.dev_image_name }}@${{ steps.release_build.outputs.digest }}
             base_image=${{ steps.prereqs.outputs.base_image }}
+            QUTIP_WHEEL=./wheelhouse/qutip*.whl
             release_version=${{ steps.prereqs.outputs.image_tag }}
           tags: ${{ steps.cudaq_metadata.outputs.tags }}
           labels: ${{ steps.cudaq_metadata.outputs.labels }}

diff --git a/docker/build/cudaq.dev.Dockerfile b/docker/build/cudaq.dev.Dockerfile
@@ -30,6 +30,16 @@ ARG destination="$CUDAQ_REPO_ROOT"
 ADD "$workspace" "$destination"
 WORKDIR "$destination"
 
+# Accept QUTIP_WHEEL as a build argument
+ARG QUTIP_WHEEL
+
+# Install qutip from the wheel if provided
+RUN if [ -n "$QUTIP_WHEEL" ]; then \
+        pip install "$QUTIP_WHEEL"; \
+    else \
+        pip install qutip; \
+    fi
+
 # mpich or openmpi
 ARG mpi=
 RUN if [ -n "$mpi" ]; \

diff --git a/docker/release/cudaq.Dockerfile b/docker/release/cudaq.Dockerfile
@@ -62,10 +62,20 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     && ln -s /bin/python3 /bin/python
 RUN apt-get update && apt-get install -y --no-install-recommends gcc g++ python3-dev \
     # Ref: https://github.com/qutip/qutip/issues/2412
-    && python3 -m pip install --no-cache-dir notebook==7.1.3 "qutip<5" matplotlib \
+    && python3 -m pip install --no-cache-dir notebook==7.1.3 matplotlib \
     && apt-get remove -y gcc g++ python3-dev \
     && apt-get autoremove -y && apt-get clean && rm -rf /var/lib/apt/lists/*
 
+# Accept QUTIP_WHEEL as a build argument
+ARG QUTIP_WHEEL
+
+# Install qutip from the wheel if provided
+RUN if [ -n "$QUTIP_WHEEL" ]; then \
+        pip install "$QUTIP_WHEEL"; \
+    else \
+        pip install qutip; \
+    fi
+
 # Copy over the CUDA-Q installation, and the necessary compiler tools.
 
 ARG release_version=