Fix S3 creation and access from MLflow (#68)

* fix s3 impl for mlflow

* allow mlflow sa to access s3

* fmt files

Author: wjayesh

aws-modular/eks.tf

@@ -95,6 +95,19 @@ resource "aws_iam_role" "ng" {
       Principal = {
         Service = "ec2.amazonaws.com"
       }
+      },
+      {
+        Action = "sts:AssumeRoleWithWebIdentity"
+        Effect = "Allow"
+        Principal = {
+          Federated = "arn:aws:iam::${data.aws_caller_identity.current.account_id}:oidc-provider/${aws_eks_cluster.cluster[0].identity[0].oidc[0].issuer}"
+        }
+        Condition = {
+          StringLike = {
+            "${aws_eks_cluster.cluster[0].identity[0].oidc[0].issuer}:aud" = "sts.amazonaws.com"
+            "${aws_eks_cluster.cluster[0].identity[0].oidc[0].issuer}:sub" = "system:serviceaccount:mlflow:*"
+          }
+        }
     }]
     Version = "2012-10-17"
   })

aws-modular/mlflow.tf

@@ -40,17 +40,50 @@ resource "aws_s3_bucket" "mlflow-bucket" {
   tags = local.tags
 }
 
-resource "aws_s3_bucket_acl" "mlflow" {
-  count  = length(aws_s3_bucket.mlflow-bucket) > 0 ? 1 : 0
+# block public access to the bucket
+resource "aws_s3_bucket_public_access_block" "mlflow" {
+  count  = (var.enable_experiment_tracker_mlflow && var.mlflow_bucket == "") ? 1 : 0
   bucket = aws_s3_bucket.mlflow-bucket[0].id
-  acl    = "private"
+
+  block_public_acls       = true
+  block_public_policy     = true
+  ignore_public_acls      = false
+  restrict_public_buckets = false
 }
 
-# block public access to the bucket
-resource "aws_s3_bucket_public_access_block" "mlflow" {
-  count  = length(aws_s3_bucket.mlflow-bucket) > 0 ? 1 : 0
+resource "aws_s3_bucket_policy" "allow_access_from_another_account_mlflow" {
+  count  = (var.enable_experiment_tracker_mlflow && var.mlflow_bucket == "") ? 1 : 0
   bucket = aws_s3_bucket.mlflow-bucket[0].id
+  policy = data.aws_iam_policy_document.allow_access_from_another_account_mlflow[0].json
+}
+
+data "aws_iam_policy_document" "allow_access_from_another_account_mlflow" {
+  count = (var.enable_experiment_tracker_mlflow && var.mlflow_bucket == "") ? 1 : 0
+  statement {
+    principals {
+      type        = "AWS"
+      identifiers = [aws_iam_role.ng[0].arn]
+    }
+    actions = [
+      "s3:*",
+    ]
+    resources = [
+      aws_s3_bucket.mlflow-bucket[0].arn,
+      "${aws_s3_bucket.mlflow-bucket[0].arn}/*",
+    ]
+  }
+}
+
+# allow the mlflow kubernetes SA to assume the IAM role
+resource "null_resource" "mlflow-iam-access" {
+
+  count = var.enable_experiment_tracker_mlflow ? 1 : 0
 
-  block_public_acls   = true
-  block_public_policy = true
+  provisioner "local-exec" {
+    command = "kubectl annotate serviceaccount -n mlflow mlflow-tracking eks.amazonaws.com/role-arn=${aws_iam_role.ng[0].arn}"
+  }
+
+  depends_on = [
+    module.mlflow,
+  ]
 }