From 32373d2ac594cb1cec99dc60450cb1c2ed4e3e19 Mon Sep 17 00:00:00 2001
From: Edison Gustavo Muenz <edisongustavo@gmail.com>
Date: Mon, 7 Feb 2022 16:13:50 +0100
Subject: [PATCH] ec2: Add GpuInfo to describe_instance_types (#4828)

---
 moto/ec2/responses/instances.py       | 17 +++++++++++
 tests/test_ec2/test_instance_types.py | 43 +++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/moto/ec2/responses/instances.py b/moto/ec2/responses/instances.py
index 93bf3761d..475c14422 100644
--- a/moto/ec2/responses/instances.py
+++ b/moto/ec2/responses/instances.py
@@ -871,6 +871,23 @@ EC2_DESCRIBE_INSTANCE_TYPES = """<?xml version="1.0" encoding="UTF-8"?>
                     {% endfor %}
                 </supportedArchitectures>
             </processorInfo>
+            {% if instance_type.get('GpuInfo', {})|length > 0 %}
+            <gpuInfo>
+                <gpus>
+                    {% for gpu in instance_type.get('GpuInfo').get('Gpus') %}
+                    <item>
+                        <count>{{ gpu['Count']|int }}</count>
+                        <manufacturer>{{ gpu['Manufacturer'] }}</manufacturer>
+                        <memoryInfo>
+                            <sizeInMiB>{{ gpu['MemoryInfo']['SizeInMiB']|int }}</sizeInMiB>
+                        </memoryInfo>
+                        <name>{{ gpu['Name'] }}</name>
+                    </item>
+                    {% endfor %}
+                </gpus>
+                <totalGpuMemoryInMiB>{{ instance_type['GpuInfo']['TotalGpuMemoryInMiB']|int }}</totalGpuMemoryInMiB>
+            </gpuInfo>
+            {% endif %}
         </item>
     {% endfor %}
     </instanceTypeSet>
diff --git a/tests/test_ec2/test_instance_types.py b/tests/test_ec2/test_instance_types.py
index 57b6897e5..5549a5b16 100644
--- a/tests/test_ec2/test_instance_types.py
+++ b/tests/test_ec2/test_instance_types.py
@@ -37,6 +37,49 @@ def test_describe_instance_types_filter_by_type():
     )
 
 
+@mock_ec2
+def test_describe_instance_types_gpu_instance_types():
+    client = boto3.client("ec2", "us-east-1")
+    instance_types = client.describe_instance_types(
+        InstanceTypes=["p3.2xlarge", "g4ad.8xlarge"]
+    )
+
+    instance_types.should.have.key("InstanceTypes")
+    instance_types["InstanceTypes"].should_not.be.empty
+    instance_types["InstanceTypes"].should.have.length_of(2)
+    instance_types["InstanceTypes"][0]["GpuInfo"].should_not.be.empty
+    instance_types["InstanceTypes"][1]["GpuInfo"].should_not.be.empty
+
+    instance_type_to_gpu_info = {
+        instance_info["InstanceType"]: instance_info["GpuInfo"]
+        for instance_info in instance_types["InstanceTypes"]
+    }
+    assert instance_type_to_gpu_info == {
+        "g4ad.8xlarge": {
+            "Gpus": [
+                {
+                    "Count": 2,
+                    "Manufacturer": "AMD",
+                    "MemoryInfo": {"SizeInMiB": 8192},
+                    "Name": "Radeon Pro V520",
+                }
+            ],
+            "TotalGpuMemoryInMiB": 16384,
+        },
+        "p3.2xlarge": {
+            "Gpus": [
+                {
+                    "Count": 1,
+                    "Manufacturer": "NVIDIA",
+                    "MemoryInfo": {"SizeInMiB": 16384},
+                    "Name": "V100",
+                }
+            ],
+            "TotalGpuMemoryInMiB": 16384,
+        },
+    }
+
+
 @mock_ec2
 def test_describe_instance_types_unknown_type():
     client = boto3.client("ec2", "us-east-1")