From 78e14570f1c0e49b36ece96bc5c5fee5f259e216 Mon Sep 17 00:00:00 2001
From: Albin Vass <albinvass@gmail.com>
Date: Wed, 27 Mar 2024 10:49:05 +0100
Subject: [PATCH] Fix issue with unknown instance family

This fixes an issue where the aws driver errors when it encounters an
unknown instances family while listing instances, making it unable to
launch new instances.

Change-Id: I5c0a6eaeebe6038806149a9c5592899db7406574
Co-Authored-By: James E. Blair <jim@acmegating.com>
---
 nodepool/driver/aws/adapter.py             | 11 +++++++----
 nodepool/tests/fixtures/aws/aws-quota.yaml |  5 +++++
 nodepool/tests/unit/test_driver_aws.py     | 16 ++++++++++++++++
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/nodepool/driver/aws/adapter.py b/nodepool/driver/aws/adapter.py
index 42b602640..b015fef95 100644
--- a/nodepool/driver/aws/adapter.py
+++ b/nodepool/driver/aws/adapter.py
@@ -459,9 +459,6 @@ class AwsAdapter(statemachine.Adapter):
                 if code in args:
                     continue
                 if not code:
-                    self.log.warning(
-                        "Unknown quota code for instance type: %s",
-                        instance_type)
                     continue
                 if code not in ec2_quotas:
                     self.log.warning(
@@ -929,7 +926,13 @@ class AwsAdapter(statemachine.Adapter):
         m = self.instance_key_re.match(instance_type)
         if m:
             key = m.group(1)
-            return QUOTA_CODES.get(key)[market_type_option]
+            code = QUOTA_CODES.get(key)
+            if code:
+                return code[market_type_option]
+            self.log.warning(
+                "Unknown quota code for instance type: %s",
+                instance_type)
+        return None
 
     def _getQuotaForInstanceType(self, instance_type, market_type_option):
         try:
diff --git a/nodepool/tests/fixtures/aws/aws-quota.yaml b/nodepool/tests/fixtures/aws/aws-quota.yaml
index bc040094e..10b202ddf 100644
--- a/nodepool/tests/fixtures/aws/aws-quota.yaml
+++ b/nodepool/tests/fixtures/aws/aws-quota.yaml
@@ -17,6 +17,7 @@ labels:
   - name: high
   - name: spot
   - name: on-demand
+  - name: unknown
 
 providers:
   - name: ec2-us-west-2
@@ -52,3 +53,7 @@ providers:
             cloud-image: ubuntu1404
             instance-type: m6i.32xlarge
             key-name: zuul
+          - name: unknown
+            cloud-image: ubuntu1404
+            instance-type: is4gen.medium
+            key-name: zuul
diff --git a/nodepool/tests/unit/test_driver_aws.py b/nodepool/tests/unit/test_driver_aws.py
index 41c4d5878..52d0f9772 100644
--- a/nodepool/tests/unit/test_driver_aws.py
+++ b/nodepool/tests/unit/test_driver_aws.py
@@ -398,6 +398,22 @@ class TestDriverAws(tests.DBTestCase):
         req3 = self.waitForNodeRequest(req3)
         self.assertSuccess(req3)
 
+    def test_aws_multi_quota_unknown(self):
+        # Test multiple instance type quotas (standard, high-mem and spot)
+        configfile = self.setup_config('aws/aws-quota.yaml')
+        pool = self.useNodepool(configfile, watermark_sleep=1)
+        pool.start()
+
+        # We don't have quota information for this node type; make
+        # sure we can still launch a node with it.
+        req1 = zk.NodeRequest()
+        req1.state = zk.REQUESTED
+        req1.node_types.append('unknown')
+        self.zk.storeNodeRequest(req1)
+        self.log.debug("Waiting for request %s", req1.id)
+        req1 = self.waitForNodeRequest(req1)
+        self.assertSuccess(req1)
+
     @ec2_quotas({
         'L-1216C47A': 1000,
         'L-43DA4232': 1000,