Add Trackio experiment tracker integration

docs/book/component-guide/experiment-trackers/README.md

@@ -37,6 +37,7 @@ Experiment Trackers are optional stack components provided by integrations:
 | [Comet](comet.md)                  | `comet`   | `comet`     | Add Comet experiment tracking and visualization capabilities to your ZenML pipelines            |
 | [MLflow](mlflow.md)                | `mlflow`  | `mlflow`    | Add MLflow experiment tracking and visualization capabilities to your ZenML pipelines           |
 | [Neptune](neptune.md)              | `neptune` | `neptune`   | Add Neptune experiment tracking and visualization capabilities to your ZenML pipelines          |
+| [Trackio](trackio.md)              | `trackio` | `trackio`   | Add Trackio experiment tracking and visualization capabilities to your ZenML pipelines          |
 | [Weights & Biases](wandb.md)       | `wandb`   | `wandb`     | Add Weights & Biases experiment tracking and visualization capabilities to your ZenML pipelines |
 | [Custom Implementation](custom.md) | _custom_  |             | _custom_                                                                                        |
 

docs/book/component-guide/experiment-trackers/trackio.md

@@ -0,0 +1,154 @@
+---
+description: Logging and visualizing experiments with Trackio.
+---
+
+# Trackio
+
+The Trackio Experiment Tracker is an [Experiment Tracker](./) flavor provided by the Trackio-ZenML integration. It connects your ZenML pipelines to [Trackio](https://github.com/gradio-app/trackio), a lightweight experiment tracking solution maintained by the Gradio team, so you can visualize and compare the artifacts and metrics generated by your runs.
+
+### When would you want to use it?
+
+Trackio focuses on offering a simple workflow for collecting model metrics and artifacts while staying close to the developer experience that Gradio users are familiar with. You should consider the Trackio Experiment Tracker if:
+
+* you already use Trackio (or plan to) to keep a history of training runs and evaluations and would like to surface the results of your automated ZenML pipelines in the same place
+* you prefer a managed experiment tracking option that integrates nicely with the Gradio/Spaces ecosystem without the complexity of heavier experiment tracking suites
+* you want to experiment locally with an easy to self-host service before committing to a larger-scale tracking solution
+
+You may want to use one of the other [Experiment Tracker flavors](./#experiment-tracker-flavors) if you rely on advanced governance features (e.g. model registry, lineage dashboards) that are better covered by tools like MLflow, Neptune, or Weights & Biases.
+
+### How do you deploy it?
+
+The Trackio Experiment Tracker flavor is provided through the Trackio ZenML integration. Install the integration to make the component available for registration:
+
+```shell
+zenml integration install trackio -y
+```
+
+Once installed, you can register the experiment tracker and add it to a stack. The Trackio tracker accepts the following configuration attributes:
+
+* `workspace` (optional): Trackio workspace/organization slug. Leave empty to use the default workspace associated with your API key.
+* `project` (optional): Project identifier inside Trackio where new runs should be stored.
+* `api_key` (optional): API token used to authenticate against Trackio. When omitted, the library falls back to the environment or CLI authentication.
+* `base_url` (optional): Override the default Trackio endpoint, e.g. when talking to a self-hosted Trackio server.
+
+### Authentication methods
+
+You can authenticate with Trackio using the same mechanisms recommended for other experiment trackers. Storing credentials inside a [ZenML Secret](https://docs.zenml.io/how-to/project-setup-and-management/interact-with-secrets) keeps your configuration secure and auditable.
+
+{% tabs %}
+{% tab title="ZenML Secret (Recommended)" %}
+
+```shell
+zenml secret create trackio_secret --api_key=<TRACKIO_API_KEY>
+
+zenml experiment-tracker register trackio_tracker \
+    --flavor=trackio \
+    --workspace=<workspace_slug> \
+    --project=<project_slug> \
+    --api_key={{trackio_secret.api_key}}
+
+zenml stack register trackio_stack -e trackio_tracker ... --set
+```
+
+{% endtab %}
+{% tab title="Basic authentication" %}
+
+You can also provide the credentials directly on the command line (useful for quick experiments):
+
+```shell
+zenml experiment-tracker register trackio_tracker \
+    --flavor=trackio \
+    --workspace=<workspace_slug> \
+    --project=<project_slug> \
+    --api_key=<TRACKIO_API_KEY>
+
+zenml stack register trackio_stack -e trackio_tracker ... --set
+```
+
+{% hint style="warning" %}
+Providing credentials inline is not recommended for production use because they will be stored in plain text in the stack configuration.
+{% endhint %}
+
+{% endtab %}
+{% endtabs %}
+
+For more information about the configuration parameters supported by the Trackio experiment tracker, refer to the [SDK docs](https://sdkdocs.zenml.io/latest/integration_code_docs/integrations-trackio.html).
+
+### How do you use it?
+
+Once the experiment tracker is part of the active stack, enable it for the steps that should log to Trackio by using the `experiment_tracker` argument of the `@step` decorator. Inside the step you can fetch the Trackio run object and use the regular Trackio API to log metrics, files or metadata:
+
+```python
+from zenml import step
+from zenml.integrations.trackio.experiment_trackers import get_trackio_run
+
+@step(experiment_tracker="trackio_tracker")
+def evaluate_model(accuracy: float, params: dict) -> None:
+    run = get_trackio_run()
+
+    # Log high level metrics / metadata using the Trackio client API
+    if hasattr(run, "log"):
+        run.log({"accuracy": accuracy})
+    elif hasattr(run, "log_metrics"):
+        run.log_metrics({"accuracy": accuracy})
+
+    if hasattr(run, "log_params"):
+        run.log_params(params)
+    elif hasattr(run, "__setitem__"):
+        run["parameters"] = params
+```
+
+Any metadata logged during the step is persisted by Trackio and the resulting dashboard URL is attached to the ZenML run metadata. You can retrieve it programmatically:
+
+```python
+from zenml.client import Client
+
+pipeline_run = Client().get_pipeline_run("<PIPELINE_RUN_NAME>")
+step = pipeline_run.steps["evaluate_model"]
+experiment_tracker_url = step.run_metadata["experiment_tracker_url"].value
+```
+
+{% hint style="info" %}
+Instead of hardcoding the experiment tracker name, you can obtain it dynamically from the active stack:
+
+```python
+from zenml.client import Client
+
+tracker_name = Client().active_stack.experiment_tracker.name
+
+@step(experiment_tracker=tracker_name)
+def my_step(...):
+    ...
+```
+{% endhint %}
+
+### Additional settings
+
+The Trackio experiment tracker exposes additional runtime settings that can be provided via the `@step` decorator:
+
+* `run_name`: Override the auto-generated run name that ZenML derives from the pipeline and step names.
+* `tags`: Attach custom tags to the Trackio run. ZenML always appends the pipeline name, pipeline run name, and step name to help with filtering.
+* `metadata`: Provide a dictionary of values that should be logged immediately after the run is created (e.g. static parameters, environment metadata).
+
+These settings can be declared when defining the step:
+
+```python
+from zenml import step
+from zenml.integrations.trackio.flavors import TrackioExperimentTrackerSettings
+
+@step(
+    experiment_tracker="trackio_tracker",
+    settings={
+        "experiment_tracker": TrackioExperimentTrackerSettings(
+            run_name="nightly-eval",
+            tags=["nightly", "evaluation"],
+            metadata={"dataset": "imagenet", "split": "validation"},
+        )
+    },
+)
+def nightly_evaluation() -> None:
+    run = get_trackio_run()
+    run.log({"f1": 0.81})
+```
+
+As with other experiment trackers, Trackio runs are automatically marked as failed when the corresponding ZenML pipeline step fails.

docs/book/component-guide/toc.md

@@ -52,6 +52,7 @@
   * [Comet](experiment-trackers/comet.md)
   * [MLflow](experiment-trackers/mlflow.md)
   * [Neptune](experiment-trackers/neptune.md)
+  * [Trackio](experiment-trackers/trackio.md)
   * [Weights & Biases](experiment-trackers/wandb.md)
   * [Google Cloud VertexAI Experiment Tracker](experiment-trackers/vertexai.md)
   * [Develop a custom experiment tracker](experiment-trackers/custom.md)

src/zenml/integrations/constants.py

@@ -70,6 +70,7 @@
 SKYPILOT_KUBERNETES = "skypilot_kubernetes"
 SLACK = "slack"
 SPARK = "spark"
+TRACKIO = "trackio"
 TEKTON = "tekton"
 TENSORBOARD = "tensorboard"
 TENSORFLOW = "tensorflow"

src/zenml/integrations/trackio/__init__.py

@@ -0,0 +1,42 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Module containing the Trackio integration."""
+
+from typing import List, Type
+
+from zenml.integrations.constants import TRACKIO
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+TRACKIO_EXPERIMENT_TRACKER_FLAVOR = "trackio"
+
+
+class TrackioIntegration(Integration):
+    """Definition of the Trackio integration for ZenML."""
+
+    NAME = TRACKIO
+    REQUIREMENTS = ["trackio>=0.1.0"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Trackio integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.trackio.flavors import (
+            TrackioExperimentTrackerFlavor,
+        )
+
+        return [TrackioExperimentTrackerFlavor]

src/zenml/integrations/trackio/experiment_trackers/__init__.py

@@ -0,0 +1,10 @@
+"""Trackio experiment tracker implementation."""
+
+from zenml.integrations.trackio.experiment_trackers.run_state import (
+    get_trackio_run,
+)
+from zenml.integrations.trackio.experiment_trackers.trackio_experiment_tracker import (
+    TrackioExperimentTracker,
+)
+
+__all__ = ["TrackioExperimentTracker", "get_trackio_run"]

src/zenml/integrations/trackio/experiment_trackers/run_state.py

@@ -0,0 +1,57 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Utilities to access the active Trackio run from pipeline steps."""
+
+from typing import Any
+
+from zenml.client import Client
+
+
+def get_trackio_run() -> Any:
+    """Return the Trackio run associated with the currently executing step.
+
+    Returns:
+        The Trackio run instance created for the active step.
+
+    Raises:
+        RuntimeError: If the active stack does not contain a Trackio experiment
+            tracker or the tracker has not yet been initialized.
+    """
+
+    from zenml.integrations.trackio.experiment_trackers import (
+        TrackioExperimentTracker,
+    )
+
+    experiment_tracker = Client().active_stack.experiment_tracker
+
+    if experiment_tracker is None:
+        raise RuntimeError(
+            "Unable to access a Trackio run: No experiment tracker is "
+            "configured for the active stack."
+        )
+
+    if not isinstance(experiment_tracker, TrackioExperimentTracker):
+        raise RuntimeError(
+            "Unable to access a Trackio run: The experiment tracker in the "
+            "active stack is not the Trackio integration."
+        )
+
+    try:
+        return experiment_tracker.active_run
+    except RuntimeError as exc:
+        raise RuntimeError(
+            "Unable to access a Trackio run: The tracker was not initialized. "
+            "Ensure that the step enabling the experiment tracker is running "
+            "when calling this helper."
+        ) from exc