build.yml

# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

# A workflow to trigger TE build on GitHub

name: 'Build'
on:
  pull_request:
  workflow_dispatch:
jobs:
  pytorch:
    name: 'PyTorch'
    runs-on: [ nv-8g-cicd-te ]
    defaults:
      run:
        shell: bash
    container:
      image: harbor.baai.ac.cn/flagscale/cuda12.8.1-torch2.7.1-python3.10-te2.9:20260209
      ports:
        - 80:80
      options: >-
        --gpus all 
        --shm-size=500g 
        --privileged 
        --ipc=host 
        --ulimit memlock=-1 
        --ulimit stack=67108864 
        --ulimit nofile=65535:65535 
        --user root
        --pull never
    steps:
      - name: Configure Git Safe Directory on Cuda
        run: /usr/bin/git config --global safe.directory '*'

      - name: 'Checkout'
        uses: actions/checkout@v4
        with:
          fetch-depth: 0
          submodules: recursive
          set-safe-directory: true

      - name: 'Setup Environment'
        run: |
          source /opt/miniconda3/etc/profile.d/conda.sh
          conda activate flagscale-train
          echo "PATH=$PATH" >> $GITHUB_ENV

      - name: 'Build'
        run: |
          pip uninstall transformer_engine transformer_engine_torch -y || true
          echo "GITHUB_WORKSPACE=$GITHUB_WORKSPACE"
          cd $GITHUB_WORKSPACE
          pip install nvdlfw-inspect
          pip install expecttest
          pip install . -v --no-deps --no-build-isolation
        env:
          NVTE_FRAMEWORK: pytorch
          TE_WITH_NCCL: '1'
          NVTE_WITH_CUDA: '1'
          CUDA_HOME: /usr/local/cuda-12.8
          NVCC: /usr/local/cuda-12.8/bin/nvcc

      - name: 'Sanity check'
        run: 
          python3 tests/pytorch/test_sanity_import.py