service-capacity-modeling 0.3.53__tar.gz → 0.3.54__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (96) hide show
  1. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/PKG-INFO +1 -1
  2. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/capacity_planner.py +10 -2
  3. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/utils.py +23 -7
  4. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling.egg-info/PKG-INFO +1 -1
  5. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling.egg-info/SOURCES.txt +1 -0
  6. service_capacity_modeling-0.3.54/tests/test_utils.py +175 -0
  7. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/LICENSE +0 -0
  8. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/README.md +0 -0
  9. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/__init__.py +0 -0
  10. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/__init__.py +0 -0
  11. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/__init__.py +0 -0
  12. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/pricing/aws/3yr-reserved_ec2.json +0 -0
  13. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/pricing/aws/3yr-reserved_zz-overrides.json +0 -0
  14. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/profiles.txt +0 -0
  15. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_c5.json +0 -0
  16. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_c5a.json +0 -0
  17. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_c5d.json +0 -0
  18. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_c5n.json +0 -0
  19. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_c6a.json +0 -0
  20. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_c6i.json +0 -0
  21. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_c6id.json +0 -0
  22. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_c7a.json +0 -0
  23. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_c7i.json +0 -0
  24. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m4.json +0 -0
  25. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m5.json +0 -0
  26. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m5n.json +0 -0
  27. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m6a.json +0 -0
  28. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m6i.json +0 -0
  29. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m6id.json +0 -0
  30. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m6idn.json +0 -0
  31. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m6in.json +0 -0
  32. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m7a.json +0 -0
  33. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_m7i.json +0 -0
  34. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r4.json +0 -0
  35. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r5.json +0 -0
  36. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r5n.json +0 -0
  37. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r6a.json +0 -0
  38. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r6i.json +0 -0
  39. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r6id.json +0 -0
  40. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r6idn.json +0 -0
  41. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r6in.json +0 -0
  42. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r7a.json +0 -0
  43. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/auto_r7i.json +0 -0
  44. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/manual_drives.json +0 -0
  45. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/manual_instances.json +0 -0
  46. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/hardware/profiles/shapes/aws/manual_services.json +0 -0
  47. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/interface.py +0 -0
  48. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/__init__.py +0 -0
  49. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/common.py +0 -0
  50. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/headroom_strategy.py +0 -0
  51. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/__init__.py +0 -0
  52. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/__init__.py +0 -0
  53. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/aurora.py +0 -0
  54. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/cassandra.py +0 -0
  55. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/counter.py +0 -0
  56. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/crdb.py +0 -0
  57. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/ddb.py +0 -0
  58. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/elasticsearch.py +0 -0
  59. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/entity.py +0 -0
  60. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/evcache.py +0 -0
  61. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/graphkv.py +0 -0
  62. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/iso_date_math.py +0 -0
  63. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/kafka.py +0 -0
  64. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/key_value.py +0 -0
  65. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/postgres.py +0 -0
  66. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/rds.py +0 -0
  67. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/stateless_java.py +0 -0
  68. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/time_series.py +0 -0
  69. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/time_series_config.py +0 -0
  70. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/wal.py +0 -0
  71. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/models/org/netflix/zookeeper.py +0 -0
  72. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/stats.py +0 -0
  73. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/tools/__init__.py +0 -0
  74. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/tools/auto_shape.py +0 -0
  75. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/tools/fetch_pricing.py +0 -0
  76. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/tools/generate_missing.py +0 -0
  77. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling/tools/instance_families.py +0 -0
  78. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling.egg-info/dependency_links.txt +0 -0
  79. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling.egg-info/entry_points.txt +0 -0
  80. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling.egg-info/requires.txt +0 -0
  81. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/service_capacity_modeling.egg-info/top_level.txt +0 -0
  82. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/setup.cfg +0 -0
  83. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/setup.py +0 -0
  84. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_arguments.py +0 -0
  85. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_buffers.py +0 -0
  86. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_common.py +0 -0
  87. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_desire_merge.py +0 -0
  88. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_generate_scenarios.py +0 -0
  89. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_hardware.py +0 -0
  90. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_hardware_shapes.py +0 -0
  91. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_headroom_strategy.py +0 -0
  92. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_io2.py +0 -0
  93. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_model_dump.py +0 -0
  94. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_reproducible.py +0 -0
  95. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_simulation.py +0 -0
  96. {service_capacity_modeling-0.3.53 → service_capacity_modeling-0.3.54}/tests/test_working_set.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: service-capacity-modeling
3
- Version: 0.3.53
3
+ Version: 0.3.54
4
4
  Summary: Contains utilities for modeling capacity for pluggable workloads
5
5
  Author: Joseph Lynch
6
6
  Author-email: josephl@netflix.com
@@ -511,6 +511,7 @@ class CapacityPlanner:
511
511
  num_results: Optional[int] = None,
512
512
  num_regions: int = 3,
513
513
  extra_model_arguments: Optional[Dict[str, Any]] = None,
514
+ max_results_per_family: int = 1,
514
515
  ) -> Sequence[CapacityPlan]:
515
516
  if model_name not in self._models:
516
517
  raise ValueError(
@@ -538,6 +539,7 @@ class CapacityPlanner:
538
539
  lifecycles=lifecycles,
539
540
  instance_families=instance_families,
540
541
  drives=drives,
542
+ max_results_per_family=max_results_per_family,
541
543
  )
542
544
  if sub_plan:
543
545
  results.append(sub_plan)
@@ -555,6 +557,7 @@ class CapacityPlanner:
555
557
  instance_families: Optional[Sequence[str]] = None,
556
558
  drives: Optional[Sequence[str]] = None,
557
559
  extra_model_arguments: Optional[Dict[str, Any]] = None,
560
+ max_results_per_family: int = 1,
558
561
  ) -> Sequence[CapacityPlan]:
559
562
  extra_model_arguments = extra_model_arguments or {}
560
563
  model = self._models[model_name]
@@ -577,7 +580,9 @@ class CapacityPlanner:
577
580
  plans.sort(key=lambda p: (p.rank, p.candidate_clusters.total_annual_cost))
578
581
 
579
582
  num_results = num_results or self._default_num_results
580
- return reduce_by_family(plans)[:num_results]
583
+ return reduce_by_family(plans, max_results_per_family=max_results_per_family)[
584
+ :num_results
585
+ ]
581
586
 
582
587
  # Calculates the minimum cpu, memory, and network requirements based on desires.
583
588
  def _per_instance_requirements(self, desires) -> Tuple[int, float]:
@@ -696,6 +701,7 @@ class CapacityPlanner:
696
701
  regret_params: Optional[CapacityRegretParameters] = None,
697
702
  extra_model_arguments: Optional[Dict[str, Any]] = None,
698
703
  explain: bool = False,
704
+ max_results_per_family: int = 1,
699
705
  ) -> UncertainCapacityPlan:
700
706
  extra_model_arguments = extra_model_arguments or {}
701
707
 
@@ -741,6 +747,7 @@ class CapacityPlanner:
741
747
  lifecycles=lifecycles,
742
748
  instance_families=instance_families,
743
749
  drives=drives,
750
+ max_results_per_family=max_results_per_family,
744
751
  ),
745
752
  )
746
753
  )
@@ -763,7 +770,8 @@ class CapacityPlanner:
763
770
  ],
764
771
  zonal_requirements,
765
772
  regional_requirements,
766
- )
773
+ ),
774
+ max_results_per_family=max_results_per_family,
767
775
  )[:num_results]
768
776
 
769
777
  low_p, high_p = sorted(percentiles)[0], sorted(percentiles)[-1]
@@ -1,19 +1,26 @@
1
1
  import math
2
+ from typing import Dict
2
3
  from typing import Iterable
3
4
  from typing import List
4
- from typing import Set
5
5
  from typing import Tuple
6
6
 
7
7
  from service_capacity_modeling.models import CapacityPlan
8
8
 
9
9
 
10
- def reduce_by_family(plans: Iterable[CapacityPlan]) -> List[CapacityPlan]:
10
+ def reduce_by_family(
11
+ plans: Iterable[CapacityPlan], max_results_per_family: int = 1
12
+ ) -> List[CapacityPlan]:
11
13
  """Groups a potential set of clusters by hardware family sorted by cost.
12
14
 
13
15
  Useful for showing different family options.
16
+
17
+ Args:
18
+ plans: Iterable of CapacityPlan objects to filter
19
+ max_results_per_family: Maximum number of results to return per
20
+ family combination
14
21
  """
15
- zonal_families: Set[Tuple[Tuple[str, str], ...]] = set()
16
- regional_families: Set[Tuple[Tuple[str, str], ...]] = set()
22
+ zonal_families: Dict[Tuple[Tuple[str, str], ...], int] = {}
23
+ regional_families: Dict[Tuple[Tuple[str, str], ...], int] = {}
17
24
 
18
25
  result: List[CapacityPlan] = []
19
26
  for plan in plans:
@@ -31,11 +38,20 @@ def reduce_by_family(plans: Iterable[CapacityPlan]) -> List[CapacityPlan]:
31
38
  sorted({(c.cluster_type, c.instance.family) for c in topo.zonal})
32
39
  )
33
40
 
34
- if not (zonal_type in zonal_families and regional_type in regional_families):
41
+ # Count how many of each family combination we've seen
42
+ zonal_count = zonal_families.get(zonal_type, 0)
43
+ regional_count = regional_families.get(regional_type, 0)
44
+
45
+ # Add the plan if we haven't reached the maximum for either family type
46
+ if (
47
+ zonal_count < max_results_per_family
48
+ or regional_count < max_results_per_family
49
+ ):
35
50
  result.append(plan)
36
51
 
37
- regional_families.add(regional_type)
38
- zonal_families.add(zonal_type)
52
+ # Update counters
53
+ zonal_families[zonal_type] = zonal_count + 1
54
+ regional_families[regional_type] = regional_count + 1
39
55
 
40
56
  return result
41
57
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: service-capacity-modeling
3
- Version: 0.3.53
3
+ Version: 0.3.54
4
4
  Summary: Contains utilities for modeling capacity for pluggable workloads
5
5
  Author: Joseph Lynch
6
6
  Author-email: josephl@netflix.com
@@ -90,4 +90,5 @@ tests/test_io2.py
90
90
  tests/test_model_dump.py
91
91
  tests/test_reproducible.py
92
92
  tests/test_simulation.py
93
+ tests/test_utils.py
93
94
  tests/test_working_set.py
@@ -0,0 +1,175 @@
1
+ from decimal import Decimal
2
+ from typing import Dict
3
+ from typing import List
4
+
5
+ from service_capacity_modeling.interface import CapacityPlan
6
+ from service_capacity_modeling.interface import Clusters
7
+ from service_capacity_modeling.interface import Instance
8
+ from service_capacity_modeling.interface import Requirements
9
+ from service_capacity_modeling.interface import ZoneClusterCapacity
10
+ from service_capacity_modeling.models.utils import reduce_by_family
11
+
12
+
13
+ # Create mock hardware instances with different families for all tests
14
+ def get_test_instances():
15
+ shape_family_a1 = Instance(
16
+ name="family_a.a1",
17
+ family_separator=".",
18
+ cpu=2,
19
+ cpu_ghz=2.4,
20
+ ram_gib=8,
21
+ net_mbps=1000,
22
+ )
23
+ shape_family_a2 = Instance(
24
+ name="family_a.a2",
25
+ family_separator=".",
26
+ cpu=4,
27
+ cpu_ghz=2.4,
28
+ ram_gib=16,
29
+ net_mbps=2000,
30
+ )
31
+ shape_family_a3 = Instance(
32
+ name="family_a.a3",
33
+ family_separator=".",
34
+ cpu=8,
35
+ cpu_ghz=2.4,
36
+ ram_gib=32,
37
+ net_mbps=4000,
38
+ )
39
+
40
+ shape_family_b1 = Instance(
41
+ name="family_b.b1",
42
+ family_separator=".",
43
+ cpu=2,
44
+ cpu_ghz=2.4,
45
+ ram_gib=8,
46
+ net_mbps=1000,
47
+ )
48
+ shape_family_b2 = Instance(
49
+ name="family_b.b2",
50
+ family_separator=".",
51
+ cpu=4,
52
+ cpu_ghz=2.4,
53
+ ram_gib=16,
54
+ net_mbps=2000,
55
+ )
56
+
57
+ return (
58
+ shape_family_a1,
59
+ shape_family_a2,
60
+ shape_family_a3,
61
+ shape_family_b1,
62
+ shape_family_b2,
63
+ )
64
+
65
+
66
+ def create_test_capacity_plans() -> List[CapacityPlan]:
67
+ """Create test capacity plans with different hardware families for testing."""
68
+ shapes = get_test_instances()
69
+ (
70
+ shape_family_a1,
71
+ shape_family_a2,
72
+ shape_family_a3,
73
+ shape_family_b1,
74
+ shape_family_b2,
75
+ ) = shapes
76
+
77
+ plans = []
78
+
79
+ # Family A plans
80
+ for i, shape in enumerate([shape_family_a1, shape_family_a2, shape_family_a3]):
81
+ annual_cost = (i + 1) * 100.0 # Different costs
82
+
83
+ cluster = ZoneClusterCapacity(
84
+ cluster_type="test_cluster",
85
+ count=i + 1,
86
+ instance=shape,
87
+ annual_cost=annual_cost,
88
+ )
89
+
90
+ annual_costs_a: Dict[str, Decimal] = {"test_cluster": Decimal(str(annual_cost))}
91
+
92
+ plans.append(
93
+ CapacityPlan(
94
+ requirements=Requirements(),
95
+ candidate_clusters=Clusters(
96
+ annual_costs=annual_costs_a, zonal=[cluster], regional=[]
97
+ ),
98
+ cost=annual_cost,
99
+ efficiency=1.0,
100
+ )
101
+ )
102
+
103
+ # Family B plans
104
+ for i, shape in enumerate([shape_family_b1, shape_family_b2]):
105
+ annual_cost = (i + 1) * 200.0 # Different costs
106
+
107
+ cluster = ZoneClusterCapacity(
108
+ cluster_type="test_cluster",
109
+ count=i + 1,
110
+ instance=shape,
111
+ annual_cost=annual_cost,
112
+ )
113
+
114
+ annual_costs_b: Dict[str, Decimal] = {"test_cluster": Decimal(str(annual_cost))}
115
+
116
+ plans.append(
117
+ CapacityPlan(
118
+ requirements=Requirements(),
119
+ candidate_clusters=Clusters(
120
+ annual_costs=annual_costs_b, zonal=[cluster], regional=[]
121
+ ),
122
+ cost=annual_cost,
123
+ efficiency=1.0,
124
+ )
125
+ )
126
+
127
+ return plans
128
+
129
+
130
+ def test_reduce_by_family_default():
131
+ """Test that reduce_by_family with default parameter returns one plan per family."""
132
+ plans = create_test_capacity_plans()
133
+ result = reduce_by_family(plans)
134
+
135
+ # Should return only 2 plans - one from family_a and one from family_b
136
+ assert len(result) == 2
137
+
138
+ # Verify we have one from each family
139
+ families = set()
140
+ for plan in result:
141
+ for cluster in plan.candidate_clusters.zonal:
142
+ families.add(cluster.instance.family)
143
+
144
+ assert families == {"family_a", "family_b"}
145
+
146
+
147
+ def test_reduce_by_family_multiple():
148
+ """Test that reduce_by_family with max_results_per_family > 1
149
+ returns multiple plans per family."""
150
+ plans = create_test_capacity_plans()
151
+ result = reduce_by_family(plans, max_results_per_family=2)
152
+
153
+ # Should return 4 plans - two from family_a and two from family_b
154
+ assert len(result) == 4
155
+
156
+ # Count plans per family
157
+ family_counts = {"family_a": 0, "family_b": 0}
158
+ for plan in result:
159
+ for cluster in plan.candidate_clusters.zonal:
160
+ family_counts[cluster.instance.family] += 1
161
+
162
+ # Verify we have exactly 2 from each family
163
+ assert family_counts["family_a"] == 2
164
+ assert family_counts["family_b"] == 2
165
+
166
+
167
+ def test_reduce_by_family_unlimited():
168
+ """Test that reduce_by_family with max_results_per_family > available plans
169
+ returns all plans."""
170
+ plans = create_test_capacity_plans()
171
+ # Set max_results_per_family higher than the number of plans we have
172
+ result = reduce_by_family(plans, max_results_per_family=10)
173
+
174
+ # Should return all 5 plans since we have 3 from family_a and 2 from family_b
175
+ assert len(result) == 5