Event Bus & Hooks

from d9d.loop.control import TrainTask, RegisterTaskEventsContext
from d9d.loop.event import (
    subscribe,
    subscribe_annotated
)
from d9d.loop.event.catalogue.train import (
    EVENT_TRAIN_MODEL_STAGES_READY,
    EVENT_TRAIN_STEP_POST,
    EventModelStagesReadyContext,
    EventStepContext,
)

class CustomTrainTask(TrainTask):
    def __init__(self):
        self._modules = []

    def register_events(self, ctx: RegisterTaskEventsContext) -> None:
        # Automatically scans this instance for @subscribe decorators
        subscribe_annotated(ctx.event_bus, self)

    @subscribe(EVENT_TRAIN_MODEL_STAGES_READY)
    def _on_model_ready(self, ctx: EventModelStagesReadyContext) -> None:
        self._modules = ctx.modules

    @subscribe(EVENT_TRAIN_STEP_POST)
    def _on_step_post(self, ctx: EventStepContext) -> None:
        for module in self._modules:
            print(f"Step {ctx.stepper.current_step} completed; MoE routing stats: {module.moe_stats}.")

    def compute_loss(self, ctx):
        ... # Task math overrides

from d9d.loop.control import TrainTask, RegisterTaskEventsContext
from d9d.loop.event import subscribe
from d9d.loop.event.catalogue.train import (
    EVENT_TRAIN_OPTIMIZER_READY
)

class CustomTrainTask(TrainTask):
    def register_events(self, ctx: RegisterTaskEventsContext) -> None:
        ctx.event_bus.subscribe(
            EVENT_TRAIN_OPTIMIZER_READY, 
            lambda event_ctx: print(f"Optimizer loaded: {event_ctx.optimizer}")
        )

from d9d.loop.event import EventBus, Event
import dataclasses
from pathlib import Path

@dataclasses.dataclass(kw_only=True)
class CheckpointContext:
    step: int
    path: Path

# Define a new event
EVENT_CHECKPOINT_SAVED = Event[CheckpointContext](id="user.checkpoint_saved")

# Trigger it from somewhere in your task
def process_something(bus: EventBus):
    bus.trigger(
        EVENT_CHECKPOINT_SAVED, 
        CheckpointContext(step=1000, path=Path("/checkpoints/step_1000"))
    )