Add wandb-compatible API for trackio (#394)

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>

Author: abidlabs

.changeset/six-crabs-type.md

@@ -0,0 +1,5 @@
+---
+"trackio": minor
+---
+
+feat:Add wandb-compatible API for trackio

docs/source/_toctree.yml

@@ -17,6 +17,8 @@
     title: Deploy and Embed Dashboards
   - local: manage
     title: Manage Projects
+  - local: python_api
+    title: Python API for Managing Runs
   - local: cli_commands
     title: CLI Commands
   - local: api_mcp_server

docs/source/python_api.md

@@ -0,0 +1,152 @@
+# Python API for Managing Runs
+
+Trackio provides a Python API class (`trackio.Api()`) that allows you to programmatically manage runs in your projects. This API is similar to `wandb.Api()` and provides methods to delete runs, move runs between projects, and access run information.
+
+**Note:** This is different from [Trackio as an API Server](api_mcp_server.md), which runs the Trackio dashboard as a web server with API endpoints. The Python API (`trackio.Api()`) is a client-side interface for managing runs in your local Trackio database.
+
+## Basic Usage
+
+```python
+import trackio
+
+# Initialize the API
+api = trackio.Api()
+
+# Get all runs in a project
+runs = api.runs("my_project")
+
+# Access individual runs
+for run in runs:
+    print(f"Run: {run.name}, Project: {run.project}")
+    print(f"Config: {run.config}")
+
+# Or access by index
+first_run = runs[0]
+```
+
+## Deleting Runs
+
+```python
+api = trackio.Api()
+runs = api.runs("my_project")
+
+# Delete a specific run
+run = runs[0]
+success = run.delete()  # Returns True if successful
+```
+
+## Moving Runs Between Projects
+
+```python
+api = trackio.Api()
+runs = api.runs("source_project")
+
+# Move a run to a different project
+run = runs[0]
+success = run.move("target_project")  # Returns True if successful
+
+# After moving, the run object's project is updated
+print(run.project)  # "target_project"
+```
+
+When you move a run, all associated data is transferred:
+- All metrics and logs
+- Run configuration
+- System metrics
+- Media files (images, videos, audio)
+
+The run is completely removed from the source project and added to the target project.
+
+## API Reference
+
+### Api
+
+Main entry point for the Trackio Python API.
+
+```python
+api = trackio.Api()
+```
+
+#### Methods
+
+- **`runs(project: str) -> Runs`**: Returns a collection of runs for the specified project. Raises `ValueError` if the project doesn't exist.
+
+### Runs
+
+A collection of runs that supports iteration and indexing.
+
+```python
+runs = api.runs("my_project")
+len(runs)  # Number of runs
+runs[0]    # First run
+for run in runs:  # Iterate over runs
+    ...
+```
+
+### Run
+
+Represents a single run in a project.
+
+#### Properties
+
+- **`id`**: The run name (same as `name`)
+- **`name`**: The run name
+- **`project`**: The project this run belongs to
+- **`config`**: The run's configuration dictionary (lazy-loaded)
+
+#### Methods
+
+- **`delete() -> bool`**: Deletes the run from its project. Returns `True` if successful, `False` otherwise.
+- **`move(new_project: str) -> bool`**: Moves the run to a different project. Returns `True` if successful, `False` otherwise. Updates the run's `project` property after a successful move.
+
+## Examples
+
+### List all runs across projects
+
+```python
+import trackio
+from trackio.sqlite_storage import SQLiteStorage
+
+api = trackio.Api()
+
+# Get all projects
+projects = SQLiteStorage.get_projects()
+
+# List runs in each project
+for project in projects:
+    print(f"\nProject: {project}")
+    runs = api.runs(project)
+    for run in runs:
+        print(f"  - {run.name}")
+```
+
+### Clean up old runs
+
+```python
+api = trackio.Api()
+runs = api.runs("my_project")
+
+# Delete runs older than a certain date
+from datetime import datetime
+cutoff_date = datetime(2024, 1, 1)
+
+for run in runs:
+    if run.config and "_Created" in run.config:
+        created = datetime.fromisoformat(run.config["_Created"])
+        if created < cutoff_date:
+            run.delete()
+            print(f"Deleted old run: {run.name}")
+```
+
+### Organize runs by moving them
+
+```python
+api = trackio.Api()
+
+# Move all runs from "experiments" to "archive"
+source_runs = api.runs("experiments")
+for run in source_runs:
+    run.move("archive")
+    print(f"Moved {run.name} to archive")
+```
+

examples/api-example.py

@@ -0,0 +1,22 @@
+import trackio
+from trackio import Api
+
+project = "api_example_project"
+
+for i in range(3):
+    run_name = f"training_run_{i}"
+    trackio.init(project=project, name=run_name)
+
+    for step in range(5):
+        trackio.log(
+            {
+                "loss": 1.0 / (step + 1),
+                "accuracy": 0.5 + step * 0.1,
+            }
+        )
+
+    trackio.finish()
+
+api = Api()
+runs = api.runs(project)
+runs[0].delete()

tests/conftest.py

@@ -23,7 +23,7 @@ def temp_dir(monkeypatch):
     with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
         for name in ["trackio.sqlite_storage"]:
             monkeypatch.setattr(f"{name}.TRACKIO_DIR", Path(tmpdir))
-        for name in ["trackio.media.media", "trackio.media.utils"]:
+        for name in ["trackio.media.media", "trackio.media.utils", "trackio.utils"]:
             monkeypatch.setattr(f"{name}.MEDIA_DIR", Path(tmpdir) / "media")
         yield tmpdir
 

tests/e2e-local/test_api.py

@@ -0,0 +1,103 @@
+import trackio
+from trackio import Api
+from trackio.sqlite_storage import SQLiteStorage
+
+
+def test_delete_run(temp_dir):
+    project = "test_delete_project"
+    run_name = "test_delete_run"
+
+    trackio.init(project=project, name=run_name)
+    trackio.log(metrics={"loss": 0.1, "accuracy": 0.9})
+    trackio.log(metrics={"loss": 0.2, "accuracy": 0.95})
+    trackio.finish()
+
+    logs = SQLiteStorage.get_logs(project=project, run=run_name)
+    assert len(logs) == 2
+    assert logs[0]["loss"] == 0.1
+    assert logs[1]["loss"] == 0.2
+
+    api = Api()
+    runs = api.runs(project)
+    run = runs[0]
+    assert run.name == run_name
+
+    success = run.delete()
+    assert success is True
+
+    logs_after = SQLiteStorage.get_logs(project=project, run=run_name)
+    assert len(logs_after) == 0
+
+    config_after = SQLiteStorage.get_run_config(project=project, run=run_name)
+    assert config_after is None
+
+    runs_after = SQLiteStorage.get_runs(project=project)
+    assert run_name not in runs_after
+
+
+def test_move_run(temp_dir, image_ndarray):
+    source_project = "test_move_source"
+    target_project = "test_move_target"
+    run_name = "test_move_run"
+
+    trackio.init(project=source_project, name=run_name)
+
+    image1 = trackio.Image(image_ndarray, caption="test_image_1")
+    image2 = trackio.Image(image_ndarray, caption="test_image_2")
+
+    trackio.log(metrics={"loss": 0.1, "acc": 0.9, "img1": image1})
+    trackio.log(metrics={"loss": 0.2, "acc": 0.95, "img2": image2})
+    trackio.finish()
+
+    source_logs = SQLiteStorage.get_logs(project=source_project, run=run_name)
+    assert len(source_logs) == 2
+    assert source_logs[0]["loss"] == 0.1
+    assert source_logs[1]["loss"] == 0.2
+
+    image1_path = source_logs[0]["img1"].get("file_path")
+    assert image1_path is not None
+    normalized_path = str(image1_path).replace("\\", "/")
+    assert normalized_path.startswith(f"{source_project}/{run_name}/")
+
+    api = Api()
+    runs = api.runs(source_project)
+    run = runs[0]
+    assert run.name == run_name
+    assert run.project == source_project
+
+    success = run.move(target_project)
+    assert success is True
+    assert run.project == target_project
+
+    target_logs = SQLiteStorage.get_logs(project=target_project, run=run_name)
+    assert len(target_logs) == 2
+    assert target_logs[0]["loss"] == 0.1
+    assert target_logs[1]["loss"] == 0.2
+
+    target_image1_path = target_logs[0]["img1"].get("file_path")
+    assert target_image1_path is not None
+    normalized_path1 = str(target_image1_path).replace("\\", "/")
+    assert normalized_path1.startswith(f"{target_project}/{run_name}/")
+
+    target_image2_path = target_logs[1]["img2"].get("file_path")
+    assert target_image2_path is not None
+    normalized_path2 = str(target_image2_path).replace("\\", "/")
+    assert normalized_path2.startswith(f"{target_project}/{run_name}/")
+
+    source_logs_after = SQLiteStorage.get_logs(project=source_project, run=run_name)
+    assert len(source_logs_after) == 0
+
+    source_runs_after = SQLiteStorage.get_runs(project=source_project)
+    assert run_name not in source_runs_after
+    assert len(source_runs_after) == 0
+
+    target_runs = SQLiteStorage.get_runs(project=target_project)
+    assert run_name in target_runs
+
+    source_config_after = SQLiteStorage.get_run_config(
+        project=source_project, run=run_name
+    )
+    assert source_config_after is None
+
+    target_config = SQLiteStorage.get_run_config(project=target_project, run=run_name)
+    assert target_config is not None

trackio/__init__.py

@@ -15,6 +15,7 @@
 from huggingface_hub.errors import LocalTokenNotFoundError
 
 from trackio import context_vars, deploy, utils
+from trackio.api import Api
 from trackio.deploy import sync
 from trackio.gpu import gpu_available, log_gpu
 from trackio.histogram import Histogram
@@ -57,6 +58,7 @@
     "Audio",
     "Table",
     "Histogram",
+    "Api",
 ]
 
 Image = TrackioImage

trackio/api.py

@@ -0,0 +1,66 @@
+from typing import Iterator
+
+from trackio.sqlite_storage import SQLiteStorage
+
+
+class Run:
+    def __init__(self, project: str, name: str):
+        self.project = project
+        self.name = name
+        self._config = None
+
+    @property
+    def id(self) -> str:
+        return self.name
+
+    @property
+    def config(self) -> dict | None:
+        if self._config is None:
+            self._config = SQLiteStorage.get_run_config(self.project, self.name)
+        return self._config
+
+    def delete(self) -> bool:
+        return SQLiteStorage.delete_run(self.project, self.name)
+
+    def move(self, new_project: str) -> bool:
+        success = SQLiteStorage.move_run(self.project, self.name, new_project)
+        if success:
+            self.project = new_project
+        return success
+
+    def __repr__(self) -> str:
+        return f"<Run {self.name} in project {self.project}>"
+
+
+class Runs:
+    def __init__(self, project: str):
+        self.project = project
+        self._runs = None
+
+    def _load_runs(self):
+        if self._runs is None:
+            run_names = SQLiteStorage.get_runs(self.project)
+            self._runs = [Run(self.project, name) for name in run_names]
+
+    def __iter__(self) -> Iterator[Run]:
+        self._load_runs()
+        return iter(self._runs)
+
+    def __getitem__(self, index: int) -> Run:
+        self._load_runs()
+        return self._runs[index]
+
+    def __len__(self) -> int:
+        self._load_runs()
+        return len(self._runs)
+
+    def __repr__(self) -> str:
+        self._load_runs()
+        return f"<Runs project={self.project} count={len(self._runs)}>"
+
+
+class Api:
+    def runs(self, project: str) -> Runs:
+        if not SQLiteStorage.get_project_db_path(project).exists():
+            raise ValueError(f"Project '{project}' does not exist")
+        return Runs(project)