Coverage for cli/nodepools.py: 89.47%
171 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""
2NodePool management utilities for GCO CLI.
4Provides functionality to create and manage Karpenter NodePools with
5support for On-Demand Capacity Reservations (ODCRs) and Capacity Blocks.
7Key Features:
8- Generate NodePool manifests for ODCR-backed capacity
9- List and describe NodePools in EKS clusters
10- Support for fallback to on-demand when ODCR is exhausted
12See: https://karpenter.sh/docs/tasks/odcrs/
13"""
15import base64
16import logging
17from dataclasses import dataclass
18from typing import Any
20import boto3
21import yaml
22from kubernetes.client import CustomObjectsApi
24logger = logging.getLogger(__name__)
26# Default vCPU count when instance type lookup fails (conservative estimate)
27DEFAULT_VCPUS_PER_NODE = 96
30def get_vcpus_for_instance_type(instance_type: str, region: str = "us-east-1") -> int:
31 """
32 Get the vCPU count for an instance type from EC2 API.
34 Args:
35 instance_type: EC2 instance type (e.g., "p4d.24xlarge")
36 region: AWS region for API calls
38 Returns:
39 Number of vCPUs for the instance type, or DEFAULT_VCPUS_PER_NODE if lookup fails
40 """
41 try:
42 ec2 = boto3.client("ec2", region_name=region)
43 response = ec2.describe_instance_types(InstanceTypes=[instance_type])
44 if response["InstanceTypes"]:
45 return int(response["InstanceTypes"][0]["VCpuInfo"]["DefaultVCpus"])
46 except Exception as e:
47 logger.debug("Failed to get vCPU count for %s: %s", instance_type, e)
49 return DEFAULT_VCPUS_PER_NODE
52def calculate_cpu_limit(
53 instance_types: list[str] | None, max_nodes: int, region: str = "us-east-1"
54) -> int:
55 """
56 Calculate the CPU limit for a NodePool based on instance types.
58 If multiple instance types are specified, uses the maximum vCPU count
59 to ensure the limit can accommodate the largest instances.
61 Args:
62 instance_types: List of instance types (None means any)
63 max_nodes: Maximum number of nodes in the pool
64 region: AWS region for API calls
66 Returns:
67 Total CPU limit (max_nodes * max_vcpus_per_instance)
68 """
69 if not instance_types:
70 # No specific instance types - use conservative default
71 return max_nodes * DEFAULT_VCPUS_PER_NODE
73 # Get vCPU count for each instance type and use the maximum
74 vcpu_counts = [get_vcpus_for_instance_type(it, region) for it in instance_types]
75 max_vcpus = max(vcpu_counts) if vcpu_counts else DEFAULT_VCPUS_PER_NODE
77 return max_nodes * max_vcpus
80@dataclass
81class NodePoolInfo:
82 """Information about a Karpenter NodePool."""
84 name: str
85 capacity_type: str # "on-demand", "spot", "reserved"
86 instance_types: list[str]
87 max_nodes: int | None
88 status: str
89 node_count: int
90 capacity_reservation_id: str | None = None
93def generate_odcr_nodepool_manifest(
94 name: str,
95 region: str,
96 capacity_reservation_id: str,
97 instance_types: list[str] | None = None,
98 max_nodes: int = 100,
99 fallback_on_demand: bool = False,
100 efa: bool = False,
101 project_name: str = "gco",
102) -> str:
103 """
104 Generate a Karpenter NodePool manifest for ODCR-backed capacity.
106 Args:
107 name: Name for the NodePool
108 region: AWS region
109 capacity_reservation_id: EC2 Capacity Reservation ID (cr-xxx) or ODCR group ARN
110 instance_types: List of instance types (if None, uses ODCR's instance type)
111 max_nodes: Maximum number of nodes
112 fallback_on_demand: Whether to fall back to on-demand if ODCR exhausted
113 efa: Whether to enable EFA support (adds EFA taint and labels)
115 Returns:
116 YAML manifest string for the NodePool and EC2NodeClass
117 """
118 # Determine capacity types based on fallback setting
119 capacity_types = ["reserved", "on-demand"] if fallback_on_demand else ["reserved"]
121 # Build the EC2NodeClass with capacity reservation selector
122 ec2_node_class = {
123 "apiVersion": "karpenter.k8s.aws/v1",
124 "kind": "EC2NodeClass",
125 "metadata": {
126 "name": f"{name}-nodeclass",
127 "labels": {
128 "app.kubernetes.io/part-of": "gco",
129 "gco.io/nodepool": name,
130 },
131 },
132 "spec": {
133 "role": f"KarpenterNodeRole-{project_name}",
134 "subnetSelectorTerms": [
135 {"tags": {"karpenter.sh/discovery": f"{project_name}-{region}"}}
136 ],
137 "securityGroupSelectorTerms": [
138 {"tags": {"karpenter.sh/discovery": f"{project_name}-{region}"}}
139 ],
140 "capacityReservationSelectorTerms": [{"id": capacity_reservation_id}],
141 "tags": {
142 "Name": f"{project_name}-{name}",
143 "gco.io/nodepool": name,
144 "gco.io/capacity-reservation": capacity_reservation_id,
145 },
146 },
147 }
149 # Build requirements list
150 requirements: list[dict[str, Any]] = [
151 {
152 "key": "karpenter.sh/capacity-type",
153 "operator": "In",
154 "values": capacity_types,
155 },
156 {
157 "key": "kubernetes.io/arch",
158 "operator": "In",
159 "values": ["amd64"],
160 },
161 ]
163 # Build the NodePool
164 nodepool: dict[str, Any] = {
165 "apiVersion": "karpenter.sh/v1",
166 "kind": "NodePool",
167 "metadata": {
168 "name": name,
169 "labels": {
170 "app.kubernetes.io/part-of": "gco",
171 },
172 },
173 "spec": {
174 "template": {
175 "metadata": {
176 "labels": {
177 "workload-type": "reserved-capacity",
178 "project": "gco",
179 "gco.io/capacity-reservation": capacity_reservation_id,
180 },
181 },
182 "spec": {
183 "nodeClassRef": {
184 "group": "karpenter.k8s.aws",
185 "kind": "EC2NodeClass",
186 "name": f"{name}-nodeclass",
187 },
188 "requirements": requirements,
189 },
190 },
191 "limits": {
192 "cpu": str(calculate_cpu_limit(instance_types, max_nodes, region)),
193 },
194 "disruption": {
195 "consolidationPolicy": "WhenEmptyOrUnderutilized",
196 "consolidateAfter": "30s",
197 "budgets": [{"nodes": "10%"}],
198 },
199 },
200 }
202 # Add instance type requirements if specified
203 if instance_types:
204 requirements.append(
205 {
206 "key": "node.kubernetes.io/instance-type",
207 "operator": "In",
208 "values": instance_types,
209 }
210 )
212 # Add GPU taints for GPU instances
213 gpu_families = ["p3", "p4", "p5", "p6", "g4", "g5", "g6"]
214 if instance_types and any(
215 any(it.startswith(fam) for fam in gpu_families) for it in instance_types
216 ):
217 taints = [
218 {
219 "key": "nvidia.com/gpu",
220 "value": "true",
221 "effect": "NoSchedule",
222 }
223 ]
224 if efa:
225 taints.append(
226 {
227 "key": "vpc.amazonaws.com/efa",
228 "value": "true",
229 "effect": "NoSchedule",
230 }
231 )
232 nodepool["spec"]["template"]["spec"]["taints"] = taints
234 # Add EFA labels
235 if efa:
236 nodepool["spec"]["template"]["metadata"]["labels"]["efa"] = "true"
237 nodepool["spec"]["template"]["metadata"]["labels"]["workload-type"] = "gpu-efa"
238 # Use WhenEmpty consolidation for EFA workloads to avoid disrupting
239 # long-running distributed training jobs
240 nodepool["spec"]["disruption"] = {
241 "consolidationPolicy": "WhenEmpty",
242 "consolidateAfter": "300s",
243 "budgets": [{"nodes": "10%"}],
244 }
246 # Generate YAML output
247 output = []
248 output.append("# ODCR-backed NodePool for GCO")
249 output.append(f"# Capacity Reservation: {capacity_reservation_id}")
250 output.append(f"# Region: {region}")
251 if fallback_on_demand:
252 output.append("# Fallback: on-demand (when ODCR exhausted)")
253 output.append("#")
254 output.append("# Apply with: kubectl apply -f <this-file>.yaml")
255 output.append("# See: https://karpenter.sh/docs/tasks/odcrs/")
256 output.append("---")
257 output.append(yaml.dump(ec2_node_class, default_flow_style=False, sort_keys=False))
258 output.append("---")
259 output.append(yaml.dump(nodepool, default_flow_style=False, sort_keys=False))
261 return "\n".join(output)
264def generate_capacity_block_nodepool_manifest(
265 name: str,
266 region: str,
267 capacity_reservation_id: str,
268 instance_types: list[str] | None = None,
269 max_nodes: int = 100,
270 fallback_on_demand: bool = False,
271 efa: bool = False,
272 project_name: str = "gco",
273) -> str:
274 """
275 Generate a Karpenter NodePool manifest for Capacity Block-backed capacity.
277 The Capacity Block counterpart to :func:`generate_odcr_nodepool_manifest`.
278 Purchasing a Capacity Block ("gco capacity reserve") yields a normal EC2
279 Capacity Reservation id (``cr-...``), which Karpenter consumes through the
280 same ``capacityReservationSelectorTerms`` / ``reserved`` capacity type as an
281 ODCR. The difference is intent: a Capacity Block is prepaid for a fixed term,
282 so the generated NodePool defaults to holding that capacity (``WhenEmpty``
283 consolidation with a long delay) rather than consolidating aggressively — you
284 have already paid for the whole block and want it available for its duration.
286 Args:
287 name: Name for the NodePool.
288 region: AWS region.
289 capacity_reservation_id: Capacity Reservation ID (cr-xxx) of the purchased
290 Capacity Block (from ``gco capacity reserve`` / reservation-check).
291 instance_types: List of instance types (if None, uses the reservation's type).
292 max_nodes: Maximum number of nodes.
293 fallback_on_demand: Whether to fall back to on-demand when the block is
294 exhausted or after it expires. Off by default — a Capacity Block is
295 fixed-term guaranteed capacity, so silent on-demand fallback can
296 surprise-bill; opt in explicitly.
297 efa: Whether to enable EFA support (adds EFA taint and labels).
299 Returns:
300 YAML manifest string for the NodePool and EC2NodeClass.
301 """
302 capacity_types = ["reserved", "on-demand"] if fallback_on_demand else ["reserved"]
304 ec2_node_class = {
305 "apiVersion": "karpenter.k8s.aws/v1",
306 "kind": "EC2NodeClass",
307 "metadata": {
308 "name": f"{name}-nodeclass",
309 "labels": {
310 "app.kubernetes.io/part-of": "gco",
311 "gco.io/nodepool": name,
312 },
313 },
314 "spec": {
315 "role": f"KarpenterNodeRole-{project_name}",
316 "subnetSelectorTerms": [
317 {"tags": {"karpenter.sh/discovery": f"{project_name}-{region}"}}
318 ],
319 "securityGroupSelectorTerms": [
320 {"tags": {"karpenter.sh/discovery": f"{project_name}-{region}"}}
321 ],
322 "capacityReservationSelectorTerms": [{"id": capacity_reservation_id}],
323 "tags": {
324 "Name": f"{project_name}-{name}",
325 "gco.io/nodepool": name,
326 "gco.io/capacity-block": capacity_reservation_id,
327 },
328 },
329 }
331 requirements: list[dict[str, Any]] = [
332 {
333 "key": "karpenter.sh/capacity-type",
334 "operator": "In",
335 "values": capacity_types,
336 },
337 {
338 "key": "kubernetes.io/arch",
339 "operator": "In",
340 "values": ["amd64"],
341 },
342 ]
344 nodepool: dict[str, Any] = {
345 "apiVersion": "karpenter.sh/v1",
346 "kind": "NodePool",
347 "metadata": {
348 "name": name,
349 "labels": {
350 "app.kubernetes.io/part-of": "gco",
351 },
352 },
353 "spec": {
354 "template": {
355 "metadata": {
356 "labels": {
357 "workload-type": "capacity-block",
358 "project": "gco",
359 "gco.io/capacity-type": "capacity-block",
360 "gco.io/capacity-block": capacity_reservation_id,
361 },
362 },
363 "spec": {
364 "nodeClassRef": {
365 "group": "karpenter.k8s.aws",
366 "kind": "EC2NodeClass",
367 "name": f"{name}-nodeclass",
368 },
369 "requirements": requirements,
370 },
371 },
372 "limits": {
373 "cpu": str(calculate_cpu_limit(instance_types, max_nodes, region)),
374 },
375 # A Capacity Block is prepaid for a fixed term — hold the nodes rather
376 # than consolidating them away, so the paid capacity stays available
377 # for the whole block.
378 "disruption": {
379 "consolidationPolicy": "WhenEmpty",
380 "consolidateAfter": "600s",
381 "budgets": [{"nodes": "10%"}],
382 },
383 },
384 }
386 if instance_types:
387 requirements.append(
388 {
389 "key": "node.kubernetes.io/instance-type",
390 "operator": "In",
391 "values": instance_types,
392 }
393 )
395 gpu_families = ["p3", "p4", "p5", "p6", "g4", "g5", "g6"]
396 if instance_types and any(
397 any(it.startswith(fam) for fam in gpu_families) for it in instance_types
398 ):
399 taints = [
400 {
401 "key": "nvidia.com/gpu",
402 "value": "true",
403 "effect": "NoSchedule",
404 }
405 ]
406 if efa:
407 taints.append(
408 {
409 "key": "vpc.amazonaws.com/efa",
410 "value": "true",
411 "effect": "NoSchedule",
412 }
413 )
414 nodepool["spec"]["template"]["spec"]["taints"] = taints
416 if efa:
417 nodepool["spec"]["template"]["metadata"]["labels"]["efa"] = "true"
418 nodepool["spec"]["template"]["metadata"]["labels"]["workload-type"] = "gpu-efa"
420 output = []
421 output.append("# Capacity Block-backed NodePool for GCO")
422 output.append(f"# Capacity Reservation (from Capacity Block): {capacity_reservation_id}")
423 output.append(f"# Region: {region}")
424 if fallback_on_demand:
425 output.append("# Fallback: on-demand (when Capacity Block exhausted/expired)")
426 output.append("#")
427 output.append("# Apply with: kubectl apply -f <this-file>.yaml")
428 output.append(
429 "# See: https://karpenter.sh/docs/concepts/nodeclasses/#speccapacityreservationselectorterms"
430 )
431 output.append("---")
432 output.append(yaml.dump(ec2_node_class, default_flow_style=False, sort_keys=False))
433 output.append("---")
434 output.append(yaml.dump(nodepool, default_flow_style=False, sort_keys=False))
436 return "\n".join(output)
439def get_eks_token(cluster_name: str, region: str) -> str:
440 """Generate EKS authentication token using STS presigned URL."""
441 from botocore.signers import RequestSigner
443 session = boto3.Session()
444 sts_client = session.client("sts", region_name=region)
445 service_id = sts_client.meta.service_model.service_id
447 signer = RequestSigner(
448 service_id, region, "sts", "v4", session.get_credentials(), session.events
449 )
451 params = {
452 "method": "GET",
453 "url": f"https://sts.{region}.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15",
454 "body": {},
455 "headers": {"x-k8s-aws-id": cluster_name},
456 "context": {},
457 }
459 url = signer.generate_presigned_url(
460 params, region_name=region, expires_in=60, operation_name=""
461 )
463 token_b64 = base64.urlsafe_b64encode(url.encode("utf-8")).decode("utf-8").rstrip("=")
464 return f"k8s-aws-v1.{token_b64}"
467def get_k8s_client(cluster_name: str, region: str) -> CustomObjectsApi:
468 """Get configured Kubernetes client for EKS cluster."""
469 from kubernetes import client
471 eks = boto3.client("eks", region_name=region)
472 cluster_info = eks.describe_cluster(name=cluster_name)
473 cluster = cluster_info["cluster"]
475 configuration = client.Configuration()
476 configuration.host = cluster["endpoint"]
477 configuration.verify_ssl = True
479 # Decode CA certificate using secure tempfile method
480 ca_cert = base64.b64decode(cluster["certificateAuthority"]["data"])
481 import os
482 import tempfile
484 fd, ca_cert_path = tempfile.mkstemp(suffix=".crt")
485 try:
486 with os.fdopen(fd, "wb") as ca_file:
487 ca_file.write(ca_cert)
488 ca_file.flush()
489 configuration.ssl_ca_cert = ca_cert_path
490 except Exception:
491 os.close(fd)
492 raise
494 # Generate EKS token
495 eks_token = get_eks_token(cluster_name, region)
496 configuration.api_key = {"authorization": f"Bearer {eks_token}"}
498 # Create API client with the configuration explicitly
499 api_client = client.ApiClient(configuration)
500 return client.CustomObjectsApi(api_client)
503def list_cluster_nodepools(cluster_name: str, region: str) -> list[dict[str, Any]]:
504 """
505 List NodePools in an EKS cluster.
507 Args:
508 cluster_name: EKS cluster name
509 region: AWS region
511 Returns:
512 List of NodePool information dictionaries
513 """
514 try:
515 custom_api = get_k8s_client(cluster_name, region)
517 nodepools = custom_api.list_cluster_custom_object(
518 group="karpenter.sh",
519 version="v1",
520 plural="nodepools",
521 )
523 result = []
524 for np in nodepools.get("items", []):
525 spec = np.get("spec", {})
526 template = spec.get("template", {}).get("spec", {})
527 requirements = template.get("requirements", [])
529 # Extract capacity types
530 capacity_types = []
531 instance_types = []
532 for req in requirements:
533 if req.get("key") == "karpenter.sh/capacity-type":
534 capacity_types = req.get("values", [])
535 elif req.get("key") == "node.kubernetes.io/instance-type": 535 ↛ 532line 535 didn't jump to line 532 because the condition on line 535 was always true
536 instance_types = req.get("values", [])
538 # Get status
539 status = np.get("status", {})
540 conditions = status.get("conditions", [])
541 ready_condition: dict[str, Any] = next(
542 (c for c in conditions if c.get("type") == "Ready"), {}
543 )
545 result.append(
546 {
547 "name": np["metadata"]["name"],
548 "capacity_types": ", ".join(capacity_types) or "on-demand",
549 "instance_types": ", ".join(instance_types[:3])
550 + ("..." if len(instance_types) > 3 else "")
551 or "any",
552 "status": "Ready" if ready_condition.get("status") == "True" else "NotReady",
553 "limits": spec.get("limits", {}),
554 }
555 )
557 return result
559 except Exception as e:
560 raise RuntimeError(f"Failed to list NodePools: {e}") from e
563def describe_cluster_nodepool(
564 cluster_name: str, region: str, nodepool_name: str
565) -> dict[str, Any] | None:
566 """
567 Describe a specific NodePool in an EKS cluster.
569 Args:
570 cluster_name: EKS cluster name
571 region: AWS region
572 nodepool_name: Name of the NodePool
574 Returns:
575 NodePool details or None if not found
576 """
577 try:
578 custom_api = get_k8s_client(cluster_name, region)
580 nodepool = custom_api.get_cluster_custom_object(
581 group="karpenter.sh",
582 version="v1",
583 plural="nodepools",
584 name=nodepool_name,
585 )
587 if isinstance(nodepool, dict): 587 ↛ 589line 587 didn't jump to line 589 because the condition on line 587 was always true
588 return nodepool
589 return None
591 except Exception as e:
592 if "404" in str(e):
593 return None
594 raise RuntimeError(f"Failed to describe NodePool: {e}") from e
597def delete_cluster_nodepool(cluster_name: str, region: str, nodepool_name: str) -> dict[str, Any]:
598 """
599 Delete a NodePool (and its paired EC2NodeClass) from an EKS cluster.
601 Deletes the Karpenter NodePool ``nodepool_name`` and, when present, the
602 EC2NodeClass named ``<nodepool_name>-nodeclass`` that the GCO manifest
603 generators create alongside it. Karpenter drains and terminates any nodes
604 the NodePool provisioned once it is removed. A missing EC2NodeClass (custom
605 name, or already deleted) is not treated as an error.
607 Args:
608 cluster_name: EKS cluster name
609 region: AWS region
610 nodepool_name: Name of the NodePool to delete
612 Returns:
613 Dict describing what was deleted: ``{"nodepool": <name>,
614 "ec2nodeclass": <name-or-None>}``.
615 """
616 try:
617 custom_api = get_k8s_client(cluster_name, region)
619 custom_api.delete_cluster_custom_object(
620 group="karpenter.sh",
621 version="v1",
622 plural="nodepools",
623 name=nodepool_name,
624 )
625 deleted: dict[str, Any] = {"nodepool": nodepool_name, "ec2nodeclass": None}
627 nodeclass_name = f"{nodepool_name}-nodeclass"
628 try:
629 custom_api.delete_cluster_custom_object(
630 group="karpenter.k8s.aws",
631 version="v1",
632 plural="ec2nodeclasses",
633 name=nodeclass_name,
634 )
635 deleted["ec2nodeclass"] = nodeclass_name
636 except Exception as e: # noqa: BLE001 - best effort; the NodePool is the primary target
637 if "404" not in str(e):
638 logger.warning("Could not delete EC2NodeClass %s: %s", nodeclass_name, e)
640 return deleted
642 except Exception as e:
643 raise RuntimeError(f"Failed to delete NodePool: {e}") from e