added reference counting on the dask client in the pool

Author: sanderegg

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_base.py

@@ -589,6 +589,12 @@ async def _process_executing_tasks(
     ) -> None:
         """process executing tasks from the 3rd party backend"""
 
+    @abstractmethod
+    async def _release_resources(
+        self, user_id: UserID, project_id: ProjectID, comp_run: CompRunsAtDB
+    ) -> None:
+        """release resources used by the scheduler for a given user and project"""
+
     async def apply(
         self,
         *,
@@ -654,6 +660,7 @@ async def apply(
 
                 # 7. Are we done scheduling that pipeline?
                 if not dag.nodes() or pipeline_result in COMPLETED_STATES:
+                    await self._release_resources(user_id, project_id, comp_run)
                     # there is nothing left, the run is completed, we're done here
                     _logger.info(
                         "pipeline %s scheduling completed with result %s",

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_dask.py

@@ -4,7 +4,7 @@
 from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Final
 
 import arrow
 from dask_task_models_library.container_tasks.errors import TaskCancelledError
@@ -23,7 +23,7 @@
 from models_library.users import UserID
 from pydantic import PositiveInt
 from servicelib.common_headers import UNDEFINED_DEFAULT_SIMCORE_USER_AGENT_VALUE
-from servicelib.logging_utils import log_catch
+from servicelib.logging_utils import log_catch, log_context
 
 from ...core.errors import (
     ComputationalBackendNotConnectedError,
@@ -56,13 +56,16 @@
 
 _logger = logging.getLogger(__name__)
 
+_DASK_CLIENT_RUN_REF: Final[str] = "{user_id}:{comp_run.run_id}"
+
 
 @asynccontextmanager
 async def _cluster_dask_client(
     user_id: UserID,
     scheduler: "DaskScheduler",
     *,
     use_on_demand_clusters: bool,
+    run_id: PositiveInt,
     run_metadata: RunMetadataDict,
 ) -> AsyncIterator[DaskClient]:
     cluster: BaseCluster = scheduler.settings.default_cluster
@@ -72,7 +75,9 @@ async def _cluster_dask_client(
             user_id=user_id,
             wallet_id=run_metadata.get("wallet_id"),
         )
-    async with scheduler.dask_clients_pool.acquire(cluster) as client:
+    async with scheduler.dask_clients_pool.acquire(
+        cluster, ref=_DASK_CLIENT_RUN_REF.format(user_id=user_id, run_id=run_id)
+    ) as client:
         yield client
 
 
@@ -101,6 +106,7 @@ async def _start_tasks(
             user_id,
             self,
             use_on_demand_clusters=comp_run.use_on_demand_clusters,
+            run_id=comp_run.run_id,
             run_metadata=comp_run.metadata,
         ) as client:
             # Change the tasks state to PENDING
@@ -151,6 +157,7 @@ async def _get_tasks_status(
                 user_id,
                 self,
                 use_on_demand_clusters=comp_run.use_on_demand_clusters,
+                run_id=comp_run.run_id,
                 run_metadata=comp_run.metadata,
             ) as client:
                 return await client.get_tasks_status([f"{t.job_id}" for t in tasks])
@@ -171,6 +178,7 @@ async def _process_executing_tasks(
                 user_id,
                 self,
                 use_on_demand_clusters=comp_run.use_on_demand_clusters,
+                run_id=comp_run.run_id,
                 run_metadata=comp_run.metadata,
             ) as client:
                 task_progresses = await client.get_tasks_progress(
@@ -217,6 +225,22 @@ async def _process_executing_tasks(
             )
         )
 
+    async def _release_resources(
+        self, user_id: UserID, project_id: ProjectID, comp_run: CompRunsAtDB
+    ) -> None:
+        """release resources used by the scheduler for a given user and project"""
+        with (
+            log_catch(_logger, reraise=False),
+            log_context(
+                _logger,
+                logging.INFO,
+                msg=f"releasing resources for {user_id=}, {project_id=}, {comp_run.run_id=}",
+            ),
+        ):
+            await self.dask_clients_pool.release_client_ref(
+                ref=_DASK_CLIENT_RUN_REF.format(user_id=user_id, run_id=comp_run.run_id)
+            )
+
     async def _stop_tasks(
         self, user_id: UserID, tasks: list[CompTaskAtDB], comp_run: CompRunsAtDB
     ) -> None:
@@ -226,6 +250,7 @@ async def _stop_tasks(
                 user_id,
                 self,
                 use_on_demand_clusters=comp_run.use_on_demand_clusters,
+                run_id=comp_run.run_id,
                 run_metadata=comp_run.metadata,
             ) as client:
                 await asyncio.gather(
@@ -259,6 +284,7 @@ async def _process_completed_tasks(
                 user_id,
                 self,
                 use_on_demand_clusters=comp_run.use_on_demand_clusters,
+                run_id=comp_run.run_id,
                 run_metadata=comp_run.metadata,
             ) as client:
                 tasks_results = await asyncio.gather(
@@ -278,6 +304,7 @@ async def _process_completed_tasks(
                 user_id,
                 self,
                 use_on_demand_clusters=comp_run.use_on_demand_clusters,
+                run_id=comp_run.run_id,
                 run_metadata=comp_run.metadata,
             ) as client:
                 await asyncio.gather(

services/director-v2/src/simcore_service_director_v2/modules/dask_clients_pool.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+from collections import defaultdict
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
@@ -24,6 +25,7 @@
 
 
 _ClusterUrl: TypeAlias = AnyUrl
+ClientRef: TypeAlias = str
 
 
 @dataclass
@@ -33,6 +35,10 @@ class DaskClientsPool:
     _client_acquisition_lock: asyncio.Lock = field(init=False)
     _cluster_to_client_map: dict[_ClusterUrl, DaskClient] = field(default_factory=dict)
     _task_handlers: TaskHandlers | None = None
+    # Track references to each client by endpoint
+    _client_refs: defaultdict[_ClusterUrl, set[str]] = field(
+        default_factory=lambda: defaultdict(set)
+    )
 
     def __post_init__(self):
         # NOTE: to ensure the correct loop is used
@@ -60,13 +66,54 @@ async def delete(self) -> None:
             *[client.delete() for client in self._cluster_to_client_map.values()],
             return_exceptions=True,
         )
+        self._cluster_to_client_map.clear()
+        self._client_refs.clear()
+
+    async def release_client_ref(self, ref: ClientRef) -> None:
+        """Release a dask client reference by its ref.
+
+        If all the references to the client are released,
+        the client will be deleted from the pool.
+        This method is thread-safe and can be called concurrently.
+        """
+        async with self._client_acquisition_lock:
+            # Find which endpoint this ref belongs to
+            endpoint_to_remove = None
+            for endpoint, refs in self._client_refs.items():
+                if ref in refs:
+                    refs.remove(ref)
+                    _logger.debug("Released reference %s for client %s", ref, endpoint)
+                    if not refs:  # No more references to this client
+                        endpoint_to_remove = endpoint
+                    break
+
+            # If we found an endpoint with no more refs, clean it up
+            if endpoint_to_remove and (
+                dask_client := self._cluster_to_client_map.pop(endpoint_to_remove, None)
+            ):
+                _logger.info(
+                    "Last reference to client %s released, deleting client",
+                    endpoint_to_remove,
+                )
+                await dask_client.delete()
+                # Clean up the empty refs set
+                del self._client_refs[endpoint_to_remove]
+                _logger.debug(
+                    "Remaining clients: %s",
+                    [f"{k}" for k in self._cluster_to_client_map],
+                )
 
     @asynccontextmanager
-    async def acquire(self, cluster: BaseCluster) -> AsyncIterator[DaskClient]:
-        """returns a dask client for the given cluster
+    async def acquire(
+        self, cluster: BaseCluster, *, ref: ClientRef
+    ) -> AsyncIterator[DaskClient]:
+        """Returns a dask client for the given cluster.
+
         This method is thread-safe and can be called concurrently.
         If the cluster is not found in the pool, it will create a new dask client for it.
 
+        The passed reference is used to track the client usage, user should call
+        `release_client_ref` to release the client reference when done.
         """
 
         async def _concurently_safe_acquire_client() -> DaskClient:
@@ -102,15 +149,21 @@ async def _concurently_safe_acquire_client() -> DaskClient:
                         if self._task_handlers:
                             dask_client.register_handlers(self._task_handlers)
 
-                        _logger.debug(
-                            "list of clients: %s", f"{self._cluster_to_client_map=}"
-                        )
+                    # Track the reference
+                    self._client_refs[cluster.endpoint].add(ref)
+
+                    _logger.debug(
+                        "Client %s now has %d references",
+                        cluster.endpoint,
+                        len(self._client_refs[cluster.endpoint]),
+                    )
 
                     assert dask_client  # nosec
                     return dask_client
 
         try:
             dask_client = await _concurently_safe_acquire_client()
+
         except Exception as exc:
             raise DaskClientAcquisisitonError(cluster=cluster, error=exc) from exc