diff --git a/package/callbacks/trackio/LICENSE b/package/callbacks/trackio/LICENSE new file mode 100644 index 00000000..c7d4dcab --- /dev/null +++ b/package/callbacks/trackio/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 y0z, neel04, nabenabe0928, and nzw0301 + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/package/callbacks/trackio/README.md b/package/callbacks/trackio/README.md new file mode 100644 index 00000000..902b9f12 --- /dev/null +++ b/package/callbacks/trackio/README.md @@ -0,0 +1,191 @@ +--- +author: ParagEkbote +title: Trackio Callback +description: A callback to track Optuna trials with Trackio. +tags: [callback, trackio, logging, built-in] +optuna_versions: [4.9.0] +license: MIT License +--- + +## Installation + +```shell +pip install trackio +``` + +## Abstract + +This callback enables tracking of Optuna studies in Trackio. By default, the study is tracked as a single experiment run, where all suggested hyperparameters and optimized metrics are logged and visualized as a function of optimizer steps. + +Trackio is offline-first and does not require authentication for local experiment tracking. Optionally, tracked experiments can be synchronized to Hugging Face Spaces for remote visualization and sharing. + +The callback also supports multi-run mode, where each Optuna trial is tracked as an independent Trackio run. This is useful for sweep-style dashboards, parameter importance analysis, and per-trial experiment inspection. + +## APIs + +- `TrackioCallback(project: str, metric_name: str | Sequence[str] = "value", as_multirun: bool = False, space_id: str | None = None, dataset_id: str | None = None, private: bool | None = None, resume: str = "allow", sync_on_finish: bool = False, sync_frequency: str = "study", sync_run_in_background: bool = False, trackio_kwargs: dict[str, Any] | None = None)` + + - `project`: + Name of the Trackio project used for local storage and optional synchronization. + + - `metric_name`: + Name assigned to the optimized metric. In case of multi-objective optimization, a list of names can be passed. These names will be assigned to objective values in the order returned by the objective function. + + If a single name is provided, it will be broadcast to multiple objectives using numerical suffixes such as `value_0`, `value_1`. + + - `as_multirun`: + Creates a new Trackio run for each Optuna trial. Useful for generating sweep-style dashboards and trial-level visualizations. + + - `space_id`: + Optional Hugging Face Space ID (`"username/space-name"`) used for synchronization and remote visualization. + + - `dataset_id`: + Optional Hugging Face Dataset ID used for exporting experiment metadata and metrics. + + - `private`: + Whether synchronized Hugging Face artifacts should be private. + + - `resume`: + Resume policy for Trackio runs. Accepted values are `"allow"`, `"must"`, and `"never"`. + + - `sync_on_finish`: + Whether to synchronize the project to Hugging Face after study completion. + + - `sync_frequency`: + Synchronization frequency strategy. + + - `"study"` synchronizes once after the study completes. + - `"trial"` synchronizes after each completed trial in multirun mode. + + - `sync_run_in_background`: + Whether synchronization should run asynchronously in a background thread. + + - `trackio_kwargs`: + Additional keyword arguments passed directly to `trackio.init()`. + +- `TrackioCallback.track_in_trackio() -> Callable` + + - Decorator for enabling Trackio logging inside the objective function. + + The decorator initializes and finalizes Trackio runs automatically. Additional metrics logged inside the objective function using `trackio.log()` are associated with the corresponding Optuna trial run. + + Use as: + + ```python + @trackioc.track_in_trackio() + ``` + +- `TrackioCallback.finish() -> None` + + - Explicitly finalizes synchronization and cleanup after `study.optimize()` completes. + +## Example + +### Add Trackio callback to Optuna optimization + +```python +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() +``` + +### Trackio logging in multirun mode + +```python +import optuna +import optunahub +import trackio + + +module = optunahub.load_module("callbacks/trackio") +TrackioCallback = module.TrackioCallback + + +trackioc = TrackioCallback( + project="my-optuna-study", + as_multirun=True, +) + + +@trackioc.track_in_trackio() +def objective(trial): + x = trial.suggest_float("x", -10, 10) + + # Additional logging inside Trackio + trackio.log( + { + "power": 2, + "base_of_metric": x - 2, + } + ) + + return (x - 2) ** 2 + + +study = optuna.create_study( + study_name="trackio-multirun", +) + +study.optimize( + objective, + n_trials=10, + callbacks=[trackioc], +) + +trackioc.finish() +``` + +## Notes + +- Trackio synchronization to Hugging Face Spaces is eventually consistent and may take time to become remotely visible. + +- For large studies or multirun experiments, it is strongly recommended to: + + 1. complete the Optuna study locally first, + 1. verify local experiment tracking, + 1. then synchronize results to Hugging Face. + +- In most cases, the recommended configuration is: + +```python +TrackioCallback( + ..., + sync_on_finish=True, + sync_frequency="study", +) +``` + +instead of per-trial synchronization. + +- Per-trial synchronization may significantly increase runtime due to repeated remote uploads and Hugging Face Space propagation delays. + +- To ensure proper Trackio lifecycle management in multi-run mode, the objective function should always be wrapped with: + +```python +@trackioc.track_in_trackio() +``` diff --git a/package/callbacks/trackio/__init__.py b/package/callbacks/trackio/__init__.py new file mode 100644 index 00000000..1f590c39 --- /dev/null +++ b/package/callbacks/trackio/__init__.py @@ -0,0 +1,4 @@ +from .callback import TrackioCallback + + +__all__ = ["TrackioCallback"] diff --git a/package/callbacks/trackio/callback.py b/package/callbacks/trackio/callback.py new file mode 100644 index 00000000..859886be --- /dev/null +++ b/package/callbacks/trackio/callback.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +from collections.abc import Callable +from collections.abc import Sequence +import functools +from typing import Any +from typing import cast +from typing import TYPE_CHECKING +import uuid + +import optuna +from optuna._imports import try_import + + +with try_import() as _imports: + import trackio + + +if TYPE_CHECKING: + from optuna.study.study import ObjectiveFuncType + + +class TrackioCallback: + """Callback to track Optuna trials with Trackio. + + This callback enables tracking of an Optuna study using Trackio. + By default, the entire study is recorded as a single experiment run, + where all suggested hyperparameters and optimized metrics are logged + and visualized as a function of optimizer steps. + + Trackio is offline-first and does not require authentication for local + usage. Optionally, results can be synchronized to Hugging Face Spaces + or exported as Hugging Face Datasets for sharing, visualization, and + reproducibility. + + .. note:: + Trackio does not require users to be logged in for local experiment + tracking. Authentication is only required when synchronizing results + to Hugging Face Hub (e.g., Spaces or Datasets). + + .. note:: + Unlike Weights & Biases, Trackio does not rely on global mutable state. + Each run is explicitly initialized and finalized, which makes this + callback safe to use in long-running processes and research pipelines. + + .. note:: + To ensure deterministic trial ordering in logged metrics, this + callback should only be used with ``study.optimize(n_jobs=1)``. + Parallel optimization may result in out-of-order steps. + + """ + + def __init__( + self, + project: str, + metric_name: str | Sequence[str] = "value", + *, + as_multirun: bool = False, + space_id: str | None = None, + dataset_id: str | None = None, + private: bool | None = None, + resume: str = "allow", + sync_on_finish: bool = False, + sync_frequency: str = "study", + sync_run_in_background: bool = True, + trackio_kwargs: dict[str, Any] | None = None, + ) -> None: + _imports.check() + + if not isinstance(metric_name, (str, Sequence)): + raise TypeError(f"metric_name must be str or Sequence[str], got {type(metric_name)}") + + if sync_frequency not in {"study", "trial"}: + raise ValueError("sync_frequency must be either 'study' or 'trial'") + + self._project: str = project + self._metric_name: str | Sequence[str] = metric_name + self._as_multirun: bool = as_multirun + self._space_id: str | None = space_id + self._dataset_id: str | None = dataset_id + self._private: bool | None = private + self._resume: str = resume + self._sync_on_finish: bool = sync_on_finish + self._sync_frequency: str = sync_frequency + self._sync_run_in_background: bool = sync_run_in_background + self._trackio_kwargs: dict[str, Any] = trackio_kwargs or {} + + # Explicit internal state + self._objective_wrapped: bool = False + self._resolved_run_name: str | None = None + self._study_instance_id: str | None = None + self._study_run_initialized: bool = False + + def __call__( + self, + study: optuna.study.Study, + trial: optuna.trial.FrozenTrial, + ) -> None: + if trial.values is None: + return + + # Multirun logging is handled entirely inside the wrapped objective lifecycle. + if self._as_multirun: + return + + self._ensure_study_run_initialized(study) + + metrics = self._build_metrics(trial) + + trackio.log( + { + **trial.params, + **metrics, + "trial_number": trial.number, + }, + step=trial.number, + ) + + def track_in_trackio(self) -> Callable: + """Decorator for enabling Trackio logging inside the objective function. + + This decorator wraps an Optuna objective function so that a Trackio run + is initialized before the objective executes and finalized afterward. + Any calls to :func:`trackio.log` inside the objective will be associated + with the correct run. + + The decorator is required when logging from inside the objective + function, since Optuna callbacks are invoked *after* a trial finishes + and therefore cannot manage per-trial runtime state. + + When ``as_multirun=True``, a separate Trackio run is created for each + Optuna trial. When ``as_multirun=False``, all trials are logged into a + single run. + + Returns: + A wrapped objective function with Trackio logging enabled. + """ + + def decorator(func: ObjectiveFuncType) -> ObjectiveFuncType: + self._objective_wrapped = True + wrapped = self._wrap_objective(func) + + @functools.wraps(func) + def wrapper(trial: optuna.trial.Trial) -> Any: + return wrapped(trial) + + return wrapper + + return decorator + + def finish(self) -> None: + """Finalize Trackio synchronization and cleanup. + + Must be called after ``study.optimize()`` completes to close the + single-run session and trigger any end-of-study synchronization. + For multirun mode, individual trial runs are closed automatically + after each trial; this method handles the optional study-level sync. + """ + if self._sync_on_finish: + if not self._as_multirun or self._sync_frequency == "study": + self._safe_sync() + + if not self._as_multirun and self._study_run_initialized: + cast(Any, trackio).finish() + + def _safe_sync(self) -> None: + try: + trackio.sync( + project=self._project, + run_in_background=self._sync_run_in_background, + ) + except TimeoutError as exc: + print( + "Trackio sync timed out while waiting for " + "remote visibility. Local experiment data " + "was successfully recorded and uploaded. " + f"Original error: {exc}" + ) + + def _initialize_study_identity( + self, + study_name: str | None, + ) -> None: + if self._resolved_run_name is not None: + return + + resolved_study_name = study_name or "optuna-study" + + if self._study_instance_id is None: + self._study_instance_id = uuid.uuid4().hex[:8] + + self._resolved_run_name = f"{resolved_study_name}-{self._study_instance_id}" + + def _ensure_study_run_initialized( + self, + study: optuna.study.Study, + ) -> None: + if self._as_multirun: + return + + if self._study_run_initialized: + return + + self._initialize_study_identity(study.study_name) + + trackio.init( + project=self._project, + name=cast(str, self._resolved_run_name), + space_id=self._space_id, + dataset_id=self._dataset_id, + private=self._private, + resume=self._resume, + **self._trackio_kwargs, + ) + + self._study_run_initialized = True + + trackio.log({"study_name": study.study_name}) + + def _wrap_objective(self, func: ObjectiveFuncType) -> ObjectiveFuncType: + @functools.wraps(func) + def wrapped(trial: optuna.trial.Trial) -> Any: + study = trial.study + + self._initialize_study_identity(study.study_name) + + base_name = cast(str, self._resolved_run_name) + + if self._as_multirun: + run_name = f"{base_name}/trial-{trial.number}" + + trackio.init( + project=self._project, + name=run_name, + space_id=self._space_id, + dataset_id=self._dataset_id, + private=self._private, + resume=self._resume, + **self._trackio_kwargs, + ) + + else: + self._ensure_study_run_initialized(study) + + trial_completed = False + + try: + result = func(trial) + + if self._as_multirun: + values = result if isinstance(result, (list, tuple)) else [result] + metrics = self._build_result_metrics(values) + + trackio.log( + { + **trial.params, + **metrics, + "trial_number": trial.number, + }, + step=trial.number, + ) + + trial_completed = True + return result + + except optuna.exceptions.TrialPruned: + trackio.log({"trial_state": "pruned"}) + raise + + except Exception as exc: + trackio.log({"trial_state": "failed", "error": str(exc)}) + raise + + finally: + if self._as_multirun: + cast(Any, trackio).finish() + + if ( + trial_completed + and self._sync_on_finish + and self._sync_frequency == "trial" + ): + self._safe_sync() + + return wrapped + + def _build_result_metrics( + self, + values: Sequence[float], + ) -> dict[str, float]: + if isinstance(self._metric_name, str): + if len(values) == 1: + names = [self._metric_name] + else: + names = [f"{self._metric_name}_{i}" for i in range(len(values))] + else: + if len(self._metric_name) != len(values): + raise ValueError( + "Metric names must match number of objectives " + f"({len(self._metric_name)} vs {len(values)})" + ) + names = list(self._metric_name) + + return dict(zip(names, values)) + + def _build_metrics( + self, + trial: optuna.trial.FrozenTrial, + ) -> dict[str, float]: + values = trial.values + assert values is not None + return self._build_result_metrics(values)