Report training metrics from Python

Goal: let atune schedule a training loop owned by another library.

The seam is atune.Trial.report and atune.Trial.should_prune. Report a validation or evaluation metric at comparable resource boundaries. report itself raises atune.Pruned when the scheduler prunes; let it leave the objective so Study.optimize records the trial's state and last intermediate value. Close environments in finally blocks so pruning cleans up too.

These executable examples use a fixed study seed, one worker, CPU execution, small budgets, and no downloads or display. They demonstrate integration; they make no learning-quality claim. Each finishes by taking an Optuna reading snapshot and serializing a temporary HTML optimization-history plot. The temporary directory is removed after the file is checked.

Install and run

Build and install the wheel as described in Install. For the reproducible Linux x86_64, CPython 3.10 example environment:

python -m pip install --require-hashes --only-binary=:all: --extra-index-url https://download.pytorch.org/whl/cpu -r .github/release/python-adoption-requirements.txt -r .github/release/python-callback-requirements.txt
python examples/python/lightning_callback.py
python examples/python/sb3_callback.py
python examples/python/cleanrl_loop.py

The callback lock selects torch==2.7.1+cpu from the PyTorch CPU index, with hashes for torch and every transitive dependency. It targets the private self-hosted Linux/X64 runner (glibc 2.28 or newer), not the wheel's full platform matrix. Other environments may install Lightning, Stable-Baselines3, Gymnasium, CPU PyTorch, and atune[optuna] with Plotly independently.

Lightning validation callback

A one-layer regression model trains for three epochs on 64 synthetic rows. The callback ignores sanity validation and reports epoch-level val_loss. Three learning rates are tried, with the Median scheduler comparing losses.

Python

class AtuneValidation(pl.Callback):
    """Report validation loss at epoch boundaries; propagate pruning."""

    def __init__(self, trial: atune.Trial) -> None:
        self.trial = trial
        self.value = float("inf")

    def on_validation_epoch_end(
        self, trainer: pl.Trainer, module: pl.LightningModule
    ) -> None:
        if trainer.sanity_checking:
            return
        self.value = float(trainer.callback_metrics["val_loss"].item())
        # report itself raises atune.Pruned when the scheduler stops the trial.
        self.trial.report(self.value, step=trainer.current_epoch + 1)
        if self.trial.should_prune():
            raise atune.Pruned()

Stable-Baselines3 evaluation callback

PPO trains on CartPole-v1 for at most 512 steps per trial, four trials total. Every 256 training steps, the callback evaluates two episodes on a separate seeded environment and reports mean reward. Environments cap episodes at 100 steps and close on completion or pruning.

Python

class AtuneEvaluation(BaseCallback):
    """Evaluate a separate, seeded environment every 256 training steps."""

    def __init__(self, trial: atune.Trial, evaluation: gym.Env, seed: int) -> None:
        super().__init__()
        self.trial = trial
        self.evaluation = evaluation
        self.seed = seed
        self.reward = 0.0

    def _on_step(self) -> bool:
        if self.num_timesteps % 256 == 0:
            self.evaluation.reset(seed=self.seed)
            self.reward, _ = evaluate_policy(
                self.model,
                self.evaluation,
                n_eval_episodes=2,
                deterministic=True,
            )
            self.trial.report(self.reward, step=self.num_timesteps)
            if self.trial.should_prune():
                raise atune.Pruned()
        return True

CleanRL-style single-file loop

A small REINFORCE loop trains six short CartPole episodes per trial, three trials total. This example needs PyTorch and Gymnasium, without importing CleanRL. Greedy evaluation uses a separate environment with the same seed at each boundary. Reports use episodes as their resource unit so the scheduler compares the same progress boundary despite varying episode lengths.

Python

# The scheduler compares curves at the same resource boundary.
# Episodes are the resource unit here; environment steps may differ.
trial.report(reward, step=episode + 1)
if trial.should_prune():
    raise atune.Pruned()

Evidence

The Rust workflow's installed-wheel adoption gate requires Optuna 5.0.0 and CPU torch, runs the adapter tests, and executes all three files end to end. The documentation workflow's Python example gate discovers and runs them too. Both environments use the committed hash locks. No GPU execution, hosted CI run, or public deployment is implied by a local pass.