feat: add general CloudProvider class

src/k8s_node_operator/npat.yaml → examples/npat.yaml

@@ -3,4 +3,4 @@ kind: NodepoolAllocationTarget
 metadata:
   name: example-npat
 spec:
-  minimumNodeCount: 0
+  minimumNodeCount: 1

helm/k8s_node_operator/crds/nodepool_allocation_target.yaml

@@ -27,18 +27,30 @@ spec:
               type: object
               x-kubernetes-preserve-unknown-fields: true
       additionalPrinterColumns:
+      - name: Nodepool
+        type: string
+        description: Nodepool name
+        jsonPath: .status.nodepool_allocation.name
+      - name: Status
+        type: string
+        description: Nodepool allocation target status
+        jsonPath: .status.nodepool_allocation.status
       - name: Node count
         type: integer
         description: Current node count
-        jsonPath: .status.create_npat.node_count
+        jsonPath: .status.nodepool_allocation.node_count
       - name: Min node count
         type: integer
         description: Minimum node count
-        jsonPath: .status.create_npat.min_node_count
+        jsonPath: .status.nodepool_allocation.min_node_count
+      - name: Target Min node count
+        type: integer
+        description: Target minimum node count
+        jsonPath: .status.nodepool_allocation.target_min_node_count
       - name: Max node count
         type: integer
         description: Maximum node count
-        jsonPath: .status.create_npat.max_node_count
+        jsonPath: .status.nodepool_allocation.max_node_count
       - name: Age
         type: date
         jsonPath: .metadata.creationTimestamp

src/k8s_node_operator/operator.py

@@ -1,107 +1,27 @@
-import google.auth
-import google.api_core
-from google.cloud import container_v1
 import kopf
-from kubernetes.aio import client, config
-from kubernetes.aio.client.api_client import ApiClient
-import logging
 import os
-from typing import Any
-
-logger = logging.getLogger(__name__)
-
-class GCPClient:
-    def __init__(self):
-        self.cluster_name = os.environ.get("GCP_CLUSTER", "")
-        self.machine_type = os.environ.get("GCP_MACHINE_TYPE", "")
-        self.nodepool = os.environ.get("GCP_NODEPOOL", "")
-        self.project_name = os.environ.get("GCP_PROJECT_ID", "")
-        self.zone = os.environ.get("GCP_ZONE", "") # TODO: add support for regional clusters
-        self.region = os.environ.get("GCP_REGION", "")
-        self.prefix =  f"projects/{self.project_name}/zones/{self.zone}" if self.zone else f"projects/{self.project_name}/region/{self.region}"
-        self.nodepool_name = self.prefix + f"/clusters/{self.cluster_name}/nodePools/{self.nodepool}"
-        self.credentials_file = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "")
-        self.credentials, self.project = google.auth.default()
-        self.client = None
-        logger.debug(self.credentials.get_cred_info())
-
-    async def __aenter__(self):
-        self.client = container_v1.ClusterManagerAsyncClient(
-            credentials=self.credentials
-        )
-        return self
-
-    async def __aexit__(self, exc_type, exc, tb):
-        if self.client:
-            await self.client.transport.close()
-
-    async def get_nodepool(self):
-        request = container_v1.GetNodePoolRequest(
-           name=self.nodepool_name
-        )
-        return await self.client.get_node_pool(request=request)
-
-    async def update_autoscaling_min_node_count(self, min_node_count: int, nodepool):
-        if min_node_count >= nodepool.autoscaling.max_node_count:
-            logger.error(f'Minimum node count {min_node_count} exceeds maximum node count.')
-            return # TODO: update Status
-        elif min_node_count != nodepool.autoscaling.min_node_count:
-            nodepool_autoscaling = container_v1.NodePoolAutoscaling(
-                enabled = nodepool.autoscaling.enabled,
-                min_node_count = min_node_count,
-                max_node_count = nodepool.autoscaling.max_node_count,
-                location_policy = nodepool.autoscaling.location_policy
-            )
-            request = container_v1.SetNodePoolAutoscalingRequest(
-                name=self.nodepool_name,
-                autoscaling=nodepool_autoscaling
-            )
-            logger.info(f'Nodepool Allocation Target requested.')
-            return await self.client.set_node_pool_autoscaling(request=request)
-        else:
-            logger.info(f'Minimum node count is already set to {min_node_count}.')
-            return # TODO: update status
-
-    async def wait_gcp_operation(self, operation_name: str):
-        """
-        Blocking call to wait until operation is completed.
-        """
-        name = '/'.join([self.prefix, "operations", operation_name])
-        logger.debug(f'{name=}')
-        request = container_v1.GetOperationRequest(name=name)
-        while True:
-            response = await self.client.get_operation(request=request)
-            if response.status != container_v1.Operation.Status.DONE:
-                logger.info(f'Operation is {container_v1.Operation.Status(response.status).name}')
-            # TODO: backoff on error
-            else:
-                return response
+from typing import Any, Dict
+from k8s_node_operator.providers import create_provider
 
 @kopf.on.create('nodepoolallocationtarget')
-async def create_npat(spec: kopf.Spec, name: str, namespace: str | None, logger: kopf.Logger, **_: Any) -> None:
+@kopf.on.update('nodepoolallocationtarget')
+async def nodepool_allocation(spec: kopf.Spec, name: str, namespace: str | None, logger: kopf.Logger, **_: Any) -> Dict:
     # Parse npat spec
-    min_node_count = spec.get('minimumNodeCount')
-    if not min_node_count:
-        min_node_count = 0
+    target_min_node_count = spec.get('minimumNodeCount')
+    if not target_min_node_count:
+        target_min_node_count = 0
     # Send nodepool scaling request to cloud provider
-    async with GCPClient() as gke:
-        nodepool = await gke.get_nodepool()
-        operation = await gke.update_autoscaling_min_node_count(min_node_count, nodepool)
-        # Block until scaling operation is completed
-        if operation:
-            response = await gke.wait_gcp_operation(operation.name)
-            logger.debug(f'{response.progress=}')
-        # Get updated nodepool
-        nodepool = await gke.get_nodepool()
-    # Block on current node count with k8s api until minimum nodepool count is reached
-    await config.load_kube_config()
-    node_count = 0
-    while node_count < min_node_count:
-        async with ApiClient() as api:
-            v1 = client.CoreV1Api(api)
-            node_list = await v1.list_node()
-        node_count = len(node_list.items)
-        logger.info(f'Node count = {node_count}.')
-    return {'status': 'SUCCESS', 'node_count': node_count, 'min_node_count': nodepool.autoscaling.min_node_count, 'max_node_count': nodepool.autoscaling.max_node_count} # type: ignore
-
-# TODO: we want to update/patch the npat over time, so change create_fn to handle first instantiation of npat, and then convert current fn logic to an update_fn to respond to @kopf.on.patch/update
+    provider_name = os.environ.get("K8S_NODE_OPERATOR_CLOUD_PROVIDER")
+    async with create_provider(name=provider_name, logger=logger) as provider:
+        nodepool = await provider.set_min_node_count(target_min_node_count=target_min_node_count)
+    # Store output in k8s npat object
+    return {'status': nodepool.status.name, 'name': nodepool.name, 'node_count': nodepool.current_node_count, 'min_node_count': nodepool.min_node_count, 'max_node_count': nodepool.max_node_count, 'target_min_node_count': nodepool.target_min_node_count}
+
+@kopf.on.delete('nodepoolallocationtarget')
+async def delete_nodepool_allocation(logger: kopf.Logger, **_: Any) -> None:
+    # Set minimum node count to zero
+    target_min_node_count = 0
+    provider_name = os.environ.get("K8S_NODE_OPERATOR_CLOUD_PROVIDER")
+    async with create_provider(name=provider_name, logger=logger) as provider:
+        nodepool = await provider.set_min_node_count(target_min_node_count=target_min_node_count)
+        logger.info(f'npat deleted: "{nodepool.name}" minimum node count set to {nodepool.min_node_count}.')

src/k8s_node_operator/providers.py

@@ -0,0 +1,177 @@
+import google.auth
+import google.api_core
+import kopf
+import os
+from dataclasses import dataclass
+from enum import Enum
+from google.cloud import container_v1
+from kubernetes.aio import client, config
+from kubernetes.aio.client.api_client import ApiClient
+from typing import Protocol
+
+class NodepoolStatus(Enum):
+    READY = 1
+    UPDATING = 2
+    ERROR = 3
+
+
+@dataclass
+class Nodepool:
+    status: NodepoolStatus
+    name: str
+    min_node_count: int
+    max_node_count: int
+    current_node_count: int
+    target_min_node_count: int
+
+class CloudProvider(Protocol):
+    """
+    Abstract class that allows structural subtyping/duck typing of all cloud providers (see https://typing.python.org/en/latest/reference/protocols.html).
+    """
+    async def get_nodepool(self):
+        ... # Note that `...` is a Python placeholder object
+
+    async def set_min_node_count(self, min_node_count: int, nodepool: Nodepool):
+        ...
+
+    async def get_k8s_current_node_count(self):
+        """
+        Get current node count with Kubernetes API. We use this as the source of truth for the number of nodes online, rather than cloud provider specific APIs.
+        """
+        await config.load_kube_config()
+        async with ApiClient() as api:
+            v1 = client.CoreV1Api(api)
+            node_list = await v1.list_node()
+        node_count = len(node_list.items)
+        return node_count
+
+class GCPProvider(CloudProvider):
+    """
+    Methods for Google Cloud Platform (GCP).
+    """
+    def __init__(self, logger: kopf.Logger):
+        self.cluster_name = os.environ.get("GCP_CLUSTER", "")
+        self.machine_type = os.environ.get("GCP_MACHINE_TYPE", "")
+        self.nodepool = os.environ.get("GCP_NODEPOOL", "") # TODO: get this from the npat spec
+        self.project_name = os.environ.get("GCP_PROJECT_ID", "")
+        self.zone = os.environ.get("GCP_ZONE", "") # TODO: add support for regional clusters
+        self.region = os.environ.get("GCP_REGION", "")
+        self.prefix =  f"projects/{self.project_name}/zones/{self.zone}" if self.zone else f"projects/{self.project_name}/region/{self.region}"
+        self.nodepool_name = self.prefix + f"/clusters/{self.cluster_name}/nodePools/{self.nodepool}"
+        self.credentials_file = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "")
+        self.credentials, self.project = google.auth.default()
+        self.client = None
+        self.log = logger
+        self.log.debug(self.credentials.get_cred_info())
+
+    async def __aenter__(self):
+        self.client = container_v1.ClusterManagerAsyncClient(
+            credentials=self.credentials
+        )
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        if self.client:
+            await self.client.transport.close()
+
+    async def _get_gcp_nodepool(self):
+        request = container_v1.GetNodePoolRequest(
+           name=self.nodepool_name
+        )
+        response = await self.client.get_node_pool(request=request)
+        return response
+
+    async def get_nodepool(self, target_min_node_count: int, gcp_nodepool: container_v1.NodePool | None = None):
+        if not gcp_nodepool:
+            gcp_nodepool = await self._get_gcp_nodepool()
+        current_node_count = await self.get_k8s_current_node_count()
+        nodepool = Nodepool(status=NodepoolStatus.READY, name=self.nodepool, min_node_count=gcp_nodepool.autoscaling.min_node_count, max_node_count=gcp_nodepool.autoscaling.max_node_count, current_node_count=current_node_count,
+        target_min_node_count=target_min_node_count)
+        return nodepool
+
+    async def wait_gcp_operation(self, operation_name: str):
+        """
+        Blocking call to wait until operation is completed.
+        """
+        name = '/'.join([self.prefix, "operations", operation_name])
+        self.log.debug(f'Operation name: {name}')
+        request = container_v1.GetOperationRequest(name=name)
+        while True:
+            response = await self.client.get_operation(request=request)
+            if response.status != container_v1.Operation.Status.DONE:
+                self.log.debug(f'Operation is {container_v1.Operation.Status(response.status).name}')
+            # TODO: backoff on error
+            else:
+                return response
+
+    async def set_min_node_count(self, target_min_node_count: int):
+        gcp_nodepool = await self._get_gcp_nodepool()
+        if target_min_node_count >= gcp_nodepool.autoscaling.max_node_count:
+            self.log.warning(f'Target minimum node count {target_min_node_count} exceeds maximum node count.')
+        elif target_min_node_count != gcp_nodepool.autoscaling.min_node_count:
+            gcp_nodepool_autoscaling = container_v1.NodePoolAutoscaling(
+                enabled = gcp_nodepool.autoscaling.enabled,
+                min_node_count = target_min_node_count,
+                max_node_count = gcp_nodepool.autoscaling.max_node_count,
+                location_policy = gcp_nodepool.autoscaling.location_policy
+            )
+            request = container_v1.SetNodePoolAutoscalingRequest(
+                name=self.nodepool_name,
+                autoscaling=gcp_nodepool_autoscaling
+            )
+            operation = await self.client.set_node_pool_autoscaling(request=request)
+            # Block until scaling operation is completed
+            if operation:
+                await self.wait_gcp_operation(operation_name = operation.name)
+            # Update with new nodepool config
+            gcp_nodepool = await self._get_gcp_nodepool()
+            self.log.info(f'Minimum node count set to {gcp_nodepool.autoscaling.min_node_count}.')
+        else:
+            self.log.warning(f'Minimum node count is already set to {target_min_node_count}.')
+        nodepool = await self.get_nodepool(gcp_nodepool=gcp_nodepool, target_min_node_count=target_min_node_count)
+        self.log.debug(f'{nodepool=}')
+        return nodepool
+
+class TestProvider(CloudProvider):
+    """
+    No-op cloud provider for testing and mocking.
+    """
+    def __init__(self, logger: kopf.Logger):
+        self._entered = False
+        self._exited = False
+        self.nodepool = "test-nodepool"
+        self.log = logger
+
+    async def __aenter__(self):
+        self._entered = True
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        self._exited = True
+
+    async def get_nodepool(self, target_min_node_count: int):
+        self.nodepool = Nodepool(
+            name="test-pool",
+            min_node_count=0, 
+            max_node_count=0,
+            current_node_count=0,
+            target_min_node_count=target_min_node_count, # target_min_node_count is the only variable we are testing
+        )
+        return self.nodepool
+
+    async def set_min_node_count(self, target_min_node_count: int):
+        self.nodepool = Nodepool(
+            name="test-nodepool",
+            min_node_count=0,
+            max_node_count=0,
+            current_node_count=0,
+            target_min_node_count=target_min_node_count, # target_min_node_count is the only variable we are testing
+        )
+        return self.nodepool
+
+
+def create_provider(name: str, logger: kopf.Logger):
+    if name == "GCP":
+        return GCPProvider(logger=logger)
+    elif name == "TEST":
+        return TestProvider(logger=logger)