#!/usr/bin/env groovy

// Map from CUDA version to URL to obtain windows installer
def cuda_version_url = [
    // CUDA drivers on test machines only available for CUDA version >= 11.0
    //  see: https://gitlab-master.nvidia.com/ipp/cloud-infra/blossom/dev/windows-gpu-pods/-/tree/master/ContainerDriverSetup
    // test machines currently are the only option, named 'gpu_tester'
    //      two machines exist, only the TITAN RTX will pass tests
    '11.1': 'http://developer.download.nvidia.com/compute/cuda/11.1.1/local_installers/cuda_11.1.1_456.81_win10.exe',
    '11.3': 'https://developer.download.nvidia.com/compute/cuda/11.3.1/local_installers/cuda_11.3.1_465.89_win10.exe',
    '11.5': 'https://developer.download.nvidia.com/compute/cuda/11.5.2/local_installers/cuda_11.5.2_496.13_windows.exe',
    '11.6': 'https://developer.download.nvidia.com/compute/cuda/11.6.2/local_installers/cuda_11.6.2_511.65_windows.exe',
    '11.7': 'https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda_11.7.1_516.94_windows.exe',
    '11.8': 'https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_522.06_windows.exe',
    '12.1': 'https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_531.14_windows.exe'
]


docker_registry_server = targetImageTag.split(':')[0..1].join(':')
// This will be the "RUN" displayed on Blue Ocean
currentBuild.displayName = targetImageTag.split(':')[2]
// This will be the "MESSAGE" displayed on Blue Ocean
currentBuild.description = sourceBranch + ": " + commitHash

gitlabCommitStatus("build-${configName}") {

podTemplate(
cloud:'sc-ipp-blossom-116',
envVars:[envVar(key:"JENKINS_URL", value:"${env.JENKINS_URL}")],
yaml:'''
apiVersion: v1
kind: Pod
spec:
  volumes:
  - name: pvc-mount
    persistentVolumeClaim:
      claimName: 'kaolin-pvc'
  containers:
  - name: jnlp
    image: urm.nvidia.com/sw-ipp-blossom-sre-docker-local/jnlp-agent:jdk11-windows
    env: 
    - name: JENKINS_AGENT_WORKDIR
      value: C:/Jenkins/agent
    - name: DOCKER_HOST
      value: "win-docker-proxy.blossom-system.svc.cluster.local"
    - name: DOCKER_TLS_CERTDIR
      value: ""
    volumeMounts:
      - mountPath: c:/mnt
        name: pvc-mount
    resources:
        requests:
          memory: 32Gi
        limits:
          memory: 32Gi
  imagePullSecrets:
  - name: gitlabcred
  nodeSelector:
    kubernetes.io/os: windows
''') 
{
  node(POD_LABEL) {
    try {
      timeout(time: 300, unit: 'MINUTES') {
        stage("Checkout") {
          checkout([
              $class: 'GitSCM',
              branches: [[name: "${commitHash}"]],
              // We need submodules
              extensions: [[
                  $class: 'SubmoduleOption',
                  disableSubmodules: false,
                  parentCredentials: false,
                  recursiveSubmodules: true,
                  reference: '',
                  trackingSubmodules: false
              ]],
              userRemoteConfigs: [[
                  credentialsId: 'kaolin-gitlab-access-token-as-password',
                  url: "${repoUrl}"
              ]]
          ])
        }
        docker.withRegistry("https://${docker_registry_server}", 'kaolin-gitlab-access-token-as-password') {
          stage("Build base") {
            cudaUrl = cuda_version_url[cudaVer]
            baseImage = docker.build(
                "${targetImageTag}-base",
                """-m 32g --no-cache -f ./tools/windows/Dockerfile.base \
                    --build-arg CUDA_VERSION=${cudaVer} \
                    --build-arg CUDA_URL=${cudaUrl} \
                    --build-arg PYTHON_VERSION=${pythonVer} \
                    --build-arg PYTORCH_VERSION=${torchVer} \
                    .
                """)
          }

          if (buildWheel.toBoolean()) {
            stage("Build with wheel") {
              targetImage = docker.build(
                  "${targetImageTag}",
                  """-m 32g --no-cache -f ./tools/windows/Dockerfile.install_wheel \
                     --build-arg BASE_IMAGE=${targetImageTag}-base \
                     .
                  """
              )
            }
          } else {
            stage("Build") {
              targetImage = docker.build(
                  "${targetImageTag}",
                  """-m 32g --no-cache -f ./tools/windows/Dockerfile.install \
                      --build-arg BASE_IMAGE=${targetImageTag}-base \
                      .
                  """)
            }
          }
          
          stage("Push") {
            targetImage.push()
          }
        }
      }
    } catch (e) {
        // In case of build failure, we need to update the following tests as we won't run them.
        if (buildWheel.toBoolean()) {
          updateGitlabCommitStatus(name: "test-${configName}", state: 'canceled')
        } else {
          for (arch in archsToTest.split(';')) {
            updateGitlabCommitStatus(name: "test-${configName}-${arch}", state: 'canceled')
          }
        }
        throw e
    }
    stage("Launch tests") {
      jobMap = [:]
      if (buildWheel.toBoolean()) {
        jobMap["test"] = prepareWindowsWheelTestJob()
      } //else { 
      //  for (arch in archsToTest.split(';')) {
      //    jobMap["${arch}"] = prepareWindowsTestJob(arch)
      //  }
      parallel jobMap
    }
  }
}

}  // gitlabCommitStatus

/*
def prepareWindowsTestJob(arch) {
  return {
    stage("Test ${arch}") {
      build job: "windows_test_template_CI",
      parameters: [
        string(name: 'sourceBranch', value: "${sourceBranch}"),
        string(name: 'configName', value: "${configName}"),
        string(name: 'imageTag', value: "${targetImageTag}"),
        string(name: 'arch', value: "${arch}"),
        string(name: 'commitHash', value: "${commitHash}"),
      ],
      // This node doesn't need to be held while tests run.
      wait: false,
      // Success of this script depends only on successful build
      // and launch of tests, not successful tests.
      propagate: false
    }
  }
}
*/

def prepareWindowsWheelTestJob() {
  return {
    stage("Test") {
      build job: "windows_wheels_template_CI",
      parameters: [
        string(name: 'sourceBranch', value: "${sourceBranch}"),
        string(name: 'configName', value: "${configName}"),
        string(name: 'imageTag', value: "${targetImageTag}"),
        string(name: 'commitHash', value: "${commitHash}"),
        string(name: 'torchVer', value: "${torchVer}"),
        string(name: 'cudaVer', value: "${cudaVer}"),
      ],
      // This node doesn't need to be held while tests run.
      wait: false,
      // Success of this script depends only on successful build
      // and launch of tests, not successful tests.
      propagate: false
    }
  }
}