We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
add EFA devices to AWS named_resources
## Description
For EFA devices to be available within Batch job container they have to be mapped via Job Definition (created by torchx's `aws_batch_scheduler`).
It has been made possible to register EFA devices when [this code](https://github.com/pytorch/torchx/blob/61533150ccd277ef002398559de54e5f578dd666/torchx/schedulers/devices.py#L26) was introduced.

All that's left is to add these devices to AWS named_resources that supports them.

## Motivation/Background
EFA significantly increases the throughput for cross-device communication. It is opt-in for the training job to use if or not (using an `env` flag), but it won't work unless it's registered in the Batch Job Definition.


## Detailed Proposal
Simple add EFA devices for AWS named_resources that have them.


## Alternatives
- Otherwise the customers have to copy-paste and register modified named_resources


## Additional context/links
* https://github.com/pytorch/torchx/blob/61533150ccd277ef002398559de54e5f578dd666/torchx/schedulers/devices.py#L26
* https://docs.aws.amazon.com/batch/latest/userguide/efa.html


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torchx/specs/named_resources_aws.py`
Content:
```
1 # Copyright (c) Meta Platforms, Inc. and affiliates.
2 # All rights reserved.
3 #
4 # This source code is licensed under the BSD-style license found in the
5 # LICENSE file in the root directory of this source tree.
6 
7 r"""
8 `torchx.specs.named_resources_aws` contains resource definitions that represent corresponding AWS instance types
9 taken from https://aws.amazon.com/ec2/instance-types/. The resources are exposed
10 via entrypoints after installing torchx lib. The mapping is stored in the `setup.py` file.
11 
12 The named resources currently do not specify AWS instance type capabilities but merely represent
13 the equvalent resource in mem, cpu and gpu numbers.
14 
15 .. note::
16     These resource definitions may change in future. It is expected for each user to
17     manage their own resources. Follow https://pytorch.org/torchx/latest/specs.html#torchx.specs.get_named_resources
18     to set up named resources.
19 
20 Usage:
21 
22     ::
23 
24      from torchx.specs import named_resources
25      print(named_resources["aws_t3.medium"])
26      print(named_resources["aws_m5.2xlarge"])
27      print(named_resources["aws_p3.2xlarge"])
28      print(named_resources["aws_p3.8xlarge"])
29 
30 """
31 
32 from typing import Callable, Mapping
33 
34 from torchx.specs.api import Resource
35 
36 # ecs and ec2 have memtax and currently AWS Batch uses hard memory limits
37 # so we have to account for mem tax when registering these resources for AWS
38 # otherwise the job will be stuck in the jobqueue forever
39 # 97% is based on empirical observation that works well for most instance types
40 # see: https://docs.aws.amazon.com/batch/latest/userguide/memory-management.html
41 MEM_TAX = 0.97
42 K8S_ITYPE = "node.kubernetes.io/instance-type"
43 GiB: int = int(1024 * MEM_TAX)
44 
45 
46 def aws_p3_2xlarge() -> Resource:
47     return Resource(
48         cpu=8, gpu=1, memMB=61 * GiB, capabilities={K8S_ITYPE: "p3.2xlarge"}
49     )
50 
51 
52 def aws_p3_8xlarge() -> Resource:
53     return Resource(
54         cpu=32, gpu=4, memMB=244 * GiB, capabilities={K8S_ITYPE: "p3.8xlarge"}
55     )
56 
57 
58 def aws_p3_16xlarge() -> Resource:
59     return Resource(
60         cpu=64, gpu=8, memMB=488 * GiB, capabilities={K8S_ITYPE: "p3.16xlarge"}
61     )
62 
63 
64 def aws_p3dn_24xlarge() -> Resource:
65     return Resource(
66         cpu=96, gpu=8, memMB=768 * GiB, capabilities={K8S_ITYPE: "p3dn.24xlarge"}
67     )
68 
69 
70 def aws_p4d_24xlarge() -> Resource:
71     return Resource(
72         cpu=96, gpu=8, memMB=1152 * GiB, capabilities={K8S_ITYPE: "p4d.24xlarge"}
73     )
74 
75 
76 def aws_p4de_24xlarge() -> Resource:
77     # p4de has same cpu, gpu, memMB as p4d but gpu memory is 2x (32GB vs 64GB per GPU)
78     return Resource(
79         cpu=96, gpu=8, memMB=1152 * GiB, capabilities={K8S_ITYPE: "p4de.24xlarge"}
80     )
81 
82 
83 def aws_t3_medium() -> Resource:
84     return Resource(cpu=2, gpu=0, memMB=4 * GiB, capabilities={K8S_ITYPE: "t3.medium"})
85 
86 
87 def aws_m5_2xlarge() -> Resource:
88     return Resource(
89         cpu=8, gpu=0, memMB=32 * GiB, capabilities={K8S_ITYPE: "m5.2xlarge"}
90     )
91 
92 
93 def aws_g4dn_xlarge() -> Resource:
94     return Resource(
95         cpu=4, gpu=1, memMB=16 * GiB, capabilities={K8S_ITYPE: "g4dn.xlarge"}
96     )
97 
98 
99 def aws_g4dn_2xlarge() -> Resource:
100     return Resource(
101         cpu=8, gpu=1, memMB=32 * GiB, capabilities={K8S_ITYPE: "g4dn.2xlarge"}
102     )
103 
104 
105 def aws_g4dn_4xlarge() -> Resource:
106     return Resource(
107         cpu=16, gpu=1, memMB=64 * GiB, capabilities={K8S_ITYPE: "g4dn.4xlarge"}
108     )
109 
110 
111 def aws_g4dn_8xlarge() -> Resource:
112     return Resource(
113         cpu=32, gpu=1, memMB=128 * GiB, capabilities={K8S_ITYPE: "g4dn.8xlarge"}
114     )
115 
116 
117 def aws_g4dn_16xlarge() -> Resource:
118     return Resource(
119         cpu=64, gpu=1, memMB=256 * GiB, capabilities={K8S_ITYPE: "g4dn.16xlarge"}
120     )
121 
122 
123 def aws_g4dn_12xlarge() -> Resource:
124     return Resource(
125         cpu=48, gpu=4, memMB=192 * GiB, capabilities={K8S_ITYPE: "g4dn.12xlarge"}
126     )
127 
128 
129 def aws_g4dn_metal() -> Resource:
130     return Resource(
131         cpu=96, gpu=8, memMB=384 * GiB, capabilities={K8S_ITYPE: "g4dn.metal"}
132     )
133 
134 
135 def aws_g5_xlarge() -> Resource:
136     return Resource(cpu=4, gpu=1, memMB=16 * GiB, capabilities={K8S_ITYPE: "g5.xlarge"})
137 
138 
139 def aws_g5_2xlarge() -> Resource:
140     return Resource(
141         cpu=8, gpu=1, memMB=32 * GiB, capabilities={K8S_ITYPE: "g5.2xlarge"}
142     )
143 
144 
145 def aws_g5_4xlarge() -> Resource:
146     return Resource(
147         cpu=16, gpu=1, memMB=64 * GiB, capabilities={K8S_ITYPE: "g5.4xlarge"}
148     )
149 
150 
151 def aws_g5_8xlarge() -> Resource:
152     return Resource(
153         cpu=32, gpu=1, memMB=128 * GiB, capabilities={K8S_ITYPE: "g5.8xlarge"}
154     )
155 
156 
157 def aws_g5_16xlarge() -> Resource:
158     return Resource(
159         cpu=64, gpu=1, memMB=256 * GiB, capabilities={K8S_ITYPE: "g5.16xlarge"}
160     )
161 
162 
163 def aws_g5_12xlarge() -> Resource:
164     return Resource(
165         cpu=48, gpu=4, memMB=192 * GiB, capabilities={K8S_ITYPE: "g5.12xlarge"}
166     )
167 
168 
169 def aws_g5_24xlarge() -> Resource:
170     return Resource(
171         cpu=96, gpu=4, memMB=384 * GiB, capabilities={K8S_ITYPE: "g5.24xlarge"}
172     )
173 
174 
175 def aws_g5_48xlarge() -> Resource:
176     return Resource(
177         cpu=192, gpu=8, memMB=768 * GiB, capabilities={K8S_ITYPE: "g5.48xlarge"}
178     )
179 
180 
181 def aws_trn1_2xl() -> Resource:
182     return Resource(cpu=8, gpu=0, memMB=32 * GiB, capabilities={K8S_ITYPE: "trn1.2xl"})
183 
184 
185 def aws_trn1_32xl() -> Resource:
186     return Resource(
187         cpu=128, gpu=0, memMB=512 * GiB, capabilities={K8S_ITYPE: "trn1.32xl"}
188     )
189 
190 
191 NAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = {
192     "aws_t3.medium": aws_t3_medium,
193     "aws_m5.2xlarge": aws_m5_2xlarge,
194     "aws_p3.2xlarge": aws_p3_2xlarge,
195     "aws_p3.8xlarge": aws_p3_8xlarge,
196     "aws_p3.16xlarge": aws_p3_16xlarge,
197     "aws_p3dn.24xlarge": aws_p3dn_24xlarge,
198     "aws_p4d.24xlarge": aws_p4d_24xlarge,
199     "aws_p4de.24xlarge": aws_p4de_24xlarge,
200     "aws_g4dn.xlarge": aws_g4dn_xlarge,
201     "aws_g4dn.2xlarge": aws_g4dn_2xlarge,
202     "aws_g4dn.4xlarge": aws_g4dn_4xlarge,
203     "aws_g4dn.8xlarge": aws_g4dn_8xlarge,
204     "aws_g4dn.16xlarge": aws_g4dn_16xlarge,
205     "aws_g4dn.12xlarge": aws_g4dn_12xlarge,
206     "aws_g4dn.metal": aws_g4dn_metal,
207     "aws_g5.xlarge": aws_g5_xlarge,
208     "aws_g5.2xlarge": aws_g5_2xlarge,
209     "aws_g5.4xlarge": aws_g5_4xlarge,
210     "aws_g5.8xlarge": aws_g5_8xlarge,
211     "aws_g5.16xlarge": aws_g5_16xlarge,
212     "aws_g5.12xlarge": aws_g5_12xlarge,
213     "aws_g5.24xlarge": aws_g5_24xlarge,
214     "aws_g5.48xlarge": aws_g5_48xlarge,
215     "aws_trn1.2xl": aws_trn1_2xl,
216     "aws_trn1.32xl": aws_trn1_32xl,
217 }
218
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torchx/specs/named_resources_aws.py b/torchx/specs/named_resources_aws.py
--- a/torchx/specs/named_resources_aws.py
+++ b/torchx/specs/named_resources_aws.py
@@ -33,6 +33,8 @@
 
 from torchx.specs.api import Resource
 
+EFA_DEVICE = "vpc.amazonaws.com/efa"
+
 # ecs and ec2 have memtax and currently AWS Batch uses hard memory limits
 # so we have to account for mem tax when registering these resources for AWS
 # otherwise the job will be stuck in the jobqueue forever
@@ -63,20 +65,32 @@
 
 def aws_p3dn_24xlarge() -> Resource:
     return Resource(
-        cpu=96, gpu=8, memMB=768 * GiB, capabilities={K8S_ITYPE: "p3dn.24xlarge"}
+        cpu=96,
+        gpu=8,
+        memMB=768 * GiB,
+        capabilities={K8S_ITYPE: "p3dn.24xlarge"},
+        devices={EFA_DEVICE: 1},
     )
 
 
 def aws_p4d_24xlarge() -> Resource:
     return Resource(
-        cpu=96, gpu=8, memMB=1152 * GiB, capabilities={K8S_ITYPE: "p4d.24xlarge"}
+        cpu=96,
+        gpu=8,
+        memMB=1152 * GiB,
+        capabilities={K8S_ITYPE: "p4d.24xlarge"},
+        devices={EFA_DEVICE: 4},
     )
 
 
 def aws_p4de_24xlarge() -> Resource:
     # p4de has same cpu, gpu, memMB as p4d but gpu memory is 2x (32GB vs 64GB per GPU)
     return Resource(
-        cpu=96, gpu=8, memMB=1152 * GiB, capabilities={K8S_ITYPE: "p4de.24xlarge"}
+        cpu=96,
+        gpu=8,
+        memMB=1152 * GiB,
+        capabilities={K8S_ITYPE: "p4de.24xlarge"},
+        devices={EFA_DEVICE: 4},
     )

{"golden_diff": "diff --git a/torchx/specs/named_resources_aws.py b/torchx/specs/named_resources_aws.py\n--- a/torchx/specs/named_resources_aws.py\n+++ b/torchx/specs/named_resources_aws.py\n@@ -33,6 +33,8 @@\n \n from torchx.specs.api import Resource\n \n+EFA_DEVICE = \"vpc.amazonaws.com/efa\"\n+\n # ecs and ec2 have memtax and currently AWS Batch uses hard memory limits\n # so we have to account for mem tax when registering these resources for AWS\n # otherwise the job will be stuck in the jobqueue forever\n@@ -63,20 +65,32 @@\n \n def aws_p3dn_24xlarge() -> Resource:\n     return Resource(\n-        cpu=96, gpu=8, memMB=768 * GiB, capabilities={K8S_ITYPE: \"p3dn.24xlarge\"}\n+        cpu=96,\n+        gpu=8,\n+        memMB=768 * GiB,\n+        capabilities={K8S_ITYPE: \"p3dn.24xlarge\"},\n+        devices={EFA_DEVICE: 1},\n     )\n \n \n def aws_p4d_24xlarge() -> Resource:\n     return Resource(\n-        cpu=96, gpu=8, memMB=1152 * GiB, capabilities={K8S_ITYPE: \"p4d.24xlarge\"}\n+        cpu=96,\n+        gpu=8,\n+        memMB=1152 * GiB,\n+        capabilities={K8S_ITYPE: \"p4d.24xlarge\"},\n+        devices={EFA_DEVICE: 4},\n     )\n \n \n def aws_p4de_24xlarge() -> Resource:\n     # p4de has same cpu, gpu, memMB as p4d but gpu memory is 2x (32GB vs 64GB per GPU)\n     return Resource(\n-        cpu=96, gpu=8, memMB=1152 * GiB, capabilities={K8S_ITYPE: \"p4de.24xlarge\"}\n+        cpu=96,\n+        gpu=8,\n+        memMB=1152 * GiB,\n+        capabilities={K8S_ITYPE: \"p4de.24xlarge\"},\n+        devices={EFA_DEVICE: 4},\n     )\n", "issue": "add EFA devices to AWS named_resources\n## Description\r\nFor EFA devices to be available within Batch job container they have to be mapped via Job Definition (created by torchx's `aws_batch_scheduler`).\r\nIt has been made possible to register EFA devices when [this code](https://github.com/pytorch/torchx/blob/61533150ccd277ef002398559de54e5f578dd666/torchx/schedulers/devices.py#L26) was introduced.\r\n\r\nAll that's left is to add these devices to AWS named_resources that supports them.\r\n\r\n## Motivation/Background\r\nEFA significantly increases the throughput for cross-device communication. It is opt-in for the training job to use if or not (using an `env` flag), but it won't work unless it's registered in the Batch Job Definition.\r\n\r\n\r\n## Detailed Proposal\r\nSimple add EFA devices for AWS named_resources that have them.\r\n\r\n\r\n## Alternatives\r\n- Otherwise the customers have to copy-paste and register modified named_resources\r\n\r\n\r\n## Additional context/links\r\n* https://github.com/pytorch/torchx/blob/61533150ccd277ef002398559de54e5f578dd666/torchx/schedulers/devices.py#L26\r\n* https://docs.aws.amazon.com/batch/latest/userguide/efa.html\r\n\n", "before_files": [{"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the BSD-style license found in the\n# LICENSE file in the root directory of this source tree.\n\nr\"\"\"\n`torchx.specs.named_resources_aws` contains resource definitions that represent corresponding AWS instance types\ntaken from https://aws.amazon.com/ec2/instance-types/. The resources are exposed\nvia entrypoints after installing torchx lib. The mapping is stored in the `setup.py` file.\n\nThe named resources currently do not specify AWS instance type capabilities but merely represent\nthe equvalent resource in mem, cpu and gpu numbers.\n\n.. note::\n    These resource definitions may change in future. It is expected for each user to\n    manage their own resources. Follow https://pytorch.org/torchx/latest/specs.html#torchx.specs.get_named_resources\n    to set up named resources.\n\nUsage:\n\n    ::\n\n     from torchx.specs import named_resources\n     print(named_resources[\"aws_t3.medium\"])\n     print(named_resources[\"aws_m5.2xlarge\"])\n     print(named_resources[\"aws_p3.2xlarge\"])\n     print(named_resources[\"aws_p3.8xlarge\"])\n\n\"\"\"\n\nfrom typing import Callable, Mapping\n\nfrom torchx.specs.api import Resource\n\n# ecs and ec2 have memtax and currently AWS Batch uses hard memory limits\n# so we have to account for mem tax when registering these resources for AWS\n# otherwise the job will be stuck in the jobqueue forever\n# 97% is based on empirical observation that works well for most instance types\n# see: https://docs.aws.amazon.com/batch/latest/userguide/memory-management.html\nMEM_TAX = 0.97\nK8S_ITYPE = \"node.kubernetes.io/instance-type\"\nGiB: int = int(1024 * MEM_TAX)\n\n\ndef aws_p3_2xlarge() -> Resource:\n    return Resource(\n        cpu=8, gpu=1, memMB=61 * GiB, capabilities={K8S_ITYPE: \"p3.2xlarge\"}\n    )\n\n\ndef aws_p3_8xlarge() -> Resource:\n    return Resource(\n        cpu=32, gpu=4, memMB=244 * GiB, capabilities={K8S_ITYPE: \"p3.8xlarge\"}\n    )\n\n\ndef aws_p3_16xlarge() -> Resource:\n    return Resource(\n        cpu=64, gpu=8, memMB=488 * GiB, capabilities={K8S_ITYPE: \"p3.16xlarge\"}\n    )\n\n\ndef aws_p3dn_24xlarge() -> Resource:\n    return Resource(\n        cpu=96, gpu=8, memMB=768 * GiB, capabilities={K8S_ITYPE: \"p3dn.24xlarge\"}\n    )\n\n\ndef aws_p4d_24xlarge() -> Resource:\n    return Resource(\n        cpu=96, gpu=8, memMB=1152 * GiB, capabilities={K8S_ITYPE: \"p4d.24xlarge\"}\n    )\n\n\ndef aws_p4de_24xlarge() -> Resource:\n    # p4de has same cpu, gpu, memMB as p4d but gpu memory is 2x (32GB vs 64GB per GPU)\n    return Resource(\n        cpu=96, gpu=8, memMB=1152 * GiB, capabilities={K8S_ITYPE: \"p4de.24xlarge\"}\n    )\n\n\ndef aws_t3_medium() -> Resource:\n    return Resource(cpu=2, gpu=0, memMB=4 * GiB, capabilities={K8S_ITYPE: \"t3.medium\"})\n\n\ndef aws_m5_2xlarge() -> Resource:\n    return Resource(\n        cpu=8, gpu=0, memMB=32 * GiB, capabilities={K8S_ITYPE: \"m5.2xlarge\"}\n    )\n\n\ndef aws_g4dn_xlarge() -> Resource:\n    return Resource(\n        cpu=4, gpu=1, memMB=16 * GiB, capabilities={K8S_ITYPE: \"g4dn.xlarge\"}\n    )\n\n\ndef aws_g4dn_2xlarge() -> Resource:\n    return Resource(\n        cpu=8, gpu=1, memMB=32 * GiB, capabilities={K8S_ITYPE: \"g4dn.2xlarge\"}\n    )\n\n\ndef aws_g4dn_4xlarge() -> Resource:\n    return Resource(\n        cpu=16, gpu=1, memMB=64 * GiB, capabilities={K8S_ITYPE: \"g4dn.4xlarge\"}\n    )\n\n\ndef aws_g4dn_8xlarge() -> Resource:\n    return Resource(\n        cpu=32, gpu=1, memMB=128 * GiB, capabilities={K8S_ITYPE: \"g4dn.8xlarge\"}\n    )\n\n\ndef aws_g4dn_16xlarge() -> Resource:\n    return Resource(\n        cpu=64, gpu=1, memMB=256 * GiB, capabilities={K8S_ITYPE: \"g4dn.16xlarge\"}\n    )\n\n\ndef aws_g4dn_12xlarge() -> Resource:\n    return Resource(\n        cpu=48, gpu=4, memMB=192 * GiB, capabilities={K8S_ITYPE: \"g4dn.12xlarge\"}\n    )\n\n\ndef aws_g4dn_metal() -> Resource:\n    return Resource(\n        cpu=96, gpu=8, memMB=384 * GiB, capabilities={K8S_ITYPE: \"g4dn.metal\"}\n    )\n\n\ndef aws_g5_xlarge() -> Resource:\n    return Resource(cpu=4, gpu=1, memMB=16 * GiB, capabilities={K8S_ITYPE: \"g5.xlarge\"})\n\n\ndef aws_g5_2xlarge() -> Resource:\n    return Resource(\n        cpu=8, gpu=1, memMB=32 * GiB, capabilities={K8S_ITYPE: \"g5.2xlarge\"}\n    )\n\n\ndef aws_g5_4xlarge() -> Resource:\n    return Resource(\n        cpu=16, gpu=1, memMB=64 * GiB, capabilities={K8S_ITYPE: \"g5.4xlarge\"}\n    )\n\n\ndef aws_g5_8xlarge() -> Resource:\n    return Resource(\n        cpu=32, gpu=1, memMB=128 * GiB, capabilities={K8S_ITYPE: \"g5.8xlarge\"}\n    )\n\n\ndef aws_g5_16xlarge() -> Resource:\n    return Resource(\n        cpu=64, gpu=1, memMB=256 * GiB, capabilities={K8S_ITYPE: \"g5.16xlarge\"}\n    )\n\n\ndef aws_g5_12xlarge() -> Resource:\n    return Resource(\n        cpu=48, gpu=4, memMB=192 * GiB, capabilities={K8S_ITYPE: \"g5.12xlarge\"}\n    )\n\n\ndef aws_g5_24xlarge() -> Resource:\n    return Resource(\n        cpu=96, gpu=4, memMB=384 * GiB, capabilities={K8S_ITYPE: \"g5.24xlarge\"}\n    )\n\n\ndef aws_g5_48xlarge() -> Resource:\n    return Resource(\n        cpu=192, gpu=8, memMB=768 * GiB, capabilities={K8S_ITYPE: \"g5.48xlarge\"}\n    )\n\n\ndef aws_trn1_2xl() -> Resource:\n    return Resource(cpu=8, gpu=0, memMB=32 * GiB, capabilities={K8S_ITYPE: \"trn1.2xl\"})\n\n\ndef aws_trn1_32xl() -> Resource:\n    return Resource(\n        cpu=128, gpu=0, memMB=512 * GiB, capabilities={K8S_ITYPE: \"trn1.32xl\"}\n    )\n\n\nNAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = {\n    \"aws_t3.medium\": aws_t3_medium,\n    \"aws_m5.2xlarge\": aws_m5_2xlarge,\n    \"aws_p3.2xlarge\": aws_p3_2xlarge,\n    \"aws_p3.8xlarge\": aws_p3_8xlarge,\n    \"aws_p3.16xlarge\": aws_p3_16xlarge,\n    \"aws_p3dn.24xlarge\": aws_p3dn_24xlarge,\n    \"aws_p4d.24xlarge\": aws_p4d_24xlarge,\n    \"aws_p4de.24xlarge\": aws_p4de_24xlarge,\n    \"aws_g4dn.xlarge\": aws_g4dn_xlarge,\n    \"aws_g4dn.2xlarge\": aws_g4dn_2xlarge,\n    \"aws_g4dn.4xlarge\": aws_g4dn_4xlarge,\n    \"aws_g4dn.8xlarge\": aws_g4dn_8xlarge,\n    \"aws_g4dn.16xlarge\": aws_g4dn_16xlarge,\n    \"aws_g4dn.12xlarge\": aws_g4dn_12xlarge,\n    \"aws_g4dn.metal\": aws_g4dn_metal,\n    \"aws_g5.xlarge\": aws_g5_xlarge,\n    \"aws_g5.2xlarge\": aws_g5_2xlarge,\n    \"aws_g5.4xlarge\": aws_g5_4xlarge,\n    \"aws_g5.8xlarge\": aws_g5_8xlarge,\n    \"aws_g5.16xlarge\": aws_g5_16xlarge,\n    \"aws_g5.12xlarge\": aws_g5_12xlarge,\n    \"aws_g5.24xlarge\": aws_g5_24xlarge,\n    \"aws_g5.48xlarge\": aws_g5_48xlarge,\n    \"aws_trn1.2xl\": aws_trn1_2xl,\n    \"aws_trn1.32xl\": aws_trn1_32xl,\n}\n", "path": "torchx/specs/named_resources_aws.py"}], "after_files": [{"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the BSD-style license found in the\n# LICENSE file in the root directory of this source tree.\n\nr\"\"\"\n`torchx.specs.named_resources_aws` contains resource definitions that represent corresponding AWS instance types\ntaken from https://aws.amazon.com/ec2/instance-types/. The resources are exposed\nvia entrypoints after installing torchx lib. The mapping is stored in the `setup.py` file.\n\nThe named resources currently do not specify AWS instance type capabilities but merely represent\nthe equvalent resource in mem, cpu and gpu numbers.\n\n.. note::\n    These resource definitions may change in future. It is expected for each user to\n    manage their own resources. Follow https://pytorch.org/torchx/latest/specs.html#torchx.specs.get_named_resources\n    to set up named resources.\n\nUsage:\n\n    ::\n\n     from torchx.specs import named_resources\n     print(named_resources[\"aws_t3.medium\"])\n     print(named_resources[\"aws_m5.2xlarge\"])\n     print(named_resources[\"aws_p3.2xlarge\"])\n     print(named_resources[\"aws_p3.8xlarge\"])\n\n\"\"\"\n\nfrom typing import Callable, Mapping\n\nfrom torchx.specs.api import Resource\n\nEFA_DEVICE = \"vpc.amazonaws.com/efa\"\n\n# ecs and ec2 have memtax and currently AWS Batch uses hard memory limits\n# so we have to account for mem tax when registering these resources for AWS\n# otherwise the job will be stuck in the jobqueue forever\n# 97% is based on empirical observation that works well for most instance types\n# see: https://docs.aws.amazon.com/batch/latest/userguide/memory-management.html\nMEM_TAX = 0.97\nK8S_ITYPE = \"node.kubernetes.io/instance-type\"\nGiB: int = int(1024 * MEM_TAX)\n\n\ndef aws_p3_2xlarge() -> Resource:\n    return Resource(\n        cpu=8, gpu=1, memMB=61 * GiB, capabilities={K8S_ITYPE: \"p3.2xlarge\"}\n    )\n\n\ndef aws_p3_8xlarge() -> Resource:\n    return Resource(\n        cpu=32, gpu=4, memMB=244 * GiB, capabilities={K8S_ITYPE: \"p3.8xlarge\"}\n    )\n\n\ndef aws_p3_16xlarge() -> Resource:\n    return Resource(\n        cpu=64, gpu=8, memMB=488 * GiB, capabilities={K8S_ITYPE: \"p3.16xlarge\"}\n    )\n\n\ndef aws_p3dn_24xlarge() -> Resource:\n    return Resource(\n        cpu=96,\n        gpu=8,\n        memMB=768 * GiB,\n        capabilities={K8S_ITYPE: \"p3dn.24xlarge\"},\n        devices={EFA_DEVICE: 1},\n    )\n\n\ndef aws_p4d_24xlarge() -> Resource:\n    return Resource(\n        cpu=96,\n        gpu=8,\n        memMB=1152 * GiB,\n        capabilities={K8S_ITYPE: \"p4d.24xlarge\"},\n        devices={EFA_DEVICE: 4},\n    )\n\n\ndef aws_p4de_24xlarge() -> Resource:\n    # p4de has same cpu, gpu, memMB as p4d but gpu memory is 2x (32GB vs 64GB per GPU)\n    return Resource(\n        cpu=96,\n        gpu=8,\n        memMB=1152 * GiB,\n        capabilities={K8S_ITYPE: \"p4de.24xlarge\"},\n        devices={EFA_DEVICE: 4},\n    )\n\n\ndef aws_t3_medium() -> Resource:\n    return Resource(cpu=2, gpu=0, memMB=4 * GiB, capabilities={K8S_ITYPE: \"t3.medium\"})\n\n\ndef aws_m5_2xlarge() -> Resource:\n    return Resource(\n        cpu=8, gpu=0, memMB=32 * GiB, capabilities={K8S_ITYPE: \"m5.2xlarge\"}\n    )\n\n\ndef aws_g4dn_xlarge() -> Resource:\n    return Resource(\n        cpu=4, gpu=1, memMB=16 * GiB, capabilities={K8S_ITYPE: \"g4dn.xlarge\"}\n    )\n\n\ndef aws_g4dn_2xlarge() -> Resource:\n    return Resource(\n        cpu=8, gpu=1, memMB=32 * GiB, capabilities={K8S_ITYPE: \"g4dn.2xlarge\"}\n    )\n\n\ndef aws_g4dn_4xlarge() -> Resource:\n    return Resource(\n        cpu=16, gpu=1, memMB=64 * GiB, capabilities={K8S_ITYPE: \"g4dn.4xlarge\"}\n    )\n\n\ndef aws_g4dn_8xlarge() -> Resource:\n    return Resource(\n        cpu=32, gpu=1, memMB=128 * GiB, capabilities={K8S_ITYPE: \"g4dn.8xlarge\"}\n    )\n\n\ndef aws_g4dn_16xlarge() -> Resource:\n    return Resource(\n        cpu=64, gpu=1, memMB=256 * GiB, capabilities={K8S_ITYPE: \"g4dn.16xlarge\"}\n    )\n\n\ndef aws_g4dn_12xlarge() -> Resource:\n    return Resource(\n        cpu=48, gpu=4, memMB=192 * GiB, capabilities={K8S_ITYPE: \"g4dn.12xlarge\"}\n    )\n\n\ndef aws_g4dn_metal() -> Resource:\n    return Resource(\n        cpu=96, gpu=8, memMB=384 * GiB, capabilities={K8S_ITYPE: \"g4dn.metal\"}\n    )\n\n\ndef aws_g5_xlarge() -> Resource:\n    return Resource(cpu=4, gpu=1, memMB=16 * GiB, capabilities={K8S_ITYPE: \"g5.xlarge\"})\n\n\ndef aws_g5_2xlarge() -> Resource:\n    return Resource(\n        cpu=8, gpu=1, memMB=32 * GiB, capabilities={K8S_ITYPE: \"g5.2xlarge\"}\n    )\n\n\ndef aws_g5_4xlarge() -> Resource:\n    return Resource(\n        cpu=16, gpu=1, memMB=64 * GiB, capabilities={K8S_ITYPE: \"g5.4xlarge\"}\n    )\n\n\ndef aws_g5_8xlarge() -> Resource:\n    return Resource(\n        cpu=32, gpu=1, memMB=128 * GiB, capabilities={K8S_ITYPE: \"g5.8xlarge\"}\n    )\n\n\ndef aws_g5_16xlarge() -> Resource:\n    return Resource(\n        cpu=64, gpu=1, memMB=256 * GiB, capabilities={K8S_ITYPE: \"g5.16xlarge\"}\n    )\n\n\ndef aws_g5_12xlarge() -> Resource:\n    return Resource(\n        cpu=48, gpu=4, memMB=192 * GiB, capabilities={K8S_ITYPE: \"g5.12xlarge\"}\n    )\n\n\ndef aws_g5_24xlarge() -> Resource:\n    return Resource(\n        cpu=96, gpu=4, memMB=384 * GiB, capabilities={K8S_ITYPE: \"g5.24xlarge\"}\n    )\n\n\ndef aws_g5_48xlarge() -> Resource:\n    return Resource(\n        cpu=192, gpu=8, memMB=768 * GiB, capabilities={K8S_ITYPE: \"g5.48xlarge\"}\n    )\n\n\ndef aws_trn1_2xl() -> Resource:\n    return Resource(cpu=8, gpu=0, memMB=32 * GiB, capabilities={K8S_ITYPE: \"trn1.2xl\"})\n\n\ndef aws_trn1_32xl() -> Resource:\n    return Resource(\n        cpu=128, gpu=0, memMB=512 * GiB, capabilities={K8S_ITYPE: \"trn1.32xl\"}\n    )\n\n\nNAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = {\n    \"aws_t3.medium\": aws_t3_medium,\n    \"aws_m5.2xlarge\": aws_m5_2xlarge,\n    \"aws_p3.2xlarge\": aws_p3_2xlarge,\n    \"aws_p3.8xlarge\": aws_p3_8xlarge,\n    \"aws_p3.16xlarge\": aws_p3_16xlarge,\n    \"aws_p3dn.24xlarge\": aws_p3dn_24xlarge,\n    \"aws_p4d.24xlarge\": aws_p4d_24xlarge,\n    \"aws_p4de.24xlarge\": aws_p4de_24xlarge,\n    \"aws_g4dn.xlarge\": aws_g4dn_xlarge,\n    \"aws_g4dn.2xlarge\": aws_g4dn_2xlarge,\n    \"aws_g4dn.4xlarge\": aws_g4dn_4xlarge,\n    \"aws_g4dn.8xlarge\": aws_g4dn_8xlarge,\n    \"aws_g4dn.16xlarge\": aws_g4dn_16xlarge,\n    \"aws_g4dn.12xlarge\": aws_g4dn_12xlarge,\n    \"aws_g4dn.metal\": aws_g4dn_metal,\n    \"aws_g5.xlarge\": aws_g5_xlarge,\n    \"aws_g5.2xlarge\": aws_g5_2xlarge,\n    \"aws_g5.4xlarge\": aws_g5_4xlarge,\n    \"aws_g5.8xlarge\": aws_g5_8xlarge,\n    \"aws_g5.16xlarge\": aws_g5_16xlarge,\n    \"aws_g5.12xlarge\": aws_g5_12xlarge,\n    \"aws_g5.24xlarge\": aws_g5_24xlarge,\n    \"aws_g5.48xlarge\": aws_g5_48xlarge,\n    \"aws_trn1.2xl\": aws_trn1_2xl,\n    \"aws_trn1.32xl\": aws_trn1_32xl,\n}\n", "path": "torchx/specs/named_resources_aws.py"}]}