Modify run_task() to use vCPU and memory in GB (facebookresearch#399)

Summary:
Pull Request resolved: facebookresearch#399

- Modify run_task() to use vCPU and memory in GB so it's more user friendly
- Add mappers in aws.py to convert cpu & memory units

Reviewed By: liliarizona

Differential Revision: D38525040

fbshipit-source-id: 2ce2cefb4e71fb8e154b4ca90edaec09c18d4bd4

fbpcp/gateway/ecs.py

@@ -21,6 +21,8 @@
     map_ecstask_to_containerinstance,
     map_ecstaskdefinition_to_containerdefinition,
     map_esccluster_to_clusterinstance,
+    map_gb_to_mb,
+    map_vcpu_to_unit,
 )
 from fbpcp.metrics.emitter import MetricsEmitter
 from fbpcp.metrics.getter import MetricsGetter
@@ -87,6 +89,7 @@ def run_task(
             "environment": environment,
         }
         if cpu and memory:
+            cpu, memory = map_vcpu_to_unit(cpu), map_gb_to_mb(memory)
             container_override.update(
                 {
                     "cpu": cpu,

fbpcp/mapper/aws.py

@@ -32,6 +32,25 @@
     get_json_values,
 )
 
+CPU_VIRTUAL_TO_UNIT = 1024
+MEMORY_GB_TO_MB = 1024
+
+
+def map_vcpu_to_unit(vcpu: int) -> int:
+    return vcpu * CPU_VIRTUAL_TO_UNIT
+
+
+def map_gb_to_mb(gb: int) -> int:
+    return gb * MEMORY_GB_TO_MB
+
+
+def map_unit_to_vcpu(cpu_unit: int) -> int:
+    return cpu_unit // CPU_VIRTUAL_TO_UNIT
+
+
+def map_mb_to_gb(mb: int) -> int:
+    return mb // MEMORY_GB_TO_MB
+
 
 def map_ecstask_to_containerinstance(task: Dict[str, Any]) -> ContainerInstance:
     container = task["containers"][0]

tests/gateway/test_ecs.py

@@ -12,6 +12,7 @@
 from fbpcp.entity.container_definition import ContainerDefinition
 from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus
 from fbpcp.gateway.ecs import ECSGateway
+from fbpcp.mapper.aws import map_gb_to_mb, map_vcpu_to_unit
 from fbpcp.util.aws import convert_list_to_dict, get_container_definition_id
 
 
@@ -42,8 +43,8 @@ class TestECSGateway(unittest.TestCase):
     TEST_CLUSTER_TAG_KEY = "test-tag-key"
     TEST_CLUSTER_TAG_VALUE = "test-tag-value"
     REGION = "us-west-2"
-    TEST_CPU_UNITS = 1024
-    TEST_MEMORY_IN_MIB = 2048
+    TEST_CPU = 4  # in vCPU
+    TEST_MEMORY = 30  # in GB
 
     @patch("boto3.client")
     def setUp(self, BotoClient) -> None:
@@ -54,6 +55,8 @@ def setUp(self, BotoClient) -> None:
 
     def test_run_task(self) -> None:
         # Arrange
+        cpu_response = map_vcpu_to_unit(self.TEST_CPU)
+        memory_response = map_gb_to_mb(self.TEST_MEMORY)
         client_return_response = {
             "tasks": [
                 {
@@ -78,12 +81,12 @@ def test_run_task(self) -> None:
                                     self.TEST_CMD,
                                 ],
                                 "environment": [],
-                                "cpu": self.TEST_CPU_UNITS,
-                                "memory": self.TEST_MEMORY_IN_MIB,
+                                "cpu": cpu_response,
+                                "memory": memory_response,
                             }
                         ],
-                        "cpu": str(self.TEST_CPU_UNITS),
-                        "memory": str(self.TEST_MEMORY_IN_MIB),
+                        "cpu": str(cpu_response),
+                        "memory": str(memory_response),
                     },
                 },
             ]
@@ -101,8 +104,8 @@ def test_run_task(self) -> None:
             self.TEST_CMD,
             self.TEST_CLUSTER,
             self.TEST_SUBNETS,
-            cpu=self.TEST_CPU_UNITS,
-            memory=self.TEST_MEMORY_IN_MIB,
+            cpu=self.TEST_CPU,
+            memory=self.TEST_MEMORY,
         )
         # Assert
         self.assertEqual(task, expected_task)
@@ -121,12 +124,12 @@ def test_run_task(self) -> None:
                         "name": self.TEST_CONTAINER,
                         "command": [self.TEST_CMD],
                         "environment": [],
-                        "cpu": self.TEST_CPU_UNITS,
-                        "memory": self.TEST_MEMORY_IN_MIB,
+                        "cpu": cpu_response,
+                        "memory": memory_response,
                     }
                 ],
-                "cpu": str(self.TEST_CPU_UNITS),
-                "memory": str(self.TEST_MEMORY_IN_MIB),
+                "cpu": str(cpu_response),
+                "memory": str(memory_response),
             },
         )
 

tests/mapper/test_aws.py

@@ -20,6 +20,8 @@
     map_ec2subnet_to_subnet,
     map_ecstask_to_containerinstance,
     map_esccluster_to_clusterinstance,
+    map_gb_to_mb,
+    map_vcpu_to_unit,
 )
 
 
@@ -34,8 +36,12 @@ class TestAWSMapper(unittest.TestCase):
     TEST_IGW_TARGET_ID = "igw-a1b2c3d000"
     TEST_ROUTE_STATE_ACTIVE = "active"
     TEST_ROUTE_STATE_INACTIVE = "blackhole"
+    TEST_CPU = 1
+    TEST_MEMORY = 2
 
     def test_map_ecstask_to_containerinstance(self):
+        cpu_response = map_vcpu_to_unit(self.TEST_CPU)
+        memory_response = map_gb_to_mb(self.TEST_MEMORY)
         ecs_task_response = {
             "tasks": [
                 {
@@ -61,6 +67,8 @@ def test_map_ecstask_to_containerinstance(self):
                         },
                     ],
                     "taskArn": self.TEST_TASK_ARN,
+                    "cpu": cpu_response,
+                    "memory": memory_response,
                 },
                 {
                     "containers": [
@@ -71,6 +79,10 @@ def test_map_ecstask_to_containerinstance(self):
                         },
                     ],
                     "taskArn": self.TEST_TASK_ARN,
+                    "overrides": {
+                        "cpu": cpu_response,
+                        "memory": memory_response,
+                    },
                 },
                 {
                     "containers": [