# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.

efa_supported_instance = set(
    [
        "c5n.18xlarge",
        "c5n.9xlarge",
        "c5n.metal",
        "c6a.48xlarge",
        "c6a.metal",
        "c6gn.16xlarge",
        "c6i.32xlarge",
        "c6i.metal",
        "c6id.32xlarge",
        "c6id.metal",
        "c6in.32xlarge",
        "c6in.metal",
        "c7a.48xlarge",
        "c7a.metal-48xl",
        "c7g.16xlarge",
        "c7g.metal",
        "c7gd.16xlarge",
        "c7gd.metal",
        "c7gn.16xlarge",
        "c7gn.metal",
        "c7i.48xlarge",
        "c7i.metal-48xl",
        "dl1.24xlarge",
        "dl2q.24xlarge",
        "g4dn.12xlarge",
        "g4dn.16xlarge",
        "g4dn.8xlarge",
        "g4dn.metal",
        "g5.12xlarge",
        "g5.16xlarge",
        "g5.24xlarge",
        "g5.48xlarge",
        "g5.8xlarge",
        "g6.12xlarge",
        "g6.16xlarge",
        "g6.24xlarge",
        "g6.48xlarge",
        "g6.8xlarge",
        "g6e.8xlarge",
        "g6e.16xlarge",
        "g6e.12xlarge",
        "g6e.24xlarge",
        "g6e.48xlarge",
        "gr6.8xlarge",
        "i3en.12xlarge",
        "i3en.24xlarge",
        "i3en.metal",
        "i4g.16xlarge",
        "i4i.32xlarge",
        "i4i.metal",
        "im4gn.16xlarge",
        "inf1.24xlarge",
        "m5dn.24xlarge",
        "m5dn.metal",
        "m5n.24xlarge",
        "m5n.metal",
        "m5zn.12xlarge",
        "m5zn.metal",
        "m6a.48xlarge",
        "m6a.metal",
        "m6i.32xlarge",
        "m6i.metal",
        "m6id.32xlarge",
        "m6id.metal",
        "m6idn.32xlarge",
        "m6idn.metal",
        "m6in.32xlarge",
        "m6in.metal",
        "m7a.48xlarge",
        "m7a.metal-48xl",
        "m7g.16xlarge",
        "m7g.metal",
        "m7gd.16xlarge",
        "m7gd.metal",
        "m7i.48xlarge",
        "m7i.metal-48xl",
        "p3dn.24xlarge",
        "p4d.24xlarge",
        "p4de.24xlarge",
        "p5.48xlarge",
        "p5e.48xlarge",
        "p5en.48xlarge",
        "r5dn.24xlarge",
        "r5dn.metal",
        "r5n.24xlarge",
        "r5n.metal",
        "r6a.48xlarge",
        "r6a.metal",
        "r6i.32xlarge",
        "r6i.metal",
        "r6id.32xlarge",
        "r6id.metal",
        "r6idn.32xlarge",
        "r6idn.metal",
        "r6in.32xlarge",
        "r6in.metal",
        "r7a.48xlarge",
        "r7a.metal-48xl",
        "r7g.16xlarge",
        "r7g.metal",
        "r7gd.16xlarge",
        "r7gd.metal",
        "r7i.48xlarge",
        "r7i.metal-48xl",
        "r7iz.32xlarge",
        "r7iz.metal-32xl",
        "r8g.24xlarge",
        "r8g.48xlarge",
        "r8g.metal-24xl",
        "r8g.metal-48xl",
        "trn1.32xlarge",
        "trn1n.32xlarge",
        "trn2.48xlarge",
        "u7i-12tb.224xlarge",
        "u7in-16tb.224xlarge",
        "u7in-24tb.224xlarge",
        "u7in-32tb.224xlarge",
        "vt1.24xlarge",
        "x2idn.32xlarge",
        "x2idn.metal",
        "x2iedn.32xlarge",
        "x2iedn.metal",
        "x2iezn.12xlarge",
        "x2iezn.metal",
    ]
)

instanceWithMultipleEFAs = {
    "p4d.24xlarge": 4,
    "p4de.24xlarge": 4,
    "p5.4xlarge": 4,
    "p5.24xlarge": 16,
    "p5.48xlarge": 32,
    "p5e.48xlarge": 32,
    "p5en.48xlarge": 16,
    "trn1.32xlarge": 8,
    "trn1n.32xlarge": 16,
    "trn2.48xlarge": 16,
    "g6e.24xlarge": 2,
    "g6e.48xlarge": 4,
}

instanceWithRDMASupport = set(
    [
        "p4d.24xlarge",
        "p4de.24xlarge",
        "p5.4xlarge",
        "p5.24xlarge",
        "p5.48xlarge",
        "p5e.48xlarge",
        "p5en.48xlarge",
        "trn1.32xlarge",
        "trn1n.32xlarge",
        "trn2.48xlarge",
    ]
)