GNN training blog

kumo-ai · Nov 8, 2022 · 19708c5 · 19708c5
1 parent da7c5ea
commit 19708c5
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 0 deletions.
diff --git a/training/provisioner.tf b/training/provisioner.tf
@@ -0,0 +1,68 @@
+resource "kubernetes_manifest" "provisioner_gpu_provisioner" {
+  manifest = {
+    "apiVersion" = "karpenter.sh/v1alpha5"
+    "kind" = "Provisioner"
+    "metadata" = {
+      "name" = "gpu-provisioner"
+    }
+    "spec" = {
+      "provider" = {
+        "blockDeviceMappings" = [
+          {
+            "deviceName" = "/dev/xvda"
+            "ebs" = {
+              "deleteOnTermination" = true
+              "iops" = 3000
+              "throughput" = 125
+              "volumeSize" = "100Gi"
+              "volumeType" = "gp3"
+            }
+          },
+        ]
+        "securityGroupSelector" = {
+          "karpenter.sh/discovery/training-cluster" = "training-cluster"
+        }
+        "subnetSelector" = {
+          "Name" = "training-cluster-vpc-private*"
+        }
+        "tags" = {
+          "karpenter.sh/discovery/training-cluster" = "training-cluster"
+        }
+      }
+      "requirements" = [
+        {
+          "key" = "node.kubernetes.io/instance-type"
+          "operator" = "In"
+          "values" = [
+            "g4dn.xlarge",
+            "g4dn.4xlarge",
+            "g4dn.8xlarge",
+            "g4dn.16xlarge",
+          ]
+        },
+        {
+          "key" = "karpenter.sh/capacity-type"
+          "operator" = "In"
+          "values" = [
+            "spot",
+            "on-demand",
+          ]
+        },
+        {
+          "key" = "kubernetes.io/arch"
+          "operator" = "In"
+          "values" = [
+            "amd64",
+          ]
+        },
+      ]
+      "taints" = [
+        {
+          "effect" = "NoSchedule"
+          "key" = "nvidia.com/gpu"
+        },
+      ]
+      "ttlSecondsAfterEmpty" = 30
+    }
+  }
+}
diff --git a/training/training-job.yaml b/training/training-job.yaml
@@ -0,0 +1,32 @@
+apiVersion: batch/v1 
+kind: Job 
+metadata:   
+  name: gnn-training-job
+  namespace: default  
+  labels:     
+    jobname: gnn-training-job 
+spec:   
+  ttlSecondsAfterFinished: 60
+  template:     
+    metadata:       
+      name: gnn-training-job-pod      
+      labels:         
+        app: gnn-training-job-pod   
+    spec:     
+      containers:       
+      - name: gnn-trainer
+        resources:
+          requests:
+            memory: "250G"
+          limits:
+            nvidia.com/gpu: 1      
+        image: <obfuscated>
+        volumeMounts:
+        - name: dshm
+          mountPath: /dev/shm 
+        imagePullPolicy: IfNotPresent     
+      restartPolicy: OnFailure
+      volumes:
+      - name: dshm
+        emptyDir:
+          medium: Memory