feat: MPI operator installation code for distributed training use case (

#362)

* MPI operator code for distributed training

* Making MPI operator optional for users

* added type string to mpi operator variable version

ai-ml/trainium-inferentia/addons.tf

@@ -499,3 +499,21 @@ resource "aws_launch_template" "trn1_lt" {
     }
   }
 }
+
+#---------------------------------------------------------------
+# MPI Operator for distributed training on Trainium
+#---------------------------------------------------------------
+data "http" "mpi_operator_yaml" {
+  url = "https://raw.githubusercontent.com/kubeflow/mpi-operator/${var.mpi_operator_version}/deploy/v2beta1/mpi-operator.yaml"
+}
+
+data "kubectl_file_documents" "mpi_operator_yaml" {
+  content = data.http.mpi_operator_yaml.response_body
+}
+
+resource "kubectl_manifest" "mpi_operator" {
+  for_each   = var.enable_mpi_operator ? data.kubectl_file_documents.mpi_operator_yaml.manifests : {}
+  yaml_body  = each.value
+  depends_on = [module.eks.eks_cluster_id]
+}
+

ai-ml/trainium-inferentia/variables.tf

@@ -37,3 +37,15 @@ variable "enable_amazon_prometheus" {
   type        = bool
   default     = true
 }
+
+variable "mpi_operator_version" {
+  description = "The version of the MPI Operator to install"
+  default     = "v0.4.0"
+  type        = string
+}
+
+variable "enable_mpi_operator" {
+  description = "Flag to enable the MPI Operator deployment"
+  type        = bool
+  default     = false
+}