"""A dashboard command changes the learning rate between steps.

The loop registers ``set_lr`` at the ``after_step`` safe point and calls
``sync`` once per step. ``sync`` returns immediately when no command is
waiting. Issue the command from the run's Controls tab while this runs.

    python change_learning_rate.py
"""

import time

import launchhelm


class Schedule:
    """The learning rate this fake loop is using."""

    def __init__(self) -> None:
        self.lr = 0.0003

    def set_lr(self, lr: float) -> None:
        """Store the requested learning rate."""
        self.lr = lr


schedule = Schedule()
with launchhelm.init("quickstart", config={"lr": schedule.lr}) as run:
    run.command("after_step")(schedule.set_lr)
    for step in range(60):
        loss = 2.0 / (step + 1)
        run.log({"loss": loss, "lr": schedule.lr}, step=step)
        run.sync("after_step", step=step)
        # Real steps take seconds; the pause leaves a window to issue
        # set_lr from the run's Controls tab.
        time.sleep(0.5)

print("run finished:", run.run_id)