feat(BA-1428): Refactor event dispatcher and handlers directory structure (#4497)

Author: HyeockJinKim

changes/4497.feature.md

@@ -0,0 +1 @@
+Refactor event dispatcher and handlers directory structure

src/ai/backend/manager/api/context.py

@@ -15,6 +15,7 @@
 )
 from ai.backend.manager.config.provider import ManagerConfigProvider
 from ai.backend.manager.plugin.network import NetworkPluginContext
+from ai.backend.manager.scheduler.dispatcher import SchedulerDispatcher
 from ai.backend.manager.service.base import ServicesContext
 from ai.backend.manager.services.processors import Processors
 
@@ -65,6 +66,7 @@ class RootContext(BaseContext):
 
     registry: AgentRegistry
     agent_cache: AgentRPCCache
+    scheduler_dispatcher: SchedulerDispatcher
 
     error_monitor: ErrorPluginContext
     stats_monitor: StatsPluginContext

src/ai/backend/manager/config/unified.py

@@ -1604,7 +1604,7 @@ class IdleCheckerConfig(BaseModel):
         Enabled idle checkers.
         Comma-separated list of checker names.
         """,
-        examples=["timeout", "utilization"],
+        examples=["network_timeout", "utilization"],
     )
     app_streaming_packet_timeout: TimeDuration = Field(
         default=_TimeDurationPydanticAnnotation.time_duration_validator("5m"),

src/ai/backend/manager/event_dispatcher/bgtask.py

@@ -8,7 +8,15 @@
     AgentTerminatedEvent,
     DoAgentResourceCheckEvent,
 )
+from ai.backend.common.events.bgtask import (
+    BgtaskCancelledEvent,
+    BgtaskDoneEvent,
+    BgtaskFailedEvent,
+    BgtaskPartialSuccessEvent,
+    BgtaskUpdatedEvent,
+)
 from ai.backend.common.events.dispatcher import (
+    CoalescingOptions,
     EventDispatcher,
 )
 from ai.backend.common.events.hub.hub import EventHub
@@ -32,8 +40,15 @@
     ModelServiceStatusEvent,
     RouteCreatedEvent,
 )
+from ai.backend.common.events.schedule import (
+    DoCheckPrecondEvent,
+    DoScaleEvent,
+    DoScheduleEvent,
+    DoStartSessionEvent,
+)
 from ai.backend.common.events.session import (
     DoTerminateSessionEvent,
+    DoUpdateSessionStatusEvent,
     SessionCancelledEvent,
     SessionEnqueuedEvent,
     SessionFailureEvent,
@@ -49,22 +64,24 @@
     VFolderDeletionSuccessEvent,
 )
 from ai.backend.common.plugin.event import EventDispatcherPluginContext
-from ai.backend.manager.event_dispatcher.propagator import PropagatorEventDispatcher
+from ai.backend.manager.event_dispatcher.handlers.propagator import PropagatorEventHandler
+from ai.backend.manager.event_dispatcher.handlers.schedule import ScheduleEventHandler
 from ai.backend.manager.registry import AgentRegistry
+from ai.backend.manager.scheduler.dispatcher import SchedulerDispatcher
 
 from ..models.utils import ExtendedAsyncSAEngine
-from .agent import AgentEventHandler
-from .bgtask import dispatch_bgtask_events
-from .image import ImageEventHandler
-from .kernel import KernelEventHandler
-from .model_serving import ModelServingEventHandler
+from .handlers.agent import AgentEventHandler
+from .handlers.image import ImageEventHandler
+from .handlers.kernel import KernelEventHandler
+from .handlers.model_serving import ModelServingEventHandler
+from .handlers.session import SessionEventHandler
+from .handlers.vfolder import VFolderEventHandler
 from .reporters import EventLogger
-from .session import SessionEventHandler
-from .vfolder import VFolderEventHandler
 
 
 @dataclass
 class DispatcherArgs:
+    scheduler_dispatcher: SchedulerDispatcher
     event_hub: EventHub
     agent_registry: AgentRegistry
     db: ExtendedAsyncSAEngine
@@ -73,10 +90,11 @@ class DispatcherArgs:
 
 class Dispatchers:
     _db: ExtendedAsyncSAEngine
-    _propagator_dispatcher: PropagatorEventDispatcher
+    _propagator_handler: PropagatorEventHandler
     _agent_event_handler: AgentEventHandler
     _image_event_handler: ImageEventHandler
     _kernel_event_handler: KernelEventHandler
+    _schedule_event_handler: ScheduleEventHandler
     _model_serving_event_handler: ModelServingEventHandler
     _session_event_handler: SessionEventHandler
     _vfolder_event_handler: VFolderEventHandler
@@ -87,10 +105,11 @@ def __init__(self, args: DispatcherArgs) -> None:
         """
         self._db = args.db
         self._event_dispatcher_plugin_ctx = args.event_dispatcher_plugin_ctx
-        self._propagator_dispatcher = PropagatorEventDispatcher(args.event_hub)
+        self._propagator_handler = PropagatorEventHandler(args.event_hub)
         self._agent_event_handler = AgentEventHandler(args.agent_registry, args.db)
         self._image_event_handler = ImageEventHandler(args.agent_registry, args.db)
         self._kernel_event_handler = KernelEventHandler(args.agent_registry, args.db)
+        self._schedule_event_handler = ScheduleEventHandler(args.scheduler_dispatcher)
         self._model_serving_event_handler = ModelServingEventHandler(args.agent_registry, args.db)
         self._session_event_handler = SessionEventHandler(args.agent_registry, args.db)
         self._vfolder_event_handler = VFolderEventHandler(args.db)
@@ -99,15 +118,39 @@ def dispatch(self, event_dispatcher: EventDispatcher) -> None:
         """
         Dispatch events to the appropriate dispatcher.
         """
-        dispatch_bgtask_events(event_dispatcher, self._propagator_dispatcher)
+        self._dispatch_bgtask_events(event_dispatcher)
         self._dispatch_agent_events(event_dispatcher)
         self._dispatch_error_monitor_events(event_dispatcher)
         self._dispatch_image_events(event_dispatcher)
         self._dispatch_kernel_events(event_dispatcher)
+        self._dispatch_schedule_events(event_dispatcher)
         self._dispatch_model_serving_events(event_dispatcher)
         self._dispatch_session_events(event_dispatcher)
         self._dispatch_vfolder_events(event_dispatcher)
 
+    def _dispatch_bgtask_events(
+        self,
+        event_dispatcher: EventDispatcher,
+    ) -> None:
+        """
+        Register event dispatchers for background task events.
+        """
+        event_dispatcher.subscribe(
+            BgtaskUpdatedEvent, None, self._propagator_handler.propagate_event
+        )
+        event_dispatcher.subscribe(BgtaskDoneEvent, None, self._propagator_handler.propagate_event)
+        event_dispatcher.subscribe(
+            BgtaskPartialSuccessEvent,
+            None,
+            self._propagator_handler.propagate_event,
+        )
+        event_dispatcher.subscribe(
+            BgtaskCancelledEvent, None, self._propagator_handler.propagate_event
+        )
+        event_dispatcher.subscribe(
+            BgtaskFailedEvent, None, self._propagator_handler.propagate_event
+        )
+
     def _dispatch_agent_events(
         self,
         event_dispatcher: EventDispatcher,
@@ -229,6 +272,47 @@ def _dispatch_model_serving_events(self, event_dispatcher: EventDispatcher) -> N
             RouteCreatedEvent, None, self._model_serving_event_handler.handle_route_creation
         )
 
+    def _dispatch_schedule_events(self, event_dispatcher: EventDispatcher) -> None:
+        coalescing_opts: CoalescingOptions = {
+            "max_wait": 0.5,
+            "max_batch_size": 32,
+        }
+        event_dispatcher.consume(
+            SessionEnqueuedEvent,
+            None,
+            self._schedule_event_handler.handle_session_enqueued,
+            coalescing_opts,
+            name="dispatcher.schedule/enqueue",
+        )
+        event_dispatcher.consume(
+            SessionTerminatedEvent,
+            None,
+            self._schedule_event_handler.handle_session_terminated,
+            coalescing_opts,
+            name="dispatcher.term",
+        )
+        event_dispatcher.consume(
+            AgentStartedEvent,
+            None,
+            self._schedule_event_handler.handle_agent_started,
+            name="dispatcher.schedule",
+        )
+        event_dispatcher.consume(
+            DoScheduleEvent, None, self._schedule_event_handler.handle_do_schedule, coalescing_opts
+        )
+        event_dispatcher.consume(
+            DoStartSessionEvent, None, self._schedule_event_handler.handle_do_start_session
+        )
+        event_dispatcher.consume(
+            DoCheckPrecondEvent, None, self._schedule_event_handler.handle_do_check_precond
+        )
+        event_dispatcher.consume(DoScaleEvent, None, self._schedule_event_handler.handle_do_scale)
+        event_dispatcher.consume(
+            DoUpdateSessionStatusEvent,
+            None,
+            self._schedule_event_handler.handle_do_update_session_status,
+        )
+
     def _dispatch_session_events(self, event_dispatcher: EventDispatcher) -> None:
         # action-trigerring events
         event_dispatcher.consume(

src/ai/backend/manager/event_dispatcher/dispatch.py

@@ -0,0 +1 @@
+python_sources(name="src")

src/ai/backend/manager/event_dispatcher/handlers/BUILD

@@ -16,8 +16,8 @@
 from ai.backend.manager.errors.exceptions import InstanceNotFound
 from ai.backend.manager.registry import AgentRegistry
 
-from ..models.agent import AgentStatus, agents
-from ..models.utils import (
+from ...models.agent import AgentStatus, agents
+from ...models.utils import (
     ExtendedAsyncSAEngine,
 )
 

src/ai/backend/manager/event_dispatcher/handlers/__init__.py

@@ -12,7 +12,7 @@
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.manager.registry import AgentRegistry
 
-from ..models.utils import (
+from ...models.utils import (
     ExtendedAsyncSAEngine,
 )
 

src/ai/backend/manager/event_dispatcher/agent.py renamed to src/ai/backend/manager/event_dispatcher/handlers/agent.py

@@ -21,8 +21,8 @@
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.manager.registry import AgentRegistry
 
-from ..models.kernel import kernels
-from ..models.utils import (
+from ...models.kernel import kernels
+from ...models.utils import (
     ExtendedAsyncSAEngine,
     execute_with_retry,
 )

src/ai/backend/manager/event_dispatcher/image.py renamed to src/ai/backend/manager/event_dispatcher/handlers/image.py

@@ -19,19 +19,19 @@
 from ai.backend.manager.errors.exceptions import SessionNotFound
 from ai.backend.manager.registry import AgentRegistry
 
-from ..models.endpoint import EndpointRow
-from ..models.image import ImageIdentifier, ImageRow
-from ..models.keypair import KeyPairRow
-from ..models.routing import RouteStatus, RoutingRow
-from ..models.session import KernelLoadingStrategy, SessionRow
-from ..models.user import UserRow
-from ..models.utils import (
+from ...models.endpoint import EndpointRow
+from ...models.image import ImageIdentifier, ImageRow
+from ...models.keypair import KeyPairRow
+from ...models.routing import RouteStatus, RoutingRow
+from ...models.session import KernelLoadingStrategy, SessionRow
+from ...models.user import UserRow
+from ...models.utils import (
     ExtendedAsyncSAEngine,
     execute_with_retry,
     is_db_retry_error,
 )
-from ..types import UserScope
-from ..utils import query_userinfo
+from ...types import UserScope
+from ...utils import query_userinfo
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 

src/ai/backend/manager/event_dispatcher/kernel.py renamed to src/ai/backend/manager/event_dispatcher/handlers/kernel.py

@@ -4,7 +4,7 @@
 from ai.backend.manager.errors.exceptions import InternalServerError
 
 
-class PropagatorEventDispatcher:
+class PropagatorEventHandler:
     _event_hub: EventHub
 
     def __init__(self, event_hub: EventHub) -> None:

src/ai/backend/manager/event_dispatcher/model_serving.py renamed to src/ai/backend/manager/event_dispatcher/handlers/model_serving.py

@@ -0,0 +1,61 @@
+import logging
+
+from ai.backend.common.events.agent import AgentStartedEvent
+from ai.backend.common.events.schedule import (
+    DoCheckPrecondEvent,
+    DoScaleEvent,
+    DoScheduleEvent,
+    DoStartSessionEvent,
+)
+from ai.backend.common.events.session import (
+    DoUpdateSessionStatusEvent,
+    SessionEnqueuedEvent,
+    SessionTerminatedEvent,
+)
+from ai.backend.logging.utils import BraceStyleAdapter
+from ai.backend.manager.scheduler.dispatcher import SchedulerDispatcher
+
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
+
+class ScheduleEventHandler:
+    _scheduler_dispatcher: SchedulerDispatcher
+
+    def __init__(self, scheduler_dispatcher: SchedulerDispatcher) -> None:
+        self._scheduler_dispatcher = scheduler_dispatcher
+
+    async def handle_session_enqueued(
+        self, context: None, agent_id: str, ev: SessionEnqueuedEvent
+    ) -> None:
+        await self._scheduler_dispatcher.schedule(ev.event_name())
+
+    async def handle_session_terminated(
+        self, context: None, agent_id: str, ev: SessionTerminatedEvent
+    ) -> None:
+        await self._scheduler_dispatcher.schedule(ev.event_name())
+
+    async def handle_agent_started(
+        self, context: None, agent_id: str, ev: AgentStartedEvent
+    ) -> None:
+        await self._scheduler_dispatcher.schedule(ev.event_name())
+
+    async def handle_do_schedule(self, context: None, agent_id: str, ev: DoScheduleEvent) -> None:
+        await self._scheduler_dispatcher.schedule(ev.event_name())
+
+    async def handle_do_start_session(
+        self, context: None, agent_id: str, ev: DoStartSessionEvent
+    ) -> None:
+        await self._scheduler_dispatcher.start(ev.event_name())
+
+    async def handle_do_check_precond(
+        self, context: None, agent_id: str, ev: DoCheckPrecondEvent
+    ) -> None:
+        await self._scheduler_dispatcher.check_precond(ev.event_name())
+
+    async def handle_do_scale(self, context: None, agent_id: str, ev: DoScaleEvent) -> None:
+        await self._scheduler_dispatcher.scale_services(ev.event_name())
+
+    async def handle_do_update_session_status(
+        self, context: None, agent_id: str, ev: DoUpdateSessionStatusEvent
+    ) -> None:
+        await self._scheduler_dispatcher.update_session_status()

src/ai/backend/manager/event_dispatcher/propagator.py renamed to src/ai/backend/manager/event_dispatcher/handlers/propagator.py

@@ -32,10 +32,10 @@
 from ai.backend.manager.errors.exceptions import SessionNotFound
 from ai.backend.manager.registry import AgentRegistry
 
-from ..models.endpoint import EndpointLifecycle, EndpointRow
-from ..models.routing import RouteStatus, RoutingRow
-from ..models.session import KernelLoadingStrategy, SessionRow
-from ..models.utils import (
+from ...models.endpoint import EndpointLifecycle, EndpointRow
+from ...models.routing import RouteStatus, RoutingRow
+from ...models.session import KernelLoadingStrategy, SessionRow
+from ...models.utils import (
     ExtendedAsyncSAEngine,
     execute_with_retry,
     is_db_retry_error,

src/ai/backend/manager/event_dispatcher/handlers/schedule.py

@@ -9,10 +9,10 @@
 )
 from ai.backend.logging import BraceStyleAdapter
 
-from ..models.utils import (
+from ...models.utils import (
     ExtendedAsyncSAEngine,
 )
-from ..models.vfolder import VFolderOperationStatus, update_vfolder_status
+from ...models.vfolder import VFolderOperationStatus, update_vfolder_status
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))