feat: add aws p5e instance type as named resources

Differential Revision: D83504664

Pull Request resolved: #1098

Author: CollapsedSilhouette

torchx/specs/named_resources_aws.py

@@ -120,6 +120,16 @@ def aws_p5_48xlarge() -> Resource:
     )
 
 
+def aws_p5e_48xlarge() -> Resource:
+    return Resource(
+        cpu=192,
+        gpu=8,
+        memMB=2048 * GiB,
+        capabilities={K8S_ITYPE: "p5e.48xlarge"},
+        devices={EFA_DEVICE: 32},
+    )
+
+
 def aws_p5en_48xlarge() -> Resource:
     return Resource(
         cpu=192,
@@ -419,6 +429,7 @@ def aws_inf2_48xlarge() -> Resource:
     "aws_p4d.24xlarge": aws_p4d_24xlarge,
     "aws_p4de.24xlarge": aws_p4de_24xlarge,
     "aws_p5.48xlarge": aws_p5_48xlarge,
+    "aws_p5e.48xlarge": aws_p5e_48xlarge,
     "aws_p5en.48xlarge": aws_p5en_48xlarge,
     "aws_g4dn.xlarge": aws_g4dn_xlarge,
     "aws_g4dn.2xlarge": aws_g4dn_2xlarge,

torchx/specs/test/named_resources_aws_test.py

@@ -44,6 +44,7 @@
     aws_p4d_24xlarge,
     aws_p4de_24xlarge,
     aws_p5_48xlarge,
+    aws_p5e_48xlarge,
     aws_p5en_48xlarge,
     aws_t3_medium,
     aws_trn1_2xlarge,
@@ -95,13 +96,19 @@ def test_aws_p4(self) -> None:
 
     def test_aws_p5(self) -> None:
         p5 = aws_p5_48xlarge()
+        p5e = aws_p5e_48xlarge()
         p5en = aws_p5en_48xlarge()
 
         self.assertEqual(192, p5.cpu)
         self.assertEqual(8, p5.gpu)
         self.assertEqual(2048 * GiB, p5.memMB)
         self.assertEqual({EFA_DEVICE: 32}, p5.devices)
 
+        self.assertEqual(192, p5e.cpu)
+        self.assertEqual(8, p5e.gpu)
+        self.assertEqual(2048 * GiB, p5e.memMB)
+        self.assertEqual({EFA_DEVICE: 32}, p5e.devices)
+
         self.assertEqual(192, p5en.cpu)
         self.assertEqual(8, p5en.gpu)
         self.assertEqual(2048 * GiB, p5en.memMB)