diff --git a/torchx/specs/named_resources_aws.py b/torchx/specs/named_resources_aws.py index 555513a06..49456a0c9 100644 --- a/torchx/specs/named_resources_aws.py +++ b/torchx/specs/named_resources_aws.py @@ -130,6 +130,12 @@ def aws_m5_2xlarge() -> Resource: ) +def aws_c5_18xlarge() -> Resource: + return Resource( + cpu=72, gpu=0, memMB=144 * GiB, capabilities={K8S_ITYPE: "c5.18xlarge"} + ) + + def aws_g4dn_xlarge() -> Resource: return Resource( cpu=4, gpu=1, memMB=16 * GiB, capabilities={K8S_ITYPE: "g4dn.xlarge"} @@ -351,6 +357,7 @@ def aws_trn1_32xlarge() -> Resource: NAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = { "aws_t3.medium": aws_t3_medium, "aws_m5.2xlarge": aws_m5_2xlarge, + "aws_c5.18xlarge": aws_c5_18xlarge, "aws_p3.2xlarge": aws_p3_2xlarge, "aws_p3.8xlarge": aws_p3_8xlarge, "aws_p3.16xlarge": aws_p3_16xlarge, diff --git a/torchx/specs/test/named_resources_aws_test.py b/torchx/specs/test/named_resources_aws_test.py index 3480fa463..fcd4526dd 100644 --- a/torchx/specs/test/named_resources_aws_test.py +++ b/torchx/specs/test/named_resources_aws_test.py @@ -8,6 +8,7 @@ import unittest from torchx.specs.named_resources_aws import ( + aws_c5_18xlarge, aws_g4dn_12xlarge, aws_g4dn_16xlarge, aws_g4dn_2xlarge, @@ -237,6 +238,12 @@ def test_aws_m5_2xlarge(self) -> None: self.assertEqual(0, resource.gpu) self.assertEqual(32 * GiB, resource.memMB) + def test_aws_c5_18xlarge(self) -> None: + resource = aws_c5_18xlarge() + self.assertEqual(72, resource.cpu) + self.assertEqual(0, resource.gpu) + self.assertEqual(144 * GiB, resource.memMB) + def test_aws_t3_medium(self) -> None: resource = aws_t3_medium() self.assertEqual(2, resource.cpu)