Add Trackio Integration for Optuna by ParagEkbote · Pull Request #259 · optuna/optuna-integration

ParagEkbote · 2025-12-18T19:57:24Z

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

…in toml.

c-bata · 2025-12-24T07:23:35Z

@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.

optuna_integration/trackio/trackio.py

Co-authored-by: c-bata <contact@c-bata.link>

github-actions · 2026-01-04T23:07:10Z

This pull request has not seen any recent activity.

ParagEkbote · 2026-01-09T09:32:09Z

Could you please review the changes?

cc: @c-bata

c-bata

I left two early feedback comments 🙏

README.md

c-bata · 2026-01-09T10:56:20Z

tests/trackio/test_trackio.py

@@ -0,0 +1,203 @@
+from __future__ import annotations


My first impression is that most of these test scenarios rely heavily on unittest.mock, which makes it difficult to see the long-term value of maintaining them.

For example, if TrackioCallback were to break due to future changes in trackio, these tests would not detect that issue. Given this, I wonder if it might make sense to remove these tests. What do you think?

I have modeled them similarly to the WeightsAndBiasesCallback tests, since the public api of trackio is modeled based on wandb. What type of alternative testing strategy do you think would work?

cc: @c-bata

One option might be to remove these tests altogether. Since they rely heavily on mocks, I’m not sure they provide much value in catching real regressions in trackio itself.

The only alternative I can think of would be integration tests against the actual trackio behavior, but those tend to be fragile and hard to maintain. From a maintainability perspective, I would personally prefer to avoid such tests and instead address issues as they are reported.

I have removed the tests and can also remove the gh-action additions if required . Could you please review?

Co-authored-by: c-bata <contact@c-bata.link>

c-bata

Thank you for the update. Changes itself look good to me. I will merge this PR after following CI issue will be resolved:
https://github.com/optuna/optuna-integration/actions/runs/20967917853/job/60551934863?pr=259

c-bata · 2026-01-16T05:42:34Z

@ParagEkbote Thank you for your investigation. It seems that the cause of the problem is not related to your changes. I’m now investigating the root cause. 🙏

c-bata · 2026-01-16T06:01:33Z

@ParagEkbote I fixed the CI issue at #264. Could you merge the latest main branch? 🙏

c-bata

Thank you for the update.

ParagEkbote added 15 commits December 18, 2025 17:05

Add the initial files and update deps for integration.

3e9dd4c

Add an initial implementation.

ade11c0

lint and fix styling.

c282c7e

register the callback in init.py

b2dab0e

also register in the callback itself.

39c5d24

update for fixing of multi-run trials.

b7e0fc7

update some more to fix errors.

fd28c4d

register the tests in gh actions

81c87ef

add the tests for trackio.

fd77a5b

make improvements in tests to mirror the callback and fix whitespace …

1ab1c39

…in toml.

add missing return types for mypy

f5f6a7b

Add ref to docs and fix mypy errors.

4e2832a

add type hints in test file and mention in trackio

3f02188

improve docstring for the decorator and fix indent in docstring.

b44be95

fix flake8 issues.

dc65e7a

ParagEkbote marked this pull request as ready for review December 19, 2025 11:11

ParagEkbote added 2 commits December 19, 2025 11:23

fix flake8 issue.

07981b2

update comments.

fd6aa49

not522 assigned not522 and c-bata and unassigned not522 Dec 24, 2025

c-bata added the feature Change that does not break compatibility, but affects the public interfaces. label Dec 24, 2025

specify python version in gh actions.

57340e4

c-bata reviewed Dec 26, 2025

View reviewed changes

optuna_integration/trackio/trackio.py Outdated Show resolved Hide resolved

ParagEkbote and others added 2 commits December 26, 2025 18:24

Update optuna_integration/trackio/trackio.py

59d25a2

Co-authored-by: c-bata <contact@c-bata.link>

add a missing param to suppress warnings, make style.

f2937fa

ParagEkbote requested a review from c-bata December 26, 2025 14:17

update type hints

70813de

github-actions bot added the stale Exempt from stale bot labeling. label Jan 4, 2026

c-bata removed the stale Exempt from stale bot labeling. label Jan 5, 2026

c-bata reviewed Jan 9, 2026

View reviewed changes

ParagEkbote and others added 2 commits January 9, 2026 19:21

Update README.md

48d0571

Co-authored-by: c-bata <contact@c-bata.link>

rm tests.

07fb081

ParagEkbote requested a review from c-bata January 13, 2026 18:18

rm files as well.

7ebd60d

c-bata reviewed Jan 16, 2026

View reviewed changes

update docstrings.

498db19

c-bata mentioned this pull request Jan 16, 2026

Remove FastAIV2PruningCallback from docs to resolve warning #264

Merged

Merge branch 'optuna:main' into add-trackio-integration

6df38fc

ParagEkbote requested a review from c-bata January 16, 2026 07:29

c-bata approved these changes Jan 16, 2026

View reviewed changes

c-bata merged commit 6bf37a5 into optuna:main Jan 16, 2026
129 checks passed

c-bata added this to the v4.8.0 milestone Jan 16, 2026

ParagEkbote deleted the add-trackio-integration branch January 16, 2026 13:26

This was referenced Mar 23, 2026

Add Trackio as a New Backend for Experiment Visualization NVIDIA-NeMo/RL#2142

Open

[Feature] Add Trackio as a New Backend for Experiment Visualization inclusionAI/AReaL#1089

Open

Add Trackio as a New Backend for Experiment Visualization alibaba/ROLL#402

Open

Conversation

ParagEkbote commented Dec 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Description of the changes

Uh oh!

c-bata commented Dec 24, 2025

Uh oh!

c-bata left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented Jan 4, 2026

Uh oh!

ParagEkbote commented Jan 9, 2026

Uh oh!

c-bata left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

c-bata Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

ParagEkbote Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

c-bata Jan 13, 2026

Choose a reason for hiding this comment

Uh oh!

ParagEkbote Jan 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

c-bata left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

c-bata commented Jan 16, 2026

Uh oh!

c-bata commented Jan 16, 2026

Uh oh!

c-bata left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ParagEkbote commented Dec 18, 2025 •

edited

Loading

ParagEkbote Jan 9, 2026 •

edited

Loading

ParagEkbote Jan 13, 2026 •

edited

Loading

c-bata left a comment •

edited

Loading