Trackio logger

[edufschmidt]

Description & Motivation

I propose adding a dedicated Trackio logger for PyTorch Lightning. Trackio is a lightweight experiment tracking library that works well for local prototyping and small-scale experiments. Trackio is very easy to work with and hack, lightweight, and is based on Gradio. It is ideal for users who want minimal setup or cannot rely on external services. Trackio's API is meant as a drop-in replacement to WandB's, but I believe that a dedicated Lightning logger might be justifiable as per my understanding implementation details make it kind of cumbersome to re-use the existing WanDB logger. Ofc I'd be happy to discuss more suitable alternatives.

Pitch

I would like to add a new TrackioLogger class, modeled after the existing WandB logger, that integrates seamlessly with PyTorch Lightning. The implementation would include proper unit tests and documentation, enabling Lightning users to track experiments with Trackio just as they would with other loggers. This would provide a native, lightweight option for local and flexible experiment tracking.

Alternatives

The closest alternative is running a local WandB server, but this approach is less flexible, especially if customizations are needed to the visualizations/graphics, harder to hack, and requires additional setup (e.g., a Docker container). Trackio offers similar functionality with less overhead, making it a more convenient solution for quick prototyping and lightweight workflows.

Additional context

No response

cc @lantiga

[GdoongMathew]

Hi @edufschmidt

I just implemented one a few months ago in another project. I’m not sure if the following fully meets your use case, but I wanted to share it here as a reference:

class TrackioLogger(Logger):

    def __init__(
        self,
        project: str,
        name: Optional[str] = None,
        resume: Literal["never", "allow", "must"] = "allow",
        **kwargs,
    ):
        if not _TRACKIO_AVAILABLE:
            raise ModuleNotFoundError(str(_TRACKIO_AVAILABLE))
        super().__init__()
        self._project = project
        self._name = name

        kwargs["resume"] = resume
        self._kwargs = kwargs
        self._experiment: Run = None

    @property
    def name(self) -> str:
        return self._project

    @property
    def version(self) -> Optional[Union[int, str]]:
        return self._name if self._experiment is None else self.experiment.name

    @property
    @rank_zero_experiment
    def experiment(self) -> "Run":
        if self._experiment is None:
            import trackio

            run = trackio.init(
                project=self._project,
                name=self._name,
                **self._kwargs,
            )
            self._experiment = run
        return self._experiment

    @override
    @rank_zero_only
    def log_hyperparams(
        self,
        params: Union[dict[str, Any], Namespace],
        *args: Any,
        **kwargs: Any,
    ) -> None:
        params = _convert_params(params)
        params = _sanitize_callable_params(params)
        params = _convert_json_serializable(params)
        self.experiment.config.update(params)

    @override
    @rank_zero_only
    def log_metrics(
        self,
        metrics: dict[str, float],
        step: Optional[int] = None,
    ) -> None:
        self.experiment.log(metrics, step=step)

    @override
    @rank_zero_only
    def finalize(self, *args, **kwargs) -> None:
        if self._experiment is not None and not self.experiment._stop_flag.is_set():
            self.experiment.finish()

    def __del__(self):
        self.finalize()

I would love to finish the remaining unittests if the Lightning team accepts this feature.

[edufschmidt]

Thanks @GdoongMathew. While I haven't tested this yet, it seems quite similar with the implementation I came up with, so I believe it should meet my requirements as well. Looking forward to hearing from the Lightning team if they agree this addition makes sense.