Merge pull request #210 from Yelp/mpiano/CLUSTERMAN-692

Author: piax93

clusterman/autoscaler/pool_manager.py

@@ -89,8 +89,7 @@ def reload_state(self, **cluster_connector_kwargs) -> None:
         ``Autoscaler.run()``.
         """
         logger.info("Reloading cluster connector state")
-        # TODO: update mypy to avoid having to ignore this error (CLUSTERMAN-692)
-        self.cluster_connector.reload_state(**cluster_connector_kwargs)  # type: ignore
+        self.cluster_connector.reload_state(**cluster_connector_kwargs)
 
         logger.info("Reloading resource groups")
         self._reload_resource_groups()

clusterman/aws/auto_scaling_resource_group.py

@@ -199,7 +199,7 @@ def _get_launch_template_and_overrides(
             template = policy["LaunchTemplate"]["LaunchTemplateSpecification"]
             overrides = policy["LaunchTemplate"]["Overrides"]
         else:
-            logger.warn(f"ASG {self.id} is not using LaunchTemplates, it will be unable to do smart scheduling")
+            logger.warning(f"ASG {self.id} is not using LaunchTemplates, it will be unable to do smart scheduling")
             return None, []
 
         launch_template_name = template["LaunchTemplateName"]

clusterman/aws/aws_resource_group.py

@@ -17,6 +17,7 @@
 from collections import defaultdict
 from socket import gethostbyaddr
 from typing import Any
+from typing import cast
 from typing import Collection
 from typing import List
 from typing import Mapping
@@ -32,6 +33,7 @@
 from clusterman.aws.client import InstanceDict
 from clusterman.aws.markets import get_instance_market
 from clusterman.aws.markets import InstanceMarket
+from clusterman.aws.markets import MarketDict
 from clusterman.interfaces.resource_group import InstanceMetadata
 from clusterman.interfaces.resource_group import ResourceGroup
 
@@ -77,7 +79,7 @@ def get_instance_metadatas(self, state_filter: Optional[Collection[str]] = None)
             if state_filter and aws_state not in state_filter:
                 continue
 
-            instance_market = get_instance_market(instance_dict)
+            instance_market = get_instance_market(cast(MarketDict, instance_dict))
             instance_ip = instance_dict.get("PrivateIpAddress")
             hostname = gethostbyaddr(instance_ip)[0] if instance_ip else None
             is_cordoned = self._is_instance_cordoned(instance_dict)
@@ -127,14 +129,16 @@ def terminate_instances_by_id(self, instance_ids: List[str], batch_size: int = 5
 
         instance_weights = {}
         for instance in ec2_describe_instances(instance_ids):
-            instance_market = get_instance_market(instance)
+            instance_market = get_instance_market(cast(MarketDict, instance))
             if not instance_market.az:
                 logger.warning(
                     f"Instance {instance['InstanceId']} missing AZ info, likely already terminated so skipping",
                 )
                 instance_ids.remove(instance["InstanceId"])
                 continue
-            instance_weights[instance["InstanceId"]] = self.market_weight(get_instance_market(instance))
+            instance_weights[instance["InstanceId"]] = self.market_weight(
+                get_instance_market(cast(MarketDict, instance))
+            )
 
         # AWS API recommends not terminating more than 1000 instances at a time, and to
         # terminate larger numbers in batches
@@ -186,15 +190,21 @@ def _get_instances_by_market(self):
         """Responses from this API call are cached to prevent hitting any AWS request limits"""
         instance_dict: Mapping[InstanceMarket, List[Mapping]] = defaultdict(list)
         for instance in ec2_describe_instances(self.instance_ids):
-            instance_dict[get_instance_market(instance)].append(instance)
+            instance_dict[get_instance_market(cast(MarketDict, instance))].append(instance)
         return instance_dict
 
     @abstractproperty
     def _target_capacity(self):  # pragma: no cover
         pass
 
     @classmethod
-    def load(cls, cluster: str, pool: str, config: Any, **kwargs: Any) -> Mapping[str, "AWSResourceGroup"]:
+    def load(  # type: ignore # (mypy errors with "incompatible signature with supertype")
+        cls,
+        cluster: str,
+        pool: str,
+        config: Any,
+        **kwargs: Any,
+    ) -> Mapping[str, "AWSResourceGroup"]:
         """Load a list of corresponding resource groups
 
         :param cluster: a cluster name

clusterman/aws/client.py

@@ -53,6 +53,7 @@
         "PrivateIpAddress": str,
         "State": InstanceStateDict,
         "LaunchTime": str,
+        "PrivateDnsName": str,
         "Tags": Sequence[Mapping[str, str]],
     },
 )

clusterman/aws/spot_fleet_resource_group.py

@@ -145,7 +145,7 @@ def _target_capacity(self) -> float:
         return self._configuration["SpotFleetRequestConfig"]["TargetCapacity"]
 
     @classmethod
-    def load(
+    def load(  # type: ignore # (mypy errors with "incompatible signature with supertype")
         cls,
         cluster: str,
         pool: str,

clusterman/batch/autoscaler_bootstrap.py

@@ -16,6 +16,9 @@
 import time
 import traceback
 import xmlrpc.client
+from typing import Any
+from typing import cast
+from typing import Mapping
 
 import colorlog
 from yelp_batch.batch import batch_command_line_arguments
@@ -56,7 +59,10 @@ def wait_for_process(
     while True:
         try:
             states = [
-                rpc.supervisor.getProcessInfo(f"{process_name}:{process_name}_{i}")["statename"]
+                cast(
+                    Mapping[str, Any],
+                    rpc.supervisor.getProcessInfo(f"{process_name}:{process_name}_{i}"),
+                )["statename"]
                 for i in range(num_procs)
             ]
         except OSError:

clusterman/config.py

@@ -61,7 +61,7 @@ def setup_config(args: argparse.Namespace) -> None:
     # different region.
     elif cluster:
         aws_region = staticconf.read_string(f"clusters.{cluster}.aws_region", default=None)
-        if pool:
+        if pool and scheduler:
             load_cluster_pool_config(cluster, pool, scheduler, signals_branch_or_tag)
 
     staticconf.DictConfiguration({"aws": {"region": aws_region}})

clusterman/kubernetes/util.py

@@ -17,7 +17,10 @@
 from enum import auto
 from enum import Enum
 from functools import partial
+from typing import Any
+from typing import Hashable
 from typing import List
+from typing import MutableMapping
 from typing import Type
 
 import colorlog
@@ -43,7 +46,9 @@
 DEFAULT_KUBERNETES_DISK_REQUEST = "0"  # Kubernetes doesn't schedule based on disk allocation right now
 KUBERNETES_API_CACHE_SIZE = 16
 KUBERNETES_API_CACHE_TTL = 60
-KUBERNETES_API_CACHE = TTLCache(maxsize=KUBERNETES_API_CACHE_SIZE, ttl=KUBERNETES_API_CACHE_TTL)
+KUBERNETES_API_CACHE: MutableMapping[Hashable, Any] = TTLCache(
+    maxsize=KUBERNETES_API_CACHE_SIZE, ttl=KUBERNETES_API_CACHE_TTL
+)
 VERSION_MATCH_EXPR = re.compile(r"(\W|^)(?P<release>\d+\.\d+(\.\d+)?)(\W|$)")
 logger = colorlog.getLogger(__name__)
 

clusterman/math/piecewise.py

@@ -287,7 +287,8 @@ def _merged_breakpoints(
     bp1 = zip_longest(fn1.breakpoints.items(), [], fillvalue=1)
     yprev0, yprev1 = fn0._initial_value, fn1._initial_value
 
-    for (x, y), fnnum in merge(bp0, bp1):
+    for xy_tuple, fnnum in merge(bp0, bp1):
+        x, y = cast(Tuple[XValue[T], float], xy_tuple)
         if fnnum == 0:
             yield x, y, yprev1
             yprev0 = y

clusterman/math/piecewise_types.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Hashable
 from typing import TypeVar
 
 from typing_extensions import Protocol
@@ -19,15 +20,15 @@
 T = TypeVar("T")
 
 
-class XValueDiff(Protocol[T]):
+class XValueDiff(Protocol[T], Hashable):
     def __mul__(self, other: int) -> "XValueDiff[T]":
         ...
 
     def __truediv__(self, other: "XValueDiff[T]") -> float:
         ...
 
 
-class XValue(Protocol[T]):
+class XValue(Protocol[T], Hashable):
     def __add__(self, other: XValueDiff[T]) -> "XValue[T]":
         ...