Force the planner to use a certain instance type

service_capacity_modeling/capacity_planner.py

@@ -21,6 +21,8 @@
 from service_capacity_modeling.interface import CapacityRequirement
 from service_capacity_modeling.interface import certain_float
 from service_capacity_modeling.interface import DataShape
+from service_capacity_modeling.interface import Drive
+from service_capacity_modeling.interface import Instance
 from service_capacity_modeling.interface import Interval
 from service_capacity_modeling.interface import interval
 from service_capacity_modeling.interface import interval_percentile
@@ -171,25 +173,43 @@ def model_desires_percentiles(
 
     return results, d
 
+def _allow_instance(
+    instance: Instance,
+    allowed_names: Sequence[str],
+    allowed_lifecycles: Sequence[Lifecycle],
+) -> bool:
+    # If the user has explicitly asked for particular families instead
+    # of all lifecycles filter based on that
+    if allowed_names:
+        for name in allowed_names:
+            exact_match = instance.family_separator in name
+            if exact_match and name == instance.name:
+                return True
+            if (not exact_match) and name == instance.family:
+                return True
+        return False
+    # Otherwise consider lifecycle (default)
+    else:
+        if instance.lifecycle not in allowed_lifecycles:
+            return False
+    return True
 
-def _allow_hardware(
-    name: str,
-    lifecycle: Lifecycle,
+def _allow_drive(
+    drive: Drive,
     allowed_names: Sequence[str],
     allowed_lifecycles: Sequence[Lifecycle],
 ) -> bool:
     # If the user has explicitly asked for particular families instead
     # of all lifecycles filter based on that
     if allowed_names:
-        if name not in allowed_names:
+        if drive.name not in allowed_names:
             return False
     # Otherwise consider lifecycle (default)
     else:
-        if lifecycle not in allowed_lifecycles:
+        if drive.lifecycle not in allowed_lifecycles:
             return False
     return True
 
-
 def _regret(
     capacity_plans: Sequence[Tuple[CapacityDesires, CapacityPlan]],
     regret_params: CapacityRegretParameters,
@@ -365,16 +385,16 @@ def _plan_certain(
 
         plans = []
         for instance in hardware.instances.values():
-            if not _allow_hardware(
-                instance.family, instance.lifecycle, instance_families, lifecycles
+            if not _allow_instance(
+                instance, instance_families, lifecycles
             ):
                 continue
 
             if per_instance_mem > instance.ram_gib:
                 continue
 
             for drive in hardware.drives.values():
-                if not _allow_hardware(drive.name, drive.lifecycle, drives, lifecycles):
+                if not _allow_drive(drive, drives, lifecycles):
                     continue
 
                 plan = self._models[model_name].capacity_plan(
@@ -578,4 +598,4 @@ def _sub_models(
 
 
 planner = CapacityPlanner()
-planner.register_group(netflix.models)
+planner.register_group(netflix.models)