Event Bus and Hooks

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


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

    def register_events(self, context: RegisterTaskEventsContext) -> None:
        # Registers every method of this instance marked with @subscribe.
        subscribe_annotated(context.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:
        print(f"Step {ctx.schedule.current_step} of {ctx.schedule.total_steps} completed.")

    def build_forward_inputs(self, ctx):
        ...  # Task logic

    def compute_loss(self, ctx):
        ...  # Task logic

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


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

    def build_forward_inputs(self, ctx):
        ...  # Task logic

    def compute_loss(self, ctx):
        ...  # Task logic

import dataclasses
from pathlib import Path

from d9d.loop.event import Event, EventBus


@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 your task.
def process_something(bus: EventBus):
    bus.trigger(
        EVENT_CHECKPOINT_SAVED,
        CheckpointContext(step=1000, path=Path("/checkpoints/step_1000")),
    )