service-capacity-modeling 0.3.76__py3-none-any.whl → 0.3.78__py3-none-any.whl

service_capacity_modeling/models/org/netflix/cassandra.py

@@ -79,7 +79,7 @@ def _write_buffer_gib_zone(
     return float(write_buffer_gib) / zones_per_region
 
 
-def _get_cores_from_desires(desires, instance):
+def _get_cores_from_desires(desires: CapacityDesires, instance: Instance) -> int:
     cpu_buffer = buffer_for_components(
         buffers=desires.buffers, components=[BACKGROUND_BUFFER]
     )
@@ -97,7 +97,7 @@ def _get_cores_from_desires(desires, instance):
     return needed_cores
 
 
-def _get_disk_from_desires(desires, copies_per_region):
+def _get_disk_from_desires(desires: CapacityDesires, copies_per_region: int) -> int:
     disk_buffer = buffer_for_components(
         buffers=desires.buffers, components=[BufferComponent.disk]
     )
@@ -116,7 +116,7 @@ def _get_min_count(
     needed_disk_gib: float,
     disk_per_node_gib: float,
     cluster_size_lambda: Callable[[int], int],
-):
+) -> int:
     """
     Compute the minimum number of nodes required for a zone.
 
@@ -158,7 +158,10 @@ def _get_min_count(
 
 
 def _zonal_requirement_for_new_cluster(
-    desires, instance, copies_per_region, zones_per_region
+    desires: CapacityDesires,
+    instance: Instance,
+    copies_per_region: int,
+    zones_per_region: int,
 ) -> CapacityRequirement:
     needed_cores = _get_cores_from_desires(desires, instance)
     needed_disk = _get_disk_from_desires(desires, copies_per_region)
@@ -325,14 +328,14 @@ def _estimate_cassandra_requirement(
     )
 
 
-def _get_current_cluster_size(desires) -> int:
+def _get_current_cluster_size(desires: CapacityDesires) -> int:
     current_capacity = _get_current_capacity(desires)
     if current_capacity is None:
         return 0
     return math.ceil(current_capacity.cluster_instance_count.mid)
 
 
-def _get_current_capacity(desires) -> Optional[CurrentClusterCapacity]:
+def _get_current_capacity(desires: CapacityDesires) -> Optional[CurrentClusterCapacity]:
     current_capacity = (
         None
         if desires.current_clusters is None
@@ -345,7 +348,7 @@ def _get_current_capacity(desires) -> Optional[CurrentClusterCapacity]:
     return current_capacity
 
 
-def _upsert_params(cluster, params):
+def _upsert_params(cluster: Any, params: Dict[str, Any]) -> None:
     if cluster.cluster_params:
         cluster.cluster_params.update(params)
     else:
@@ -374,7 +377,7 @@ def _estimate_cassandra_cluster_zonal(  # pylint: disable=too-many-positional-ar
     desires: CapacityDesires,
     zones_per_region: int = 3,
     copies_per_region: int = 3,
-    require_local_disks: bool = True,
+    require_local_disks: bool = False,
     require_attached_disks: bool = False,
     required_cluster_size: Optional[int] = None,
     max_rps_to_disk: int = 500,
@@ -575,7 +578,7 @@ def _estimate_cassandra_cluster_zonal(  # pylint: disable=too-many-positional-ar
 
 
 # C* LCS has 160 MiB sstables by default and 10 sstables per level
-def _cass_io_per_read(node_size_gib, sstable_size_mb=160):
+def _cass_io_per_read(node_size_gib: float, sstable_size_mb: int = 160) -> int:
     gb = node_size_gib * 1024
     sstables = max(1, gb // sstable_size_mb)
     # 10 sstables per level, plus 1 for L0 (avg)
@@ -585,7 +588,7 @@ def _cass_io_per_read(node_size_gib, sstable_size_mb=160):
     return 2 * levels
 
 
-def _get_base_memory(desires: CapacityDesires):
+def _get_base_memory(desires: CapacityDesires) -> float:
     return (
         desires.data_shape.reserved_instance_app_mem_gib
         + desires.data_shape.reserved_instance_system_mem_gib
@@ -641,7 +644,7 @@ class NflxCassandraArguments(BaseModel):
         " this will be deduced from durability and consistency desires",
     )
     require_local_disks: bool = Field(
-        default=True,
+        default=False,
         description="If local (ephemeral) drives are required",
     )
     require_attached_disks: bool = Field(
@@ -679,8 +682,13 @@ class NflxCassandraArguments(BaseModel):
 
 
 class NflxCassandraCapacityModel(CapacityModel):
+    def __init__(self) -> None:
+        pass
+
     @staticmethod
-    def get_required_cluster_size(tier, extra_model_arguments):
+    def get_required_cluster_size(
+        tier: int, extra_model_arguments: Dict[str, Any]
+    ) -> Optional[int]:
         required_cluster_size: Optional[int] = (
             math.ceil(extra_model_arguments["required_cluster_size"])
             if "required_cluster_size" in extra_model_arguments
@@ -724,7 +732,7 @@ class NflxCassandraCapacityModel(CapacityModel):
             desires, extra_model_arguments.get("copies_per_region", None)
         )
         require_local_disks: bool = extra_model_arguments.get(
-            "require_local_disks", True
+            "require_local_disks", False
         )
         require_attached_disks: bool = extra_model_arguments.get(
             "require_attached_disks", False
@@ -775,7 +783,7 @@ class NflxCassandraCapacityModel(CapacityModel):
         )
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming Cassandra Model"
 
     @staticmethod
@@ -803,7 +811,9 @@ class NflxCassandraCapacityModel(CapacityModel):
         )
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments: Dict[str, Any]):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         acceptable_consistency = {
             None,
             AccessConsistency.best_effort,

service_capacity_modeling/models/org/netflix/counter.py

@@ -81,7 +81,7 @@ class NflxCounterCapacityModel(CapacityModel):
         return counter_app
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming Counter Model"
 
     @staticmethod
@@ -129,7 +129,9 @@ class NflxCounterCapacityModel(CapacityModel):
         return tuple(stores)
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         if user_desires.query_pattern.access_pattern == AccessPattern.latency:
             return CapacityDesires(
                 query_pattern=QueryPattern(

service_capacity_modeling/models/org/netflix/crdb.py

@@ -29,6 +29,7 @@ from service_capacity_modeling.interface import Interval
 from service_capacity_modeling.interface import QueryPattern
 from service_capacity_modeling.interface import RegionContext
 from service_capacity_modeling.interface import Requirements
+from service_capacity_modeling.interface import ZoneClusterCapacity
 from service_capacity_modeling.models import CapacityModel
 from service_capacity_modeling.models.common import buffer_for_components
 from service_capacity_modeling.models.common import compute_stateful_zone
@@ -44,7 +45,7 @@ logger = logging.getLogger(__name__)
 
 # Pebble does Leveled compaction with tieres of size??
 # (FIXME) What does pebble actually do
-def _crdb_io_per_read(node_size_gib, sstable_size_mb=1000):
+def _crdb_io_per_read(node_size_gib: float, sstable_size_mb: int = 1000) -> int:
     gb = node_size_gib * 1024
     sstables = max(1, gb // sstable_size_mb)
     # 10 sstables per level, plus 1 for L0 (avg)
@@ -128,7 +129,7 @@ def _estimate_cockroachdb_requirement(  # noqa=E501 pylint: disable=too-many-pos
     )
 
 
-def _upsert_params(cluster, params):
+def _upsert_params(cluster: ZoneClusterCapacity, params: Dict[str, Any]) -> None:
     if cluster.cluster_params:
         cluster.cluster_params.update(params)
     else:
@@ -329,7 +330,7 @@ class NflxCockroachDBCapacityModel(CapacityModel):
         )
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming CockroachDB Model"
 
     @staticmethod
@@ -337,7 +338,9 @@ class NflxCockroachDBCapacityModel(CapacityModel):
         return NflxCockroachDBArguments.model_json_schema()
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments: Dict[str, Any]):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         acceptable_consistency = {
             None,
             AccessConsistency.linearizable,

service_capacity_modeling/models/org/netflix/ddb.py

@@ -139,11 +139,12 @@ def _get_read_consistency_percentages(
     )
     if total_percent == 0:
         access_consistency = desires.query_pattern.access_consistency.same_region
-        if access_consistency == AccessConsistency.serializable:
+        target_consistency = access_consistency.target_consistency
+        if target_consistency == AccessConsistency.serializable:
             transactional_read_percent = 1.0
             eventual_read_percent = 0.0
             strong_read_percent = 0.0
-        elif access_consistency in (
+        elif target_consistency in (
             AccessConsistency.read_your_writes,
             AccessConsistency.linearizable,
         ):
@@ -184,7 +185,8 @@ def _get_write_consistency_percentages(
     total_percent = transactional_write_percent + non_transactional_write_percent
     if total_percent == 0:
         access_consistency = desires.query_pattern.access_consistency.same_region
-        if access_consistency == AccessConsistency.serializable:
+        target_consistency = access_consistency.target_consistency
+        if target_consistency == AccessConsistency.serializable:
             transactional_write_percent = 1.0
             non_transactional_write_percent = 0.0
         else:
@@ -547,7 +549,7 @@ class NflxDynamoDBCapacityModel(CapacityModel):
         )
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming DynamoDB Model"
 
     @staticmethod
@@ -559,7 +561,9 @@ class NflxDynamoDBCapacityModel(CapacityModel):
         return False
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments: Dict[str, Any]):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         acceptable_consistency = {
             "same_region": {
                 None,

service_capacity_modeling/models/org/netflix/elasticsearch.py

@@ -58,7 +58,7 @@ def _target_rf(desires: CapacityDesires, user_copies: Optional[int]) -> int:
 # segments of 512 megs per
 # https://lucene.apache.org/core/8_1_0/core/org/apache/lucene/index/TieredMergePolicy.html#setSegmentsPerTier(double)
 # (FIXME) Verify what elastic merge actually does
-def _es_io_per_read(node_size_gib, segment_size_mb=512):
+def _es_io_per_read(node_size_gib: float, segment_size_mb: int = 512) -> int:
     size_mib = node_size_gib * 1024
     segments = max(1, size_mib // segment_size_mb)
     # 10 segments per tier, plus 1 for L0 (avg)
@@ -75,7 +75,7 @@ def _estimate_elasticsearch_requirement(  # noqa: E501 pylint: disable=too-many-
     max_rps_to_disk: int,
     zones_in_region: int = 3,
     copies_per_region: int = 3,
-    jvm_memory_overhead=1.2,
+    jvm_memory_overhead: float = 1.2,
 ) -> CapacityRequirement:
     """Estimate the capacity required for one zone given a regional desire
 
@@ -153,7 +153,7 @@ def _estimate_elasticsearch_requirement(  # noqa: E501 pylint: disable=too-many-
     )
 
 
-def _upsert_params(cluster, params):
+def _upsert_params(cluster: ZoneClusterCapacity, params: Dict[str, Any]) -> None:
     if cluster.cluster_params:
         cluster.cluster_params.update(params)
     else:
@@ -194,7 +194,7 @@ class NflxElasticsearchDataCapacityModel(CapacityModel):
 
     @staticmethod
     def default_desires(
-        user_desires, extra_model_arguments: Dict[str, Any]
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
     ) -> CapacityDesires:
         desires = CapacityModel.default_desires(user_desires, extra_model_arguments)
         desires.buffers = NflxElasticsearchDataCapacityModel.default_buffers()
@@ -452,7 +452,7 @@ class NflxElasticsearchCapacityModel(CapacityModel):
         return None
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming Elasticsearch Model"
 
     @staticmethod
@@ -482,7 +482,9 @@ class NflxElasticsearchCapacityModel(CapacityModel):
         return NflxElasticsearchArguments.model_json_schema()
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments: Dict[str, Any]):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         acceptable_consistency = {
             AccessConsistency.best_effort,
             AccessConsistency.eventual,

service_capacity_modeling/models/org/netflix/entity.py

@@ -48,7 +48,7 @@ class NflxEntityCapacityModel(CapacityModel):
         return entity_app
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming Entity Model"
 
     @staticmethod
@@ -76,7 +76,7 @@ class NflxEntityCapacityModel(CapacityModel):
                         / 1024**3
                     )
                 else:
-                    item_size_gib = 10 / 1024**2
+                    item_size_gib = 10 / 1024**2  # type: ignore[unreachable]
                 item_count = user_desires.data_shape.estimated_state_size_gib.scale(
                     1 / item_size_gib
                 )
@@ -102,7 +102,9 @@ class NflxEntityCapacityModel(CapacityModel):
         )
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         if user_desires.query_pattern.access_pattern == AccessPattern.latency:
             return CapacityDesires(
                 query_pattern=QueryPattern(

service_capacity_modeling/models/org/netflix/evcache.py

@@ -69,7 +69,9 @@ def calculate_read_cpu_time_evcache_ms(read_size_bytes: float) -> float:
     return max(read_latency_ms, 0.005)
 
 
-def calculate_spread_cost(cluster_size: int, max_cost=100000, min_cost=0.0) -> float:
+def calculate_spread_cost(
+    cluster_size: int, max_cost: float = 100000, min_cost: float = 0.0
+) -> float:
     if cluster_size > 10:
         return min_cost
     if cluster_size < 2:
@@ -82,14 +84,14 @@ def calculate_vitals_for_capacity_planner(
     instance: Instance,
     current_memory_gib: float,
     current_disk_gib: float,
-):
+) -> Tuple[float, float, float, float]:
     # First calculate assuming new deployment
     needed_cores = normalize_cores(
         core_count=sqrt_staffed_cores(desires),
         target_shape=instance,
         reference_shape=desires.reference_shape,
     )
-    needed_network_mbps = simple_network_mbps(desires)
+    needed_network_mbps = float(simple_network_mbps(desires))
     needed_memory_gib = current_memory_gib
     needed_disk_gib = current_disk_gib
 
@@ -110,8 +112,8 @@ def calculate_vitals_for_capacity_planner(
         target_shape=instance,
         reference_shape=current_capacity.cluster_instance,
     )
-    needed_network_mbps = requirements.network_mbps
-    needed_disk_gib = (
+    needed_network_mbps = float(requirements.network_mbps)
+    needed_disk_gib = float(
         requirements.disk_gib if current_capacity.cluster_drive is not None else 0.0
     )
     needed_memory_gib = requirements.mem_gib
@@ -196,7 +198,7 @@ def _estimate_evcache_requirement(
     )
 
 
-def _upsert_params(cluster, params):
+def _upsert_params(cluster: Any, params: Dict[str, Any]) -> None:
     if cluster.cluster_params:
         cluster.cluster_params.update(params)
     else:
@@ -263,7 +265,7 @@ def _estimate_evcache_cluster_zonal(  # noqa: C901,E501 pylint: disable=too-many
     # larger to account for additional overhead. Note that the
     # reserved_instance_system_mem_gib has a base of 1 GiB OSMEM so this
     # just represents the variable component
-    def reserve_memory(instance_mem_gib):
+    def reserve_memory(instance_mem_gib: float) -> float:
         # (Joey) From the chart it appears to be about a 3% overhead for
         # OS memory.
         variable_os = int(instance_mem_gib * 0.03)
@@ -439,7 +441,7 @@ class NflxEVCacheCapacityModel(CapacityModel):
         )
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming EVCache (memcached) Model"
 
     @staticmethod
@@ -447,7 +449,9 @@ class NflxEVCacheCapacityModel(CapacityModel):
         return NflxEVCacheArguments.model_json_schema()
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments: Dict[str, Any]):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         acceptable_consistency = {
             AccessConsistency.best_effort,
             AccessConsistency.never,

service_capacity_modeling/models/org/netflix/graphkv.py

@@ -45,7 +45,7 @@ class NflxGraphKVCapacityModel(CapacityModel):
         return graphkv_app
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming Graph Abstraction"
 
     @staticmethod
@@ -89,7 +89,7 @@ class NflxGraphKVCapacityModel(CapacityModel):
                         / 1024**3
                     )
                 else:
-                    item_size_gib = 1 / 1024**2
+                    item_size_gib = 1 / 1024**2  # type: ignore[unreachable]
                 item_count = user_desires.data_shape.estimated_state_size_gib.scale(
                     1 / item_size_gib
                 )
@@ -102,7 +102,9 @@ class NflxGraphKVCapacityModel(CapacityModel):
         return (("org.netflix.key-value", _modify_kv_desires),)
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         if user_desires.query_pattern.access_pattern == AccessPattern.latency:
             return CapacityDesires(
                 query_pattern=QueryPattern(

service_capacity_modeling/models/org/netflix/iso_date_math.py

@@ -1,8 +1,9 @@
 import math
 from datetime import timedelta
 from decimal import Decimal
+from typing import Union
 
-import isodate
+import isodate  # type: ignore
 from isodate import parse_duration
 
 DURATION_1M = 2592000
@@ -11,24 +12,26 @@ D_YEAR = Decimal(DURATION_1Y)
 D_MONTH = Decimal(DURATION_1M)
 
 
-def _iso_to_proto_duration(iso_duration: str):
+def _iso_to_proto_duration(iso_duration: str) -> str:
     parsed = parse_duration(iso_duration)
     return f"{int(parsed.total_seconds())}s"
 
 
-def iso_to_seconds(iso_duration: str, unlimited=0) -> int:
+def iso_to_seconds(iso_duration: str, unlimited: int = 0) -> int:
     if iso_duration == "unlimited":
         return unlimited
-    parsed = isodate.parse_duration(iso_duration)
+    parsed: Union[timedelta, isodate.Duration] = isodate.parse_duration(iso_duration)
     # isodate package parse_duration returns either timedelta OR Duration
     if isinstance(parsed, isodate.Duration):
-        return math.ceil(
-            parsed.years * D_YEAR
-            + parsed.months * D_MONTH
-            + Decimal(parsed.tdelta.total_seconds())
+        return int(
+            math.ceil(
+                parsed.years * D_YEAR
+                + parsed.months * D_MONTH
+                + Decimal(parsed.tdelta.total_seconds())
+            )
         )
     return int(parsed.total_seconds())
 
 
-def _iso_to_timedelta(iso_duration: str, unlimited=0) -> timedelta:
+def _iso_to_timedelta(iso_duration: str, unlimited: int = 0) -> timedelta:
     return timedelta(seconds=iso_to_seconds(iso_duration, unlimited))

service_capacity_modeling/models/org/netflix/kafka.py

@@ -64,7 +64,9 @@ def _get_current_zonal_cluster(
     )
 
 
-def _is_same_instance_family(cluster, target_family):
+def _is_same_instance_family(
+    cluster: Optional[CurrentZoneClusterCapacity], target_family: str
+) -> bool:
     """Check if cluster has a different instance family than the target."""
     if not cluster or not cluster.cluster_instance:
         return False
@@ -213,14 +215,16 @@ def _estimate_kafka_requirement(  # pylint: disable=too-many-positional-argument
     )
 
 
-def _upsert_params(cluster, params):
+def _upsert_params(cluster: Any, params: Dict[str, Any]) -> None:
     if cluster.cluster_params:
         cluster.cluster_params.update(params)
     else:
         cluster.cluster_params = params
 
 
-def _kafka_read_io(rps, io_size_kib, size_gib, recovery_seconds: int) -> float:
+def _kafka_read_io(
+    rps: float, io_size_kib: float, size_gib: float, recovery_seconds: int
+) -> float:
     # Get enough disk read IO capacity for some reads
     # In practice we have cache reducing this by 99% or more
     read_ios = rps * 0.05
@@ -229,7 +233,7 @@ def _kafka_read_io(rps, io_size_kib, size_gib, recovery_seconds: int) -> float:
     size_kib = size_gib * 0.5 * (1024 * 1024)
     recovery_ios = max(1, size_kib / io_size_kib) / recovery_seconds
     # Leave 50% headroom for read IOs since generally we will hit cache
-    return (read_ios + int(round(recovery_ios))) * 1.5
+    return float((read_ios + int(round(recovery_ios))) * 1.5)
 
 
 # pylint: disable=too-many-locals
@@ -239,7 +243,7 @@ def _estimate_kafka_cluster_zonal(  # noqa: C901
     instance: Instance,
     drive: Drive,
     desires: CapacityDesires,
-    hot_retention_seconds,
+    hot_retention_seconds: float,
     zones_per_region: int = 3,
     copies_per_region: int = 2,
     require_local_disks: bool = False,
@@ -532,7 +536,7 @@ class NflxKafkaCapacityModel(CapacityModel):
         )
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming Kafka Model"
 
     @staticmethod
@@ -545,7 +549,9 @@ class NflxKafkaCapacityModel(CapacityModel):
         return ("gp3",)
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments: Dict[str, Any]):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         # Default to 10MiB/s and a single reader
         concurrent_readers = max(
             1, int(user_desires.query_pattern.estimated_read_per_second.mid)

service_capacity_modeling/models/org/netflix/key_value.py

@@ -48,7 +48,7 @@ class NflxKeyValueCapacityModel(CapacityModel):
         return kv_app
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming Key-Value Model"
 
     @staticmethod
@@ -119,7 +119,9 @@ class NflxKeyValueCapacityModel(CapacityModel):
             return (("org.netflix.cassandra", lambda x: x),)
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         if user_desires.query_pattern.access_pattern == AccessPattern.latency:
             return CapacityDesires(
                 query_pattern=QueryPattern(

service_capacity_modeling/models/org/netflix/postgres.py

@@ -45,7 +45,7 @@ class NflxPostgresCapacityModel(CapacityModel):
         return plan
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Postgres Model"
 
     @staticmethod
@@ -53,7 +53,9 @@ class NflxPostgresCapacityModel(CapacityModel):
         return Platform.aurora_postgres, Platform.amd64
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         return CapacityDesires(
             query_pattern=QueryPattern(
                 access_pattern=AccessPattern.latency,

service_capacity_modeling/models/org/netflix/rds.py

@@ -1,6 +1,7 @@
 import logging
 import math
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import Optional
 
@@ -83,7 +84,9 @@ def _estimate_rds_requirement(
 
 # MySQL default block size is 16KiB, PostGreSQL is 8KiB Number of reads for B-Tree
 # are given by log of total pages to the base of B-Tree fan out factor
-def _rds_required_disk_ios(disk_size_gib: int, db_type: str, btree_fan_out: int = 100):
+def _rds_required_disk_ios(
+    disk_size_gib: int, db_type: str, btree_fan_out: int = 100
+) -> float:
     disk_size_kb = disk_size_gib * 1024 * 1024
     if db_type == "postgres":
         default_block_size = 8  # KiB
@@ -101,8 +104,8 @@ def _compute_rds_region(  # pylint: disable=too-many-positional-arguments
     needed_disk_gib: int,
     needed_memory_gib: int,
     needed_network_mbps: float,
-    required_disk_ios,
-    required_disk_space,
+    required_disk_ios: Callable[[int], float],
+    required_disk_space: Callable[[int], int],
     reference_shape: Instance,
 ) -> Optional[RegionClusterCapacity]:
     """Computes a regional cluster of a RDS service
@@ -239,7 +242,7 @@ class NflxRDSCapacityModel(CapacityModel):
         )
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix RDS Cluster Model"
 
     @staticmethod
@@ -247,7 +250,9 @@ class NflxRDSCapacityModel(CapacityModel):
         return NflxRDSArguments.model_json_schema()
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         return CapacityDesires(
             query_pattern=QueryPattern(
                 access_pattern=AccessPattern.latency,

service_capacity_modeling/models/org/netflix/stateless_java.py

@@ -176,7 +176,7 @@ class NflxJavaAppCapacityModel(CapacityModel):
         )
 
     @staticmethod
-    def description():
+    def description() -> str:
         return "Netflix Streaming Java App Model"
 
     @staticmethod
@@ -196,7 +196,9 @@ class NflxJavaAppCapacityModel(CapacityModel):
         return regret
 
     @staticmethod
-    def default_desires(user_desires, extra_model_arguments):
+    def default_desires(
+        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
+    ) -> CapacityDesires:
         if user_desires.query_pattern.access_pattern == AccessPattern.latency:
             return CapacityDesires(
                 query_pattern=QueryPattern(