feat: add aws_p5_48xlarge (#905)

Co-authored-by: Alexander Jipa <[email protected]>

Author: Alexander Jipa

torchx/specs/named_resources_aws.py

@@ -109,6 +109,16 @@ def aws_p4de_24xlarge() -> Resource:
     )
 
 
+def aws_p5_48xlarge() -> Resource:
+    return Resource(
+        cpu=192,
+        gpu=8,
+        memMB=2048 * GiB,
+        capabilities={K8S_ITYPE: "p5.48xlarge"},
+        devices={EFA_DEVICE: 32},
+    )
+
+
 def aws_t3_medium() -> Resource:
     return Resource(cpu=2, gpu=0, memMB=4 * GiB, capabilities={K8S_ITYPE: "t3.medium"})
 
@@ -268,6 +278,7 @@ def aws_trn1_32xlarge() -> Resource:
     "aws_p3dn.24xlarge": aws_p3dn_24xlarge,
     "aws_p4d.24xlarge": aws_p4d_24xlarge,
     "aws_p4de.24xlarge": aws_p4de_24xlarge,
+    "aws_p5.48xlarge": aws_p5_48xlarge,
     "aws_g4dn.xlarge": aws_g4dn_xlarge,
     "aws_g4dn.2xlarge": aws_g4dn_2xlarge,
     "aws_g4dn.4xlarge": aws_g4dn_4xlarge,

torchx/specs/test/named_resources_aws_test.py

@@ -30,6 +30,7 @@
     aws_p3dn_24xlarge,
     aws_p4d_24xlarge,
     aws_p4de_24xlarge,
+    aws_p5_48xlarge,
     aws_t3_medium,
     aws_trn1_2xlarge,
     aws_trn1_32xlarge,
@@ -77,6 +78,14 @@ def test_aws_p4(self) -> None:
         self.assertEqual(p4de.memMB, p4d.memMB)
         self.assertEqual({EFA_DEVICE: 4}, p4de.devices)
 
+    def test_aws_p5(self) -> None:
+        p5 = aws_p5_48xlarge()
+
+        self.assertEqual(192, p5.cpu)
+        self.assertEqual(8, p5.gpu)
+        self.assertEqual(2048 * GiB, p5.memMB)
+        self.assertEqual({EFA_DEVICE: 32}, p5.devices)
+
     def test_aws_g4dn(self) -> None:
         g4d = aws_g4dn_xlarge()
         self.assertEqual(4, g4d.cpu)