From fd193c08d342376a25ea9ed80de25fc0b3c33ebb Mon Sep 17 00:00:00 2001
From: Mike Bentley Mills <mike@valarlabs.com>
Date: Thu, 25 Apr 2024 11:00:48 -0700
Subject: [PATCH] Valar-Labs version

---
 examples/eks_argo/eks.tf      | 109 ++++++++++++++++++++++++++++++++++
 examples/eks_argo/metaflow.tf |   6 +-
 2 files changed, 112 insertions(+), 3 deletions(-)

diff --git a/examples/eks_argo/eks.tf b/examples/eks_argo/eks.tf
index 7cf4d13..14c6c55 100644
--- a/examples/eks_argo/eks.tf
+++ b/examples/eks_argo/eks.tf
@@ -27,7 +27,52 @@ module "eks" {
       update_config = {
         max_unavailable_percentage = 50
       }
+      kubelet_extra_args = "--node-labels=nvidia.com/gpu=false"
     }
+
+    multi_cpu = {
+      desired_capacity = 0
+      max_capacity     = 5
+      min_capacity     = 0
+
+      instance_types = ["r5.xlarge"]
+      update_config = {
+        max_unavailable_percentage = 50
+      }
+      kubelet_extra_args = "--node-labels=nvidia.com/gpu=false"
+    }
+
+    multi_cpu_X = {
+      desired_capacity = 0
+      max_capacity     = 5
+      min_capacity     = 0
+
+      instance_types = ["r5.2xlarge"]
+
+      update_config = {
+        max_unavailable_percentage = 50
+      }
+      kubelet_extra_args = "--node-labels=nvidia.com/gpu=false"
+
+    }
+
+    gpu_group = {
+      ami_type  = "AL2_x86_64_GPU"
+
+      desired_capacity = 0
+      max_capacity     = 1
+      min_capacity     = 0
+
+      instance_types = ["g4dn.xlarge"]
+
+      update_config = {
+        max_unavailable_percentage = 50
+      }
+      kubelet_extra_args = "--node-labels=nvidia.com/gpu=true"
+
+
+    }
+
   }
 
   workers_additional_policies = [
@@ -105,6 +150,70 @@ data "aws_iam_policy_document" "cluster_autoscaler" {
   }
 }
 
+resource "aws_autoscaling_group_tag" "cluster_autoscaler_resource_tags" {
+  for_each               = module.eks.node_groups
+  autoscaling_group_name = each.value.resources[0].autoscaling_groups[0].name
+
+  tag {
+    key   = "k8s.io/cluster-autoscaler/node-template/resources/ephemeral-storage"
+    value = "50G"
+    propagate_at_launch = true
+  }
+}
+
+# Deploy NVIDIA device plugin
+resource "kubernetes_daemonset" "nvidia_device_plugin" {
+  metadata {
+    name      = "nvidia-device-plugin-daemonset"
+    namespace = "kube-system"
+  }
+
+  spec {
+    selector {
+      match_labels = {
+        name = "nvidia-device-plugin-ds"
+      }
+    }
+
+    template {
+      metadata {
+        labels = {
+          name = "nvidia-device-plugin-ds"
+        }
+      }
+
+      spec {
+        toleration {
+          key      = "nvidia.com/gpu"
+          operator = "Exists"
+          effect   = "NoSchedule"
+        }
+
+        container {
+          image = "nvidia/k8s-device-plugin:v0.9.0"
+          name  = "nvidia-device-plugin-ctr"
+
+          security_context {
+            allow_privilege_escalation = false
+          }
+
+          volume_mount {
+            name       = "device-plugin"
+            mount_path = "/var/lib/kubelet/device-plugins"
+          }
+        }
+
+        volume {
+          name = "device-plugin"
+          host_path {
+            path = "/var/lib/kubelet/device-plugins"
+          }
+        }
+      }
+    }
+  }
+}
+
 
 data "aws_eks_cluster" "cluster" {
   name = module.eks.cluster_id
diff --git a/examples/eks_argo/metaflow.tf b/examples/eks_argo/metaflow.tf
index 58805bf..1c5bce0 100644
--- a/examples/eks_argo/metaflow.tf
+++ b/examples/eks_argo/metaflow.tf
@@ -20,7 +20,7 @@ data "aws_availability_zones" "available" {
 
 module "metaflow-datastore" {
   source  = "outerbounds/metaflow/aws//modules/datastore"
-  version = "0.10.0"
+  version = "0.12.0"
 
   force_destroy_s3_bucket = true
 
@@ -37,12 +37,12 @@ module "metaflow-datastore" {
 
 module "metaflow-common" {
   source  = "outerbounds/metaflow/aws//modules/common"
-  version = "0.10.0"
+  version = "0.12.0"
 }
 
 module "metaflow-metadata-service" {
   source  = "outerbounds/metaflow/aws//modules/metadata-service"
-  version = "0.10.0"
+  version = "0.12.0"
 
   resource_prefix = local.resource_prefix
   resource_suffix = local.resource_suffix
-- 
2.39.3 (Apple Git-146)