python-build-test.yml

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

name: Python Build & Test

on:
  push:
    branches:
      - main
  pull_request:
  workflow_dispatch:

permissions:
  contents: read

concurrency:
  group: python-build-test-${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true

jobs:
  build-and-test:
    name: py${{ matrix.python }} rdkit${{ matrix.rdkit }} cuda${{ matrix.cuda }}
    runs-on: colossus
    container:
      image: nvcr.io/nvidia/cuda:12.6.3-devel-ubuntu22.04
      options: --gpus all
    timeout-minutes: 120
    strategy:
      fail-fast: false
      matrix:
        include:
          # Exercise the latest RDKit release that supports CUDA 12 in the
          # conda-build matrix. Python 3.12 keeps this job compatible with the
          # minimum supported Volta GPU while still testing a current Python.
          - python: "3.12"
            rdkit: "2025.9.6"
            cuda: "12.6"

    steps:
      - name: Verify GPU access and select CUDA architectures
        shell: bash
        run: |
          set -xeuo pipefail
          nvidia-smi

          mapfile -t compute_caps < <(
            nvidia-smi --query-gpu=compute_cap --format=csv,noheader |
              tr -d ' .' |
              sort -u
          )
          if (( ${#compute_caps[@]} == 0 )); then
            echo "No GPU compute capabilities were reported" >&2
            exit 1
          fi
          for compute_cap in "${compute_caps[@]}"; do
            if [[ ! "${compute_cap}" =~ ^[0-9]+$ ]] || (( compute_cap < 70 )); then
              echo "Unsupported GPU compute capability: ${compute_cap}" >&2
              exit 1
            fi
          done
          cuda_architectures=$(IFS=';'; echo "${compute_caps[*]}")
          echo "CUDA_ARCHITECTURES=${cuda_architectures}" | tee -a "${GITHUB_ENV}"

      - name: Check out source tree
        uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5

      - name: Install conda and native dependencies
        shell: bash
        run: bash admin/ci/setup_dependencies.sh "${{ matrix.python }}" "${{ matrix.rdkit }}"

      - name: Install Python runtime and test dependencies
        shell: bash
        run: |
          set -xeuo pipefail
          . /usr/local/anaconda/etc/profile.d/conda.sh
          conda activate base
          # Installing the native build dependencies into base can replace
          # conda-libmamba-solver. Use the standalone solver retained by the
          # setup step so this transaction cannot fall back to classic conda.
          mamba install -q -y \
            "cuda-version=${{ matrix.cuda }}" \
            pytorch-gpu \
            "scikit-build>=0.18" \
            "setuptools>=42" \
            ninja \
            wheel \
            pandas \
            psutil \
            optuna

      - name: Verify PyTorch GPU support
        shell: bash
        working-directory: /tmp
        run: |
          set -xeuo pipefail
          . /usr/local/anaconda/etc/profile.d/conda.sh
          conda activate base
          python - <<'PY'
          import torch

          assert torch.cuda.is_available(), "PyTorch cannot access CUDA"
          print(f"PyTorch {torch.__version__}, CUDA {torch.version.cuda}")
          print(f"PyTorch architectures: {torch.cuda.get_arch_list()}")
          for device in range(torch.cuda.device_count()):
              print(f"GPU {device}: {torch.cuda.get_device_name(device)}")
              torch.zeros(1, device=f"cuda:{device}")
          PY

      - name: Build and install nvMolKit
        shell: bash
        working-directory: /tmp
        env:
          CUDA_HOME: /usr/local/cuda
          NVMOLKIT_CUDA_TARGET_MODE: default
        run: |
          set -xeuo pipefail
          . /usr/local/anaconda/etc/profile.d/conda.sh
          conda activate base
          export CMAKE_ARGS="-DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES}"
          CMAKE_BUILD_PARALLEL_LEVEL="$(nproc)"
          export CMAKE_BUILD_PARALLEL_LEVEL
          python -m pip install \
            --no-deps \
            --no-build-isolation \
            --verbose \
            "${GITHUB_WORKSPACE}"

      - name: Run pytest
        shell: bash
        working-directory: /tmp
        run: |
          set -xeuo pipefail
          . /usr/local/anaconda/etc/profile.d/conda.sh
          conda activate base
          python -m pytest \
            "${GITHUB_WORKSPACE}/nvmolkit/tests" \
            --tb=short \
            -m "not long"