diff --git a/torchx/specs/named_resources_aws.py b/torchx/specs/named_resources_aws.py index a7480e587..88fb1a4ac 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_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, diff --git a/torchx/specs/test/named_resources_aws_test.py b/torchx/specs/test/named_resources_aws_test.py index e0e5dc770..e7a36ddc6 100644 --- a/torchx/specs/test/named_resources_aws_test.py +++ b/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,6 +96,7 @@ 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) @@ -102,6 +104,11 @@ def test_aws_p5(self) -> None: 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)