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

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 

9@dataclass 

10class InstanceTypeInfo: 

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

12 

13 instance_type: str 

14 vcpus: int 

15 memory_gib: float 

16 gpu_count: int = 0 

17 gpu_type: str | None = None 

18 gpu_memory_gib: float = 0 

19 architecture: str = "x86_64" 

20 

21 @property 

22 def is_gpu(self) -> bool: 

23 return self.gpu_count > 0 

24 

25 

26@dataclass 

27class SpotPriceInfo: 

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

29 

30 instance_type: str 

31 availability_zone: str 

32 current_price: float 

33 avg_price_7d: float 

34 min_price_7d: float 

35 max_price_7d: float 

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

37 

38 

39@dataclass 

40class CapacityEstimate: 

41 """Capacity availability estimate.""" 

42 

43 instance_type: str 

44 region: str 

45 availability_zone: str | None 

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

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

48 confidence: float # 0-1 

49 estimated_wait_time: str | None = None 

50 price_per_hour: float | None = None 

51 recommendation: str = "" 

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

53 

54 

55# GPU instance type specifications (common types) 

56GPU_INSTANCE_SPECS = { 

57 # G4dn - NVIDIA T4 

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

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

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

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

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

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

64 # G5 - NVIDIA A10G 

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

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

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

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

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

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

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

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

73 # P3 - NVIDIA V100 

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

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

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

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

78 # P4d - NVIDIA A100 

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

80 # P5 - NVIDIA H100 

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

82}