Merge branch 'master' into enh/user-notifications

Author: odeimaiz

packages/service-library/src/servicelib/background_task.py

@@ -25,15 +25,36 @@ class PeriodicTaskCancellationError(PydanticErrorMixin, Exception):
     msg_template: str = "Could not cancel task '{task_name}'"
 
 
+class SleepUsingAsyncioEvent:
+    """Sleep strategy that waits on an event to be set."""
+
+    def __init__(self, event: "asyncio.Event") -> None:
+        self.event = event
+
+    async def __call__(self, timeout: float | None) -> None:
+        with contextlib.suppress(TimeoutError):
+            await asyncio.wait_for(self.event.wait(), timeout=timeout)
+            self.event.clear()
+
+
 async def _periodic_scheduled_task(
     task: Callable[..., Awaitable[None]],
     *,
     interval: datetime.timedelta,
     task_name: str,
+    early_wake_up_event: asyncio.Event | None,
     **task_kwargs,
 ) -> None:
     # NOTE: This retries forever unless cancelled
-    async for attempt in AsyncRetrying(wait=wait_fixed(interval.total_seconds())):
+    nap = (
+        asyncio.sleep
+        if early_wake_up_event is None
+        else SleepUsingAsyncioEvent(early_wake_up_event)
+    )
+    async for attempt in AsyncRetrying(
+        sleep=nap,
+        wait=wait_fixed(interval.total_seconds()),
+    ):
         with attempt:
             with log_context(
                 _logger,
@@ -51,6 +72,7 @@ def start_periodic_task(
     interval: datetime.timedelta,
     task_name: str,
     wait_before_running: datetime.timedelta = datetime.timedelta(0),
+    early_wake_up_event: asyncio.Event | None = None,
     **kwargs,
 ) -> asyncio.Task:
     with log_context(
@@ -64,6 +86,7 @@ def start_periodic_task(
                 task,
                 interval=interval,
                 task_name=task_name,
+                early_wake_up_event=early_wake_up_event,
                 **kwargs,
             ),
             name=task_name,

packages/service-library/tests/test_background_task.py

@@ -6,7 +6,8 @@
 
 import asyncio
 import datetime
-from typing import AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
+from typing import Final
 from unittest import mock
 
 import pytest
@@ -18,38 +19,50 @@
     stop_periodic_task,
 )
 
-_FAST_POLL_INTERVAL = 1
+_FAST_POLL_INTERVAL: Final[int] = 1
+_VERY_SLOW_POLL_INTERVAL: Final[int] = 100
 
 
 @pytest.fixture
 def mock_background_task(mocker: MockerFixture) -> mock.AsyncMock:
-    mocked_task = mocker.AsyncMock(return_value=None)
-    return mocked_task
+    return mocker.AsyncMock(return_value=None)
 
 
 @pytest.fixture
 def task_interval() -> datetime.timedelta:
     return datetime.timedelta(seconds=_FAST_POLL_INTERVAL)
 
 
-@pytest.fixture(params=[None, 1])
+@pytest.fixture
+def very_long_task_interval() -> datetime.timedelta:
+    return datetime.timedelta(seconds=_VERY_SLOW_POLL_INTERVAL)
+
+
+@pytest.fixture(params=[None, 1], ids=lambda x: f"stop-timeout={x}")
 def stop_task_timeout(request: pytest.FixtureRequest) -> float | None:
     return request.param
 
 
 @pytest.fixture
 async def create_background_task(
     faker: Faker, stop_task_timeout: float | None
-) -> AsyncIterator[Callable[[datetime.timedelta, Callable], Awaitable[asyncio.Task]]]:
+) -> AsyncIterator[
+    Callable[
+        [datetime.timedelta, Callable, asyncio.Event | None], Awaitable[asyncio.Task]
+    ]
+]:
     created_tasks = []
 
     async def _creator(
-        interval: datetime.timedelta, task: Callable[..., Awaitable]
+        interval: datetime.timedelta,
+        task: Callable[..., Awaitable],
+        early_wake_up_event: asyncio.Event | None,
     ) -> asyncio.Task:
         background_task = start_periodic_task(
             task,
             interval=interval,
             task_name=faker.pystr(),
+            early_wake_up_event=early_wake_up_event,
         )
         assert background_task
         created_tasks.append(background_task)
@@ -62,33 +75,69 @@ async def _creator(
     )
 
 
+@pytest.mark.parametrize(
+    "wake_up_event", [None, asyncio.Event], ids=lambda x: f"wake-up-event: {x}"
+)
 async def test_background_task_created_and_deleted(
     mock_background_task: mock.AsyncMock,
     task_interval: datetime.timedelta,
     create_background_task: Callable[
-        [datetime.timedelta, Callable], Awaitable[asyncio.Task]
+        [datetime.timedelta, Callable, asyncio.Event | None], Awaitable[asyncio.Task]
     ],
+    wake_up_event: Callable | None,
 ):
-    task = await create_background_task(
+    event = wake_up_event() if wake_up_event else None
+    _task = await create_background_task(
         task_interval,
         mock_background_task,
+        event,
     )
     await asyncio.sleep(5 * task_interval.total_seconds())
     mock_background_task.assert_called()
-    assert mock_background_task.call_count > 1
+    assert mock_background_task.call_count > 2
+
+
+async def test_background_task_wakes_up_early(
+    mock_background_task: mock.AsyncMock,
+    very_long_task_interval: datetime.timedelta,
+    create_background_task: Callable[
+        [datetime.timedelta, Callable, asyncio.Event | None], Awaitable[asyncio.Task]
+    ],
+):
+    wake_up_event = asyncio.Event()
+    _task = await create_background_task(
+        very_long_task_interval,
+        mock_background_task,
+        wake_up_event,
+    )
+    await asyncio.sleep(5 * _FAST_POLL_INTERVAL)
+    # now the task should have run only once
+    mock_background_task.assert_called_once()
+    await asyncio.sleep(5 * _FAST_POLL_INTERVAL)
+    mock_background_task.assert_called_once()
+    # this should wake up the task
+    wake_up_event.set()
+    await asyncio.sleep(5 * _FAST_POLL_INTERVAL)
+    mock_background_task.assert_called()
+    assert mock_background_task.call_count == 2
+    # no change this now waits again a very long time
+    await asyncio.sleep(5 * _FAST_POLL_INTERVAL)
+    mock_background_task.assert_called()
+    assert mock_background_task.call_count == 2
 
 
 async def test_background_task_raises_restarts(
     mock_background_task: mock.AsyncMock,
     task_interval: datetime.timedelta,
     create_background_task: Callable[
-        [datetime.timedelta, Callable], Awaitable[asyncio.Task]
+        [datetime.timedelta, Callable, asyncio.Event | None], Awaitable[asyncio.Task]
     ],
 ):
     mock_background_task.side_effect = RuntimeError("pytest faked runtime error")
-    task = await create_background_task(
+    _task = await create_background_task(
         task_interval,
         mock_background_task,
+        None,
     )
     await asyncio.sleep(5 * task_interval.total_seconds())
     mock_background_task.assert_called()
@@ -99,13 +148,14 @@ async def test_background_task_correctly_cancels(
     mock_background_task: mock.AsyncMock,
     task_interval: datetime.timedelta,
     create_background_task: Callable[
-        [datetime.timedelta, Callable], Awaitable[asyncio.Task]
+        [datetime.timedelta, Callable, asyncio.Event | None], Awaitable[asyncio.Task]
     ],
 ):
     mock_background_task.side_effect = asyncio.CancelledError
-    task = await create_background_task(
+    _task = await create_background_task(
         task_interval,
         mock_background_task,
+        None,
     )
     await asyncio.sleep(5 * task_interval.total_seconds())
     # the task will be called once, and then stop

services/director-v2/src/simcore_service_director_v2/api/dependencies/scheduler.py

@@ -1,7 +1,7 @@
 from fastapi import Depends, FastAPI, Request
 
 from ...core.settings import ComputationalBackendSettings
-from ...modules.comp_scheduler.base_scheduler import BaseCompScheduler
+from ...modules.comp_scheduler import BaseCompScheduler
 from . import get_app
 
 

services/director-v2/src/simcore_service_director_v2/api/routes/computations.py

@@ -63,7 +63,7 @@
 from ...models.comp_runs import CompRunsAtDB, ProjectMetadataDict, RunMetadataDict
 from ...models.comp_tasks import CompTaskAtDB
 from ...modules.catalog import CatalogClient
-from ...modules.comp_scheduler.base_scheduler import BaseCompScheduler
+from ...modules.comp_scheduler import BaseCompScheduler
 from ...modules.db.repositories.clusters import ClustersRepository
 from ...modules.db.repositories.comp_pipelines import CompPipelinesRepository
 from ...modules.db.repositories.comp_runs import CompRunsRepository
@@ -288,7 +288,7 @@ async def _try_start_pipeline(
 )
 # NOTE: in case of a burst of calls to that endpoint, we might end up in a weird state.
 @run_sequentially_in_context(target_args=["computation.project_id"])
-async def create_computation(  # noqa: PLR0913  # pylint:disable=too-many-positional-arguments
+async def create_computation(  # noqa: PLR0913 # pylint: disable=too-many-positional-arguments
     computation: ComputationCreate,
     request: Request,
     project_repo: Annotated[

services/director-v2/src/simcore_service_director_v2/core/application.py

@@ -173,9 +173,9 @@ def init_app(settings: AppSettings | None = None) -> FastAPI:
     )
     if dynamic_scheduler_enabled or computational_backend_enabled:
         rabbitmq.setup(app)
+        redis.setup(app)
 
     if dynamic_scheduler_enabled:
-        redis.setup(app)
         dynamic_sidecar.setup(app)
         socketio.setup(app)
         notifier.setup(app)

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

@@ -1,3 +1,15 @@
-from .background_task import setup
+from fastapi import FastAPI
 
-__all__: tuple[str, ...] = ("setup",)
+from ._base_scheduler import BaseCompScheduler
+from ._task import on_app_shutdown, on_app_startup
+
+
+def setup(app: FastAPI):
+    app.add_event_handler("startup", on_app_startup(app))
+    app.add_event_handler("shutdown", on_app_shutdown(app))
+
+
+__all__: tuple[str, ...] = (
+    "setup",
+    "BaseCompScheduler",
+)

…modules/comp_scheduler/base_scheduler.py …odules/comp_scheduler/_base_scheduler.pyservices/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/base_scheduler.py renamed to services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_base_scheduler.py services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/base_scheduler.py renamed to services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_base_scheduler.py

@@ -30,10 +30,6 @@
 from ...models.comp_runs import RunMetadataDict
 from ...models.comp_tasks import CompTaskAtDB
 from ...models.dask_subsystem import DaskClientTaskState
-from ...modules.dask_client import DaskClient, PublishedComputationTask
-from ...modules.dask_clients_pool import DaskClientsPool
-from ...modules.db.repositories.clusters import ClustersRepository
-from ...modules.db.repositories.comp_runs import CompRunsRepository
 from ...utils.comp_scheduler import Iteration, get_resource_tracking_run_id
 from ...utils.dask import (
     clean_task_output_and_log_files_if_invalid,
@@ -48,8 +44,12 @@
     publish_service_stopped_metrics,
 )
 from ..clusters_keeper import get_or_create_on_demand_cluster
+from ..dask_client import DaskClient, PublishedComputationTask
+from ..dask_clients_pool import DaskClientsPool
+from ..db.repositories.clusters import ClustersRepository
+from ..db.repositories.comp_runs import CompRunsRepository
 from ..db.repositories.comp_tasks import CompTasksRepository
-from .base_scheduler import BaseCompScheduler, ScheduledPipelineParams
+from ._base_scheduler import BaseCompScheduler, ScheduledPipelineParams
 
 _logger = logging.getLogger(__name__)
 
@@ -158,9 +158,11 @@ async def _get_tasks_status(
             for dask_task_state, task in zip(tasks_statuses, tasks, strict=True):
                 if dask_task_state is DaskClientTaskState.PENDING_OR_STARTED:
                     running_states += [
-                        RunningState.STARTED
-                        if task.progress is not None
-                        else RunningState.PENDING
+                        (
+                            RunningState.STARTED
+                            if task.progress is not None
+                            else RunningState.PENDING
+                        )
                     ]
                 else:
                     running_states += [

…modules/comp_scheduler/dask_scheduler.py …odules/comp_scheduler/_dask_scheduler.pyservices/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/dask_scheduler.py renamed to services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_dask_scheduler.py services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/dask_scheduler.py renamed to services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_dask_scheduler.py

@@ -2,16 +2,16 @@
 
 from fastapi import FastAPI
 from models_library.clusters import DEFAULT_CLUSTER_ID
-from simcore_service_director_v2.core.settings import AppSettings
 
 from ...core.errors import ConfigurationError
+from ...core.settings import AppSettings
 from ...models.comp_runs import CompRunsAtDB
-from ...modules.dask_clients_pool import DaskClientsPool
-from ...modules.rabbitmq import get_rabbitmq_client, get_rabbitmq_rpc_client
 from ...utils.comp_scheduler import SCHEDULED_STATES
+from ..dask_clients_pool import DaskClientsPool
 from ..db.repositories.comp_runs import CompRunsRepository
-from .base_scheduler import BaseCompScheduler, ScheduledPipelineParams
-from .dask_scheduler import DaskScheduler
+from ..rabbitmq import get_rabbitmq_client, get_rabbitmq_rpc_client
+from ._base_scheduler import BaseCompScheduler, ScheduledPipelineParams
+from ._dask_scheduler import DaskScheduler
 
 logger = logging.getLogger(__name__)
 
@@ -43,9 +43,9 @@ async def create_from_db(app: FastAPI) -> BaseCompScheduler:
         db_engine=db_engine,
         scheduled_pipelines={
             (r.user_id, r.project_uuid, r.iteration): ScheduledPipelineParams(
-                cluster_id=r.cluster_id
-                if r.cluster_id is not None
-                else DEFAULT_CLUSTER_ID,
+                cluster_id=(
+                    r.cluster_id if r.cluster_id is not None else DEFAULT_CLUSTER_ID
+                ),
                 run_metadata=r.metadata,
                 mark_for_cancellation=False,
                 use_on_demand_clusters=r.use_on_demand_clusters,

…tor_v2/modules/comp_scheduler/factory.py …les/comp_scheduler/_scheduler_factory.pyservices/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/factory.py renamed to services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_factory.py services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/factory.py renamed to services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_factory.py

@@ -0,0 +1,50 @@
+import datetime
+import logging
+from collections.abc import Callable, Coroutine
+from typing import Any, Final
+
+from fastapi import FastAPI
+from servicelib.background_task import start_periodic_task, stop_periodic_task
+from servicelib.logging_utils import log_context
+from servicelib.redis import RedisClientsManager
+from servicelib.redis_utils import exclusive
+from settings_library.redis import RedisDatabase
+
+from . import _scheduler_factory
+
+_logger = logging.getLogger(__name__)
+
+_COMPUTATIONAL_SCHEDULER_INTERVAL: Final[datetime.timedelta] = datetime.timedelta(
+    seconds=5
+)
+_TASK_NAME: Final[str] = "computational services scheduler"
+
+
+def on_app_startup(app: FastAPI) -> Callable[[], Coroutine[Any, Any, None]]:
+    async def start_scheduler() -> None:
+        with log_context(
+            _logger, level=logging.INFO, msg="starting computational scheduler"
+        ):
+            redis_clients_manager: RedisClientsManager = app.state.redis_clients_manager
+            lock_key = f"{app.title}:computational_scheduler"
+            app.state.scheduler = scheduler = await _scheduler_factory.create_from_db(
+                app
+            )
+            app.state.computational_scheduler_task = start_periodic_task(
+                exclusive(
+                    redis_clients_manager.client(RedisDatabase.LOCKS),
+                    lock_key=lock_key,
+                )(scheduler.schedule_all_pipelines),
+                interval=_COMPUTATIONAL_SCHEDULER_INTERVAL,
+                task_name=_TASK_NAME,
+                early_wake_up_event=scheduler.wake_up_event,
+            )
+
+    return start_scheduler
+
+
+def on_app_shutdown(app: FastAPI) -> Callable[[], Coroutine[Any, Any, None]]:
+    async def stop_scheduler() -> None:
+        await stop_periodic_task(app.state.computational_scheduler_task)
+
+    return stop_scheduler