refactor: clean up unused functions and improve node state checks

Author: sanderegg

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

@@ -17,7 +17,7 @@
 )
 from aws_library.ec2._errors import EC2TooManyInstancesError
 from fastapi import FastAPI
-from models_library.generated_models.docker_rest_api import Node, NodeState
+from models_library.generated_models.docker_rest_api import Node
 from models_library.rabbitmq_messages import ProgressType
 from servicelib.logging_utils import log_catch, log_context
 from servicelib.utils import limited_gather
@@ -66,11 +66,6 @@
 _logger = logging.getLogger(__name__)
 
 
-def _node_not_ready(node: Node) -> bool:
-    assert node.status  # nosec
-    return bool(node.status.state != NodeState.ready)
-
-
 async def _analyze_current_cluster(
     app: FastAPI,
     auto_scaling_mode: BaseAutoscaling,
@@ -141,7 +136,7 @@ async def _analyze_current_cluster(
                     - node_used_resources,
                 )
             )
-        elif auto_scaling_mode.is_instance_drained(instance):
+        elif utils_docker.is_instance_drained(instance):
             all_drained_nodes.append(instance)
         elif await auto_scaling_mode.is_instance_retired(app, instance):
             # it should be drained, but it is not, so we force it to be drained such that it might be re-used if needed
@@ -166,7 +161,9 @@ async def _analyze_current_cluster(
         terminated_instances=[
             NonAssociatedInstance(ec2_instance=i) for i in terminated_ec2_instances
         ],
-        disconnected_nodes=[n for n in docker_nodes if _node_not_ready(n)],
+        disconnected_nodes=[
+            n for n in docker_nodes if not utils_docker.is_node_ready(n)
+        ],
         retired_nodes=retired_nodes,
     )
     _logger.info("current state: %s", f"{cluster!r}")
@@ -343,10 +340,10 @@ async def _sorted_allowed_instance_types(app: FastAPI) -> list[EC2InstanceType]:
         allowed_instance_type_names
     ), "EC2_INSTANCES_ALLOWED_TYPES cannot be empty!"
 
-    allowed_instance_types: list[
-        EC2InstanceType
-    ] = await ec2_client.get_ec2_instance_capabilities(
-        cast(set[InstanceTypeType], set(allowed_instance_type_names))
+    allowed_instance_types: list[EC2InstanceType] = (
+        await ec2_client.get_ec2_instance_capabilities(
+            cast(set[InstanceTypeType], set(allowed_instance_type_names))
+        )
     )
 
     def _as_selection(instance_type: EC2InstanceType) -> int:
@@ -1078,9 +1075,9 @@ async def _notify_based_on_machine_type(
     launch_time_to_tasks: dict[datetime.datetime, list] = collections.defaultdict(list)
     now = datetime.datetime.now(datetime.UTC)
     for instance in instances:
-        launch_time_to_tasks[instance.ec2_instance.launch_time] += (
-            instance.assigned_tasks
-        )
+        launch_time_to_tasks[
+            instance.ec2_instance.launch_time
+        ] += instance.assigned_tasks
 
     for launch_time, tasks in launch_time_to_tasks.items():
         time_since_launch = now - launch_time

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

@@ -8,7 +8,6 @@
 from types_aiobotocore_ec2.literals import InstanceTypeType
 
 from ..models import AssociatedInstance
-from ..utils import utils_docker
 
 
 @dataclass
@@ -71,10 +70,6 @@ async def is_instance_retired(
         app: FastAPI, instance: AssociatedInstance
     ) -> bool: ...
 
-    @staticmethod
-    def is_instance_drained(instance: AssociatedInstance) -> bool:
-        return not utils_docker.is_node_osparc_ready(instance.node)
-
     @staticmethod
     @abstractmethod
     async def try_retire_nodes(app: FastAPI) -> None: ...

services/autoscaling/src/simcore_service_autoscaling/utils/utils_docker.py

@@ -35,6 +35,7 @@
 from types_aiobotocore_ec2.literals import InstanceTypeType
 
 from ..core.settings import ApplicationSettings
+from ..models import AssociatedInstance
 from ..modules.docker import AutoscalingDocker
 
 logger = logging.getLogger(__name__)
@@ -278,24 +279,32 @@ def get_max_resources_from_docker_task(task: Task) -> Resources:
         return Resources(
             cpus=max(
                 (
-                    task.spec.resources.reservations
-                    and task.spec.resources.reservations.nano_cp_us
+                    (
+                        task.spec.resources.reservations
+                        and task.spec.resources.reservations.nano_cp_us
+                    )
                     or 0
                 ),
                 (
-                    task.spec.resources.limits
-                    and task.spec.resources.limits.nano_cp_us
+                    (
+                        task.spec.resources.limits
+                        and task.spec.resources.limits.nano_cp_us
+                    )
                     or 0
                 ),
             )
             / _NANO_CPU,
             ram=TypeAdapter(ByteSize).validate_python(
                 max(
-                    task.spec.resources.reservations
-                    and task.spec.resources.reservations.memory_bytes
+                    (
+                        task.spec.resources.reservations
+                        and task.spec.resources.reservations.memory_bytes
+                    )
                     or 0,
-                    task.spec.resources.limits
-                    and task.spec.resources.limits.memory_bytes
+                    (
+                        task.spec.resources.limits
+                        and task.spec.resources.limits.memory_bytes
+                    )
                     or 0,
                 )
             ),
@@ -382,7 +391,7 @@ async def compute_cluster_used_resources(
     list_of_used_resources = await logged_gather(
         *(compute_node_used_resources(docker_client, node) for node in nodes)
     )
-    counter = collections.Counter({k: 0 for k in list(Resources.model_fields)})
+    counter = collections.Counter(dict.fromkeys(list(Resources.model_fields), 0))
     for result in list_of_used_resources:
         counter.update(result.model_dump())
 
@@ -570,14 +579,14 @@ def get_new_node_docker_tags(
 ) -> dict[DockerLabelKey, str]:
     assert app_settings.AUTOSCALING_NODES_MONITORING  # nosec
     return (
-        {
-            tag_key: "true"
-            for tag_key in app_settings.AUTOSCALING_NODES_MONITORING.NODES_MONITORING_NODE_LABELS
-        }
-        | {
-            tag_key: "true"
-            for tag_key in app_settings.AUTOSCALING_NODES_MONITORING.NODES_MONITORING_NEW_NODES_LABELS
-        }
+        dict.fromkeys(
+            app_settings.AUTOSCALING_NODES_MONITORING.NODES_MONITORING_NODE_LABELS,
+            "true",
+        )
+        | dict.fromkeys(
+            app_settings.AUTOSCALING_NODES_MONITORING.NODES_MONITORING_NEW_NODES_LABELS,
+            "true",
+        )
         | {DOCKER_TASK_EC2_INSTANCE_TYPE_PLACEMENT_CONSTRAINT_KEY: ec2_instance.type}
     )
 
@@ -601,6 +610,10 @@ def is_node_osparc_ready(node: Node) -> bool:
     )
 
 
+def is_instance_drained(instance: AssociatedInstance) -> bool:
+    return not is_node_osparc_ready(instance.node)
+
+
 async def set_node_osparc_ready(
     app_settings: ApplicationSettings,
     docker_client: AutoscalingDocker,
@@ -702,3 +715,8 @@ async def attach_node(
         tags=new_tags,
         available=app_settings.AUTOSCALING_DRAIN_NODES_WITH_LABELS,  # NOTE: full drain sometimes impede on performance
     )
+
+
+def is_node_ready(node: Node) -> bool:
+    assert node.status  # nosec
+    return bool(node.status.state is NodeState.ready)