From b0eff9900e08e7eb349771630b0e5c6a45373443 Mon Sep 17 00:00:00 2001
From: Alexander Zhipa <azzhipa@amazon.com>
Date: Fri, 27 Dec 2024 10:07:52 -0500
Subject: [PATCH] feat: add aws_p5en.48xlarge

---
 torchx/specs/named_resources_aws.py           | 11 +++++++++++
 torchx/specs/test/named_resources_aws_test.py |  7 +++++++
 2 files changed, 18 insertions(+)

diff --git a/torchx/specs/named_resources_aws.py b/torchx/specs/named_resources_aws.py
index 49456a0c9..c290c3c38 100644
--- 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
index fcd4526dd..cdc99573a 100644
--- 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()