Coverage for cli/capacity/models.py: 100.00%

1"""Data classes and GPU instance specifications for capacity checking.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass, field 

6from typing import Any 

7 

8 

9class CapacityCheckError(Exception): 

10 """A primary capacity availability check failed at the AWS API level. 

11 

12 Raised when an availability lookup (e.g. DescribeInstanceTypeOfferings) 

13 fails due to throttling, expired/invalid credentials, denied permissions, 

14 or a region that isn't opted in. This is distinct from a *successful* lookup 

15 that reports an instance type as genuinely not offered — masking such 

16 failures as "unavailable" hides real, actionable errors from the caller. 

17 """ 

18 

19 

20@dataclass 

21class InstanceTypeInfo: 

22 """Information about an EC2 instance type.""" 

23 

24 instance_type: str 

25 vcpus: int 

26 memory_gib: float 

27 gpu_count: int = 0 

28 gpu_type: str | None = None 

29 gpu_memory_gib: float = 0 

30 architecture: str = "x86_64" 

31 

32 @property 

33 def is_gpu(self) -> bool: 

34 return self.gpu_count > 0 

35 

36 

37@dataclass 

38class SpotPriceInfo: 

39 """Spot price information for an instance type.""" 

40 

41 instance_type: str 

42 availability_zone: str 

43 current_price: float 

44 avg_price_7d: float 

45 min_price_7d: float 

46 max_price_7d: float 

47 price_stability: float # 0-1, higher is more stable 

48 

49 

50@dataclass 

51class CapacityEstimate: 

52 """Capacity availability estimate.""" 

53 

54 instance_type: str 

55 region: str 

56 availability_zone: str | None 

57 capacity_type: str # "spot" or "on-demand" 

58 availability: str # "high", "medium", "low", "unavailable", "unknown" 

59 confidence: float # 0-1 

60 estimated_wait_time: str | None = None 

61 price_per_hour: float | None = None 

62 recommendation: str = "" 

63 details: dict[str, Any] = field(default_factory=dict) 

64 error: str | None = None # Explicit error string for partial/degraded results 

65 

66 

67# GPU instance type specifications (common types) 

68GPU_INSTANCE_SPECS = { 

69 # G4dn - NVIDIA T4 

70 "g4dn.xlarge": InstanceTypeInfo("g4dn.xlarge", 4, 16, 1, "T4", 16, "x86_64"), 

71 "g4dn.2xlarge": InstanceTypeInfo("g4dn.2xlarge", 8, 32, 1, "T4", 16, "x86_64"), 

72 "g4dn.4xlarge": InstanceTypeInfo("g4dn.4xlarge", 16, 64, 1, "T4", 16, "x86_64"), 

73 "g4dn.8xlarge": InstanceTypeInfo("g4dn.8xlarge", 32, 128, 1, "T4", 16, "x86_64"), 

74 "g4dn.12xlarge": InstanceTypeInfo("g4dn.12xlarge", 48, 192, 4, "T4", 64, "x86_64"), 

75 "g4dn.16xlarge": InstanceTypeInfo("g4dn.16xlarge", 64, 256, 1, "T4", 16, "x86_64"), 

76 # G5 - NVIDIA A10G 

77 "g5.xlarge": InstanceTypeInfo("g5.xlarge", 4, 16, 1, "A10G", 24, "x86_64"), 

78 "g5.2xlarge": InstanceTypeInfo("g5.2xlarge", 8, 32, 1, "A10G", 24, "x86_64"), 

79 "g5.4xlarge": InstanceTypeInfo("g5.4xlarge", 16, 64, 1, "A10G", 24, "x86_64"), 

80 "g5.8xlarge": InstanceTypeInfo("g5.8xlarge", 32, 128, 1, "A10G", 24, "x86_64"), 

81 "g5.12xlarge": InstanceTypeInfo("g5.12xlarge", 48, 192, 4, "A10G", 96, "x86_64"), 

82 "g5.16xlarge": InstanceTypeInfo("g5.16xlarge", 64, 256, 1, "A10G", 24, "x86_64"), 

83 "g5.24xlarge": InstanceTypeInfo("g5.24xlarge", 96, 384, 4, "A10G", 96, "x86_64"), 

84 "g5.48xlarge": InstanceTypeInfo("g5.48xlarge", 192, 768, 8, "A10G", 192, "x86_64"), 

85 # P3 - NVIDIA V100 

86 "p3.2xlarge": InstanceTypeInfo("p3.2xlarge", 8, 61, 1, "V100", 16, "x86_64"), 

87 "p3.8xlarge": InstanceTypeInfo("p3.8xlarge", 32, 244, 4, "V100", 64, "x86_64"), 

88 "p3.16xlarge": InstanceTypeInfo("p3.16xlarge", 64, 488, 8, "V100", 128, "x86_64"), 

89 "p3dn.24xlarge": InstanceTypeInfo("p3dn.24xlarge", 96, 768, 8, "V100", 256, "x86_64"), 

90 # P4d / P4de - NVIDIA A100 

91 "p4d.24xlarge": InstanceTypeInfo("p4d.24xlarge", 96, 1152, 8, "A100", 320, "x86_64"), 

92 "p4de.24xlarge": InstanceTypeInfo("p4de.24xlarge", 96, 1152, 8, "A100", 640, "x86_64"), 

93 # P5 / P5e / P5en - NVIDIA H100 / H200 

94 "p5.48xlarge": InstanceTypeInfo("p5.48xlarge", 192, 2048, 8, "H100", 640, "x86_64"), 

95 "p5e.48xlarge": InstanceTypeInfo("p5e.48xlarge", 192, 2048, 8, "H200", 1128, "x86_64"), 

96 "p5en.48xlarge": InstanceTypeInfo("p5en.48xlarge", 192, 2048, 8, "H200", 1128, "x86_64"), 

97 # P6 - NVIDIA Blackwell B200 / B300 Ultra (standalone EC2 types). The 

98 # Grace-Blackwell GB200/GB300 NVL72 superchips ship only as P6e-GB 

99 # UltraServers, so those are intentionally absent from this standalone table. 

100 "p6-b200.48xlarge": InstanceTypeInfo("p6-b200.48xlarge", 192, 2048, 8, "B200", 1440, "x86_64"), 

101 "p6-b300.48xlarge": InstanceTypeInfo("p6-b300.48xlarge", 192, 4096, 8, "B300", 2144, "x86_64"), 

102}