k8s_template.yaml

apiVersion: batch/v1
kind: Job
spec:
  template:
    spec:
      restartPolicy: OnFailure
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
              - matchExpressions:
                  # Require compatibility with CUDA 12.4
                  # This version needs to be specified in the dockerfile `FROM` and
                  # `RUN micromamba install` lines too
                  # Check the version used by NRP nodes with:
                  # kubectl get nodes -L nvidia.com/gpu.product,nvidia.com/cuda.runtime.major,nvidia.com/cuda.runtime.minor -l nvidia.com/gpu.product
                  - key: nvidia.com/cuda.runtime.major
                    operator: In
                    values:
                      - "12"
                  - key: nvidia.com/cuda.runtime.minor
                    operator: In
                    values:
                      - "4"
      initContainers:
        - name: init-git
          image: alpine/git
          resources:
            limits:
              memory: "100Mi"
              cpu: "100m"
            requests:
              memory: "100Mi"
              cpu: "100m"
          command:
            - "/bin/sh"
            - "-c"
            - |
              set -euf
              cd /opt/repo
              git clone https://github.com/openforcefield/proteinbenchmark-nrp.git /opt/repo || true
              git checkout $PROTBENCH_SCRIPT_COMMIT
          volumeMounts:
            - mountPath: /opt/repo
              name: repo-vol
        - name: init-rclone
          image: rclone/rclone
          resources:
            limits:
              memory: "1Gi"
              cpu: "200m"
              ephemeral-storage: "20Gi"
            requests:
              memory: "1Gi"
              cpu: "200m"
              ephemeral-storage: "20Gi"
          env:
            - name: RCLONE_CONFIG
              value: /secrets/rclone.conf
          command:
            - "/bin/sh"
            - "-c"
            - |
              set -euf
              rclone --version
              echo "about to clone"
              IFS=$'\n'
              for FILE in ${PROTBENCH_REQUIRED_FILES}; do
                echo "cloning to /results/$FILE"
                rclone copyto --progress nrp-internal:proteinbenchmark-jm-bucket/results/$FILE /results/$FILE
              done
              unset IFS
              echo "done cloning"
              ls /results
              chmod -R 777 /results
              echo "done chmoding"
          volumeMounts:
            - mountPath: /results
              name: results-vol
            - name: rclone-config
              mountPath: /secrets/rclone.conf
              subPath: rclone.conf
      containers:
        - name: main
          image: ghcr.io/openforcefield/proteinbenchmark-nrp:latest
          resources:
            limits:
              memory: "5Gi"
              cpu: "1"
              nvidia.com/gpu: 1
              ephemeral-storage: "20Gi"
            requests:
              memory: "5Gi"
              cpu: "1"
              nvidia.com/gpu: 1
              ephemeral-storage: "20Gi"
          volumeMounts:
            - name: rclone-config
              mountPath: /secrets/rclone.conf
              subPath: rclone.conf
            - name: openeye-license
              mountPath: /secrets/oe_license.txt
              subPath: oe_license.txt
            - mountPath: /opt/repo
              name: repo-vol
            - mountPath: /results
              name: results-vol
          env:
            - name: OE_LICENSE
              value: /secrets/oe_license.txt
            - name: RCLONE_CONFIG
              value: /secrets/rclone.conf
            - name: THIS_POD_NAME
              valueFrom:
                fieldRef:
                  fieldPath: metadata.name
          command:
            - /bin/sh
            - -c
            - |
              micromamba run -n base pip install git+https://github.com/openforcefield/proteinbenchmark.git@nagl
              micromamba run -n base python /opt/repo/$(PROTBENCH_SCRIPT_PATH) -o/results
              rclone copy --update /results nrp-internal:proteinbenchmark-jm-bucket/results
          lifecycle:
            preStop:
              exec:
                command:
                  - /bin/sh
                  - -c
                  - rclone copy --update /results nrp-internal:proteinbenchmark-jm-bucket/results 2>&1 > /proc/1/fd/1
      volumes:
        - name: rclone-config
          secret:
            secretName: jm-rclone-config
        - name: openeye-license
          secret:
            secretName: oe-license-feb-2024
        - name: repo-vol
          emptyDir: {}
        - name: results-vol
          emptyDir: {}
  backoffLimit: 10