Coverage for cli / commands / nodepools_cmd.py: 100%

1"""NodePool commands.""" 

2 

3import sys 

4from typing import Any 

5 

6import click 

7 

8from ..config import GCOConfig 

9from ..output import get_output_formatter 

10 

11pass_config = click.make_pass_decorator(GCOConfig, ensure=True) 

12 

13 

14@click.group() 

15@pass_config 

16def nodepools(config: Any) -> None: 

17 """Manage Karpenter NodePools with ODCR/Capacity Reservation support.""" 

18 pass 

19 

20 

21@nodepools.command("create-odcr") 

22@click.option("--name", "-n", required=True, help="NodePool name") 

23@click.option("--region", "-r", required=True, help="AWS region") 

24@click.option( 

25 "--capacity-reservation-id", 

26 "-c", 

27 required=True, 

28 help="EC2 Capacity Reservation ID (cr-xxx) or ODCR group ARN", 

29) 

30@click.option( 

31 "--instance-type", 

32 "-i", 

33 multiple=True, 

34 help="Instance types (can specify multiple)", 

35) 

36@click.option("--max-nodes", default=100, help="Maximum nodes in pool") 

37@click.option("--fallback-on-demand", is_flag=True, help="Fall back to on-demand if ODCR exhausted") 

38@click.option( 

39 "--efa", 

40 is_flag=True, 

41 help="Enable EFA support (adds EFA taint and labels for p4d/p5 instances)", 

42) 

43@click.option("--output-file", "-o", help="Output manifest to file instead of applying") 

44@pass_config 

45def create_odcr_nodepool( 

46 config: Any, 

47 name: Any, 

48 region: Any, 

49 capacity_reservation_id: Any, 

50 instance_type: Any, 

51 max_nodes: Any, 

52 fallback_on_demand: Any, 

53 efa: Any, 

54 output_file: Any, 

55) -> None: 

56 """Create a NodePool backed by an On-Demand Capacity Reservation (ODCR). 

57 

58 This creates a Karpenter NodePool and EC2NodeClass configured to use 

59 a specific capacity reservation for guaranteed capacity. 

60 

61 Examples: 

62 gco nodepools create-odcr -n gpu-reserved -r us-east-1 \\ 

63 -c cr-0123456789abcdef0 -i p4d.24xlarge 

64 

65 gco nodepools create-odcr -n ml-training -r us-west-2 \\ 

66 -c cr-0123456789abcdef0 -i p5.48xlarge --fallback-on-demand 

67 

68 gco nodepools create-odcr -n efa-training -r us-east-1 \\ 

69 -c cr-0123456789abcdef0 -i p4d.24xlarge --efa 

70 

71 See: https://karpenter.sh/docs/tasks/odcrs/ 

72 """ 

73 from ..nodepools import generate_odcr_nodepool_manifest 

74 

75 formatter = get_output_formatter(config) 

76 

77 try: 

78 manifest = generate_odcr_nodepool_manifest( 

79 name=name, 

80 region=region, 

81 capacity_reservation_id=capacity_reservation_id, 

82 instance_types=list(instance_type) if instance_type else None, 

83 max_nodes=max_nodes, 

84 fallback_on_demand=fallback_on_demand, 

85 efa=efa, 

86 ) 

87 

88 if output_file: 

89 with open(output_file, "w", encoding="utf-8") as f: 

90 f.write(manifest) 

91 formatter.print_success(f"NodePool manifest written to {output_file}") 

92 formatter.print_info(f"Apply with: kubectl apply -f {output_file}") 

93 else: 

94 # Print the manifest for review 

95 print(manifest) 

96 formatter.print_info("\nTo apply this manifest, save it to a file and run:") 

97 formatter.print_info(" kubectl apply -f <filename>.yaml") 

98 

99 except Exception as e: 

100 formatter.print_error(f"Failed to create ODCR NodePool: {e}") 

101 sys.exit(1) 

102 

103 

104@nodepools.command("list") 

105@click.option("--region", "-r", help="Filter by region") 

106@click.option("--cluster", help="EKS cluster name (defaults to gco-<region>)") 

107@pass_config 

108def list_nodepools(config: Any, region: Any, cluster: Any) -> None: 

109 """List NodePools in the cluster. 

110 

111 Examples: 

112 gco nodepools list --region us-east-1 

113 gco nodepools list --cluster my-cluster 

114 """ 

115 from ..nodepools import list_cluster_nodepools 

116 

117 formatter = get_output_formatter(config) 

118 

119 try: 

120 if not region and not cluster: 

121 formatter.print_error("Either --region or --cluster is required") 

122 sys.exit(1) 

123 

124 cluster_name = cluster or f"gco-{region}" 

125 nodepools_list = list_cluster_nodepools(cluster_name, region or config.default_region) 

126 

127 if not nodepools_list: 

128 formatter.print_info("No NodePools found") 

129 return 

130 

131 formatter.print(nodepools_list) 

132 

133 except Exception as e: 

134 formatter.print_error(f"Failed to list NodePools: {e}") 

135 sys.exit(1) 

136 

137 

138@nodepools.command("describe") 

139@click.argument("nodepool_name") 

140@click.option("--region", "-r", required=True, help="AWS region") 

141@click.option("--cluster", help="EKS cluster name (defaults to gco-<region>)") 

142@pass_config 

143def describe_nodepool(config: Any, nodepool_name: Any, region: Any, cluster: Any) -> None: 

144 """Describe a NodePool. 

145 

146 Examples: 

147 gco nodepools describe gpu-x86-pool --region us-east-1 

148 """ 

149 from ..nodepools import describe_cluster_nodepool 

150 

151 formatter = get_output_formatter(config) 

152 

153 try: 

154 cluster_name = cluster or f"gco-{region}" 

155 nodepool = describe_cluster_nodepool(cluster_name, region, nodepool_name) 

156 

157 if not nodepool: 

158 formatter.print_error(f"NodePool {nodepool_name} not found") 

159 sys.exit(1) 

160 

161 formatter.print(nodepool) 

162 

163 except Exception as e: 

164 formatter.print_error(f"Failed to describe NodePool: {e}") 

165 sys.exit(1)