Add Trackio Integration for Optuna

[ParagEkbote]

Fixes #248

Motivation

As described in the issue, this PR implements an TrackioCallback for optuna-integration. Currently, trackio has a min python version of 3.10 and since python 3.9 has reached EOL, we will be required to bump the min. version in this repo to ensure that CI will be green.

Description of the changes

The implementation of the callback is similar to the wandb callback, as the library is designed to be an open-source and offline alternative to wandb. The main difference between the wandb and trackio is that trackio does not rely on a global mutable state, each run is explicitly created and each run has a fixed lifetime. The API docs have also helped to define the scope of the callback.

The runs can be visualized as follows:

We can also display the results as an HF Space and an dataset. The example is as follows:

import optuna
import trackio

from optuna_integration.trackio import TrackioCallback


def main():
    callback = TrackioCallback(
        project="toy-optuna-trackio-multirun",
        metric_name="loss",
        as_multirun=True,
        resume="allow",


        dataset_id="AINovice2005/toy-optuna-hpo-dataset",
        space_id="AINovice2005/optuna-dashboard",
        private=False,  # set True if you want private
    )

    @callback.track_in_trackio()
    def objective(trial: optuna.trial.Trial) -> float:
        x = trial.suggest_float("x", -10.0, 10.0)
        loss = (x - 2.0) ** 2

        trackio.log(
            {
                "x": x,
                "loss_squared": loss,
            }
        )

        return loss

    study = optuna.create_study(direction="minimize")
    study.optimize(
        objective,
        n_trials=5,
        callbacks=[callback],
        n_jobs=1,
    )

    print("Done multi-run test.")


if __name__ == "__main__":
    main()

If the results are not synced, you can do so with the following command:

trackio sync \
  --project toy-optuna-trackio-multirun \
  --space-id AINovice2005/optuna-dashboard

Could you please review?

cc: @nzw0301

[c-bata]

@ParagEkbote Thank you for your pull request. Could you fix the CI error? I think you can skip Python 3.9 test by specifying python_matrix like below:

optuna-integration/.github/workflows/checks-and-tests.yml

Line 121 in 0b6a416

python_matrix: "['3.9', '3.10', '3.11', '3.12', '3.13']"

[c-bata]

Thank you the updates! I left an early feedback comment.

Due to the year-end and New Year holidays, my review may take a bit longer than usual. I’ll get back to you once I’m able to resume the review.

[github-actions]

This pull request has not seen any recent activity.