♻️ Autoscaling: one docker client (ITISFoundation#3654)

Author: sanderegg

services/autoscaling/src/simcore_service_autoscaling/api/health.py

@@ -10,6 +10,7 @@
 from fastapi.responses import PlainTextResponse
 from pydantic import BaseModel
 
+from ..modules.docker import get_docker_client
 from ..modules.rabbitmq import get_rabbitmq_client
 from .dependencies.application import get_app
 
@@ -30,6 +31,7 @@ class _ComponentStatus(BaseModel):
 class _StatusGet(BaseModel):
     rabbitmq: _ComponentStatus
     ec2: _ComponentStatus
+    docker: _ComponentStatus
 
 
 @router.get("/status", include_in_schema=True, response_model=_StatusGet)
@@ -48,4 +50,8 @@ async def get_status(app: FastAPI = Depends(get_app)) -> _StatusGet:
             if app.state.ec2_client
             else False,
         ),
+        docker=_ComponentStatus(
+            is_enabled=bool(app.state.docker_client),
+            is_responsive=await get_docker_client(app).ping(),
+        ),
     )

services/autoscaling/src/simcore_service_autoscaling/core/application.py

@@ -11,6 +11,7 @@
 )
 from ..api.routes import setup_api_routes
 from ..dynamic_scaling import setup as setup_background_task
+from ..modules.docker import setup as setup_docker
 from ..modules.ec2 import setup as setup_ec2
 from ..modules.rabbitmq import setup as setup_rabbitmq
 from .settings import ApplicationSettings
@@ -37,6 +38,7 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
 
     # PLUGINS SETUP
     setup_api_routes(app)
+    setup_docker(app)
     setup_rabbitmq(app)
     setup_ec2(app)
     # autoscaler background task

services/autoscaling/src/simcore_service_autoscaling/dynamic_scaling_core.py

@@ -7,12 +7,12 @@
 from pydantic import parse_obj_as
 from types_aiobotocore_ec2.literals import InstanceTypeType
 
-from . import utils_docker
 from ._meta import VERSION
 from .core.errors import Ec2InstanceNotFoundError
 from .core.settings import ApplicationSettings
+from .modules.docker import get_docker_client
 from .modules.ec2 import get_ec2_client
-from .utils import ec2, rabbitmq
+from .utils import ec2, rabbitmq, utils_docker
 
 logger = logging.getLogger(__name__)
 
@@ -28,25 +28,28 @@ async def check_dynamic_resources(app: FastAPI) -> None:
     assert app_settings.AUTOSCALING_NODES_MONITORING  # nosec
 
     # 1. get monitored nodes information and resources
+    docker_client = get_docker_client(app)
     monitored_nodes = await utils_docker.get_monitored_nodes(
-        node_labels=app_settings.AUTOSCALING_NODES_MONITORING.NODES_MONITORING_NODE_LABELS
+        docker_client,
+        node_labels=app_settings.AUTOSCALING_NODES_MONITORING.NODES_MONITORING_NODE_LABELS,
     )
 
     cluster_total_resources = await utils_docker.compute_cluster_total_resources(
         monitored_nodes
     )
     logger.info("%s", f"{cluster_total_resources=}")
     cluster_used_resources = await utils_docker.compute_cluster_used_resources(
-        monitored_nodes
+        docker_client, monitored_nodes
     )
     logger.info("%s", f"{cluster_used_resources=}")
 
     # 2. Remove nodes that are gone
-    await utils_docker.remove_monitored_down_nodes(monitored_nodes)
+    await utils_docker.remove_monitored_down_nodes(docker_client, monitored_nodes)
 
     # 3. Scale up nodes if there are pending tasks
     pending_tasks = await utils_docker.pending_service_tasks_with_insufficient_resources(
-        service_labels=app_settings.AUTOSCALING_NODES_MONITORING.NODES_MONITORING_SERVICE_LABELS
+        docker_client,
+        service_labels=app_settings.AUTOSCALING_NODES_MONITORING.NODES_MONITORING_SERVICE_LABELS,
     )
     await rabbitmq.post_state_message(
         app,
@@ -106,8 +109,11 @@ async def check_dynamic_resources(app: FastAPI) -> None:
             # NOTE: new_instance_dns_name is of type ip-123-23-23-3.ec2.internal and we need only the first part
             if match := re.match(_EC2_INTERNAL_DNS_RE, new_instance_dns_name):
                 new_instance_dns_name = match.group(1)
-                new_node = await utils_docker.wait_for_node(new_instance_dns_name)
+                new_node = await utils_docker.wait_for_node(
+                    docker_client, new_instance_dns_name
+                )
                 await utils_docker.tag_node(
+                    docker_client,
                     new_node,
                     tags={
                         tag_key: "true"

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

@@ -0,0 +1,46 @@
+import logging
+from typing import cast
+
+import aiodocker
+from fastapi import FastAPI
+from tenacity._asyncio import AsyncRetrying
+from tenacity.before_sleep import before_sleep_log
+from tenacity.stop import stop_after_delay
+from tenacity.wait import wait_random_exponential
+
+logger = logging.getLogger(__name__)
+
+
+class AutoscalingDocker(aiodocker.Docker):
+    async def ping(self) -> bool:
+        try:
+            await self.version()
+            return True
+        except Exception:  # pylint: disable=broad-except
+            return False
+
+
+def setup(app: FastAPI) -> None:
+    async def on_startup() -> None:
+        app.state.docker_client = client = AutoscalingDocker()
+
+        async for attempt in AsyncRetrying(
+            reraise=True,
+            stop=stop_after_delay(120),
+            wait=wait_random_exponential(max=30),
+            before_sleep=before_sleep_log(logger, logging.WARNING),
+        ):
+            with attempt:
+                # this will raise if the connection is not working
+                await client.version()
+
+    async def on_shutdown() -> None:
+        if app.state.docker_client:
+            await cast(AutoscalingDocker, app.state.docker_client).close()
+
+    app.add_event_handler("startup", on_startup)
+    app.add_event_handler("shutdown", on_shutdown)
+
+
+def get_docker_client(app: FastAPI) -> AutoscalingDocker:
+    return cast(AutoscalingDocker, app.state.docker_client)

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

@@ -8,7 +8,6 @@
 import re
 from typing import Final
 
-import aiodocker
 from models_library.docker import DockerLabelKey
 from models_library.generated_models.docker_rest_api import (
     Node,
@@ -24,7 +23,8 @@
 from tenacity.stop import stop_after_delay
 from tenacity.wait import wait_fixed
 
-from .models import Resources
+from ..models import Resources
+from ..modules.docker import AutoscalingDocker
 
 logger = logging.getLogger(__name__)
 _NANO_CPU: Final[float] = 10**9
@@ -40,18 +40,21 @@
 _TIMEOUT_WAITING_FOR_NODES_S: Final[int] = 5 * _MINUTE
 
 
-async def get_monitored_nodes(node_labels: list[DockerLabelKey]) -> list[Node]:
-    async with aiodocker.Docker() as docker:
-        nodes = parse_obj_as(
-            list[Node],
-            await docker.nodes.list(
-                filters={"node.label": [f"{label}=true" for label in node_labels]}
-            ),
-        )
+async def get_monitored_nodes(
+    docker_client: AutoscalingDocker, node_labels: list[DockerLabelKey]
+) -> list[Node]:
+    nodes = parse_obj_as(
+        list[Node],
+        await docker_client.nodes.list(
+            filters={"node.label": [f"{label}=true" for label in node_labels]}
+        ),
+    )
     return nodes
 
 
-async def remove_monitored_down_nodes(nodes: list[Node]) -> list[Node]:
+async def remove_monitored_down_nodes(
+    docker_client: AutoscalingDocker, nodes: list[Node]
+) -> list[Node]:
     """removes docker nodes that are in the down state"""
 
     def _check_if_node_is_removable(node: Node) -> bool:
@@ -69,15 +72,15 @@ def _check_if_node_is_removable(node: Node) -> bool:
         return False
 
     nodes_that_need_removal = [n for n in nodes if _check_if_node_is_removable(n)]
-    async with aiodocker.Docker() as docker:
-        for node in nodes_that_need_removal:
-            assert node.ID  # nosec
-            with log_context(logger, logging.INFO, msg=f"remove {node.ID=}"):
-                await docker.nodes.remove(node_id=node.ID)
+    for node in nodes_that_need_removal:
+        assert node.ID  # nosec
+        with log_context(logger, logging.INFO, msg=f"remove {node.ID=}"):
+            await docker_client.nodes.remove(node_id=node.ID)
     return nodes_that_need_removal
 
 
 async def pending_service_tasks_with_insufficient_resources(
+    docker_client: AutoscalingDocker,
     service_labels: list[DockerLabelKey],
 ) -> list[Task]:
     """
@@ -88,16 +91,15 @@ async def pending_service_tasks_with_insufficient_resources(
     - have an error message with "insufficient resources"
     - are not scheduled on any node
     """
-    async with aiodocker.Docker() as docker:
-        tasks = parse_obj_as(
-            list[Task],
-            await docker.tasks.list(
-                filters={
-                    "desired-state": "running",
-                    "label": service_labels,
-                }
-            ),
-        )
+    tasks = parse_obj_as(
+        list[Task],
+        await docker_client.tasks.list(
+            filters={
+                "desired-state": "running",
+                "label": service_labels,
+            }
+        ),
+    )
 
     def _is_task_waiting_for_resources(task: Task) -> bool:
         # NOTE: https://docs.docker.com/engine/swarm/how-swarm-mode-works/swarm-task-states/
@@ -171,36 +173,38 @@ def get_max_resources_from_docker_task(task: Task) -> Resources:
     return Resources(cpus=0, ram=ByteSize(0))
 
 
-async def compute_node_used_resources(node: Node) -> Resources:
+async def compute_node_used_resources(
+    docker_client: AutoscalingDocker,
+    node: Node,
+) -> Resources:
     cluster_resources_counter = collections.Counter({"ram": 0, "cpus": 0})
-    async with aiodocker.Docker() as docker:
-        all_tasks_on_node = parse_obj_as(
-            list[Task], await docker.tasks.list(filters={"node": node.ID})
-        )
-        for task in all_tasks_on_node:
-            assert task.Status  # nosec
-            if (
-                task.Status.State in _TASK_STATUS_WITH_ASSIGNED_RESOURCES
-                and task.Spec
-                and task.Spec.Resources
-                and task.Spec.Resources.Reservations
-            ):
-                task_reservations = task.Spec.Resources.Reservations.dict(
-                    exclude_none=True
-                )
-                cluster_resources_counter.update(
-                    {
-                        "ram": task_reservations.get("MemoryBytes", 0),
-                        "cpus": task_reservations.get("NanoCPUs", 0) / _NANO_CPU,
-                    }
-                )
+    all_tasks_on_node = parse_obj_as(
+        list[Task], await docker_client.tasks.list(filters={"node": node.ID})
+    )
+    for task in all_tasks_on_node:
+        assert task.Status  # nosec
+        if (
+            task.Status.State in _TASK_STATUS_WITH_ASSIGNED_RESOURCES
+            and task.Spec
+            and task.Spec.Resources
+            and task.Spec.Resources.Reservations
+        ):
+            task_reservations = task.Spec.Resources.Reservations.dict(exclude_none=True)
+            cluster_resources_counter.update(
+                {
+                    "ram": task_reservations.get("MemoryBytes", 0),
+                    "cpus": task_reservations.get("NanoCPUs", 0) / _NANO_CPU,
+                }
+            )
     return Resources.parse_obj(dict(cluster_resources_counter))
 
 
-async def compute_cluster_used_resources(nodes: list[Node]) -> Resources:
+async def compute_cluster_used_resources(
+    docker_client: AutoscalingDocker, nodes: list[Node]
+) -> Resources:
     """Returns the total amount of resources (reservations) used on each of the given nodes"""
     list_of_used_resources = await logged_gather(
-        *(compute_node_used_resources(node) for node in nodes)
+        *(compute_node_used_resources(docker_client, node) for node in nodes)
     )
     counter = collections.Counter({k: 0 for k in Resources.__fields__.keys()})
     for result in list_of_used_resources:
@@ -243,32 +247,37 @@ async def get_docker_swarm_join_bash_command() -> str:
     before_sleep=before_sleep_log(logger, logging.WARNING),
     wait=wait_fixed(5),
 )
-async def wait_for_node(node_name: str) -> Node:
-    async with aiodocker.Docker() as docker:
-        list_of_nodes = await docker.nodes.list(filters={"name": node_name})
+async def wait_for_node(
+    docker_client: AutoscalingDocker,
+    node_name: str,
+) -> Node:
+    list_of_nodes = await docker_client.nodes.list(filters={"name": node_name})
     if not list_of_nodes:
         raise TryAgain
     return parse_obj_as(Node, list_of_nodes[0])
 
 
 async def tag_node(
-    node: Node, *, tags: dict[DockerLabelKey, str], available: bool
+    docker_client: AutoscalingDocker,
+    node: Node,
+    *,
+    tags: dict[DockerLabelKey, str],
+    available: bool,
 ) -> None:
     with log_context(
         logger, logging.DEBUG, msg=f"tagging {node.ID=} with {tags=} and {available=}"
     ):
-        async with aiodocker.Docker() as docker:
-            assert node.ID  # nosec
-            assert node.Version  # nosec
-            assert node.Version.Index  # nosec
-            assert node.Spec  # nosec
-            assert node.Spec.Role  # nosec
-            await docker.nodes.update(
-                node_id=node.ID,
-                version=node.Version.Index,
-                spec={
-                    "Availability": "active" if available else "drain",
-                    "Labels": tags,
-                    "Role": node.Spec.Role.value,
-                },
-            )
+        assert node.ID  # nosec
+        assert node.Version  # nosec
+        assert node.Version.Index  # nosec
+        assert node.Spec  # nosec
+        assert node.Spec.Role  # nosec
+        await docker_client.nodes.update(
+            node_id=node.ID,
+            version=node.Version.Index,
+            spec={
+                "Availability": "active" if available else "drain",
+                "Labels": tags,
+                "Role": node.Spec.Role.value,
+            },
+        )

services/autoscaling/tests/unit/conftest.py

@@ -16,6 +16,7 @@
     Mapping,
     Optional,
     Union,
+    cast,
 )
 
 import aiodocker
@@ -39,6 +40,7 @@
 from simcore_service_autoscaling.core.application import create_app
 from simcore_service_autoscaling.core.settings import ApplicationSettings, EC2Settings
 from simcore_service_autoscaling.models import SimcoreServiceDockerLabelKeys
+from simcore_service_autoscaling.modules.docker import AutoscalingDocker
 from simcore_service_autoscaling.modules.ec2 import AutoscalingEC2
 from tenacity import retry
 from tenacity._asyncio import AsyncRetrying
@@ -165,6 +167,12 @@ async def async_client(initialized_app: FastAPI) -> AsyncIterator[httpx.AsyncCli
         yield client
 
 
+@pytest.fixture
+async def autoscaling_docker() -> AsyncIterator[AutoscalingDocker]:
+    async with AutoscalingDocker() as docker_client:
+        yield cast(AutoscalingDocker, docker_client)
+
+
 @pytest.fixture
 async def async_docker_client() -> AsyncIterator[aiodocker.Docker]:
     async with aiodocker.Docker() as docker_client: