From 0255cf98dac2179946caa900e09db8e1b181c292 Mon Sep 17 00:00:00 2001
From: Ziqi Huang <ziqihuang@fb.com>
Date: Tue, 9 Aug 2022 16:47:28 -0700
Subject: [PATCH] Add CPU & memory to ContainerInstance

Summary:
## Context
We want to enable customized containers in our infra ([design doc](https://docs.google.com/document/d/1UqaK6VXEoIVUHSJnYs_Ydju8)). Therefore we need to add cpu & memory specs to ContainerInstance;
## This commit
- Add CPU & Memory to ContainerInstance
- Parse them in map_ecstask_to_containerinstance
## How this effect ```pcs_container_instance.py```?
Since ```PCSContainerInstance``` inherits ```ContainerInstance```, we change it to use keyword argument for ```log_url``` to prevent breaks;

Differential Revision: D38553232

fbshipit-source-id: 32f83534c6185409209ace601b9f3a7a61c21938
---
 fbpcp/entity/container_instance.py |  2 ++
 fbpcp/mapper/aws.py                | 16 ++++++++++++++--
 tests/gateway/test_ecs.py          |  2 ++
 tests/mapper/test_aws.py           |  4 ++++
 4 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/fbpcp/entity/container_instance.py b/fbpcp/entity/container_instance.py
index 9e34a6d1..ea0373ed 100644
--- a/fbpcp/entity/container_instance.py
+++ b/fbpcp/entity/container_instance.py
@@ -26,3 +26,5 @@ class ContainerInstance:
     instance_id: str
     ip_address: Optional[str] = None
     status: ContainerInstanceStatus = ContainerInstanceStatus.UNKNOWN
+    cpu: Optional[int] = None  # Number of vCPU
+    memory: Optional[int] = None  # Memory in GB
diff --git a/fbpcp/mapper/aws.py b/fbpcp/mapper/aws.py
index 25246d88..18f10040 100644
--- a/fbpcp/mapper/aws.py
+++ b/fbpcp/mapper/aws.py
@@ -70,8 +70,20 @@ def map_ecstask_to_containerinstance(task: Dict[str, Any]) -> ContainerInstance:
             status = ContainerInstanceStatus.FAILED
     else:
         status = ContainerInstanceStatus.UNKNOWN
-
-    return ContainerInstance(task["taskArn"], ip_v4, status)
+    task_overrides = None
+    if (
+        "overrides" in task
+        and "cpu" in task["overrides"]
+        and "memory" in task["overrides"]
+    ):
+        task_overrides = task["overrides"]
+    elif "cpu" in task and "memory" in task:
+        task_overrides = task
+    vcpu = map_unit_to_vcpu(int(task_overrides["cpu"])) if task_overrides else None
+    memory_in_gb = (
+        map_mb_to_gb(int(task_overrides["memory"])) if task_overrides else None
+    )
+    return ContainerInstance(task["taskArn"], ip_v4, status, vcpu, memory_in_gb)
 
 
 def map_esccluster_to_clusterinstance(cluster: Dict[str, Any]) -> Cluster:
diff --git a/tests/gateway/test_ecs.py b/tests/gateway/test_ecs.py
index b4532143..0ef801d3 100644
--- a/tests/gateway/test_ecs.py
+++ b/tests/gateway/test_ecs.py
@@ -96,6 +96,8 @@ def test_run_task(self) -> None:
             self.TEST_TASK_ARN,
             self.TEST_IP_ADDRESS,
             ContainerInstanceStatus.STARTED,
+            self.TEST_CPU,
+            self.TEST_MEMORY,
         )
         # Act
         task = self.gw.run_task(
diff --git a/tests/mapper/test_aws.py b/tests/mapper/test_aws.py
index af94d4e8..c9b75d9c 100644
--- a/tests/mapper/test_aws.py
+++ b/tests/mapper/test_aws.py
@@ -107,11 +107,15 @@ def test_map_ecstask_to_containerinstance(self):
                 self.TEST_TASK_ARN,
                 None,
                 ContainerInstanceStatus.COMPLETED,
+                self.TEST_CPU,
+                self.TEST_MEMORY,
             ),
             ContainerInstance(
                 self.TEST_TASK_ARN,
                 None,
                 ContainerInstanceStatus.FAILED,
+                self.TEST_CPU,
+                self.TEST_MEMORY,
             ),
             ContainerInstance(
                 self.TEST_TASK_ARN,