removed Subs from dv-2

Author: sanderegg

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

@@ -4,6 +4,7 @@
 from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
+from typing import Any
 
 import arrow
 from dask_task_models_library.container_tasks.errors import TaskCancelledError
@@ -35,7 +36,7 @@
     clean_task_output_and_log_files_if_invalid,
     parse_output_data,
 )
-from ...utils.dask_client_utils import TaskHandlers
+from ...utils.dask_client_utils import TaskHandlers, UnixTimestamp
 from ...utils.rabbitmq import (
     publish_service_progress,
     publish_service_resource_tracking_stopped,
@@ -344,9 +345,11 @@ async def _process_task_result(
             optional_stopped=arrow.utcnow().datetime,
         )
 
-    async def _task_progress_change_handler(self, event: str) -> None:
+    async def _task_progress_change_handler(
+        self, event: tuple[UnixTimestamp, Any]
+    ) -> None:
         with log_catch(_logger, reraise=False):
-            task_progress_event = TaskProgressEvent.model_validate_json(event)
+            task_progress_event = TaskProgressEvent.model_validate_json(event[1])
             _logger.debug("received task progress update: %s", task_progress_event)
             user_id = task_progress_event.task_owner.user_id
             project_id = task_progress_event.task_owner.project_id

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

@@ -8,12 +8,11 @@
 
 """
 
-import asyncio
 import logging
 import traceback
 from collections.abc import Callable
 from copy import deepcopy
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from http.client import HTTPException
 from typing import Any, Final, cast
 
@@ -23,6 +22,7 @@
 from common_library.json_serialization import json_dumps
 from dask_task_models_library.container_tasks.docker import DockerBasicAuth
 from dask_task_models_library.container_tasks.errors import TaskCancelledError
+from dask_task_models_library.container_tasks.events import TaskProgressEvent
 from dask_task_models_library.container_tasks.io import (
     TaskCancelEventName,
     TaskInputData,
@@ -53,7 +53,7 @@
 from models_library.users import UserID
 from pydantic import TypeAdapter, ValidationError
 from pydantic.networks import AnyUrl
-from servicelib.logging_utils import log_catch
+from servicelib.logging_utils import log_catch, log_context
 from settings_library.s3 import S3Settings
 from simcore_sdk.node_ports_common.exceptions import NodeportsException
 from simcore_sdk.node_ports_v2 import FileLinkType
@@ -123,8 +123,6 @@ class DaskClient:
     tasks_file_link_type: FileLinkType
     cluster_type: ClusterTypeInModel
 
-    _subscribed_tasks: list[asyncio.Task] = field(default_factory=list)
-
     @classmethod
     async def create(
         cls,
@@ -177,24 +175,15 @@ async def create(
         raise ValueError(err_msg)
 
     async def delete(self) -> None:
-        _logger.debug("closing dask client...")
-        for task in self._subscribed_tasks:
-            task.cancel()
-        await asyncio.gather(*self._subscribed_tasks, return_exceptions=True)
-        await self.backend.close()
-        _logger.info("dask client properly closed")
+        with log_context(_logger, logging.INFO, msg="close dask client"):
+            await self.backend.close()
 
     def register_handlers(self, task_handlers: TaskHandlers) -> None:
         _event_consumer_map = [
-            (self.backend.progress_sub, task_handlers.task_progress_handler),
-        ]
-        self._subscribed_tasks = [
-            asyncio.create_task(
-                dask_utils.dask_sub_consumer_task(dask_sub, handler),
-                name=f"{dask_sub.name}_dask_sub_consumer_task",
-            )
-            for dask_sub, handler in _event_consumer_map
+            (TaskProgressEvent.topic_name(), task_handlers.task_progress_handler),
         ]
+        for topic_name, handler in _event_consumer_map:
+            self.backend.client.subscribe_topic(topic_name, handler)
 
     async def _publish_in_dask(  # noqa: PLR0913 # pylint: disable=too-many-arguments
         self,

services/director-v2/src/simcore_service_director_v2/utils/dask.py

@@ -1,8 +1,7 @@
-import asyncio
 import collections
 import logging
-from collections.abc import Awaitable, Callable, Coroutine, Generator
-from typing import Any, Final, NoReturn, ParamSpec, TypeVar, cast
+from collections.abc import Coroutine, Generator
+from typing import Any, ParamSpec, TypeVar, cast
 
 import distributed
 from common_library.json_serialization import json_dumps
@@ -30,7 +29,6 @@
 from models_library.users import UserID
 from models_library.wallets import WalletID
 from pydantic import AnyUrl, ByteSize, TypeAdapter, ValidationError
-from servicelib.logging_utils import log_catch, log_context
 from simcore_sdk import node_ports_v2
 from simcore_sdk.node_ports_common.exceptions import (
     NodeportsException,
@@ -438,40 +436,6 @@ async def clean_task_output_and_log_files_if_invalid(
         )
 
 
-async def _dask_sub_consumer(
-    dask_sub: distributed.Sub,
-    handler: Callable[[str], Awaitable[None]],
-) -> None:
-    async for dask_event in dask_sub:
-        _logger.debug(
-            "received dask event '%s' of topic %s",
-            dask_event,
-            dask_sub.name,
-        )
-        await handler(dask_event)
-
-
-_REST_TIMEOUT_S: Final[int] = 1
-
-
-async def dask_sub_consumer_task(
-    dask_sub: distributed.Sub,
-    handler: Callable[[str], Awaitable[None]],
-) -> NoReturn:
-    while True:
-        with (
-            log_catch(_logger, reraise=False),
-            log_context(
-                _logger,
-                level=logging.DEBUG,
-                msg=f"dask sub task for topic {dask_sub.name}",
-            ),
-        ):
-            await _dask_sub_consumer(dask_sub, handler)
-        # we sleep a bit before restarting
-        await asyncio.sleep(_REST_TIMEOUT_S)
-
-
 def from_node_reqs_to_dask_resources(
     node_reqs: NodeRequirements,
 ) -> dict[str, int | float]:

services/director-v2/src/simcore_service_director_v2/utils/dask_client_utils.py

@@ -3,6 +3,7 @@
 import socket
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
+from typing import Any, TypeAlias
 
 import distributed
 from models_library.clusters import ClusterAuthentication, TLSAuthentication
@@ -11,10 +12,12 @@
 from ..core.errors import ConfigurationError
 from .dask import wrap_client_async_routine
 
+UnixTimestamp: TypeAlias = float
+
 
 @dataclass
 class TaskHandlers:
-    task_progress_handler: Callable[[str], Awaitable[None]]
+    task_progress_handler: Callable[[tuple[UnixTimestamp, Any]], Awaitable[None]]
 
 
 logger = logging.getLogger(__name__)

services/director-v2/tests/unit/test_modules_dask_client.py

@@ -163,7 +163,6 @@ async def factory() -> DaskClient:
             client.settings
             == minimal_app.state.settings.DIRECTOR_V2_COMPUTATIONAL_BACKEND
         )
-        assert not client._subscribed_tasks  # noqa: SLF001
 
         assert client.backend.client
         scheduler_infos = client.backend.client.scheduler_info()  # type: ignore
@@ -181,7 +180,7 @@ async def factory() -> DaskClient:
 @pytest.fixture(params=["create_dask_client_from_scheduler"])
 async def dask_client(
     create_dask_client_from_scheduler: Callable[[], Awaitable[DaskClient]],
-    request,
+    request: pytest.FixtureRequest,
 ) -> DaskClient:
     client: DaskClient = await {
         "create_dask_client_from_scheduler": create_dask_client_from_scheduler,
@@ -402,7 +401,7 @@ def comp_run_metadata(faker: Faker) -> RunMetadataDict:
     return RunMetadataDict(
         product_name=faker.pystr(),
         simcore_user_agent=faker.pystr(),
-    ) | cast(dict[str, str], faker.pydict(allowed_types=(str,)))
+    ) | cast(RunMetadataDict, faker.pydict(allowed_types=(str,)))
 
 
 @pytest.fixture
@@ -1102,8 +1101,7 @@ def fake_remote_fct(
         log_file_url: LogFileUploadURL,
         s3_settings: S3Settings | None,
     ) -> TaskOutputData:
-        progress_pub = distributed.Pub(TaskProgressEvent.topic_name())
-        progress_pub.put("my name is progress")
+        get_worker().log_event(TaskProgressEvent.topic_name(), "my name is progress")
         # tell the client we are done
         published_event = Event(name=_DASK_START_EVENT)
         published_event.set()
@@ -1147,7 +1145,7 @@ def fake_remote_fct(
             )
             # we should have received data in our TaskHandlers
             fake_task_handlers.task_progress_handler.assert_called_with(
-                "my name is progress"
+                (mock.ANY, "my name is progress")
             )
     await _assert_wait_for_cb_call(mocked_user_completed_cb)