Add Trackio Integration for Ray

[ParagEkbote]

Description

As described in the issue, this PR adds the integration of trackio with ray. Trackio has a drop-in API similar to W&B, 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 pushed to HF Hub as a dataset or as a Space as seen below, you can view them locally.

Could you please review?

Related issues

Fixes #60708

Additional information

You can the test the same with the following script:

import random
import time

import numpy as np
import trackio

import ray
from ray import tune
from ray.air.integrations.trackio import (
    TrackioLoggerCallback,
    setup_trackio,
)
from ray.train import RunConfig, ScalingConfig
from ray.train.torch import TorchTrainer


PROJECT_NAME = "trackio-ray-demo"

HF_DATASET_ID = (
    "AINovice2005/ray-trackio-experiments"  
)
HF_SPACE_ID = (
    "AINovice2005/ray-trackio-dashboard" 
)

NUM_STEPS = 15


def tune_trainable(config):

    for step in range(NUM_STEPS):

        loss = (config["lr"] * 10) / (step + 1) + random.random()
        accuracy = 1 / (loss + 1e-3)

        # Example artifact
        image = np.random.rand(64, 64, 3)

        tune.report(
            {
                "loss": loss,
                "accuracy": accuracy,
                "image_mean": float(image.mean()),
                "step": step,
            }
        )

        time.sleep(0.2)


def run_tune_example():

    tuner = tune.Tuner(
        tune_trainable,
        param_space={
            "lr": tune.grid_search([0.001, 0.01, 0.1]),
        },
        run_config=tune.RunConfig(
            name="trackio-ray-tune-demo",
            callbacks=[
                TrackioLoggerCallback(
                    project=PROJECT_NAME,
                    # Trackio capabilities
                    auto_log_gpu=True,
                    gpu_log_interval=5,
                    dataset_id=HF_DATASET_ID,
                    space_id=HF_SPACE_ID,
                )
            ],
        ),
    )

    results = tuner.fit()

    print("\nTune finished\n")
    print(results)


def train_loop(config):

    run = setup_trackio(
        config=config,
        project=PROJECT_NAME,
        auto_log_gpu=True,
        gpu_log_interval=5,
        dataset_id=HF_DATASET_ID,
        space_id=HF_SPACE_ID,
    )

    for step in range(NUM_STEPS):

        loss = 5 / (step + 1) + random.random()
        throughput = random.uniform(50, 150)

        # Log metrics
        if run:
            run.log(
                {
                    "loss": loss,
                    "throughput": throughput,
                    "step": step,
                },
                step=step,
            )

        # Example artifact
        sample_image = np.random.rand(64, 64, 3)

        if run:
            run.log(
                {
                    "image_mean": float(sample_image.mean()),
                    "image_std": float(sample_image.std()),
                }
            )

        time.sleep(0.2)

    if run:
        run.finish()


def run_train_example():

    trainer = TorchTrainer(
        train_loop_per_worker=train_loop,
        train_loop_config={"lr": 0.01},
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(name="trackio-ray-train-demo"),
    )

    trainer.fit()


def launch_dashboard():

    print("\nLaunching Trackio dashboard...\n")

    trackio.show(
        project=PROJECT_NAME,
        open_browser=True,
    )


if __name__ == "__main__":

    ray.init()

    print("\nRunning Ray Tune experiment\n")
    run_tune_example()

    print("\nRunning Ray Train experiment\n")
    run_train_example()

    print("\nOpening dashboard\n")
    print("Run dashboard manually with:")
    print('trackio show --project "trackio-ray-demo"')

    trackio.finish()
    ray.shutdown()

    print("\nDemo completed\n")

Local usage

HF Space: https://huggingface.co/spaces/AINovice2005/ray-trackio-dashboard

HF Dataset: https://huggingface.co/datasets/AINovice2005/ray-trackio-experiments

[gemini-code-assist]

Code Review

This pull request introduces a new integration for Trackio with Ray, providing both a setup_trackio function for Ray Train and a TrackioLoggerCallback for Ray Tune. The implementation is clean and follows existing patterns for logger integrations in Ray.

I've identified a few areas for improvement:

• The example in the setup_trackio docstring can lead to an AttributeError and should be corrected.
• The error handling in TrackioLoggerCallback can be improved by logging exceptions instead of silently passing, which will make debugging easier for users.

Overall, this is a great addition. My comments are focused on improving robustness and user experience.

[ParagEkbote]

Could you please review the changes?

cc: @alexeykudinkin

[ParagEkbote]

Which maintainer do I have to reach out to get a first review for this integration?

A gentle ping to @alexeykudinkin

[pseudo-rnd-thoughts]

Thanks for the PR and integration.

I've made a couple of edits based on personal style, reducing the number of new line between code (IMO new line should show a group of code that works together), changed Dict[Trial, object] to Dict[Trial, trackio.Run].
Ray uses Google docstring styling, so I've moved the docsrings to start on the """ rather than next line.

The PR is missing two important things

• Documentation - The other experiment trackers have a page to explain how to use it. For example comet. To do this, you need to create a file in doc/source/tune/examples/tune-trackio.ipynb
• Testing - Comet, WandB, etc have their own testing to ensure that the implementation continues working. Create a couple of tests in ray/python/air/tests/test_integration_trackio.py

I've also put some comments on the rest of the file if you could answer them

[pseudo-rnd-thoughts]

Whats the difference between the excludes and exclude_results? Could we combine the two?

[ParagEkbote]

That's a good point, I have updated the integration to combine them.

[pseudo-rnd-thoughts]

This will silently ignore values to be logged without telling the user why.
I think we should log a warning to the user and why.

[ParagEkbote]

I have added a logging warning about the same.

[pseudo-rnd-thoughts]

This function is now deprecated, is it necessary / can we find a replacement?

[ParagEkbote]

I have removed the import and added a helper function for it, the function is needed since trackio callback needs flat keys for the trials.

[pseudo-rnd-thoughts]

When can the run be None? Looking through the code I can't see anywhere that a trial's run is set to None? If the run is never None, then could you remove all the checking

[ParagEkbote]

There are several edge cases where the trial runs do not get recorded for the callback, especially if a trial errors or partially initializes, the run can be None. Secondly, if trials are resumed after a fault, log_trial_result may be called before log_trial_start, leading to a race condition. This lazy initialization ensures that metrics are not dropped in these cases.

[pseudo-rnd-thoughts]

Ahh ok thanks

[ParagEkbote]

Looking at the CI, the error is

[2026-04-02T09:42:19Z] #7 111.7 The conflict is caused by:
--
[2026-04-02T09:42:19Z] #7 111.7     The user requested gradio<7.0.0 and >=6.8.0
[2026-04-02T09:42:19Z] #7 111.7     The user requested (constraint) gradio==3.50.2

Therefore, it looks like we need to wait for gradio to be updated to support v6 before we can merge this PR

In this repo, a separate integration of ray serve has gradio v3 used. We could update this module with gradio v6 public API and bump the gradio deps in requirements.txt

Can I do this in a separate PR?

[pseudo-rnd-thoughts]

@ParagEkbote One of my colleagues are in the process of updating Gradio to v5 and will look at updating it to v6 afterwards.

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[ParagEkbote]

Not Stale.

[pseudo-rnd-thoughts]

@ParagEkbote We've updated gradio to 5.50.0 however gradio==6.x will require a breaking gradio-client 1.x -> 2.x change. Could you research to understand what changes are necessary both on the packages and implementations for us?

[ParagEkbote]

@ParagEkbote We've updated gradio to 5.50.0 however gradio==6.x will require a breaking gradio-client 1.x -> 2.x change. Could you research to understand what changes are necessary both on the packages and implementations for us?

Hi, I'd like to thank you for continuing to reach out and supporting this PR with your inputs.

I've rechecked the main branch and the gradio version bump is not visible in the integration module and the requirements.txt file. Could you let me know if the gradio dependency has been updated elsewhere, since I'm only aware about these files since they are visible in the buildkite logs.

[ParagEkbote]

As for upgrading the existing gradio integration, could you please point me to the module or file which uses gradio-client as a dependency, I'm unable to locate the imports used within this repo.

cc: @pseudo-rnd-thoughts

[pseudo-rnd-thoughts]

Gradio-client is one of our requirements code however I can't see its use in Python.
@elliot-barn Could you clarify where we needed to use gradio-client and what needs updating

[ParagEkbote]

Gradio-client is one of our requirements code however I can't see its use in Python. @elliot-barn Could you clarify where we needed to use gradio-client and what needs updating

As far as I know, trackio pulls in gradio-client based on their fixed version pin, similarly gradio has a fixed version of gradio-client for their repo. Does this clarification help?

[ParagEkbote]

Do you think that we can enable the tests in CI in a separate PR until the gradio dependency issue is resolved, I'll add a TODO note in BUILD.baze and merge this working implementation of trackio?

cc: @pseudo-rnd-thoughts

[pseudo-rnd-thoughts]

@ParagEkbote Reviewing the CI, your PR is failing due to

[2026-05-08T09:51:39Z]     from ray.air.integrations.trackio import (
--
  | [2026-05-08T09:51:39Z]   File "/rayci/python/ray/air/integrations/trackio.py", line 7, in <module>
  | [2026-05-08T09:51:39Z]     import trackio
  | [2026-05-08T09:51:39Z]   File "/opt/miniforge/lib/python3.10/site-packages/trackio/__init__.py", line 8, in <module>
  | [2026-05-08T09:51:39Z]     from trackio.ui import demo
  | [2026-05-08T09:51:39Z]   File "/opt/miniforge/lib/python3.10/site-packages/trackio/ui.py", line 68, in <module>
  | [2026-05-08T09:51:39Z]     with gr.Sidebar() as sidebar:
  | [2026-05-08T09:51:39Z] AttributeError: module 'gradio' has no attribute 'Sidebar'

How important is Sidebar? Do you know what the minimal gradio version contains this?

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Reviewed by Cursor Bugbot for commit 5deb5ae. Configure here.}

[ParagEkbote]

@ParagEkbote Reviewing the CI, your PR is failing due to

[2026-05-08T09:51:39Z]     from ray.air.integrations.trackio import (
--
  | [2026-05-08T09:51:39Z]   File "/rayci/python/ray/air/integrations/trackio.py", line 7, in <module>
  | [2026-05-08T09:51:39Z]     import trackio
  | [2026-05-08T09:51:39Z]   File "/opt/miniforge/lib/python3.10/site-packages/trackio/__init__.py", line 8, in <module>
  | [2026-05-08T09:51:39Z]     from trackio.ui import demo
  | [2026-05-08T09:51:39Z]   File "/opt/miniforge/lib/python3.10/site-packages/trackio/ui.py", line 68, in <module>
  | [2026-05-08T09:51:39Z]     with gr.Sidebar() as sidebar:
  | [2026-05-08T09:51:39Z] AttributeError: module 'gradio' has no attribute 'Sidebar'

How important is Sidebar? Do you know what the minimal gradio version contains this?

As per the latest version of trackio release, gradio as a dependency has been removed and we need to only set gradio-client>=2.0.0

So, I've pinned the latest version of trackio, but we still need to update gradio-client to the required version.

cc: @pseudo-rnd-thoughts