apiVersion: batch/v1
kind: Job
metadata: 
  name: tabnet-model-training-64-32
spec:
  backoffLimit: 0
  completions: 1
  template:
    metadata:
      labels:
        k8s-app: tabnet-model-training-64-32
    spec:
      restartPolicy: Never
      containers:
      - name: tabnet-model-training-64-32
        image: jmlehrer/cell-exploration:latest
        command: ["stdbuf", "-i0", "-o0", "-e0", "bash", "-c"]
        env:
        - name: PYTHONUNBUFFERED
          value: "1"
        args: [
          "wandb login --relogin $(cat wandbcreds) && \
          python src/models/train_neural_network.py --num-workers=64 --batch-size=32"
        ]
        resources:
          limits:
            memory: 256Gi
            cpu: "32"
            nvidia.com/gpu: "1"
          requests:
            memory: 128Gi
            cpu: "32"
            nvidia.com/gpu: "1"
        volumeMounts:
        - name: ephemeral
          mountPath: "/data"
        # SystemV shared memory
        - name: "dshm"
          mountPath: "/dev/shm"
      volumes:
      # Ephemeral storage
      - name: ephemeral
        emptyDir: {}
      - name: dshm
        emptyDir:
          medium: Memory
      tolerations:
        - key: "nautilus.io/chase-ci"
          operator: "Equal"
          value: "true"
          effect: "NoSchedule"
