.gitlab-ci.yml

workflow:
  rules:
    - if: ($CI_PIPELINE_SOURCE == "merge_request_event" && $CI_MERGE_REQUEST_LABELS =~ /Run tests and nightly/) || ($CI_PIPELINE_SOURCE == "schedule")
      variables:
        JET_CUSTOM_FILTER: "type == 'build' or 'mr' in spec.scope or 'nightly' in spec.scope"
    - if: $CI_PIPELINE_SOURCE == "merge_request_event" && $CI_MERGE_REQUEST_LABELS =~ /Run tests/
      variables:
        JET_CUSTOM_FILTER: "type == 'build' or 'mr' in spec.scope"
    - if: $CI_PIPELINE_SOURCE == "merge_request_event" && $CI_MERGE_REQUEST_LABELS =~ /Build only/
      variables:
        JET_CUSTOM_FILTER: "type == 'build'"
    # always run MR pipelines
    - if: $CI_PIPELINE_SOURCE == "merge_request_event"
    # always run web pipelines
    - if: $CI_PIPELINE_SOURCE == "web"
    # do not run branch pipelines if open MR exists
    - if: $CI_COMMIT_BRANCH && $CI_OPEN_MERGE_REQUESTS
      when: never
    # run branch pipeline if no open MR
    - if: $CI_COMMIT_BRANCH
    

stages:
  - build
  - test
  - jet

variables:
  SELENE_ADLR_CI_PATH: "/lustre/fsw/adlr/adlr-nlp/adlr_ci/megatron"
  DATA_DIR: "/lustre/fsw/adlr/adlr-nlp/adlr_ci/megatron/data"
  PYTHON_VIRTUAL_ENV: /lustre/fsw/adlr/adlr-nlp/adlr_ci/cienv/bin/activate
  JET_CUSTOM_FILTER:
    description: |
      Selects what functional tests to run. For mr tests: "type == 'build' or 'mr' in spec.scope". For nightly tests: "type == 'build' or 'nightly' in spec.scope"
    value: ""
  DISPLAY_OUTPUT: "True" # Set to true for new tests to copy the logs for creating golden truth file
  TIME_LIMIT: "10:00" # Default time limit for all jobs
  MOE_GROUPED_GEMM: 0 # Set to 1 to enable grouped gemm for MoE
  JET_CLUSTER_BRANCH:
    value: "mcore/draco-oci"
    options:
      - "mcore/draco-oci"
      - "mcore/eos"
    description: '"mcore/draco-oci" for OCI-IAD, "mcore/eos" for EOS'
  CI_IMAGE: gitlab-master.nvidia.com:5005/adlr/megatron-lm/mcore_ci
  LINTING_IMAGE: gitlab-master.nvidia.com:5005/adlr/megatron-lm/mcore_linting

build_image:
  tags:
    - 8xL40S
  image: docker:26.1.4-dind
  stage: build
  parallel:
    matrix:
      - IMAGE: CI_IMAGE
        FILE: Dockerfile.ci
        BASE_IMAGE: nvcr.io/nvidia/pytorch:24.01-py3
      - IMAGE: LINTING_IMAGE
        FILE: Dockerfile.linting
        BASE_IMAGE: python:3.10
  before_script:
    - echo "$CI_REGISTRY_PASSWORD" | docker login $CI_REGISTRY -u $CI_REGISTRY_USER --password-stdin
  script:
    - |
      eval "IMAGE=\$$IMAGE"

      docker build \
        -f $FILE \
        -t ${IMAGE}:${CI_PIPELINE_ID}  \
        --cache-to type=inline \
        --build-arg FROM_IMAGE_NAME=$BASE_IMAGE \
        --cache-from type=registry,ref=${IMAGE}:buildcache .

      docker push ${IMAGE}:${CI_PIPELINE_ID}  

      if [[ "$CI_COMMIT_BRANCH" = "main" ]]; then
        docker tag ${IMAGE}:${CI_PIPELINE_ID} ${IMAGE}:buildcache
        docker push ${IMAGE}:buildcache
      fi
  interruptible: true

unit_tests:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s --cov-report=term --cov-report=html --cov=megatron/core --no-cov-on-fail tests/unit_tests
  coverage: '/(?i)total.*? (100(?:\.0+)?\%|[1-9]?\d(?:\.\d+)?\%)$/'
  artifacts:
    paths:
      - coverage
    expire_in: 30 days
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event" && $CI_MERGE_REQUEST_LABELS =~ /Run tests/'
      when: never
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

unit_tests-data:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s tests/unit_tests/data
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event"'
      when: always
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

unit_tests-dist-checkpointing:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s tests/unit_tests/dist_checkpointing
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event" && $CI_MERGE_REQUEST_LABELS =~ /Run tests/'
      when: always
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

unit_tests-fusions:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s tests/unit_tests/fusions
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event"'
      when: always
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

unit_tests-inference:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s tests/unit_tests/inference
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event"'
      when: always
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

unit_tests-models:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s tests/unit_tests/models
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event"'
      when: always
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

unit_tests-pipeline-parallel:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s tests/unit_tests/pipeline_parallel
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event"'
      when: always
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

unit_tests-tensor-parallel:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s tests/unit_tests/tensor_parallel
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event"'
      when: always
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

unit_tests-transformer:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s tests/unit_tests/transformer
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event"'
      when: always
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

unit_tests-top-py:
  image: ${CI_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - 8xL40S
  stage: test
  script:
    - torchrun --nproc_per_node=8 -m pytest -x -v -s tests/unit_tests/*.py
  rules:
    - if: '$CI_PIPELINE_SOURCE == "merge_request_event"'
      when: always
    - if: $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH
  interruptible: true

docs_build_test:
  image: gitlab-master.nvidia.com:5005/adlr/megatron-lm/python-format:0.0.1
  stage: test
  tags:
    - os/linux
  script:
    - cd ..
    - rm -rf documentation && git clone https://gitlab-ci-token:${CI_JOB_TOKEN}@gitlab-master.nvidia.com/nemo-megatron-core-tme/documentation.git
    - mv megatron-lm/ documentation/
    - cd documentation/
    - ./repo docs
  allow_failure: true
  except:
    - main
  interruptible: true

formatting:
  image: ${LINTING_IMAGE}:${CI_PIPELINE_ID} 
  tags:
    - os/linux
  stage: test
  before_script:
    - git fetch origin main
  script:
    - |
      set -x
      CHANGED_FILES=$(git diff --name-only origin/main | grep '^megatron/core' || true)

      if [ -n "$CHANGED_FILES" ]; then
        black --check --verbose --diff $CHANGED_FILES
      fi

      if [ -n "$CHANGED_FILES" ]; then
        isort --check $CHANGED_FILES    
      fi
  rules:
    - when: always
  interruptible: true

include:
  - jet-tests.yml