Add AWS Inf2 instances support for aws_batch scheduler

Author: shixianc

torchx/specs/named_resources_aws.py

@@ -354,6 +354,46 @@ def aws_trn1_32xlarge() -> Resource:
     )
 
 
+def aws_inf2_xlarge() -> Resource:
+    return Resource(
+        cpu=4,
+        gpu=0,
+        memMB=16 * GiB,
+        capabilities={K8S_ITYPE: "inf2.xlarge"},
+        devices={NEURON_DEVICE: 1},
+    )
+
+
+def aws_inf2_8xlarge() -> Resource:
+    return Resource(
+        cpu=32,
+        gpu=0,
+        memMB=128 * GiB,
+        capabilities={K8S_ITYPE: "inf2.8xlarge"},
+        devices={NEURON_DEVICE: 1},
+    )
+
+
+def aws_inf2_24xlarge() -> Resource:
+    return Resource(
+        cpu=96,
+        gpu=0,
+        memMB=384 * GiB,
+        capabilities={K8S_ITYPE: "inf2.24xlarge"},
+        devices={NEURON_DEVICE: 6},
+    )
+
+
+def aws_inf2_48xlarge() -> Resource:
+    return Resource(
+        cpu=192,
+        gpu=0,
+        memMB=768 * GiB,
+        capabilities={K8S_ITYPE: "inf2.48xlarge"},
+        devices={NEURON_DEVICE: 12},
+    )
+
+
 NAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = {
     "aws_t3.medium": aws_t3_medium,
     "aws_m5.2xlarge": aws_m5_2xlarge,
@@ -390,4 +430,8 @@ def aws_trn1_32xlarge() -> Resource:
     "aws_g6e.48xlarge": aws_g6e_48xlarge,
     "aws_trn1.2xlarge": aws_trn1_2xlarge,
     "aws_trn1.32xlarge": aws_trn1_32xlarge,
+    "aws_inf2.xlarge": aws_inf2_xlarge,
+    "aws_inf2.8xlarge": aws_inf2_8xlarge,
+    "aws_inf2.24xlarge": aws_inf2_24xlarge,
+    "aws_inf2.48xlarge": aws_inf2_48xlarge,
 }

torchx/specs/test/named_resources_aws_test.py

@@ -32,6 +32,10 @@
     aws_g6e_4xlarge,
     aws_g6e_8xlarge,
     aws_g6e_xlarge,
+    aws_inf2_24xlarge,
+    aws_inf2_48xlarge,
+    aws_inf2_8xlarge,
+    aws_inf2_xlarge,
     aws_m5_2xlarge,
     aws_p3_16xlarge,
     aws_p3_2xlarge,
@@ -232,6 +236,31 @@ def test_aws_trn1(self) -> None:
         self.assertEqual(trn1_32.memMB, trn1_2.memMB * 16)
         self.assertEqual({EFA_DEVICE: 8, NEURON_DEVICE: 16}, trn1_32.devices)
 
+    def test_aws_inf2(self) -> None:
+        inf2_1 = aws_inf2_xlarge()
+        self.assertEqual(4, inf2_1.cpu)
+        self.assertEqual(0, inf2_1.gpu)
+        self.assertEqual(16 * GiB, inf2_1.memMB)
+        self.assertEqual({NEURON_DEVICE: 1}, inf2_1.devices)
+
+        inf2_8 = aws_inf2_8xlarge()
+        self.assertEqual(32, inf2_8.cpu)
+        self.assertEqual(0, inf2_8.gpu)
+        self.assertEqual(128 * GiB, inf2_8.memMB)
+        self.assertEqual({NEURON_DEVICE: 1}, inf2_8.devices)
+
+        inf2_24 = aws_inf2_24xlarge()
+        self.assertEqual(96, inf2_24.cpu)
+        self.assertEqual(0, inf2_24.gpu)
+        self.assertEqual(384 * GiB, inf2_24.memMB)
+        self.assertEqual({NEURON_DEVICE: 6}, inf2_24.devices)
+
+        inf2_48 = aws_inf2_48xlarge()
+        self.assertEqual(192, inf2_48.cpu)
+        self.assertEqual(0, inf2_48.gpu)
+        self.assertEqual(768 * GiB, inf2_48.memMB)
+        self.assertEqual({NEURON_DEVICE: 12}, inf2_48.devices)
+
     def test_aws_m5_2xlarge(self) -> None:
         resource = aws_m5_2xlarge()
         self.assertEqual(8, resource.cpu)