Coverage for cli / nodepools.py: 100.00%
192 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +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
17import os
18import tempfile
19import weakref
20from contextlib import suppress
21from dataclasses import dataclass
22from typing import Any
24import boto3
25import yaml
26from kubernetes.client import CustomObjectsApi
28logger = logging.getLogger(__name__)
30# Default vCPU count when instance type lookup fails (conservative estimate)
31DEFAULT_VCPUS_PER_NODE = 96
34def _unlink_temp_ca_cert(path: str) -> None:
35 """Best-effort cleanup for a CA file owned by a Kubernetes API client."""
36 try:
37 os.unlink(path)
38 except FileNotFoundError:
39 # Already gone (e.g. removed by another cleanup path) — the goal is an
40 # absent file, so this is success, not a failure to log or retry.
41 pass
42 except OSError as exc:
43 logger.debug("Failed to remove temporary Kubernetes CA certificate %s: %s", path, exc)
46def get_vcpus_for_instance_type(instance_type: str, region: str = "us-east-1") -> int:
47 """
48 Get the vCPU count for an instance type from EC2 API.
50 Args:
51 instance_type: EC2 instance type (e.g., "p4d.24xlarge")
52 region: AWS region for API calls
54 Returns:
55 Number of vCPUs for the instance type, or DEFAULT_VCPUS_PER_NODE if lookup fails
56 """
57 try:
58 ec2 = boto3.client("ec2", region_name=region)
59 response = ec2.describe_instance_types(InstanceTypes=[instance_type])
60 if response["InstanceTypes"]:
61 return int(response["InstanceTypes"][0]["VCpuInfo"]["DefaultVCpus"])
62 except Exception as e:
63 logger.debug("Failed to get vCPU count for %s: %s", instance_type, e)
65 return DEFAULT_VCPUS_PER_NODE
68def calculate_cpu_limit(
69 instance_types: list[str] | None, max_nodes: int, region: str = "us-east-1"
70) -> int:
71 """
72 Calculate the CPU limit for a NodePool based on instance types.
74 If multiple instance types are specified, uses the maximum vCPU count
75 to ensure the limit can accommodate the largest instances.
77 Args:
78 instance_types: List of instance types (None means any)
79 max_nodes: Maximum number of nodes in the pool
80 region: AWS region for API calls
82 Returns:
83 Total CPU limit (max_nodes * max_vcpus_per_instance)
84 """
85 if not instance_types:
86 # No specific instance types - use conservative default
87 return max_nodes * DEFAULT_VCPUS_PER_NODE
89 # Get vCPU count for each instance type and use the maximum
90 vcpu_counts = [get_vcpus_for_instance_type(it, region) for it in instance_types]
91 max_vcpus = max(vcpu_counts) if vcpu_counts else DEFAULT_VCPUS_PER_NODE
93 return max_nodes * max_vcpus
96@dataclass
97class NodePoolInfo:
98 """Information about a Karpenter NodePool."""
100 name: str
101 capacity_type: str # "on-demand", "spot", "reserved"
102 instance_types: list[str]
103 max_nodes: int | None
104 status: str
105 node_count: int
106 capacity_reservation_id: str | None = None
109def generate_odcr_nodepool_manifest(
110 name: str,
111 region: str,
112 capacity_reservation_id: str,
113 instance_types: list[str] | None = None,
114 max_nodes: int = 100,
115 fallback_on_demand: bool = False,
116 efa: bool = False,
117 project_name: str = "gco",
118) -> str:
119 """
120 Generate a Karpenter NodePool manifest for ODCR-backed capacity.
122 Args:
123 name: Name for the NodePool
124 region: AWS region
125 capacity_reservation_id: EC2 Capacity Reservation ID (cr-xxx) or ODCR group ARN
126 instance_types: List of instance types (if None, uses ODCR's instance type)
127 max_nodes: Maximum number of nodes
128 fallback_on_demand: Whether to fall back to on-demand if ODCR exhausted
129 efa: Whether to enable EFA support (adds EFA taint and labels)
131 Returns:
132 YAML manifest string for the NodePool and EC2NodeClass
133 """
134 # Determine capacity types based on fallback setting
135 capacity_types = ["reserved", "on-demand"] if fallback_on_demand else ["reserved"]
137 # Build the EC2NodeClass with capacity reservation selector
138 ec2_node_class = {
139 "apiVersion": "karpenter.k8s.aws/v1",
140 "kind": "EC2NodeClass",
141 "metadata": {
142 "name": f"{name}-nodeclass",
143 "labels": {
144 "app.kubernetes.io/part-of": "gco",
145 "gco.io/nodepool": name,
146 },
147 },
148 "spec": {
149 "role": f"KarpenterNodeRole-{project_name}",
150 "subnetSelectorTerms": [
151 {"tags": {"karpenter.sh/discovery": f"{project_name}-{region}"}}
152 ],
153 "securityGroupSelectorTerms": [
154 {"tags": {"karpenter.sh/discovery": f"{project_name}-{region}"}}
155 ],
156 "capacityReservationSelectorTerms": [{"id": capacity_reservation_id}],
157 "tags": {
158 "Name": f"{project_name}-{name}",
159 "gco.io/nodepool": name,
160 "gco.io/capacity-reservation": capacity_reservation_id,
161 },
162 },
163 }
165 # Build requirements list
166 requirements: list[dict[str, Any]] = [
167 {
168 "key": "karpenter.sh/capacity-type",
169 "operator": "In",
170 "values": capacity_types,
171 },
172 {
173 "key": "kubernetes.io/arch",
174 "operator": "In",
175 "values": ["amd64"],
176 },
177 ]
179 # Build the NodePool
180 nodepool: dict[str, Any] = {
181 "apiVersion": "karpenter.sh/v1",
182 "kind": "NodePool",
183 "metadata": {
184 "name": name,
185 "labels": {
186 "app.kubernetes.io/part-of": "gco",
187 },
188 },
189 "spec": {
190 "template": {
191 "metadata": {
192 "labels": {
193 "workload-type": "reserved-capacity",
194 "project": "gco",
195 "gco.io/capacity-reservation": capacity_reservation_id,
196 },
197 },
198 "spec": {
199 "nodeClassRef": {
200 "group": "karpenter.k8s.aws",
201 "kind": "EC2NodeClass",
202 "name": f"{name}-nodeclass",
203 },
204 "requirements": requirements,
205 },
206 },
207 "limits": {
208 "cpu": str(calculate_cpu_limit(instance_types, max_nodes, region)),
209 },
210 "disruption": {
211 "consolidationPolicy": "WhenEmptyOrUnderutilized",
212 "consolidateAfter": "30s",
213 "budgets": [{"nodes": "10%"}],
214 },
215 },
216 }
218 # Add instance type requirements if specified
219 if instance_types:
220 requirements.append(
221 {
222 "key": "node.kubernetes.io/instance-type",
223 "operator": "In",
224 "values": instance_types,
225 }
226 )
228 # Add GPU taints for GPU instances
229 gpu_families = ["p3", "p4", "p5", "p6", "g4", "g5", "g6"]
230 if instance_types and any(
231 any(it.startswith(fam) for fam in gpu_families) for it in instance_types
232 ):
233 taints = [
234 {
235 "key": "nvidia.com/gpu",
236 "value": "true",
237 "effect": "NoSchedule",
238 }
239 ]
240 if efa:
241 taints.append(
242 {
243 "key": "vpc.amazonaws.com/efa",
244 "value": "true",
245 "effect": "NoSchedule",
246 }
247 )
248 nodepool["spec"]["template"]["spec"]["taints"] = taints
250 # Add EFA labels
251 if efa:
252 nodepool["spec"]["template"]["metadata"]["labels"]["efa"] = "true"
253 nodepool["spec"]["template"]["metadata"]["labels"]["workload-type"] = "gpu-efa"
254 # Use WhenEmpty consolidation for EFA workloads to avoid disrupting
255 # long-running distributed training jobs
256 nodepool["spec"]["disruption"] = {
257 "consolidationPolicy": "WhenEmpty",
258 "consolidateAfter": "300s",
259 "budgets": [{"nodes": "10%"}],
260 }
262 # Generate YAML output
263 output = []
264 output.append("# ODCR-backed NodePool for GCO")
265 output.append(f"# Capacity Reservation: {capacity_reservation_id}")
266 output.append(f"# Region: {region}")
267 if fallback_on_demand:
268 output.append("# Fallback: on-demand (when ODCR exhausted)")
269 output.append("#")
270 output.append("# Apply with: kubectl apply -f <this-file>.yaml")
271 output.append("# See: https://karpenter.sh/docs/tasks/odcrs/")
272 output.append("---")
273 output.append(yaml.dump(ec2_node_class, default_flow_style=False, sort_keys=False))
274 output.append("---")
275 output.append(yaml.dump(nodepool, default_flow_style=False, sort_keys=False))
277 return "\n".join(output)
280def generate_capacity_block_nodepool_manifest(
281 name: str,
282 region: str,
283 capacity_reservation_id: str,
284 instance_types: list[str] | None = None,
285 max_nodes: int = 100,
286 fallback_on_demand: bool = False,
287 efa: bool = False,
288 project_name: str = "gco",
289) -> str:
290 """
291 Generate a Karpenter NodePool manifest for Capacity Block-backed capacity.
293 The Capacity Block counterpart to :func:`generate_odcr_nodepool_manifest`.
294 Purchasing a Capacity Block ("gco capacity reserve") yields a normal EC2
295 Capacity Reservation id (``cr-...``), which Karpenter consumes through the
296 same ``capacityReservationSelectorTerms`` / ``reserved`` capacity type as an
297 ODCR. The difference is intent: a Capacity Block is prepaid for a fixed term,
298 so the generated NodePool defaults to holding that capacity (``WhenEmpty``
299 consolidation with a long delay) rather than consolidating aggressively — you
300 have already paid for the whole block and want it available for its duration.
302 Args:
303 name: Name for the NodePool.
304 region: AWS region.
305 capacity_reservation_id: Capacity Reservation ID (cr-xxx) of the purchased
306 Capacity Block (from ``gco capacity reserve`` / reservation-check).
307 instance_types: List of instance types (if None, uses the reservation's type).
308 max_nodes: Maximum number of nodes.
309 fallback_on_demand: Whether to fall back to on-demand when the block is
310 exhausted or after it expires. Off by default — a Capacity Block is
311 fixed-term guaranteed capacity, so silent on-demand fallback can
312 surprise-bill; opt in explicitly.
313 efa: Whether to enable EFA support (adds EFA taint and labels).
315 Returns:
316 YAML manifest string for the NodePool and EC2NodeClass.
317 """
318 capacity_types = ["reserved", "on-demand"] if fallback_on_demand else ["reserved"]
320 ec2_node_class = {
321 "apiVersion": "karpenter.k8s.aws/v1",
322 "kind": "EC2NodeClass",
323 "metadata": {
324 "name": f"{name}-nodeclass",
325 "labels": {
326 "app.kubernetes.io/part-of": "gco",
327 "gco.io/nodepool": name,
328 },
329 },
330 "spec": {
331 "role": f"KarpenterNodeRole-{project_name}",
332 "subnetSelectorTerms": [
333 {"tags": {"karpenter.sh/discovery": f"{project_name}-{region}"}}
334 ],
335 "securityGroupSelectorTerms": [
336 {"tags": {"karpenter.sh/discovery": f"{project_name}-{region}"}}
337 ],
338 "capacityReservationSelectorTerms": [{"id": capacity_reservation_id}],
339 "tags": {
340 "Name": f"{project_name}-{name}",
341 "gco.io/nodepool": name,
342 "gco.io/capacity-block": capacity_reservation_id,
343 },
344 },
345 }
347 requirements: list[dict[str, Any]] = [
348 {
349 "key": "karpenter.sh/capacity-type",
350 "operator": "In",
351 "values": capacity_types,
352 },
353 {
354 "key": "kubernetes.io/arch",
355 "operator": "In",
356 "values": ["amd64"],
357 },
358 ]
360 nodepool: dict[str, Any] = {
361 "apiVersion": "karpenter.sh/v1",
362 "kind": "NodePool",
363 "metadata": {
364 "name": name,
365 "labels": {
366 "app.kubernetes.io/part-of": "gco",
367 },
368 },
369 "spec": {
370 "template": {
371 "metadata": {
372 "labels": {
373 "workload-type": "capacity-block",
374 "project": "gco",
375 "gco.io/capacity-type": "capacity-block",
376 "gco.io/capacity-block": capacity_reservation_id,
377 },
378 },
379 "spec": {
380 "nodeClassRef": {
381 "group": "karpenter.k8s.aws",
382 "kind": "EC2NodeClass",
383 "name": f"{name}-nodeclass",
384 },
385 "requirements": requirements,
386 },
387 },
388 "limits": {
389 "cpu": str(calculate_cpu_limit(instance_types, max_nodes, region)),
390 },
391 # A Capacity Block is prepaid for a fixed term — hold the nodes rather
392 # than consolidating them away, so the paid capacity stays available
393 # for the whole block.
394 "disruption": {
395 "consolidationPolicy": "WhenEmpty",
396 "consolidateAfter": "600s",
397 "budgets": [{"nodes": "10%"}],
398 },
399 },
400 }
402 if instance_types:
403 requirements.append(
404 {
405 "key": "node.kubernetes.io/instance-type",
406 "operator": "In",
407 "values": instance_types,
408 }
409 )
411 gpu_families = ["p3", "p4", "p5", "p6", "g4", "g5", "g6"]
412 if instance_types and any(
413 any(it.startswith(fam) for fam in gpu_families) for it in instance_types
414 ):
415 taints = [
416 {
417 "key": "nvidia.com/gpu",
418 "value": "true",
419 "effect": "NoSchedule",
420 }
421 ]
422 if efa:
423 taints.append(
424 {
425 "key": "vpc.amazonaws.com/efa",
426 "value": "true",
427 "effect": "NoSchedule",
428 }
429 )
430 nodepool["spec"]["template"]["spec"]["taints"] = taints
432 if efa:
433 nodepool["spec"]["template"]["metadata"]["labels"]["efa"] = "true"
434 nodepool["spec"]["template"]["metadata"]["labels"]["workload-type"] = "gpu-efa"
436 output = []
437 output.append("# Capacity Block-backed NodePool for GCO")
438 output.append(f"# Capacity Reservation (from Capacity Block): {capacity_reservation_id}")
439 output.append(f"# Region: {region}")
440 if fallback_on_demand:
441 output.append("# Fallback: on-demand (when Capacity Block exhausted/expired)")
442 output.append("#")
443 output.append("# Apply with: kubectl apply -f <this-file>.yaml")
444 output.append(
445 "# See: https://karpenter.sh/docs/concepts/nodeclasses/#speccapacityreservationselectorterms"
446 )
447 output.append("---")
448 output.append(yaml.dump(ec2_node_class, default_flow_style=False, sort_keys=False))
449 output.append("---")
450 output.append(yaml.dump(nodepool, default_flow_style=False, sort_keys=False))
452 return "\n".join(output)
455def get_eks_token(cluster_name: str, region: str) -> str:
456 """Generate EKS authentication token using STS presigned URL."""
457 from botocore.signers import RequestSigner
459 session = boto3.Session()
460 sts_client = session.client("sts", region_name=region)
461 service_id = sts_client.meta.service_model.service_id
463 signer = RequestSigner(
464 service_id, region, "sts", "v4", session.get_credentials(), session.events
465 )
467 params = {
468 "method": "GET",
469 "url": f"https://sts.{region}.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15",
470 "body": {},
471 "headers": {"x-k8s-aws-id": cluster_name},
472 "context": {},
473 }
475 url = signer.generate_presigned_url(
476 params, region_name=region, expires_in=60, operation_name=""
477 )
479 token_b64 = base64.urlsafe_b64encode(url.encode("utf-8")).decode("utf-8").rstrip("=")
480 return f"k8s-aws-v1.{token_b64}"
483def get_k8s_client(cluster_name: str, region: str) -> CustomObjectsApi:
484 """Get configured Kubernetes client for EKS cluster."""
485 from kubernetes import client
487 eks = boto3.client("eks", region_name=region)
488 cluster_info = eks.describe_cluster(name=cluster_name)
489 cluster = cluster_info["cluster"]
491 configuration = client.Configuration()
492 configuration.host = cluster["endpoint"]
493 configuration.verify_ssl = True
495 # The ApiClient reads this path throughout its lifetime. Keep the file until
496 # that client is collected, while cleaning every partially constructed path.
497 ca_cert = base64.b64decode(cluster["certificateAuthority"]["data"])
498 fd, ca_cert_path = tempfile.mkstemp(suffix=".crt")
499 fd_owned = True
500 api_client = None
501 try:
502 ca_file = os.fdopen(fd, "wb")
503 fd_owned = False # os.fdopen transferred descriptor ownership to ca_file.
504 with ca_file:
505 ca_file.write(ca_cert)
506 ca_file.flush()
507 configuration.ssl_ca_cert = ca_cert_path
509 # Generate EKS token
510 eks_token = get_eks_token(cluster_name, region)
511 configuration.api_key = {"authorization": f"Bearer {eks_token}"}
513 # Create API client with the configuration explicitly
514 api_client = client.ApiClient(configuration)
515 custom_api = client.CustomObjectsApi(api_client)
516 except BaseException:
517 if fd_owned:
518 with suppress(OSError):
519 os.close(fd)
520 if api_client is not None:
521 with suppress(Exception):
522 api_client.close()
523 _unlink_temp_ca_cert(ca_cert_path)
524 raise
526 # CustomObjectsApi retains api_client, so the certificate remains available
527 # for lazy TLS setup and is removed when the owning client is released.
528 weakref.finalize(api_client, _unlink_temp_ca_cert, ca_cert_path)
529 return custom_api
532def list_cluster_nodepools(cluster_name: str, region: str) -> list[dict[str, Any]]:
533 """
534 List NodePools in an EKS cluster.
536 Args:
537 cluster_name: EKS cluster name
538 region: AWS region
540 Returns:
541 List of NodePool information dictionaries
542 """
543 try:
544 custom_api = get_k8s_client(cluster_name, region)
546 nodepools = custom_api.list_cluster_custom_object(
547 group="karpenter.sh",
548 version="v1",
549 plural="nodepools",
550 )
552 result = []
553 for np in nodepools.get("items", []):
554 spec = np.get("spec", {})
555 template = spec.get("template", {}).get("spec", {})
556 requirements = template.get("requirements", [])
558 # Extract capacity types
559 capacity_types = []
560 instance_types = []
561 for req in requirements:
562 if req.get("key") == "karpenter.sh/capacity-type":
563 capacity_types = req.get("values", [])
564 elif req.get("key") == "node.kubernetes.io/instance-type":
565 instance_types = req.get("values", [])
567 # Get status
568 status = np.get("status", {})
569 conditions = status.get("conditions", [])
570 ready_condition: dict[str, Any] = next(
571 (c for c in conditions if c.get("type") == "Ready"), {}
572 )
574 result.append(
575 {
576 "name": np["metadata"]["name"],
577 "capacity_types": ", ".join(capacity_types) or "on-demand",
578 "instance_types": ", ".join(instance_types[:3])
579 + ("..." if len(instance_types) > 3 else "")
580 or "any",
581 "status": "Ready" if ready_condition.get("status") == "True" else "NotReady",
582 "limits": spec.get("limits", {}),
583 }
584 )
586 return result
588 except Exception as e:
589 raise RuntimeError(f"Failed to list NodePools: {e}") from e
592def describe_cluster_nodepool(
593 cluster_name: str, region: str, nodepool_name: str
594) -> dict[str, Any] | None:
595 """
596 Describe a specific NodePool in an EKS cluster.
598 Args:
599 cluster_name: EKS cluster name
600 region: AWS region
601 nodepool_name: Name of the NodePool
603 Returns:
604 NodePool details or None if not found
605 """
606 try:
607 custom_api = get_k8s_client(cluster_name, region)
609 nodepool = custom_api.get_cluster_custom_object(
610 group="karpenter.sh",
611 version="v1",
612 plural="nodepools",
613 name=nodepool_name,
614 )
616 if isinstance(nodepool, dict):
617 return nodepool
618 return None
620 except Exception as e:
621 if "404" in str(e):
622 return None
623 raise RuntimeError(f"Failed to describe NodePool: {e}") from e
626def delete_cluster_nodepool(cluster_name: str, region: str, nodepool_name: str) -> dict[str, Any]:
627 """
628 Delete a NodePool (and its paired EC2NodeClass) from an EKS cluster.
630 Deletes the Karpenter NodePool ``nodepool_name`` and, when present, the
631 EC2NodeClass named ``<nodepool_name>-nodeclass`` that the GCO manifest
632 generators create alongside it. Karpenter drains and terminates any nodes
633 the NodePool provisioned once it is removed. A missing EC2NodeClass (custom
634 name, or already deleted) is not treated as an error.
636 Args:
637 cluster_name: EKS cluster name
638 region: AWS region
639 nodepool_name: Name of the NodePool to delete
641 Returns:
642 Dict describing what was deleted: ``{"nodepool": <name>,
643 "ec2nodeclass": <name-or-None>}``.
644 """
645 try:
646 custom_api = get_k8s_client(cluster_name, region)
648 custom_api.delete_cluster_custom_object(
649 group="karpenter.sh",
650 version="v1",
651 plural="nodepools",
652 name=nodepool_name,
653 )
654 deleted: dict[str, Any] = {"nodepool": nodepool_name, "ec2nodeclass": None}
656 nodeclass_name = f"{nodepool_name}-nodeclass"
657 try:
658 custom_api.delete_cluster_custom_object(
659 group="karpenter.k8s.aws",
660 version="v1",
661 plural="ec2nodeclasses",
662 name=nodeclass_name,
663 )
664 deleted["ec2nodeclass"] = nodeclass_name
665 except Exception as e: # noqa: BLE001 - best effort; the NodePool is the primary target
666 if "404" not in str(e):
667 logger.warning("Could not delete EC2NodeClass %s: %s", nodeclass_name, e)
669 return deleted
671 except Exception as e:
672 raise RuntimeError(f"Failed to delete NodePool: {e}") from e