Specify vcpu and memory as resourceRequirements in AWS Batch executor (…

…#95)
insitro · Jun 6, 2024 · fd9479d · fd9479d
1 parent 367033e
commit fd9479d
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 11 deletions.
diff --git 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
@@ -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"},
+                        ],
                     },
                 },
             ]