From fd9479d13a8d94274fd8e1def14f7d30db1f9572 Mon Sep 17 00:00:00 2001
From: Alden <alden.timme@gmail.com>
Date: Thu, 6 Jun 2024 13:51:53 -0500
Subject: [PATCH] Specify vcpu and memory as resourceRequirements in AWS Batch
 executor (#95)

---
 redun/executors/aws_batch.py  | 15 +++++++++++++--
 redun/tests/test_aws_batch.py | 24 +++++++++++++++---------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/redun/executors/aws_batch.py b/redun/executors/aws_batch.py
index 091123e6..7bd62db9 100644
--- a/redun/executors/aws_batch.py
+++ b/redun/executors/aws_batch.py
@@ -425,9 +425,20 @@ def apply_resources(container_properties: Dict) -> None:
         # Switch units, redun configs are in GB but AWS uses MB.
         memory_mb = int(memory * 1024)
 
-        container_properties.update({"vcpus": vcpus, "memory": memory_mb})
+        container_properties["resourceRequirements"] = [
+            {
+                "type": "VCPU",
+                "value": str(vcpus),
+            },
+            {
+                "type": "MEMORY",
+                "value": str(memory_mb),
+            },
+        ]
         if gpus > 0:
-            container_properties["resourceRequirements"] = [{"type": "GPU", "value": str(gpus)}]
+            container_properties["resourceRequirements"].append(
+                {"type": "GPU", "value": str(gpus)}
+            )
 
     if "containerOverrides" in batch_job_args:
         apply_resources(batch_job_args["containerOverrides"])
diff --git a/redun/tests/test_aws_batch.py b/redun/tests/test_aws_batch.py
index 7a2030ef..8ae11e00 100644
--- a/redun/tests/test_aws_batch.py
+++ b/redun/tests/test_aws_batch.py
@@ -484,9 +484,11 @@ def test_job_submit_resources(get_or_create_job_definition_mock, get_aws_client_
         retryStrategy={"attempts": 1},
         containerOverrides={
             "command": ["test_command"],
-            "vcpus": 4,
-            "memory": 8192,
-            "resourceRequirements": [{"type": "GPU", "value": "2"}],
+            "resourceRequirements": [
+                {"type": "VCPU", "value": "4"},
+                {"type": "MEMORY", "value": "8192"},
+                {"type": "GPU", "value": "2"},
+            ],
         },
         propagateTags=True,
         user="user-test",
@@ -529,18 +531,22 @@ def test_job_submit_resources(get_or_create_job_definition_mock, get_aws_client_
                     "targetNodes": "0",
                     "containerOverrides": {
                         "command": ["command_rank_0"],
-                        "vcpus": 4,
-                        "memory": 8192,
-                        "resourceRequirements": [{"type": "GPU", "value": "2"}],
+                        "resourceRequirements": [
+                            {"type": "VCPU", "value": "4"},
+                            {"type": "MEMORY", "value": "8192"},
+                            {"type": "GPU", "value": "2"},
+                        ],
                     },
                 },
                 {
                     "targetNodes": "1:",
                     "containerOverrides": {
                         "command": ["command_rank_1"],
-                        "vcpus": 4,
-                        "memory": 8192,
-                        "resourceRequirements": [{"type": "GPU", "value": "2"}],
+                        "resourceRequirements": [
+                            {"type": "VCPU", "value": "4"},
+                            {"type": "MEMORY", "value": "8192"},
+                            {"type": "GPU", "value": "2"},
+                        ],
                     },
                 },
             ]