feat: add aws_p5en.48xlarge

pytorch · Dec 27, 2024 · c799ea8 · c799ea8
1 parent 5e44f20
commit c799ea8
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/torchx/specs/named_resources_aws.py b/torchx/specs/named_resources_aws.py
@@ -120,6 +120,16 @@ def aws_p5_48xlarge() -> Resource:
     )
 
 
+def aws_p5en_48xlarge() -> Resource:
+    return Resource(
+        cpu=192,
+        gpu=8,
+        memMB=2048 * GiB,
+        capabilities={K8S_ITYPE: "p5en.48xlarge"},
+        devices={EFA_DEVICE: 16},
+    )
+
+
 def aws_t3_medium() -> Resource:
     return Resource(cpu=2, gpu=0, memMB=4 * GiB, capabilities={K8S_ITYPE: "t3.medium"})
 
@@ -365,6 +375,7 @@ def aws_trn1_32xlarge() -> Resource:
     "aws_p4d.24xlarge": aws_p4d_24xlarge,
     "aws_p4de.24xlarge": aws_p4de_24xlarge,
     "aws_p5.48xlarge": aws_p5_48xlarge,
+    "aws_p5en.48xlarge": aws_p5en_48xlarge,
     "aws_g4dn.xlarge": aws_g4dn_xlarge,
     "aws_g4dn.2xlarge": aws_g4dn_2xlarge,
     "aws_g4dn.4xlarge": aws_g4dn_4xlarge,

diff --git a/torchx/specs/test/named_resources_aws_test.py b/torchx/specs/test/named_resources_aws_test.py
@@ -40,6 +40,7 @@
     aws_p4d_24xlarge,
     aws_p4de_24xlarge,
     aws_p5_48xlarge,
+    aws_p5en_48xlarge,
     aws_t3_medium,
     aws_trn1_2xlarge,
     aws_trn1_32xlarge,
@@ -90,12 +91,18 @@ def test_aws_p4(self) -> None:
 
     def test_aws_p5(self) -> None:
         p5 = aws_p5_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, p5en.cpu)
+        self.assertEqual(8, p5en.gpu)
+        self.assertEqual(2048 * GiB, p5en.memMB)
+        self.assertEqual({EFA_DEVICE: 16}, p5en.devices)
+
     def test_aws_g6e(self) -> None:
         g6e = aws_g6e_xlarge()
         g6e_2 = aws_g6e_2xlarge()