tests are passing

Author: sanderegg

services/dask-sidecar/docker/boot.sh

@@ -117,7 +117,7 @@ else
   fi
 
   # GPUs
-  num_gpus=$(python -c "from simcore_service_dask_sidecar.utils import num_available_gpus; print(num_available_gpus());")
+  num_gpus=$(python -c "from simcore_service_dask_sidecar.utils.gpus import num_available_gpus; print(num_available_gpus());")
 
   # RAM (is computed similarly as the default dask-sidecar computation)
   _value=$(python -c "import psutil; print(int(psutil.virtual_memory().total * $num_cpus/$(nproc)))")
@@ -128,7 +128,7 @@ else
 
   # add the GPUs if there are any
   if [ "$num_gpus" -gt 0 ]; then
-    total_vram=$(python -c "from simcore_service_dask_sidecar.utils import video_memory; print(video_memory());")
+    total_vram=$(python -c "from simcore_service_dask_sidecar.utils.gpus import video_memory; print(video_memory());")
     resources="$resources,GPU=$num_gpus,VRAM=$total_vram"
   fi
 

services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_plugin.py

@@ -1,8 +1,10 @@
+import asyncio
 import logging
+from asyncio import AbstractEventLoop
 from collections.abc import Awaitable
 
 import distributed
-from servicelib.logging_utils import log_context
+from servicelib.logging_utils import log_catch, log_context
 from servicelib.rabbitmq import RabbitMQClient, wait_till_rabbitmq_responsive
 from settings_library.rabbit import RabbitSettings
 
@@ -15,6 +17,7 @@ class RabbitMQPlugin(distributed.WorkerPlugin):
     """Dask Worker Plugin for RabbitMQ integration"""
 
     name = "rabbitmq_plugin"
+    _loop: AbstractEventLoop | None = None
     _client: RabbitMQClient | None = None
     _settings: RabbitSettings | None = None
 
@@ -36,6 +39,7 @@ async def _() -> None:
                 logging.INFO,
                 f"RabbitMQ client initialization for worker {worker.address}",
             ):
+                self._loop = asyncio.get_event_loop()
                 await wait_till_rabbitmq_responsive(self._settings.dsn)
                 self._client = RabbitMQClient(
                     client_name="dask-sidecar", settings=self._settings
@@ -53,7 +57,15 @@ async def _() -> None:
                 f"RabbitMQ client teardown for worker {worker.address}",
             ):
                 if self._client:
-                    await self._client.close()
+                    current_loop = asyncio.get_event_loop()
+                    if self._loop != current_loop:
+                        _logger.warning(
+                            "RabbitMQ client is de-activated (loop mismatch)"
+                        )
+                    assert self._loop  # nosec
+                    with log_catch(_logger, reraise=False):
+                        await asyncio.wait_for(self._client.close(), timeout=5.0)
+
                     self._client = None
 
         return _()

services/dask-sidecar/src/simcore_service_dask_sidecar/utils.py renamed to services/dask-sidecar/src/simcore_service_dask_sidecar/utils/gpus.py

@@ -13,6 +13,10 @@
 from simcore_service_dask_sidecar.settings import ApplicationSettings
 from typer.testing import CliRunner
 
+pytest_simcore_core_services_selection = [
+    "rabbit",
+]
+
 
 def test_cli_help_and_version(cli_runner: CliRunner):
     # invitations-maker --help

services/dask-sidecar/tests/unit/test_cli.py

@@ -6,15 +6,14 @@
 
 import asyncio
 import concurrent.futures
-import logging
 import time
 from collections.abc import AsyncIterator, Callable, Coroutine
 from typing import Any
 
 import distributed
 import pytest
 from dask_task_models_library.container_tasks.errors import TaskCancelledError
-from dask_task_models_library.container_tasks.events import TaskLogEvent
+from dask_task_models_library.container_tasks.events import TaskProgressEvent
 from dask_task_models_library.container_tasks.io import TaskCancelEventName
 from dask_task_models_library.container_tasks.protocol import TaskOwner
 from simcore_service_dask_sidecar.dask_utils import (
@@ -33,16 +32,20 @@
 DASK_TASK_STARTED_EVENT = "task_started"
 DASK_TESTING_TIMEOUT_S = 25
 
+pytest_simcore_core_services_selection = [
+    "rabbit",
+]
+
 
 def test_publish_event(
     dask_client: distributed.Client, job_id: str, task_owner: TaskOwner
 ):
     dask_pub = distributed.Pub("some_topic", client=dask_client)
     dask_sub = distributed.Sub("some_topic", client=dask_client)
-    event_to_publish = TaskLogEvent(
+    event_to_publish = TaskProgressEvent(
         job_id=job_id,
-        log="the log",
-        log_level=logging.INFO,
+        msg="the log",
+        progress=1,
         task_owner=task_owner,
     )
     publish_event(dask_pub=dask_pub, event=event_to_publish)
@@ -53,7 +56,7 @@ def test_publish_event(
     message = dask_sub.get(timeout=DASK_TESTING_TIMEOUT_S)
     assert message is not None
     assert isinstance(message, str)
-    received_task_log_event = TaskLogEvent.model_validate_json(message)
+    received_task_log_event = TaskProgressEvent.model_validate_json(message)
     assert received_task_log_event == event_to_publish
 
 
@@ -62,8 +65,8 @@ async def test_publish_event_async(
 ):
     dask_pub = distributed.Pub("some_topic", client=async_dask_client)
     dask_sub = distributed.Sub("some_topic", client=async_dask_client)
-    event_to_publish = TaskLogEvent(
-        job_id=job_id, log="the log", log_level=logging.INFO, task_owner=task_owner
+    event_to_publish = TaskProgressEvent(
+        job_id=job_id, msg="the log", progress=2, task_owner=task_owner
     )
     publish_event(dask_pub=dask_pub, event=event_to_publish)
 
@@ -74,7 +77,7 @@ async def test_publish_event_async(
     assert isinstance(message, Coroutine)
     message = await message
     assert message is not None
-    received_task_log_event = TaskLogEvent.model_validate_json(message)
+    received_task_log_event = TaskProgressEvent.model_validate_json(message)
     assert received_task_log_event == event_to_publish
 
 
@@ -117,11 +120,10 @@ async def _dask_sub_consumer_task(sub: distributed.Sub) -> None:
 
     async def _dask_publisher_task(pub: distributed.Pub) -> None:
         print("--> starting publisher task")
-        for n in range(NUMBER_OF_MESSAGES):
-            event_to_publish = TaskLogEvent(
+        for _ in range(NUMBER_OF_MESSAGES):
+            event_to_publish = TaskProgressEvent(
                 job_id=job_id,
-                log=f"the log {n}",
-                log_level=logging.INFO,
+                progress=0.5,
                 task_owner=task_owner,
             )
             publish_event(dask_pub=pub, event=event_to_publish)

services/dask-sidecar/tests/unit/test_dask_utils.py

@@ -9,7 +9,7 @@
 import logging
 import re
 import threading
-from collections.abc import AsyncIterator, Callable, Coroutine, Iterable
+from collections.abc import AsyncIterator, Callable, Iterable
 
 # copied out from dask
 from dataclasses import dataclass
@@ -613,7 +613,7 @@ def test_run_multiple_computational_sidecar_dask(
 
     results = dask_client.gather(futures)
     assert results
-    assert not isinstance(results, Coroutine)
+    assert isinstance(results, list)
     # for result in results:
     # check that the task produce the expected data, not less not more
     for output_data in results:
@@ -708,6 +708,7 @@ async def test_run_computational_sidecar_dask_does_not_lose_messages_with_pubsub
     progress_sub: distributed.Sub,
     mocked_get_image_labels: mock.Mock,
     log_rabbit_client_parser: mock.AsyncMock,
+    task_owner: TaskOwner,
 ):
     mocked_get_image_labels.assert_not_called()
     NUMBER_OF_LOGS = 200
@@ -745,14 +746,19 @@ async def test_run_computational_sidecar_dask_does_not_lose_messages_with_pubsub
 
     await asyncio.sleep(5)
     assert log_rabbit_client_parser.called
+
     worker_logs = [
         message
         for msg in log_rabbit_client_parser.call_args_list
         for message in LoggerRabbitMessage.model_validate_json(msg.args[0]).messages
     ]
     # check all the awaited logs are in there
     filtered_worker_logs = filter(lambda log: "This is iteration" in log, worker_logs)
-    assert len(list(filtered_worker_logs)) == NUMBER_OF_LOGS
+    assert (
+        len(list(filtered_worker_logs)) == (2 * NUMBER_OF_LOGS)
+        if task_owner.has_parent
+        else NUMBER_OF_LOGS
+    )
     mocked_get_image_labels.assert_called()
 
 

services/dask-sidecar/tests/unit/test_tasks.py

@@ -10,13 +10,17 @@
 import pytest
 from pytest_mock.plugin import MockerFixture
 from pytest_simcore.helpers.typing_env import EnvVarsDict
-from simcore_service_dask_sidecar.utils import num_available_gpus
+from simcore_service_dask_sidecar.utils.gpus import num_available_gpus
+
+pytest_simcore_core_services_selection = [
+    "rabbit",
+]
 
 
 @pytest.fixture
 def mock_aiodocker(mocker: MockerFixture) -> mock.MagicMock:
     return mocker.patch(
-        "simcore_service_dask_sidecar.utils.aiodocker.Docker", autospec=True
+        "simcore_service_dask_sidecar.utils.gpus.aiodocker.Docker", autospec=True
     )