refactor

Author: sanderegg

packages/aws-library/src/aws_library/ec2/__init__.py

@@ -17,6 +17,7 @@
     EC2InstanceData,
     EC2InstanceType,
     EC2Tags,
+    GenericResourceValueType,
     Resources,
 )
 
@@ -36,6 +37,7 @@
     "EC2NotConnectedError",
     "EC2RuntimeError",
     "EC2Tags",
+    "GenericResourceValueType",
     "Resources",
     "SimcoreEC2API",
 )

packages/aws-library/src/aws_library/ec2/_models.py

@@ -22,14 +22,14 @@
 from pydantic.config import JsonDict
 from types_aiobotocore_ec2.literals import InstanceStateNameType, InstanceTypeType
 
-GenericResourceValue: TypeAlias = StrictInt | StrictFloat | str
+GenericResourceValueType: TypeAlias = StrictInt | StrictFloat | str
 
 
 class Resources(BaseModel, frozen=True):
     cpus: NonNegativeFloat
     ram: ByteSize
     generic_resources: Annotated[
-        dict[str, GenericResourceValue],
+        dict[str, GenericResourceValueType],
         Field(
             default_factory=dict,
             description=(
@@ -83,7 +83,7 @@ def __add__(self, other: "Resources") -> "Resources":
         Note that only numeric generic resources are added
         Non-numeric generic resources are ignored
         """
-        merged: dict[str, GenericResourceValue] = {}
+        merged: dict[str, GenericResourceValueType] = {}
         keys = set(self.generic_resources) | set(other.generic_resources)
         for k in keys:
             a = self.generic_resources.get(k)
@@ -107,7 +107,7 @@ def __sub__(self, other: "Resources") -> "Resources":
         Note that only numeric generic resources are subtracted
         Non-numeric generic resources are ignored
         """
-        merged: dict[str, GenericResourceValue] = {}
+        merged: dict[str, GenericResourceValueType] = {}
         keys = set(self.generic_resources) | set(other.generic_resources)
         for k in keys:
             a = self.generic_resources.get(k)
@@ -129,7 +129,7 @@ def __sub__(self, other: "Resources") -> "Resources":
     def __hash__(self) -> int:
         """Deterministic hash including cpus, ram (in bytes) and generic_resources."""
         # sort generic_resources items to ensure order-independent hashing
-        generic_items: tuple[tuple[str, GenericResourceValue], ...] = tuple(
+        generic_items: tuple[tuple[str, GenericResourceValueType], ...] = tuple(
             sorted(self.generic_resources.items())
         )
         return hash((self.cpus, self.ram, generic_items))

services/autoscaling/src/simcore_service_autoscaling/modules/cluster_scaling/_utils_computational.py

@@ -1,7 +1,7 @@
 import logging
-from typing import Final
+from typing import Final, cast
 
-from aws_library.ec2 import Resources
+from aws_library.ec2 import GenericResourceValueType, Resources
 from dask_task_models_library.resource_constraints import (
     DaskTaskResources,
     get_ec2_instance_type_from_resources,
@@ -30,7 +30,8 @@ def resources_from_dask_task(task: DaskTask) -> Resources:
     )  # merge with defaults to ensure there is always some minimal resource defined
 
     return Resources.from_flat_dict(
-        task_resources, mapping=DASK_TO_RESOURCE_NAME_MAPPING
+        cast(dict[str, GenericResourceValueType], task_resources),
+        mapping=DASK_TO_RESOURCE_NAME_MAPPING,
     )
 
 

services/autoscaling/src/simcore_service_autoscaling/modules/dask.py

@@ -135,7 +135,6 @@ async def _list_cluster_known_tasks(
     def _list_on_scheduler(
         dask_scheduler: distributed.Scheduler,
     ) -> _DaskClusterTasks:
-
         worker_to_processing_tasks = defaultdict(list)
         unrunnable_tasks = {}
         for task_key, task_state in dask_scheduler.tasks.items():
@@ -319,13 +318,13 @@ async def compute_cluster_total_resources(
                 continue
             # get dask information about resources
             worker_dask_resources = worker_details["resources"]
-            worker_threads = worker_details["nthreads"]
-            worker_dask_resources: dict[str, int | float | str] = {
-                **worker_dask_resources,
-                DASK_WORKER_THREAD_RESOURCE_NAME: worker_threads,
-            }
+            worker_dask_nthreads = worker_details["nthreads"]
             cluster_resources += Resources.from_flat_dict(
-                worker_dask_resources, mapping=DASK_TO_RESOURCE_NAME_MAPPING
+                {
+                    **worker_dask_resources,
+                    DASK_WORKER_THREAD_RESOURCE_NAME: worker_dask_nthreads,
+                },
+                mapping=DASK_TO_RESOURCE_NAME_MAPPING,
             )
 
         return cluster_resources