service-capacity-modeling 0.3.75__tar.gz → 0.3.77__tar.gz

{service_capacity_modeling-0.3.75 → service_capacity_modeling-0.3.77}/PKG-INFO

@@ -1,19 +1,22 @@
 Metadata-Version: 2.4
 Name: service-capacity-modeling
-Version: 0.3.75
+Version: 0.3.77
 Summary: Contains utilities for modeling capacity for pluggable workloads
 Author: Joseph Lynch
 Author-email: josephl@netflix.com
 License: Apache 2.0
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9,<3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: pydantic>2.0
 Requires-Dist: scipy
 Requires-Dist: numpy
-Requires-Dist: importlib_resources; python_version < "3.7"
 Requires-Dist: isodate
 Provides-Extra: aws
 Requires-Dist: boto3; extra == "aws"
@@ -26,6 +29,7 @@ Dynamic: license
 Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
+Dynamic: requires-python
 Dynamic: summary
 
 # Service Capacity Modeling
@@ -43,10 +47,10 @@ remember this repository is public when making changes to it.
 Run the tests:
 ```bash
 # Test the capacity planner on included netflix models
-$ tox -e py38
+$ tox -e py39
 
 # Run a single test with a debugger attached if the test fails
-$ .tox/py38/bin/pytest -n0 -k test_java_heap_heavy --pdb --pdbcls=IPython.terminal.debugger:Pdb
+$ .tox/py39/bin/pytest -n0 -k test_java_heap_heavy --pdb --pdbcls=IPython.terminal.debugger:Pdb
 
 # Verify all type contracts
 $ tox -e mypy
@@ -247,7 +251,7 @@ To contribute to this project:
 2. Write a unit test using `pytest` in the `tests` folder.
 3. Ensure your tests pass via `tox` or debug them with:
 ```
-tox -e py38 -- -k test_<your_functionality> --pdb --pdbcls=IPython.terminal.debugger:Pdb
+tox -e py39 -- -k test_<your_functionality> --pdb --pdbcls=IPython.terminal.debugger:Pdb
 ```
 
 ### Pre-commit / Linting

{service_capacity_modeling-0.3.75 → service_capacity_modeling-0.3.77}/README.md

@@ -13,10 +13,10 @@ remember this repository is public when making changes to it.
 Run the tests:
 ```bash
 # Test the capacity planner on included netflix models
-$ tox -e py38
+$ tox -e py39
 
 # Run a single test with a debugger attached if the test fails
-$ .tox/py38/bin/pytest -n0 -k test_java_heap_heavy --pdb --pdbcls=IPython.terminal.debugger:Pdb
+$ .tox/py39/bin/pytest -n0 -k test_java_heap_heavy --pdb --pdbcls=IPython.terminal.debugger:Pdb
 
 # Verify all type contracts
 $ tox -e mypy
@@ -217,7 +217,7 @@ To contribute to this project:
 2. Write a unit test using `pytest` in the `tests` folder.
 3. Ensure your tests pass via `tox` or debug them with:
 ```
-tox -e py38 -- -k test_<your_functionality> --pdb --pdbcls=IPython.terminal.debugger:Pdb
+tox -e py39 -- -k test_<your_functionality> --pdb --pdbcls=IPython.terminal.debugger:Pdb
 ```
 
 ### Pre-commit / Linting

{service_capacity_modeling-0.3.75 → service_capacity_modeling-0.3.77}/service_capacity_modeling/capacity_planner.py

@@ -190,7 +190,7 @@ def model_desires_percentiles(
 def _set_instance_objects(
     desires: CapacityDesires,
     hardware: Hardware,
-):
+) -> None:
     if desires.current_clusters:
         for zonal_cluster_capacity in desires.current_clusters.zonal:
             if zonal_cluster_capacity.cluster_instance_name in hardware.instances:
@@ -290,7 +290,9 @@ def _regret(
     return plans_by_regret
 
 
-def _add_requirement(requirement, accum):
+def _add_requirement(
+    requirement: CapacityRequirement, accum: Dict[str, Dict[str, List[Interval]]]
+) -> None:
     if requirement.requirement_type not in accum:
         accum[requirement.requirement_type] = {}
 
@@ -305,7 +307,11 @@ def _add_requirement(requirement, accum):
                 requirements[field].append(d)
 
 
-def _merge_models(plans_by_model, zonal_requirements, regional_requirements):
+def _merge_models(
+    plans_by_model: List[List[CapacityPlan]],
+    zonal_requirements: Dict[str, Dict[str, List[Interval]]],
+    regional_requirements: Dict[str, Dict[str, List[Interval]]],
+) -> List[CapacityPlan]:
     capacity_plans = []
     for composed in zip(*filter(lambda x: x, plans_by_model)):
         merged_plans = [functools.reduce(merge_plan, composed)]
@@ -331,10 +337,10 @@ def _in_allowed(inp: str, allowed: Sequence[str]) -> bool:
 class CapacityPlanner:
     def __init__(
         self,
-        default_num_simulations=128,
-        default_num_results=2,
-        default_lifecycles=(Lifecycle.stable, Lifecycle.beta),
-    ):
+        default_num_simulations: int = 128,
+        default_num_results: int = 2,
+        default_lifecycles: Tuple[Lifecycle, ...] = (Lifecycle.stable, Lifecycle.beta),
+    ) -> None:
         self._shapes: HardwareShapes = shapes
         self._models: Dict[str, CapacityModel] = {}
 
@@ -343,11 +349,11 @@ class CapacityPlanner:
         self._default_regret_params = CapacityRegretParameters()
         self._default_lifecycles = default_lifecycles
 
-    def register_group(self, group: Callable[[], Dict[str, CapacityModel]]):
+    def register_group(self, group: Callable[[], Dict[str, CapacityModel]]) -> None:
         for name, model in group().items():
             self.register_model(name, model)
 
-    def register_model(self, name: str, capacity_model: CapacityModel):
+    def register_model(self, name: str, capacity_model: CapacityModel) -> None:
         self._models[name] = capacity_model
 
     @property
@@ -429,16 +435,16 @@ class CapacityPlanner:
 
     def _group_plans_by_percentile(  # pylint: disable=too-many-positional-arguments
         self,
-        drives,
-        extra_model_arguments,
-        instance_families,
-        lifecycles,
-        num_regions,
-        num_results,
-        region,
-        model_percentile_desires,
-        sorted_percentiles,
-    ):
+        drives: Optional[Sequence[str]],
+        extra_model_arguments: Dict[str, Any],
+        instance_families: Optional[Sequence[str]],
+        lifecycles: Sequence[Lifecycle],
+        num_regions: int,
+        num_results: Optional[int],
+        region: str,
+        model_percentile_desires: Any,
+        sorted_percentiles: List[int],
+    ) -> Dict[int, Sequence[CapacityPlan]]:
         percentile_plans = {}
         for index, percentile in enumerate(sorted_percentiles):
             percentile_plan = []
@@ -470,15 +476,15 @@ class CapacityPlanner:
 
     def _mean_plan(  # pylint: disable=too-many-positional-arguments
         self,
-        drives,
-        extra_model_arguments,
-        instance_families,
-        lifecycles,
-        num_regions,
-        num_results,
-        region,
-        model_mean_desires,
-    ):
+        drives: Optional[Sequence[str]],
+        extra_model_arguments: Dict[str, Any],
+        instance_families: Optional[Sequence[str]],
+        lifecycles: Sequence[Lifecycle],
+        num_regions: int,
+        num_results: Optional[int],
+        region: str,
+        model_mean_desires: Dict[str, CapacityDesires],
+    ) -> Sequence[CapacityPlan]:
         mean_plans = []
         for mean_sub_model, mean_sub_desire in model_mean_desires.items():
             mean_sub_plan = self._plan_certain(
@@ -585,7 +591,7 @@ class CapacityPlanner:
         ]
 
     # Calculates the minimum cpu, memory, and network requirements based on desires.
-    def _per_instance_requirements(self, desires) -> Tuple[int, float]:
+    def _per_instance_requirements(self, desires: CapacityDesires) -> Tuple[int, float]:
         # Applications often set fixed reservations of heap or OS memory
         per_instance_mem = (
             desires.data_shape.reserved_instance_app_mem_gib
@@ -621,14 +627,14 @@ class CapacityPlanner:
 
     def generate_scenarios(  # pylint: disable=too-many-positional-arguments
         self,
-        model,
-        region,
-        desires,
-        num_regions,
-        lifecycles,
-        instance_families,
-        drives,
-    ):
+        model: CapacityModel,
+        region: str,
+        desires: CapacityDesires,
+        num_regions: int,
+        lifecycles: Optional[Sequence[Lifecycle]],
+        instance_families: Optional[Sequence[str]],
+        drives: Optional[Sequence[str]],
+    ) -> Generator[Tuple[Instance, Drive, RegionContext], None, None]:
         lifecycles = lifecycles or self._default_lifecycles
         instance_families = instance_families or []
         drives = drives or []
@@ -718,8 +724,8 @@ class CapacityPlanner:
         lifecycles = lifecycles or self._default_lifecycles
 
         # requirement types -> values
-        zonal_requirements: Dict[str, Dict] = {}
-        regional_requirements: Dict[str, Dict] = {}
+        zonal_requirements: Dict[str, Dict[str, List[Interval]]] = {}
+        regional_requirements: Dict[str, Dict[str, List[Interval]]] = {}
 
         regret_clusters_by_model: Dict[
             str, Sequence[Tuple[CapacityPlan, CapacityDesires, float]]
@@ -836,7 +842,7 @@ class CapacityPlanner:
         model_name: str,
         desires: CapacityDesires,
         extra_model_arguments: Dict[str, Any],
-    ):
+    ) -> Generator[Tuple[str, CapacityDesires], None, None]:
         queue: List[Tuple[CapacityDesires, str]] = [(desires, model_name)]
         models_used = []
 

{service_capacity_modeling-0.3.75 → service_capacity_modeling-0.3.77}/service_capacity_modeling/hardware/__init__.py

@@ -5,6 +5,8 @@ import logging
 import os
 from functools import reduce
 from pathlib import Path
+from typing import Any
+from typing import cast
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -20,11 +22,11 @@ from service_capacity_modeling.interface import Service
 logger = logging.getLogger(__name__)
 
 
-def load_pricing(pricing: Dict) -> Pricing:
+def load_pricing(pricing: Dict[str, Any]) -> Pricing:
     return Pricing(regions=pricing)
 
 
-def load_hardware(hardware: Dict) -> Hardware:
+def load_hardware(hardware: Dict[str, Any]) -> Hardware:
     return Hardware(**hardware)
 
 
@@ -103,7 +105,9 @@ def merge_hardware(existing: Hardware, override: Hardware) -> Hardware:
             merged[key] = override_obj.get(key)
         elif isinstance(existing_field, Dict):
             override_field = override_obj.get(key)
-            merged_field = merged.setdefault(key, {})
+            if override_field is None:
+                override_field = {}
+            merged_field = cast(Dict[str, Any], merged.setdefault(key, {}))
 
             existing_keys = existing_field.keys()
             override_keys = override_obj.get(key, {}).keys()
@@ -119,7 +123,7 @@ def merge_hardware(existing: Hardware, override: Hardware) -> Hardware:
     return Hardware(**merged)
 
 
-def merge_pricing(existing: Dict, override: Dict) -> Dict:
+def merge_pricing(existing: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
     merged = existing.copy()
     for region, override_pricing in override.items():
         if region not in merged:
@@ -165,7 +169,7 @@ def load_hardware_from_disk(
     if shape_paths is None:
         shape_paths = []
 
-    combined_pricing: Dict = {}
+    combined_pricing: Dict[str, Any] = {}
 
     logger.debug("Loading pricing from: %s", price_paths)
     for price_path in price_paths:
@@ -185,7 +189,7 @@ def load_hardware_from_disk(
     return price_hardware(hardware=hardware, pricing=pricing)
 
 
-def load_hardware_from_s3(bucket, path) -> GlobalHardware:
+def load_hardware_from_s3(bucket: str, path: str) -> GlobalHardware:
     try:
         # boto is a heavy dependency so we only want to take it if
         # someone will be using it ...
@@ -205,7 +209,7 @@ def load_hardware_from_s3(bucket, path) -> GlobalHardware:
 
 
 class HardwareShapes:
-    def __init__(self):
+    def __init__(self) -> None:
         self._hardware: Optional[GlobalHardware] = None
 
     def load(self, new_hardware: GlobalHardware) -> None:

{service_capacity_modeling-0.3.75 → service_capacity_modeling-0.3.77}/service_capacity_modeling/interface.py

@@ -27,12 +27,12 @@ MEGABIT_IN_BYTES = (1000 * 1000) / 8
 
 
 class ExcludeUnsetModel(BaseModel):
-    def model_dump(self, *args, **kwargs):
+    def model_dump(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
         if "exclude_unset" not in kwargs:
             kwargs["exclude_unset"] = True
         return super().model_dump(*args, **kwargs)
 
-    def model_dump_json(self, *args, **kwargs):
+    def model_dump_json(self, *args: Any, **kwargs: Any) -> str:
         if "exclude_unset" not in kwargs:
             kwargs["exclude_unset"] = True
         return super().model_dump_json(*args, **kwargs)
@@ -44,10 +44,10 @@ class ExcludeUnsetModel(BaseModel):
 
 
 class IntervalModel(str, Enum):
-    def __str__(self):
+    def __str__(self) -> str:
         return str(self.value)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"D({self.value})"
 
     gamma = "gamma"
@@ -71,11 +71,11 @@ class Interval(ExcludeUnsetModel):
     model_config = ConfigDict(frozen=True, protected_namespaces=())
 
     @property
-    def can_simulate(self):
+    def can_simulate(self) -> bool:
         return self.confidence <= 0.99 and self.allow_simulate
 
     @property
-    def minimum(self):
+    def minimum(self) -> float:
         if self.minimum_value is None:
             if self.confidence == 1.0:
                 return self.low * 0.999
@@ -84,17 +84,19 @@ class Interval(ExcludeUnsetModel):
         return self.minimum_value
 
     @property
-    def maximum(self):
+    def maximum(self) -> float:
         if self.maximum_value is None:
             if self.confidence == 1.0:
                 return self.high * 1.001
             return self.high * 2
         return self.maximum_value
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash((type(self),) + tuple(self.__dict__.values()))
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Interval):
+            return False
         return self.__hash__() == other.__hash__()
 
     def scale(self, factor: float) -> Interval:
@@ -264,7 +266,7 @@ class Drive(ExcludeUnsetModel):
         return max(self.block_size_kib, self.group_size_kib)
 
     @property
-    def max_size_gib(self) -> int:
+    def max_size_gib(self) -> float:
         if self.max_scale_size_gib != 0:
             return self.max_scale_size_gib
         else:
@@ -279,7 +281,7 @@ class Drive(ExcludeUnsetModel):
 
     @computed_field(return_type=float)  # type: ignore
     @property
-    def annual_cost(self):
+    def annual_cost(self) -> float:
         size = self.size_gib or 0
         r_ios = self.read_io_per_s or 0
         w_ios = self.write_io_per_s or 0
@@ -382,15 +384,15 @@ class Instance(ExcludeUnsetModel):
     family_separator: str = "."
 
     @property
-    def family(self):
+    def family(self) -> str:
         return self.name.rsplit(self.family_separator, 1)[0]
 
     @property
-    def size(self):
+    def size(self) -> str:
         return self.name.rsplit(self.family_separator, 1)[1]
 
     @property
-    def cores(self):
+    def cores(self) -> int:
         if self.cpu_cores is not None:
             return self.cpu_cores
         return self.cpu // 2
@@ -456,7 +458,7 @@ class Service(ExcludeUnsetModel):
         low=1, mid=10, high=50, confidence=0.9
     )
 
-    def annual_cost_gib(self, data_gib: float = 0):
+    def annual_cost_gib(self, data_gib: float = 0) -> float:
         if isinstance(self.annual_cost_per_gib, float):
             return self.annual_cost_per_gib * data_gib
         else:
@@ -979,7 +981,7 @@ class CapacityRequirement(ExcludeUnsetModel):
     network_mbps: Interval = certain_int(0)
     disk_gib: Interval = certain_int(0)
 
-    context: Dict = {}
+    context: Dict[str, Any] = {}
 
 
 class ClusterCapacity(ExcludeUnsetModel):
@@ -992,7 +994,7 @@ class ClusterCapacity(ExcludeUnsetModel):
     # When provisioning services we might need to signal they
     # should have certain configuration, for example flags that
     # affect durability shut off
-    cluster_params: Dict = {}
+    cluster_params: Dict[str, Any] = {}
 
 
 class ServiceCapacity(ExcludeUnsetModel):
@@ -1003,7 +1005,7 @@ class ServiceCapacity(ExcludeUnsetModel):
     regret_cost: bool = False
     # Often while provisioning cloud services we need to represent
     # parameters to the cloud APIs, use this to inject those from models
-    service_params: Dict = {}
+    service_params: Dict[str, Any] = {}
 
 
 # For services that are provisioned by zone (e.g. Cassandra, EVCache)

{service_capacity_modeling-0.3.75 → service_capacity_modeling-0.3.77}/service_capacity_modeling/models/__init__.py

@@ -20,6 +20,25 @@ from service_capacity_modeling.interface import Platform
 from service_capacity_modeling.interface import QueryPattern
 from service_capacity_modeling.interface import RegionContext
 
+__all__ = [
+    "AccessConsistency",
+    "AccessPattern",
+    "CapacityDesires",
+    "CapacityPlan",
+    "CapacityRegretParameters",
+    "certain_float",
+    "Consistency",
+    "DataShape",
+    "Drive",
+    "FixedInterval",
+    "GlobalConsistency",
+    "Instance",
+    "Platform",
+    "QueryPattern",
+    "RegionContext",
+    "CapacityModel",
+]
+
 __common_regrets__ = frozenset(("spend", "disk", "mem"))
 
 
@@ -85,7 +104,7 @@ class CapacityModel:
 
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
     @staticmethod
@@ -270,7 +289,7 @@ class CapacityModel:
     @staticmethod
     def default_desires(
         user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
-    ):
+    ) -> CapacityDesires:
         """Optional defaults to apply given a user desires
 
         Often users do not know what the on-cpu time of their queries

{service_capacity_modeling-0.3.75 → service_capacity_modeling-0.3.77}/service_capacity_modeling/models/common.py

@@ -3,6 +3,7 @@ import logging
 import math
 import random
 from decimal import Decimal
+from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import List
@@ -552,7 +553,7 @@ def compute_stateful_zone(  # pylint: disable=too-many-positional-arguments
             ratio = ebs_gib / max_size
             count = max(cluster_size(math.ceil(count * ratio)), min_count)
             cost = count * instance.annual_cost
-            ebs_gib = max_size
+            ebs_gib = int(max_size)
 
         read_io, write_io = required_disk_ios(space_gib, count)
         read_io, write_io = (
@@ -597,27 +598,27 @@ def compute_stateful_zone(  # pylint: disable=too-many-positional-arguments
 
 
 # AWS GP2 gives 3 IOS / gb stored.
-def gp2_gib_for_io(read_ios) -> int:
+def gp2_gib_for_io(read_ios: float) -> int:
     return int(max(1, read_ios // 3))
 
 
-def cloud_gib_for_io(drive, total_ios, space_gib) -> int:
+def cloud_gib_for_io(drive: Drive, total_ios: float, space_gib: float) -> int:
     if drive.name == "gp2":
         return gp2_gib_for_io(total_ios)
     else:
-        return space_gib
+        return int(space_gib)
 
 
 class WorkingSetEstimator:
-    def __init__(self):
-        self._cache = {}
+    def __init__(self) -> None:
+        self._cache: Dict[Any, Interval] = {}
 
     def working_set_percent(
         self,
         # latency distributions of the read SLOs versus the drives
         # expressed as scipy rv_continuous objects
-        drive_read_latency_dist,
-        read_slo_latency_dist,
+        drive_read_latency_dist: Any,
+        read_slo_latency_dist: Any,
         # what percentile of disk latency should we target for keeping in
         # memory. Not as this is _increased_ more memory will be reserved
         target_percentile: float = 0.90,
@@ -655,8 +656,8 @@ _working_set_estimator = WorkingSetEstimator()
 def working_set_from_drive_and_slo(
     # latency distributions of the read SLOs versus the drives
     # expressed as scipy rv_continuous objects
-    drive_read_latency_dist,
-    read_slo_latency_dist,
+    drive_read_latency_dist: Any,
+    read_slo_latency_dist: Any,
     estimated_working_set: Optional[Interval] = None,
     # what percentile of disk latency should we target for keeping in
     # memory. Not as this is _increased_ more memory will be reserved
@@ -797,7 +798,7 @@ class DerivedBuffers(BaseModel):
         buffer: Dict[str, Buffer],
         components: List[str],
         component_fallbacks: Optional[Dict[str, List[str]]] = None,
-    ):
+    ) -> "DerivedBuffers":
         expanded_components = _expand_components(components, component_fallbacks)
 
         scale = 1.0
@@ -856,7 +857,7 @@ class RequirementFromCurrentCapacity(BaseModel):
 
     def cpu(self, instance_candidate: Instance) -> int:
         current_cpu_util = self.current_capacity.cpu_utilization.mid / 100
-        current_total_cpu = (
+        current_total_cpu = float(
             self.current_instance.cpu * self.current_capacity.cluster_instance_count.mid
         )
 
@@ -881,11 +882,11 @@ class RequirementFromCurrentCapacity(BaseModel):
 
     @property
     def mem_gib(self) -> float:
-        current_memory_utilization = (
+        current_memory_utilization = float(
             self.current_capacity.memory_utilization_gib.mid
             * self.current_capacity.cluster_instance_count.mid
         )
-        zonal_ram_allocated = (
+        zonal_ram_allocated = float(
             self.current_instance.ram_gib
             * self.current_capacity.cluster_instance_count.mid
         )
@@ -905,11 +906,11 @@ class RequirementFromCurrentCapacity(BaseModel):
 
     @property
     def disk_gib(self) -> int:
-        current_cluster_disk_util_gib = (
+        current_cluster_disk_util_gib = float(
             self.current_capacity.disk_utilization_gib.mid
             * self.current_capacity.cluster_instance_count.mid
         )
-        current_node_disk_gib = (
+        current_node_disk_gib = float(
             self.current_instance.drive.max_size_gib
             if self.current_instance.drive is not None
             else (
@@ -919,7 +920,7 @@ class RequirementFromCurrentCapacity(BaseModel):
             )
         )
 
-        zonal_disk_allocated = (
+        zonal_disk_allocated = float(
             current_node_disk_gib * self.current_capacity.cluster_instance_count.mid
         )
         # These are the desired buffers
@@ -939,11 +940,11 @@ class RequirementFromCurrentCapacity(BaseModel):
 
     @property
     def network_mbps(self) -> int:
-        current_network_utilization = (
+        current_network_utilization = float(
             self.current_capacity.network_utilization_mbps.mid
             * self.current_capacity.cluster_instance_count.mid
         )
-        zonal_network_allocated = (
+        zonal_network_allocated = float(
             self.current_instance.net_mbps
             * self.current_capacity.cluster_instance_count.mid
         )

{service_capacity_modeling-0.3.75 → service_capacity_modeling-0.3.77}/service_capacity_modeling/models/headroom_strategy.py

@@ -19,4 +19,5 @@ class QueuingBasedHeadroomStrategy(HeadroomStrategy):
     """
 
     def calculate_reserved_headroom(self, effective_cpu: float) -> float:
-        return 0.712 / (effective_cpu**0.448)
+        result: float = 0.712 / (effective_cpu**0.448)
+        return result

{service_capacity_modeling-0.3.75 → service_capacity_modeling-0.3.77}/service_capacity_modeling/models/org/netflix/__init__.py

@@ -1,3 +1,6 @@
+from typing import Any
+from typing import Dict
+
 from .aurora import nflx_aurora_capacity_model
 from .cassandra import nflx_cassandra_capacity_model
 from .counter import nflx_counter_capacity_model
@@ -20,7 +23,7 @@ from .wal import nflx_wal_capacity_model
 from .zookeeper import nflx_zookeeper_capacity_model
 
 
-def models():
+def models() -> Dict[str, Any]:
     return {
         "org.netflix.cassandra": nflx_cassandra_capacity_model,
         "org.netflix.stateless-java": nflx_java_app_capacity_model,