diff --git a/fbpcp/gateway/ecs.py b/fbpcp/gateway/ecs.py index 9b69e7fc..2efed7f3 100644 --- a/fbpcp/gateway/ecs.py +++ b/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, diff --git a/fbpcp/mapper/aws.py b/fbpcp/mapper/aws.py index d6954a7d..25246d88 100644 --- a/fbpcp/mapper/aws.py +++ b/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] diff --git a/tests/gateway/test_ecs.py b/tests/gateway/test_ecs.py index c52d3272..b4532143 100644 --- a/tests/gateway/test_ecs.py +++ b/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), }, ) diff --git a/tests/mapper/test_aws.py b/tests/mapper/test_aws.py index b99c4d9b..af94d4e8 100644 --- a/tests/mapper/test_aws.py +++ b/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": [