fix for mlflow missing s3 kms key permissions

MetroStar · May 17, 2024 · 056cb84 · 056cb84
1 parent 499013f
commit 056cb84
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 29 deletions.
diff --git a/src/nebari_plugin_mlflow_aws/__about__.py b/src/nebari_plugin_mlflow_aws/__about__.py
@@ -1 +1 @@
-__version__ = "0.0.13"
+__version__ = "0.0.14"
diff --git a/src/nebari_plugin_mlflow_aws/terraform/main.tf b/src/nebari_plugin_mlflow_aws/terraform/main.tf
@@ -28,8 +28,8 @@ resource "aws_s3_bucket" "artifact_storage" {
 
 # If enable_s3_encryption is true, create a key and apply Server Side Encryption to S3 bucket
 resource "aws_kms_key" "mlflow_kms_key" {
-  count                   = var.enable_s3_encryption ? 1 : 0
-  description             = "This key is used to encrypt bucket objects for the AWS MLflow extension"
+  count       = var.enable_s3_encryption ? 1 : 0
+  description = "This key is used to encrypt bucket objects for the AWS MLflow extension"
 }
 
 resource "aws_s3_bucket_server_side_encryption_configuration" "mlflow_s3_encryption" {
@@ -65,32 +65,41 @@ module "iam_assumable_role_admin" {
 resource "aws_iam_policy" "mlflow_s3" {
   name_prefix = "${var.project_name}-s3-mlflow-bucket-access"
   description = "Grants workloads full access to S3 bucket for MLflow artifact storage"
-  policy      = <<EOT
-{
-  "Version": "2012-10-17",
-  "Statement": [
-        {
-            "Sid": "ListAllBuckets",
-            "Effect": "Allow",
-            "Action": "s3:ListAllMyBuckets",
-            "Resource": "*"
-        },
-        {
-            "Sid": "ListObjectsInBucket",
-            "Effect": "Allow",
-            "Action": "s3:ListBucket",
-            "Resource": "${aws_s3_bucket.artifact_storage.arn}"
-        },
-        {
-            "Sid": "AllObjectActions",
-            "Effect": "Allow",
-            "Action": "s3:*Object",
-            "Resource": "${aws_s3_bucket.artifact_storage.arn}/*"
-        }
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Sid      = "ListAllBuckets"
+        Effect   = "Allow"
+        Action   = "s3:ListAllMyBuckets"
+        Resource = "*"
+      },
+      {
+        Sid      = "ListObjectsInBucket"
+        Effect   = "Allow"
+        Action   = "s3:ListBucket"
+        Resource = aws_s3_bucket.artifact_storage.arn
+      },
+      {
+        Sid      = "AllObjectActions"
+        Effect   = "Allow"
+        Action   = "s3:*Object"
+        Resource = "${aws_s3_bucket.artifact_storage.arn}/*"
+      },
+      {
+        Sid    = "KMS"
+        Effect = "Allow"
+        Action = [
+          "kms:Decrypt",
+          "kms:Encrypt",
+          "kms:ReEncrypt*",
+          "kms:DescribeKey",
+          "kms:GenerateDataKey*"
+        ]
+        Resource = aws_kms_key.mlflow_kms_key[0].arn
+      },
     ]
-}
-
-  EOT
+  })
 }
 
 module "keycloak" {
@@ -116,4 +125,4 @@ module "mlflow" {
   s3_bucket_name         = aws_s3_bucket.artifact_storage.id
   keycloak_config        = module.keycloak.config
   overrides              = var.overrides
-}
+}