.gitlab-ci.yml

include:
  - project: "dl/devops/gitlab-ci-slurm"
    ref: master
    file: "/.gitlab-ci.yml"
  - /ci/common.yml
  - /ci/template.yml
  - /ci/benchmark.yml

nightly_build_train:
  extends: .build_nightly
  variables:
    REMOTE_REPO: https://github.com/NVIDIA-Merlin/Merlin.git
    DST_IMAGE: ${IMAGE_TRAIN}
    DOCKER_FILE: ./training/dockerfile.ctr
    BUILD_ARGS: --build-arg HUGECTR_DEV_MODE=true

nightly_build_train_with_hdfs:
  extends: .build_nightly
  variables:
    REMOTE_REPO: https://github.com/NVIDIA-Merlin/Merlin.git
    DST_IMAGE: ${IMAGE_TRAIN_WITH_HDFS}
    DOCKER_FILE: ./training/dockerfile.ctr
    BUILD_ARGS: --build-arg HUGECTR_DEV_MODE=true --build-arg INSTALL_HDFS=true
  
nightly_build_inference:
  extends: .build_nightly
  variables:
    REMOTE_REPO: https://github.com/NVIDIA-Merlin/Merlin.git
    DST_IMAGE: ${IMAGE_INFER}
    DOCKER_FILE: ./inference/dockerfile.ctr
    BUILD_ARGS: --build-arg HUGECTR_DEV_MODE=true

nightly_build_sok_tf2:
  extends: .build_nightly
  variables:
    REMOTE_REPO: https://github.com/NVIDIA-Merlin/Merlin.git
    DST_IMAGE: ${IMAGE_SOK_TF2}
    DOCKER_FILE: ./training/dockerfile.tf
    BUILD_ARGS: --build-arg HUGECTR_DEV_MODE=true
  
nightly_build_sok_tf1:
  extends: .build_nightly
  variables:
    REMOTE_REPO: https://github.com/NVIDIA-Merlin/Merlin.git
    DST_IMAGE: ${IMAGE_SOK_TF1}
    DOCKER_FILE: ./training/dockerfile.tf
    BUILD_ARGS: --build-arg HUGECTR_DEV_MODE=true --build-arg IMAGE=nvcr.io/nvidia/tensorflow:21.11-tf1-py3 --build-arg INSTALL_NVT=false
  
nightly_build_unified_container.ctr:
  extends: .build_nightly
  variables:
    REMOTE_REPO: https://github.com/NVIDIA-Merlin/Merlin.git
    DST_IMAGE: ${CI_REGISTRY}/dl/hugectr/hugectr:unified.ctr.latest
    DOCKER_FILE: ./training/dockerfile.ctr
    BUILD_ARGS: --build-arg _CI_JOB_TOKEN=${RD_CI_JOB_TOKEN} --build-arg _HUGECTR_REPO=gitlab-master.nvidia.com/dl/hugectr/hugectr.git --build-arg HUGECTR_VER=${CI_COMMIT_SHA}
  
nightly_build_unified_container.tf:
  extends: .build_nightly
  variables:
    REMOTE_REPO: https://github.com/NVIDIA-Merlin/Merlin.git
    DST_IMAGE: ${CI_REGISTRY}/dl/hugectr/hugectr:unified.tf.latest
    DOCKER_FILE: ./training/dockerfile.tf
    BUILD_ARGS: --build-arg _CI_JOB_TOKEN=${RD_CI_JOB_TOKEN} --build-arg _HUGECTR_REPO=gitlab-master.nvidia.com/dl/hugectr/hugectr.git --build-arg HUGECTR_VER=${CI_COMMIT_SHA}

nightly_build_unified_container.tri:
  extends: .build_nightly
  variables:
    REMOTE_REPO: https://github.com/NVIDIA-Merlin/Merlin.git
    DST_IMAGE: ${CI_REGISTRY}/dl/hugectr/hugectr:unified.tri.latest
    DOCKER_FILE: ./inference/dockerfile.ctr
    BUILD_ARGS: --build-arg _CI_JOB_TOKEN=${RD_CI_JOB_TOKEN} --build-arg _HUGECTR_REPO=gitlab-master.nvidia.com/dl/hugectr/hugectr.git --build-arg HUGECTR_VER=${CI_COMMIT_SHA} --build-arg _HUGECTR_BACKEND_REPO=gitlab-master.nvidia.com/dl/hugectr/hugectr_inference_backend.git --build-arg HUGECTR_BACKEND_VER=hugectr_performance_test

### Stage: build
build_train_single_node:
  extends: .build
  variables:
    FROM_IMAGE: ${IMAGE_TRAIN}
    DST_IMAGE: $TRAIN_IMAGE_VERSIONED
    CMAKE_OPTION: "-DCMAKE_BUILD_TYPE=Release -DKEY_HIT_RATIO=ON -DSM=\"60;61;70;75;80\" -DCLANGFORMAT=OFF"
    BUILD_HUGECTR: 1
    BUILD_HUGECTR2ONNX: 1

build_train_single_node_with_hdfs:
  extends: .build
  variables:
    FROM_IMAGE: ${IMAGE_TRAIN_WITH_HDFS}
    DST_IMAGE: $TRAIN_IMAGE_VERSIONED_WITH_HDFS
    CMAKE_OPTION: "-DCMAKE_BUILD_TYPE=Release -DKEY_HIT_RATIO=ON -DSM=\"60;61;70;75;80\" -DCLANGFORMAT=OFF -DENABLE_HDFS=ON"
    BUILD_HUGECTR: 1
    BUILD_HUGECTR2ONNX: 1


build_train_multi_node:
  extends: .build
  variables:
    FROM_IMAGE: ${IMAGE_TRAIN}
    DST_IMAGE: $TRAIN_IMAGE_MULTINODE_VERSIONED
    CMAKE_OPTION: "-DCMAKE_BUILD_TYPE=Release -DKEY_HIT_RATIO=ON -DENABLE_MULTINODES=ON -DSM=\"60;61;70;75;80\" -DCLANGFORMAT=OFF"
    BUILD_HUGECTR: 1

build_train_inference:
  extends: .build
  variables:
    FROM_IMAGE: ${IMAGE_TRAIN}
    DST_IMAGE: $TRAIN_INFER_IMAGE_VERSIONED
    CMAKE_OPTION: "-DENABLE_INFERENCE=ON -DCMAKE_BUILD_TYPE=Release -DSM=\"60;61;70;75;80\" -DCLANGFORMAT=OFF"
    BUILD_HUGECTR: 1

### Stage: test
build_inference:
  extends: .build
  variables:
    FROM_IMAGE: ${IMAGE_INFER}
    DST_IMAGE: $INFER_IMAGE_VERSIONED
    CMAKE_OPTION: "-DENABLE_INFERENCE=ON -DCMAKE_BUILD_TYPE=Release -DSM=\"70;75;80\" -DCLANGFORMAT=OFF"
    BUILD_HUGECTR: 1
    BUILD_HUGECTR_BACKEND: 1
    TRITON_BRANCH: r21.06

build_sok_tf2:
  extends: .build
  variables:
    FROM_IMAGE: ${IMAGE_SOK_TF2}
    DST_IMAGE: $SOK_IMAGE_VERSIONED_TF2
    CMAKE_OPTION: "-DSM=\"60;61;70;75;80\""
    BUILD_SOK: 1

build_sok_tf1:
  extends: .build
  variables:
    FROM_IMAGE: ${IMAGE_SOK_TF1}
    DST_IMAGE: $SOK_IMAGE_VERSIONED_TF1
    CMAKE_OPTION: "-DSM=\"60;61;70;75;80\""
    BUILD_SOK: 1

## Stage: test
# unit test
utests:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/utests
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT},/lustre/fsw/devtech/hpc-hugectr/inference/:/hugectr/test/utest/
    SLURM_ACCOUNT: devtech
    WALLTIME: "01:00:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/utest/utest.sub
 

utests_embedding:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/utests_embedding
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT},/lustre/fsw/devtech/hpc-hugectr/inference/:/hugectr/test/utest/
    SLURM_ACCOUNT: devtech
    WALLTIME: "01:00:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/utest/utest_embedding.sub


utests_inference:
  extends: .cluster_test_job
  needs:
    - build_train_inference
  variables:
    GPFSFOLDER: $LOGDIR/utests_inference
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_INFER_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT},/lustre/fsw/devtech/hpc-hugectr/inference/:/hugectr/test/utest/
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:30:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/utest/utest_inference.sub

# single node
criteo:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/criteo
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/criteo/criteo.sub

# criteo_dgx1:
#   extends: .dlcluster_test_job
#   script:
#     - srun -N 1 -p dgx1v,dgx1v16g,dgx1v32g
#       --container-image $TRAIN_IMAGE_VERSIONED
#       --container-mounts ${DATASET_CRITEO_DLGROUP}:${DATASET_MOUNT},$(pwd):${CI_PROJECT_DIR}
#       bash -cx "
#       export PYTHONPATH=${WORK_DIR}/build/build_single/lib/ &&
#       cd ${DATASET_MOUNT}/criteo &&
#       python3 ${WORK_DIR}/test/pybind_test/single_node_test.py --json-file=${WORK_DIR}/test/scripts/criteo_8gpu.json &&
#       python3 ${WORK_DIR}/test/pybind_test/single_node_test.py --json-file=${WORK_DIR}/test/scripts/criteo_multi_slots_8gpu.json"

criteo_multi_node:
  extends: .cluster_test_job
  needs:
    - build_train_multi_node
  variables:
    GPFSFOLDER: $LOGDIR/criteo_multi_node
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_MULTINODE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 2
    TEST_CMD: ./ci/integration_test/criteo/criteo_multi_node.sub

dcn:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/dcn
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "02:00:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/dcn/dcn.sub

dlrm_benchmark_1node:
  extends: .cluster_test_job
  needs:
    - build_train_multi_node
  variables:
    GPFSFOLDER: $LOGDIR/dlrm_benchmark_1node
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_MULTINODE_VERSIONED
    MOUNTS: /raid:/raid
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/dlrm/benchmark_1node.sub

dlrm_benchmark_14node:
  extends: .cluster_test_job
  needs:
    - build_train_multi_node
  variables:
    GPFSFOLDER: $LOGDIR/dlrm_benchmark_14node
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_MULTINODE_VERSIONED
    MOUNTS: /raid:/raid
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    SBATCH_OTHER_PARAMS: --network sharp
    DGXNNODES: 14
    TEST_CMD: ./ci/integration_test/dlrm/benchmark_14node.sub

wdl:
  extends: .cluster_test_job                                                     # test on selene needs to extend .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/wdl                                                      # log dir, usually $LOGDIR + job name
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}                                     # should not change
    CONT: $TRAIN_IMAGE_VERSIONED                                                 # image name
    MOUNTS: ${DATASET}:${DATASET_MOUNT}                                          # mount
    SLURM_ACCOUNT: devtech                                                       # account, do not need change
    WALLTIME: "00:15:00"                                                         # estimate job time. Less time, higher priority
    DGXNNODES: 1                                                                 # node num
    TEST_CMD: ./ci/integration_test/wdl/wdl.sub                                  # test script

deepfm:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/deepfm
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/deepfm/deepfm.sub

dlrm:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/dlrm
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/dlrm/dlrm.sub

dcn_multi_node:
  extends: .cluster_test_job
  needs:
    - build_train_multi_node
  variables:
    GPFSFOLDER: $LOGDIR/dcn_multi_node
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_MULTINODE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "01:00:00"
    DGXNNODES: 4                                                                              # using 4 node
    TEST_CMD: ./ci/integration_test/dcn/dcn_multi_node.sub


group_fc:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/group_fc
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET_NEW_CRITEO_SELENE}:${NEW_CRITEO_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/group_fc/group_fc.sub

# python interface inference
py_inference:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/py_inference
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT},/lustre/fsw/devtech/hpc-hugectr/inference/:/hugectr/test/utest/
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:30:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/inference/inference.sub

py_low_level:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/py_low_level
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT},/raid:/raid
    SLURM_ACCOUNT: devtech
    WALLTIME: "01:00:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/py_interface/py_low_level.sub

din_single_node:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/din_single_node
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DIN_DATASET}:${DIN_DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/din/din.sub

etc_single_node:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/etc_single_node 
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:45:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/etc/etc_single_node.sub

etc_multi_node:
  extends: .cluster_test_job
  needs:
    - build_train_multi_node
  variables:
    GPFSFOLDER: $LOGDIR/etc_multi_node
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_MULTINODE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT},${LOGDIR}:/tmp_dir
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:45:00"
    DGXNNODES: 2
    TEST_CMD: ./ci/integration_test/etc/etc_multi_node.sub

# data generator single node
data_generator_single_node:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/data_generator_single_node
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    SLURM_ACCOUNT: devtech
    WALLTIME: "01:00:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/data_generator/data_generator.sub

# # python interface single node
py_single_node:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/py_single_node
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/py_interface/py_single_node.sub

py_multi_node:
  extends: .cluster_test_job
  needs:
    - build_train_multi_node
  variables:
    GPFSFOLDER: $LOGDIR/py_multi_node
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_MULTINODE_VERSIONED
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 4
    TEST_CMD: ./ci/integration_test/py_interface/py_multi_node.sub

# hugectr inference correctness test
inference_correctness:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/inference_correctness
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET_NEW_CRITEO_SELENE}:${NEW_CRITEO_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:30:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/inference/inference_correctness.sub

# hugectr to onnx converter test
hugectr2onnx:
  extends: .cluster_test_job
  needs:
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/hugectr2onnx
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: ${DATASET_NEW_CRITEO_SELENE}:${NEW_CRITEO_MOUNT},${DIN_DATASET}:${DIN_DATASET_MOUNT},${NCF_DATASET}:${NCF_DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:30:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/hugectr2onnx/hugectr2onnx.sub

# embedding_plugin
sparse_operation_kit-TF2:
  extends: .cluster_test_job
  needs:
    - build_sok_tf2
  variables:
    GPFSFOLDER: $LOGDIR/sparse_operation_kit
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $SOK_IMAGE_VERSIONED_TF2
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:45:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/sok/sok.sub

sparse_operation_kit-TF1:
  extends: .cluster_test_job
  needs:
    - build_sok_tf1
  variables:
    GPFSFOLDER: $LOGDIR/sparse_operation_kit
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $SOK_IMAGE_VERSIONED_TF1
    MOUNTS: ${DATASET}:${DATASET_MOUNT}
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:45:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/integration_test/sok/sok.sub

inference_benchmark:
  extends: .cluster_test_job
  needs:
    - build_inference
  before_script:
    - export BZ=1
    - export MIXED_PRECISION=FP32
  variables:
    GPFSFOLDER: $LOGDIR/inference_benchmark
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $INFER_IMAGE_VERSIONED
    MOUNTS: /lustre/fsw/devtech/hpc-hugectr/inference/dlrm_regression/dlrm/1:/model/dlrm/1,/lustre/fsw/devtech/hpc-hugectr/keynote_inference/perf_data:/perf_data
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/benchmark/inference_benchmark/run.sub

inference_ps_test:
  extends: .dlcluster_job
  allow_failure: false
  stage: test
  needs:
    - build_inference
  script:
    - export CONT=${INFER_IMAGE_VERSIONED}
    - srun -N 1 -p dgx1v,dgx1v16g,dgx1v32g bash ./ci/integration_test/inference/ps_test.sh
  rules:
    - if: $CI_PIPELINE_SOURCE =~ /^(web|merge_request_event)$/
      when: always
    - if: $TEST_NEW_IMAGE == "1"
      when: always
    - when: never

inference_embedding_cache_update_test:
  extends: .dlcluster_job
  allow_failure: false
  stage: test
  needs:
    - build_inference
  script:
    - export CONT=${INFER_IMAGE_VERSIONED}
    - srun -N 1 -p dgx1v,dgx1v16g,dgx1v32g bash ./ci/integration_test/inference/embedding_cache_update_test.sh
  rules:
    - if: $CI_PIPELINE_SOURCE =~ /^(web|merge_request_event)$/
      when: always
    - if: $TEST_NEW_IMAGE == "1"
      when: always
    - when: never
    
#hdfs backend test
hdfs_backend_test:
  extends: .dlcluster_job
  allow_failure: false
  stage: test
  needs:
    - build_train_single_node_with_hdfs
  script:
    - export CONT=${TRAIN_IMAGE_VERSIONED_WITH_HDFS}
    - srun -N 1 -p dgx1v32g,v100pcie32g bash ./ci/integration_test/hdfs/hdfs_backend_test.sh
  rules:
    - if: $CI_PIPELINE_SOURCE =~ /^(web|merge_request_event)$/
      when: always
    - if: $TEST_NEW_IMAGE == "1"
      when: always
    - when: never

wdl_check:
  # Push logs to gitlab
  extends: .cluster_post_test_job
  needs:
    - wdl
  variables:
    GPFSFOLDER: $LOGDIR/wdl_check
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: $LOGDIR/wdl:/logs
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/post_test/check_wdl.sub

inference_benchmark_check:
  extends: .cluster_post_test_job
  needs:
    - inference_benchmark
    - build_train_single_node
  variables:
    GPFSFOLDER: $LOGDIR/inference_benchmark_check
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: $LOGDIR/inference_benchmark:/logs
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/post_test/check_inference_benchmark.sub

dlrm_1node_check:
  # Push logs to gitlab
  extends: .cluster_post_test_job
  needs:
    - dlrm_benchmark_1node
  variables:
    GPFSFOLDER: $LOGDIR/dlrm_1node_check
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: $LOGDIR/dlrm_benchmark_1node:/logs
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/post_test/check_dlrm_1node.sub

dlrm_14node_check:
  # Push logs to gitlab
  extends: .cluster_post_test_job
  needs:
    - dlrm_benchmark_14node
  variables:
    GPFSFOLDER: $LOGDIR/dlrm_14node_check
    GIT_CLONE_PATH: ${GIT_CLONE_PATH_SELENE}
    CONT: $TRAIN_IMAGE_VERSIONED
    MOUNTS: $LOGDIR/dlrm_benchmark_14node:/logs
    SLURM_ACCOUNT: devtech
    WALLTIME: "00:15:00"
    DGXNNODES: 1
    TEST_CMD: ./ci/post_test/check_dlrm_14node.sub

# rm_logs:
#   extends: .cluster_test_job
#   variables:
#     GPFSFOLDER: "$LOGDIR"
#     GIT_CLONE_PATH: /lustre/fsw/devtech/hpc-hugectr/hugectr-ci/$CI_CONCURRENT_ID/$CI_PROJECT_NAME
#     CONT: $TRAIN_IMAGE_VERSIONED
#     MOUNTS: /lustre/fsw/devtech:/logs
#     SLURM_ACCOUNT: devtech
#     WALLTIME: "00:15:00"
#     DGXNNODES: 1
#     TEST_CMD: ./ci/common/clean_logs.sub