From b889f6ac836c6e64f39c93ce651d330dcced650e Mon Sep 17 00:00:00 2001
From: Sakari Ikonen <64256562+saikonen@users.noreply.github.com>
Date: Thu, 13 Jul 2023 19:04:39 +0300
Subject: [PATCH] fix: validate required configuration for Batch (#1483)

* add validation for presence of iam_role in batch.

* polish messaging
---
 metaflow/plugins/aws/batch/batch_decorator.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/metaflow/plugins/aws/batch/batch_decorator.py b/metaflow/plugins/aws/batch/batch_decorator.py
index 6dcb6e8e996..4c17cc080d5 100644
--- a/metaflow/plugins/aws/batch/batch_decorator.py
+++ b/metaflow/plugins/aws/batch/batch_decorator.py
@@ -91,8 +91,8 @@ class BatchDecorator(StepDecorator):
         "gpu": None,
         "memory": None,
         "image": None,
-        "queue": BATCH_JOB_QUEUE,
-        "iam_role": ECS_S3_ACCESS_IAM_ROLE,
+        "queue": BATCH_JOB_QUEUE,  # Will be replaced with an available queue if not provided.
+        "iam_role": ECS_S3_ACCESS_IAM_ROLE,  # Required
         "execution_role": ECS_FARGATE_EXECUTION_ROLE,
         "shared_memory": None,
         "max_swap": None,
@@ -149,6 +149,13 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
         if flow_datastore.TYPE != "s3":
             raise BatchException("The *@batch* decorator requires --datastore=s3.")
 
+        # Require iam_role
+        if self.attributes["iam_role"] is None:
+            # TODO: Unify messaging on various configuration options.
+            raise BatchException(
+                "The *@batch* decorator requires an IAM Role that allows AWS Batch job to communicate with Amazon S3 datastore.\n"
+                'You can specify it either with @batch(iam_role="role-name") or by setting METAFLOW_ECS_S3_ACCESS_IAM_ROLE in your configuration.'
+            )
         # Set internal state.
         self.logger = logger
         self.environment = environment