feat: add aws_c5_18xlarge

pytorch · Nov 27, 2024 · 5d689a8 · 5d689a8
1 parent 648913c
commit 5d689a8
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git 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
@@ -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)