Skip to content

Add Trackio Callback for Optuna Registry#377

Open
ParagEkbote wants to merge 7 commits into
optuna:mainfrom
ParagEkbote:add-trackio-integration
Open

Add Trackio Callback for Optuna Registry#377
ParagEkbote wants to merge 7 commits into
optuna:mainfrom
ParagEkbote:add-trackio-integration

Conversation

@ParagEkbote
Copy link
Copy Markdown

@ParagEkbote ParagEkbote commented May 21, 2026

Contributor Agreements

Please read the contributor agreements and if you agree, please click the checkbox below.

  • I agree to the contributor agreements.

Tip

Please follow the Quick TODO list to smoothly merge your PR.

Motivation

As described in the PR, we port the trackio callback from optuna-integration to optuna-registry. A couple of new features have been introduced:

  • Automatic syncing of optuna trials with trackio-sync

  • Better resume semantics for trials.

Could you please review the changes?

cc: @c-bata

Description of the changes

Included a init.py and callback.py which are the core files of the packages, also included the relevant license and readme file to describe the working of the callback.

TODO List towards PR Merge

Please remove this section if this PR is not an addition of a new package.
Otherwise, please check the following TODO list:

  • Copy ./template/ to create your package
  • Replace <COPYRIGHT HOLDER> in LICENSE of your package with your name
  • Fill out README.md in your package
  • Add import statements of your function or class names to be used in __init__.py
  • (Optional) Add from __future__ import annotations at the head of any Python files that include typing to support older Python versions
  • Apply the formatter based on the tips in README.md
  • Check whether your module works as intended based on the tips in README.md

@c-bata
Copy link
Copy Markdown
Member

c-bata commented May 22, 2026

@gen740 Could you review this PR?

@c-bata c-bata added the new-package New packages label May 22, 2026
Comment thread package/callbacks/trackio/README.md Outdated
Comment thread package/callbacks/trackio/LICENSE Outdated
Comment on lines +93 to +120
import optuna
import optunahub


module = optunahub.load_module("callbacks/trackio")
TrackioCallback = module.TrackioCallback


def objective(trial):
x = trial.suggest_float("x", -10, 10)
return (x - 2) ** 2


study = optuna.create_study(
study_name="trackio-demo",
)

trackioc = TrackioCallback(
project="my-optuna-study",
)

study.optimize(
objective,
n_trials=10,
callbacks=[trackioc],
)

trackioc.finish()
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried running the example locally, and got:

RuntimeError: Call trackio.init() before trackio.log().

It works fine after adding trackio.init(), but maybe it would be more user-friendly to include that setup in the example?

Copy link
Copy Markdown
Author

@ParagEkbote ParagEkbote May 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There was an issue in the implementation of the callback where __call__() did not correctly initialize the trials for trackio. I have updated the callback implementation to fix this and re-tested with the examples which now work correctly.

@ParagEkbote ParagEkbote requested a review from gen740 May 31, 2026 15:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

new-package New packages

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants