Spaces:
Sleeping
Sleeping
Upload folder using huggingface_hub
Browse filesThis view is limited to 50 files because it contains too many changes.
See raw diff
- .gitattributes +1 -0
- CHANGELOG.md +34 -0
- __init__.py +337 -0
- __pycache__/__init__.cpython-312.pyc +0 -0
- __pycache__/cli.cpython-312.pyc +0 -0
- __pycache__/commit_scheduler.cpython-312.pyc +0 -0
- __pycache__/context_vars.cpython-312.pyc +0 -0
- __pycache__/deploy.cpython-312.pyc +0 -0
- __pycache__/dummy_commit_scheduler.cpython-312.pyc +0 -0
- __pycache__/file_storage.cpython-312.pyc +0 -0
- __pycache__/histogram.cpython-312.pyc +0 -0
- __pycache__/imports.cpython-312.pyc +0 -0
- __pycache__/media.cpython-312.pyc +0 -0
- __pycache__/run.cpython-312.pyc +0 -0
- __pycache__/sqlite_storage.cpython-312.pyc +0 -0
- __pycache__/table.cpython-312.pyc +0 -0
- __pycache__/typehints.cpython-312.pyc +0 -0
- __pycache__/utils.cpython-312.pyc +0 -0
- __pycache__/video_writer.cpython-312.pyc +0 -0
- assets/trackio_logo_dark.png +0 -0
- assets/trackio_logo_light.png +0 -0
- assets/trackio_logo_old.png +3 -0
- assets/trackio_logo_type_dark.png +0 -0
- assets/trackio_logo_type_dark_transparent.png +0 -0
- assets/trackio_logo_type_light.png +0 -0
- assets/trackio_logo_type_light_transparent.png +0 -0
- cli.py +37 -0
- commit_scheduler.py +391 -0
- context_vars.py +18 -0
- deploy.py +258 -0
- dummy_commit_scheduler.py +12 -0
- file_storage.py +37 -0
- histogram.py +68 -0
- imports.py +302 -0
- media.py +299 -0
- package.json +6 -0
- py.typed +0 -0
- run.py +192 -0
- sqlite_storage.py +677 -0
- table.py +53 -0
- typehints.py +18 -0
- ui/__init__.py +10 -0
- ui/__pycache__/__init__.cpython-312.pyc +0 -0
- ui/__pycache__/fns.cpython-312.pyc +0 -0
- ui/__pycache__/main.cpython-312.pyc +0 -0
- ui/__pycache__/run_detail.cpython-312.pyc +0 -0
- ui/__pycache__/runs.cpython-312.pyc +0 -0
- ui/fns.py +241 -0
- ui/helpers/__pycache__/run_selection.cpython-312.pyc +0 -0
- ui/helpers/run_selection.py +46 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
assets/trackio_logo_old.png filter=lfs diff=lfs merge=lfs -text
|
CHANGELOG.md
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# trackio
|
| 2 |
+
|
| 3 |
+
## 0.6.0
|
| 4 |
+
|
| 5 |
+
### Features
|
| 6 |
+
|
| 7 |
+
- [#309](https://github.com/gradio-app/trackio/pull/309) [`1df2353`](https://github.com/gradio-app/trackio/commit/1df23534d6c01938c8db9c0f584ffa23e8d6021d) - Add histogram support with wandb-compatible API. Thanks @abidlabs!
|
| 8 |
+
- [#315](https://github.com/gradio-app/trackio/pull/315) [`76ba060`](https://github.com/gradio-app/trackio/commit/76ba06055dc43ca8f03b79f3e72d761949bd19a8) - Add guards to avoid silent fails. Thanks @Xmaster6y!
|
| 9 |
+
- [#313](https://github.com/gradio-app/trackio/pull/313) [`a606b3e`](https://github.com/gradio-app/trackio/commit/a606b3e1c5edf3d4cf9f31bd50605226a5a1c5d0) - No longer prevent certain keys from being used. Instead, dunderify them to prevent collisions with internal usage. Thanks @abidlabs!
|
| 10 |
+
- [#317](https://github.com/gradio-app/trackio/pull/317) [`27370a5`](https://github.com/gradio-app/trackio/commit/27370a595d0dbdf7eebbe7159d2ba778f039da44) - quick fixes for trackio.histogram. Thanks @abidlabs!
|
| 11 |
+
- [#312](https://github.com/gradio-app/trackio/pull/312) [`aa0f3bf`](https://github.com/gradio-app/trackio/commit/aa0f3bf372e7a0dd592a38af699c998363830eeb) - Fix video logging by adding TRACKIO_DIR to allowed_paths. Thanks @abidlabs!
|
| 12 |
+
|
| 13 |
+
## 0.5.3
|
| 14 |
+
|
| 15 |
+
### Features
|
| 16 |
+
|
| 17 |
+
- [#300](https://github.com/gradio-app/trackio/pull/300) [`5e4cacf`](https://github.com/gradio-app/trackio/commit/5e4cacf2e7ce527b4ce60de3a5bc05d2c02c77fb) - Adds more environment variables to allow customization of Trackio dashboard. Thanks @abidlabs!
|
| 18 |
+
|
| 19 |
+
## 0.5.2
|
| 20 |
+
|
| 21 |
+
### Features
|
| 22 |
+
|
| 23 |
+
- [#293](https://github.com/gradio-app/trackio/pull/293) [`64afc28`](https://github.com/gradio-app/trackio/commit/64afc28d3ea1dfd821472dc6bf0b8ed35a9b74be) - Ensures that the TRACKIO_DIR environment variable is respected. Thanks @abidlabs!
|
| 24 |
+
- [#287](https://github.com/gradio-app/trackio/pull/287) [`cd3e929`](https://github.com/gradio-app/trackio/commit/cd3e9294320949e6b8b829239069a43d5d7ff4c1) - fix(sqlite): unify .sqlite extension, allow export when DBs exist, clean WAL sidecars on import. Thanks @vaibhav-research!
|
| 25 |
+
|
| 26 |
+
### Fixes
|
| 27 |
+
|
| 28 |
+
- [#291](https://github.com/gradio-app/trackio/pull/291) [`3b5adc3`](https://github.com/gradio-app/trackio/commit/3b5adc3d1f452dbab7a714d235f4974782f93730) - Fix the wheel build. Thanks @pngwn!
|
| 29 |
+
|
| 30 |
+
## 0.5.1
|
| 31 |
+
|
| 32 |
+
### Fixes
|
| 33 |
+
|
| 34 |
+
- [#278](https://github.com/gradio-app/trackio/pull/278) [`314c054`](https://github.com/gradio-app/trackio/commit/314c05438007ddfea3383e06fd19143e27468e2d) - Fix row orientation of metrics plots. Thanks @abidlabs!
|
__init__.py
ADDED
|
@@ -0,0 +1,337 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import hashlib
|
| 2 |
+
import json
|
| 3 |
+
import logging
|
| 4 |
+
import os
|
| 5 |
+
import warnings
|
| 6 |
+
import webbrowser
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
from typing import Any
|
| 9 |
+
|
| 10 |
+
from gradio.blocks import BUILT_IN_THEMES
|
| 11 |
+
from gradio.themes import Default as DefaultTheme
|
| 12 |
+
from gradio.themes import ThemeClass
|
| 13 |
+
from gradio_client import Client
|
| 14 |
+
from huggingface_hub import SpaceStorage
|
| 15 |
+
|
| 16 |
+
from trackio import context_vars, deploy, utils
|
| 17 |
+
from trackio.histogram import Histogram
|
| 18 |
+
from trackio.imports import import_csv, import_tf_events
|
| 19 |
+
from trackio.media import TrackioImage, TrackioVideo
|
| 20 |
+
from trackio.run import Run
|
| 21 |
+
from trackio.sqlite_storage import SQLiteStorage
|
| 22 |
+
from trackio.table import Table
|
| 23 |
+
from trackio.ui.main import demo
|
| 24 |
+
from trackio.utils import TRACKIO_DIR, TRACKIO_LOGO_DIR
|
| 25 |
+
|
| 26 |
+
logging.getLogger("httpx").setLevel(logging.WARNING)
|
| 27 |
+
|
| 28 |
+
warnings.filterwarnings(
|
| 29 |
+
"ignore",
|
| 30 |
+
message="Empty session being created. Install gradio\\[oauth\\]",
|
| 31 |
+
category=UserWarning,
|
| 32 |
+
module="gradio.helpers",
|
| 33 |
+
)
|
| 34 |
+
|
| 35 |
+
__version__ = json.loads(Path(__file__).parent.joinpath("package.json").read_text())[
|
| 36 |
+
"version"
|
| 37 |
+
]
|
| 38 |
+
|
| 39 |
+
__all__ = [
|
| 40 |
+
"init",
|
| 41 |
+
"log",
|
| 42 |
+
"finish",
|
| 43 |
+
"show",
|
| 44 |
+
"import_csv",
|
| 45 |
+
"import_tf_events",
|
| 46 |
+
"Image",
|
| 47 |
+
"Video",
|
| 48 |
+
"Table",
|
| 49 |
+
"Histogram",
|
| 50 |
+
]
|
| 51 |
+
|
| 52 |
+
Image = TrackioImage
|
| 53 |
+
Video = TrackioVideo
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
config = {}
|
| 57 |
+
|
| 58 |
+
DEFAULT_THEME = "default"
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def init(
|
| 62 |
+
project: str,
|
| 63 |
+
name: str | None = None,
|
| 64 |
+
group: str | None = None,
|
| 65 |
+
space_id: str | None = None,
|
| 66 |
+
space_storage: SpaceStorage | None = None,
|
| 67 |
+
dataset_id: str | None = None,
|
| 68 |
+
config: dict | None = None,
|
| 69 |
+
resume: str = "never",
|
| 70 |
+
settings: Any = None,
|
| 71 |
+
private: bool | None = None,
|
| 72 |
+
embed: bool = True,
|
| 73 |
+
) -> Run:
|
| 74 |
+
"""
|
| 75 |
+
Creates a new Trackio project and returns a [`Run`] object.
|
| 76 |
+
|
| 77 |
+
Args:
|
| 78 |
+
project (`str`):
|
| 79 |
+
The name of the project (can be an existing project to continue tracking or
|
| 80 |
+
a new project to start tracking from scratch).
|
| 81 |
+
name (`str`, *optional*):
|
| 82 |
+
The name of the run (if not provided, a default name will be generated).
|
| 83 |
+
group (`str`, *optional*):
|
| 84 |
+
The name of the group which this run belongs to in order to help organize
|
| 85 |
+
related runs together. You can toggle the entire group's visibilitiy in the
|
| 86 |
+
dashboard.
|
| 87 |
+
space_id (`str`, *optional*):
|
| 88 |
+
If provided, the project will be logged to a Hugging Face Space instead of
|
| 89 |
+
a local directory. Should be a complete Space name like
|
| 90 |
+
`"username/reponame"` or `"orgname/reponame"`, or just `"reponame"` in which
|
| 91 |
+
case the Space will be created in the currently-logged-in Hugging Face
|
| 92 |
+
user's namespace. If the Space does not exist, it will be created. If the
|
| 93 |
+
Space already exists, the project will be logged to it.
|
| 94 |
+
space_storage ([`~huggingface_hub.SpaceStorage`], *optional*):
|
| 95 |
+
Choice of persistent storage tier.
|
| 96 |
+
dataset_id (`str`, *optional*):
|
| 97 |
+
If a `space_id` is provided, a persistent Hugging Face Dataset will be
|
| 98 |
+
created and the metrics will be synced to it every 5 minutes. Specify a
|
| 99 |
+
Dataset with name like `"username/datasetname"` or `"orgname/datasetname"`,
|
| 100 |
+
or `"datasetname"` (uses currently-logged-in Hugging Face user's namespace),
|
| 101 |
+
or `None` (uses the same name as the Space but with the `"_dataset"`
|
| 102 |
+
suffix). If the Dataset does not exist, it will be created. If the Dataset
|
| 103 |
+
already exists, the project will be appended to it.
|
| 104 |
+
config (`dict`, *optional*):
|
| 105 |
+
A dictionary of configuration options. Provided for compatibility with
|
| 106 |
+
`wandb.init()`.
|
| 107 |
+
resume (`str`, *optional*, defaults to `"never"`):
|
| 108 |
+
Controls how to handle resuming a run. Can be one of:
|
| 109 |
+
|
| 110 |
+
- `"must"`: Must resume the run with the given name, raises error if run
|
| 111 |
+
doesn't exist
|
| 112 |
+
- `"allow"`: Resume the run if it exists, otherwise create a new run
|
| 113 |
+
- `"never"`: Never resume a run, always create a new one
|
| 114 |
+
private (`bool`, *optional*):
|
| 115 |
+
Whether to make the Space private. If None (default), the repo will be
|
| 116 |
+
public unless the organization's default is private. This value is ignored
|
| 117 |
+
if the repo already exists.
|
| 118 |
+
settings (`Any`, *optional*):
|
| 119 |
+
Not used. Provided for compatibility with `wandb.init()`.
|
| 120 |
+
embed (`bool`, *optional*, defaults to `True`):
|
| 121 |
+
If running inside a jupyter/Colab notebook, whether the dashboard should
|
| 122 |
+
automatically be embedded in the cell when trackio.init() is called.
|
| 123 |
+
|
| 124 |
+
Returns:
|
| 125 |
+
`Run`: A [`Run`] object that can be used to log metrics and finish the run.
|
| 126 |
+
"""
|
| 127 |
+
if settings is not None:
|
| 128 |
+
warnings.warn(
|
| 129 |
+
"* Warning: settings is not used. Provided for compatibility with wandb.init(). Please create an issue at: https://github.com/gradio-app/trackio/issues if you need a specific feature implemented."
|
| 130 |
+
)
|
| 131 |
+
|
| 132 |
+
if space_id is None and dataset_id is not None:
|
| 133 |
+
raise ValueError("Must provide a `space_id` when `dataset_id` is provided.")
|
| 134 |
+
space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)
|
| 135 |
+
url = context_vars.current_server.get()
|
| 136 |
+
share_url = context_vars.current_share_server.get()
|
| 137 |
+
|
| 138 |
+
if url is None:
|
| 139 |
+
if space_id is None:
|
| 140 |
+
_, url, share_url = demo.launch(
|
| 141 |
+
show_api=False,
|
| 142 |
+
inline=False,
|
| 143 |
+
quiet=True,
|
| 144 |
+
prevent_thread_lock=True,
|
| 145 |
+
show_error=True,
|
| 146 |
+
favicon_path=TRACKIO_LOGO_DIR / "trackio_logo_light.png",
|
| 147 |
+
allowed_paths=[TRACKIO_LOGO_DIR, TRACKIO_DIR],
|
| 148 |
+
)
|
| 149 |
+
else:
|
| 150 |
+
url = space_id
|
| 151 |
+
share_url = None
|
| 152 |
+
context_vars.current_server.set(url)
|
| 153 |
+
context_vars.current_share_server.set(share_url)
|
| 154 |
+
if (
|
| 155 |
+
context_vars.current_project.get() is None
|
| 156 |
+
or context_vars.current_project.get() != project
|
| 157 |
+
):
|
| 158 |
+
print(f"* Trackio project initialized: {project}")
|
| 159 |
+
|
| 160 |
+
if dataset_id is not None:
|
| 161 |
+
os.environ["TRACKIO_DATASET_ID"] = dataset_id
|
| 162 |
+
print(
|
| 163 |
+
f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}"
|
| 164 |
+
)
|
| 165 |
+
if space_id is None:
|
| 166 |
+
print(f"* Trackio metrics logged to: {TRACKIO_DIR}")
|
| 167 |
+
if utils.is_in_notebook() and embed:
|
| 168 |
+
base_url = share_url + "/" if share_url else url
|
| 169 |
+
full_url = utils.get_full_url(
|
| 170 |
+
base_url, project=project, write_token=demo.write_token
|
| 171 |
+
)
|
| 172 |
+
utils.embed_url_in_notebook(full_url)
|
| 173 |
+
else:
|
| 174 |
+
utils.print_dashboard_instructions(project)
|
| 175 |
+
else:
|
| 176 |
+
deploy.create_space_if_not_exists(
|
| 177 |
+
space_id, space_storage, dataset_id, private
|
| 178 |
+
)
|
| 179 |
+
user_name, space_name = space_id.split("/")
|
| 180 |
+
space_url = deploy.SPACE_HOST_URL.format(
|
| 181 |
+
user_name=user_name, space_name=space_name
|
| 182 |
+
)
|
| 183 |
+
print(f"* View dashboard by going to: {space_url}")
|
| 184 |
+
if utils.is_in_notebook() and embed:
|
| 185 |
+
utils.embed_url_in_notebook(space_url)
|
| 186 |
+
context_vars.current_project.set(project)
|
| 187 |
+
|
| 188 |
+
client = None
|
| 189 |
+
if not space_id:
|
| 190 |
+
client = Client(url, verbose=False)
|
| 191 |
+
|
| 192 |
+
if resume == "must":
|
| 193 |
+
if name is None:
|
| 194 |
+
raise ValueError("Must provide a run name when resume='must'")
|
| 195 |
+
if name not in SQLiteStorage.get_runs(project):
|
| 196 |
+
raise ValueError(f"Run '{name}' does not exist in project '{project}'")
|
| 197 |
+
resumed = True
|
| 198 |
+
elif resume == "allow":
|
| 199 |
+
resumed = name is not None and name in SQLiteStorage.get_runs(project)
|
| 200 |
+
elif resume == "never":
|
| 201 |
+
if name is not None and name in SQLiteStorage.get_runs(project):
|
| 202 |
+
warnings.warn(
|
| 203 |
+
f"* Warning: resume='never' but a run '{name}' already exists in "
|
| 204 |
+
f"project '{project}'. Generating a new name and instead. If you want "
|
| 205 |
+
"to resume this run, call init() with resume='must' or resume='allow'."
|
| 206 |
+
)
|
| 207 |
+
name = None
|
| 208 |
+
resumed = False
|
| 209 |
+
else:
|
| 210 |
+
raise ValueError("resume must be one of: 'must', 'allow', or 'never'")
|
| 211 |
+
|
| 212 |
+
run = Run(
|
| 213 |
+
url=url,
|
| 214 |
+
project=project,
|
| 215 |
+
client=client,
|
| 216 |
+
name=name,
|
| 217 |
+
group=group,
|
| 218 |
+
config=config,
|
| 219 |
+
space_id=space_id,
|
| 220 |
+
)
|
| 221 |
+
|
| 222 |
+
if resumed:
|
| 223 |
+
print(f"* Resumed existing run: {run.name}")
|
| 224 |
+
else:
|
| 225 |
+
print(f"* Created new run: {run.name}")
|
| 226 |
+
|
| 227 |
+
context_vars.current_run.set(run)
|
| 228 |
+
globals()["config"] = run.config
|
| 229 |
+
return run
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
def log(metrics: dict, step: int | None = None) -> None:
|
| 233 |
+
"""
|
| 234 |
+
Logs metrics to the current run.
|
| 235 |
+
|
| 236 |
+
Args:
|
| 237 |
+
metrics (`dict`):
|
| 238 |
+
A dictionary of metrics to log.
|
| 239 |
+
step (`int`, *optional*):
|
| 240 |
+
The step number. If not provided, the step will be incremented
|
| 241 |
+
automatically.
|
| 242 |
+
"""
|
| 243 |
+
run = context_vars.current_run.get()
|
| 244 |
+
if run is None:
|
| 245 |
+
raise RuntimeError("Call trackio.init() before trackio.log().")
|
| 246 |
+
run.log(
|
| 247 |
+
metrics=metrics,
|
| 248 |
+
step=step,
|
| 249 |
+
)
|
| 250 |
+
|
| 251 |
+
|
| 252 |
+
def finish():
|
| 253 |
+
"""
|
| 254 |
+
Finishes the current run.
|
| 255 |
+
"""
|
| 256 |
+
run = context_vars.current_run.get()
|
| 257 |
+
if run is None:
|
| 258 |
+
raise RuntimeError("Call trackio.init() before trackio.finish().")
|
| 259 |
+
run.finish()
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
def show(
|
| 263 |
+
project: str | None = None,
|
| 264 |
+
theme: str | ThemeClass | None = None,
|
| 265 |
+
mcp_server: bool | None = None,
|
| 266 |
+
):
|
| 267 |
+
"""
|
| 268 |
+
Launches the Trackio dashboard.
|
| 269 |
+
|
| 270 |
+
Args:
|
| 271 |
+
project (`str`, *optional*):
|
| 272 |
+
The name of the project whose runs to show. If not provided, all projects
|
| 273 |
+
will be shown and the user can select one.
|
| 274 |
+
theme (`str` or `ThemeClass`, *optional*):
|
| 275 |
+
A Gradio Theme to use for the dashboard instead of the default Gradio theme,
|
| 276 |
+
can be a built-in theme (e.g. `'soft'`, `'citrus'`), a theme from the Hub
|
| 277 |
+
(e.g. `"gstaff/xkcd"`), or a custom Theme class. If not provided, the
|
| 278 |
+
`TRACKIO_THEME` environment variable will be used, or if that is not set, the
|
| 279 |
+
default Gradio theme will be used.
|
| 280 |
+
mcp_server (`bool`, *optional*):
|
| 281 |
+
If `True`, the Trackio dashboard will be set up as an MCP server and certain
|
| 282 |
+
functions will be added as MCP tools. If `None` (default behavior), then the
|
| 283 |
+
`GRADIO_MCP_SERVER` environment variable will be used to determine if the
|
| 284 |
+
MCP server should be enabled (which is `"True"` on Hugging Face Spaces).
|
| 285 |
+
"""
|
| 286 |
+
theme = theme or os.environ.get("TRACKIO_THEME", DEFAULT_THEME)
|
| 287 |
+
|
| 288 |
+
if theme != DEFAULT_THEME:
|
| 289 |
+
# TODO: It's a little hacky to reproduce this theme-setting logic from Gradio Blocks,
|
| 290 |
+
# but in Gradio 6.0, the theme will be set in `launch()` instead, which means that we
|
| 291 |
+
# will be able to remove this code.
|
| 292 |
+
if isinstance(theme, str):
|
| 293 |
+
if theme.lower() in BUILT_IN_THEMES:
|
| 294 |
+
theme = BUILT_IN_THEMES[theme.lower()]
|
| 295 |
+
else:
|
| 296 |
+
try:
|
| 297 |
+
theme = ThemeClass.from_hub(theme)
|
| 298 |
+
except Exception as e:
|
| 299 |
+
warnings.warn(f"Cannot load {theme}. Caught Exception: {str(e)}")
|
| 300 |
+
theme = DefaultTheme()
|
| 301 |
+
if not isinstance(theme, ThemeClass):
|
| 302 |
+
warnings.warn("Theme should be a class loaded from gradio.themes")
|
| 303 |
+
theme = DefaultTheme()
|
| 304 |
+
demo.theme: ThemeClass = theme
|
| 305 |
+
demo.theme_css = theme._get_theme_css()
|
| 306 |
+
demo.stylesheets = theme._stylesheets
|
| 307 |
+
theme_hasher = hashlib.sha256()
|
| 308 |
+
theme_hasher.update(demo.theme_css.encode("utf-8"))
|
| 309 |
+
demo.theme_hash = theme_hasher.hexdigest()
|
| 310 |
+
|
| 311 |
+
_mcp_server = (
|
| 312 |
+
mcp_server
|
| 313 |
+
if mcp_server is not None
|
| 314 |
+
else os.environ.get("GRADIO_MCP_SERVER", "False") == "True"
|
| 315 |
+
)
|
| 316 |
+
|
| 317 |
+
_, url, share_url = demo.launch(
|
| 318 |
+
show_api=_mcp_server,
|
| 319 |
+
quiet=True,
|
| 320 |
+
inline=False,
|
| 321 |
+
prevent_thread_lock=True,
|
| 322 |
+
favicon_path=TRACKIO_LOGO_DIR / "trackio_logo_light.png",
|
| 323 |
+
allowed_paths=[TRACKIO_LOGO_DIR, TRACKIO_DIR],
|
| 324 |
+
mcp_server=_mcp_server,
|
| 325 |
+
)
|
| 326 |
+
|
| 327 |
+
base_url = share_url + "/" if share_url else url
|
| 328 |
+
full_url = utils.get_full_url(
|
| 329 |
+
base_url, project=project, write_token=demo.write_token
|
| 330 |
+
)
|
| 331 |
+
|
| 332 |
+
if not utils.is_in_notebook():
|
| 333 |
+
print(f"* Trackio UI launched at: {full_url}")
|
| 334 |
+
webbrowser.open(full_url)
|
| 335 |
+
utils.block_main_thread_until_keyboard_interrupt()
|
| 336 |
+
else:
|
| 337 |
+
utils.embed_url_in_notebook(full_url)
|
__pycache__/__init__.cpython-312.pyc
ADDED
|
Binary file (14.8 kB). View file
|
|
|
__pycache__/cli.cpython-312.pyc
ADDED
|
Binary file (1.73 kB). View file
|
|
|
__pycache__/commit_scheduler.cpython-312.pyc
ADDED
|
Binary file (18.8 kB). View file
|
|
|
__pycache__/context_vars.cpython-312.pyc
ADDED
|
Binary file (933 Bytes). View file
|
|
|
__pycache__/deploy.cpython-312.pyc
ADDED
|
Binary file (10.3 kB). View file
|
|
|
__pycache__/dummy_commit_scheduler.cpython-312.pyc
ADDED
|
Binary file (1.03 kB). View file
|
|
|
__pycache__/file_storage.cpython-312.pyc
ADDED
|
Binary file (1.65 kB). View file
|
|
|
__pycache__/histogram.cpython-312.pyc
ADDED
|
Binary file (2.98 kB). View file
|
|
|
__pycache__/imports.cpython-312.pyc
ADDED
|
Binary file (13.2 kB). View file
|
|
|
__pycache__/media.cpython-312.pyc
ADDED
|
Binary file (15 kB). View file
|
|
|
__pycache__/run.cpython-312.pyc
ADDED
|
Binary file (9.3 kB). View file
|
|
|
__pycache__/sqlite_storage.cpython-312.pyc
ADDED
|
Binary file (31.4 kB). View file
|
|
|
__pycache__/table.cpython-312.pyc
ADDED
|
Binary file (2.34 kB). View file
|
|
|
__pycache__/typehints.cpython-312.pyc
ADDED
|
Binary file (920 Bytes). View file
|
|
|
__pycache__/utils.cpython-312.pyc
ADDED
|
Binary file (26.8 kB). View file
|
|
|
__pycache__/video_writer.cpython-312.pyc
ADDED
|
Binary file (5.34 kB). View file
|
|
|
assets/trackio_logo_dark.png
ADDED
|
assets/trackio_logo_light.png
ADDED
|
assets/trackio_logo_old.png
ADDED
|
Git LFS Details
|
assets/trackio_logo_type_dark.png
ADDED
|
assets/trackio_logo_type_dark_transparent.png
ADDED
|
assets/trackio_logo_type_light.png
ADDED
|
assets/trackio_logo_type_light_transparent.png
ADDED
|
cli.py
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import argparse
|
| 2 |
+
|
| 3 |
+
from trackio import show
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
def main():
|
| 7 |
+
parser = argparse.ArgumentParser(description="Trackio CLI")
|
| 8 |
+
subparsers = parser.add_subparsers(dest="command")
|
| 9 |
+
|
| 10 |
+
ui_parser = subparsers.add_parser(
|
| 11 |
+
"show", help="Show the Trackio dashboard UI for a project"
|
| 12 |
+
)
|
| 13 |
+
ui_parser.add_argument(
|
| 14 |
+
"--project", required=False, help="Project name to show in the dashboard"
|
| 15 |
+
)
|
| 16 |
+
ui_parser.add_argument(
|
| 17 |
+
"--theme",
|
| 18 |
+
required=False,
|
| 19 |
+
default="citrus",
|
| 20 |
+
help="A Gradio Theme to use for the dashboard instead of the default, can be a built-in theme (e.g. 'soft', 'citrus'), or a theme from the Hub (e.g. 'gstaff/xkcd').",
|
| 21 |
+
)
|
| 22 |
+
ui_parser.add_argument(
|
| 23 |
+
"--mcp-server",
|
| 24 |
+
action="store_true",
|
| 25 |
+
help="Enable MCP server functionality. The Trackio dashboard will be set up as an MCP server and certain functions will be exposed as MCP tools.",
|
| 26 |
+
)
|
| 27 |
+
|
| 28 |
+
args = parser.parse_args()
|
| 29 |
+
|
| 30 |
+
if args.command == "show":
|
| 31 |
+
show(args.project, args.theme, args.mcp_server)
|
| 32 |
+
else:
|
| 33 |
+
parser.print_help()
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
if __name__ == "__main__":
|
| 37 |
+
main()
|
commit_scheduler.py
ADDED
|
@@ -0,0 +1,391 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Originally copied from https://github.com/huggingface/huggingface_hub/blob/d0a948fc2a32ed6e557042a95ef3e4af97ec4a7c/src/huggingface_hub/_commit_scheduler.py
|
| 2 |
+
|
| 3 |
+
import atexit
|
| 4 |
+
import logging
|
| 5 |
+
import os
|
| 6 |
+
import time
|
| 7 |
+
from concurrent.futures import Future
|
| 8 |
+
from dataclasses import dataclass
|
| 9 |
+
from io import SEEK_END, SEEK_SET, BytesIO
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
from threading import Lock, Thread
|
| 12 |
+
from typing import Callable, Dict, List, Union
|
| 13 |
+
|
| 14 |
+
from huggingface_hub.hf_api import (
|
| 15 |
+
DEFAULT_IGNORE_PATTERNS,
|
| 16 |
+
CommitInfo,
|
| 17 |
+
CommitOperationAdd,
|
| 18 |
+
HfApi,
|
| 19 |
+
)
|
| 20 |
+
from huggingface_hub.utils import filter_repo_objects
|
| 21 |
+
|
| 22 |
+
logger = logging.getLogger(__name__)
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
@dataclass(frozen=True)
|
| 26 |
+
class _FileToUpload:
|
| 27 |
+
"""Temporary dataclass to store info about files to upload. Not meant to be used directly."""
|
| 28 |
+
|
| 29 |
+
local_path: Path
|
| 30 |
+
path_in_repo: str
|
| 31 |
+
size_limit: int
|
| 32 |
+
last_modified: float
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
class CommitScheduler:
|
| 36 |
+
"""
|
| 37 |
+
Scheduler to upload a local folder to the Hub at regular intervals (e.g. push to hub every 5 minutes).
|
| 38 |
+
|
| 39 |
+
The recommended way to use the scheduler is to use it as a context manager. This ensures that the scheduler is
|
| 40 |
+
properly stopped and the last commit is triggered when the script ends. The scheduler can also be stopped manually
|
| 41 |
+
with the `stop` method. Checkout the [upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#scheduled-uploads)
|
| 42 |
+
to learn more about how to use it.
|
| 43 |
+
|
| 44 |
+
Args:
|
| 45 |
+
repo_id (`str`):
|
| 46 |
+
The id of the repo to commit to.
|
| 47 |
+
folder_path (`str` or `Path`):
|
| 48 |
+
Path to the local folder to upload regularly.
|
| 49 |
+
every (`int` or `float`, *optional*):
|
| 50 |
+
The number of minutes between each commit. Defaults to 5 minutes.
|
| 51 |
+
path_in_repo (`str`, *optional*):
|
| 52 |
+
Relative path of the directory in the repo, for example: `"checkpoints/"`. Defaults to the root folder
|
| 53 |
+
of the repository.
|
| 54 |
+
repo_type (`str`, *optional*):
|
| 55 |
+
The type of the repo to commit to. Defaults to `model`.
|
| 56 |
+
revision (`str`, *optional*):
|
| 57 |
+
The revision of the repo to commit to. Defaults to `main`.
|
| 58 |
+
private (`bool`, *optional*):
|
| 59 |
+
Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.
|
| 60 |
+
token (`str`, *optional*):
|
| 61 |
+
The token to use to commit to the repo. Defaults to the token saved on the machine.
|
| 62 |
+
allow_patterns (`List[str]` or `str`, *optional*):
|
| 63 |
+
If provided, only files matching at least one pattern are uploaded.
|
| 64 |
+
ignore_patterns (`List[str]` or `str`, *optional*):
|
| 65 |
+
If provided, files matching any of the patterns are not uploaded.
|
| 66 |
+
squash_history (`bool`, *optional*):
|
| 67 |
+
Whether to squash the history of the repo after each commit. Defaults to `False`. Squashing commits is
|
| 68 |
+
useful to avoid degraded performances on the repo when it grows too large.
|
| 69 |
+
hf_api (`HfApi`, *optional*):
|
| 70 |
+
The [`HfApi`] client to use to commit to the Hub. Can be set with custom settings (user agent, token,...).
|
| 71 |
+
on_before_commit (`Callable[[], None]`, *optional*):
|
| 72 |
+
If specified, a function that will be called before the CommitScheduler lists files to create a commit.
|
| 73 |
+
|
| 74 |
+
Example:
|
| 75 |
+
```py
|
| 76 |
+
>>> from pathlib import Path
|
| 77 |
+
>>> from huggingface_hub import CommitScheduler
|
| 78 |
+
|
| 79 |
+
# Scheduler uploads every 10 minutes
|
| 80 |
+
>>> csv_path = Path("watched_folder/data.csv")
|
| 81 |
+
>>> CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path=csv_path.parent, every=10)
|
| 82 |
+
|
| 83 |
+
>>> with csv_path.open("a") as f:
|
| 84 |
+
... f.write("first line")
|
| 85 |
+
|
| 86 |
+
# Some time later (...)
|
| 87 |
+
>>> with csv_path.open("a") as f:
|
| 88 |
+
... f.write("second line")
|
| 89 |
+
```
|
| 90 |
+
|
| 91 |
+
Example using a context manager:
|
| 92 |
+
```py
|
| 93 |
+
>>> from pathlib import Path
|
| 94 |
+
>>> from huggingface_hub import CommitScheduler
|
| 95 |
+
|
| 96 |
+
>>> with CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path="watched_folder", every=10) as scheduler:
|
| 97 |
+
... csv_path = Path("watched_folder/data.csv")
|
| 98 |
+
... with csv_path.open("a") as f:
|
| 99 |
+
... f.write("first line")
|
| 100 |
+
... (...)
|
| 101 |
+
... with csv_path.open("a") as f:
|
| 102 |
+
... f.write("second line")
|
| 103 |
+
|
| 104 |
+
# Scheduler is now stopped and last commit have been triggered
|
| 105 |
+
```
|
| 106 |
+
"""
|
| 107 |
+
|
| 108 |
+
def __init__(
|
| 109 |
+
self,
|
| 110 |
+
*,
|
| 111 |
+
repo_id: str,
|
| 112 |
+
folder_path: Union[str, Path],
|
| 113 |
+
every: Union[int, float] = 5,
|
| 114 |
+
path_in_repo: str | None = None,
|
| 115 |
+
repo_type: str | None = None,
|
| 116 |
+
revision: str | None = None,
|
| 117 |
+
private: bool | None = None,
|
| 118 |
+
token: str | None = None,
|
| 119 |
+
allow_patterns: list[str] | str | None = None,
|
| 120 |
+
ignore_patterns: list[str] | str | None = None,
|
| 121 |
+
squash_history: bool = False,
|
| 122 |
+
hf_api: HfApi | None = None,
|
| 123 |
+
on_before_commit: Callable[[], None] | None = None,
|
| 124 |
+
) -> None:
|
| 125 |
+
self.api = hf_api or HfApi(token=token)
|
| 126 |
+
self.on_before_commit = on_before_commit
|
| 127 |
+
|
| 128 |
+
# Folder
|
| 129 |
+
self.folder_path = Path(folder_path).expanduser().resolve()
|
| 130 |
+
self.path_in_repo = path_in_repo or ""
|
| 131 |
+
self.allow_patterns = allow_patterns
|
| 132 |
+
|
| 133 |
+
if ignore_patterns is None:
|
| 134 |
+
ignore_patterns = []
|
| 135 |
+
elif isinstance(ignore_patterns, str):
|
| 136 |
+
ignore_patterns = [ignore_patterns]
|
| 137 |
+
self.ignore_patterns = ignore_patterns + DEFAULT_IGNORE_PATTERNS
|
| 138 |
+
|
| 139 |
+
if self.folder_path.is_file():
|
| 140 |
+
raise ValueError(
|
| 141 |
+
f"'folder_path' must be a directory, not a file: '{self.folder_path}'."
|
| 142 |
+
)
|
| 143 |
+
self.folder_path.mkdir(parents=True, exist_ok=True)
|
| 144 |
+
|
| 145 |
+
# Repository
|
| 146 |
+
repo_url = self.api.create_repo(
|
| 147 |
+
repo_id=repo_id, private=private, repo_type=repo_type, exist_ok=True
|
| 148 |
+
)
|
| 149 |
+
self.repo_id = repo_url.repo_id
|
| 150 |
+
self.repo_type = repo_type
|
| 151 |
+
self.revision = revision
|
| 152 |
+
self.token = token
|
| 153 |
+
|
| 154 |
+
self.last_uploaded: Dict[Path, float] = {}
|
| 155 |
+
self.last_push_time: float | None = None
|
| 156 |
+
|
| 157 |
+
if not every > 0:
|
| 158 |
+
raise ValueError(f"'every' must be a positive integer, not '{every}'.")
|
| 159 |
+
self.lock = Lock()
|
| 160 |
+
self.every = every
|
| 161 |
+
self.squash_history = squash_history
|
| 162 |
+
|
| 163 |
+
logger.info(
|
| 164 |
+
f"Scheduled job to push '{self.folder_path}' to '{self.repo_id}' every {self.every} minutes."
|
| 165 |
+
)
|
| 166 |
+
self._scheduler_thread = Thread(target=self._run_scheduler, daemon=True)
|
| 167 |
+
self._scheduler_thread.start()
|
| 168 |
+
atexit.register(self._push_to_hub)
|
| 169 |
+
|
| 170 |
+
self.__stopped = False
|
| 171 |
+
|
| 172 |
+
def stop(self) -> None:
|
| 173 |
+
"""Stop the scheduler.
|
| 174 |
+
|
| 175 |
+
A stopped scheduler cannot be restarted. Mostly for tests purposes.
|
| 176 |
+
"""
|
| 177 |
+
self.__stopped = True
|
| 178 |
+
|
| 179 |
+
def __enter__(self) -> "CommitScheduler":
|
| 180 |
+
return self
|
| 181 |
+
|
| 182 |
+
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
| 183 |
+
# Upload last changes before exiting
|
| 184 |
+
self.trigger().result()
|
| 185 |
+
self.stop()
|
| 186 |
+
return
|
| 187 |
+
|
| 188 |
+
def _run_scheduler(self) -> None:
|
| 189 |
+
"""Dumb thread waiting between each scheduled push to Hub."""
|
| 190 |
+
while True:
|
| 191 |
+
self.last_future = self.trigger()
|
| 192 |
+
time.sleep(self.every * 60)
|
| 193 |
+
if self.__stopped:
|
| 194 |
+
break
|
| 195 |
+
|
| 196 |
+
def trigger(self) -> Future:
|
| 197 |
+
"""Trigger a `push_to_hub` and return a future.
|
| 198 |
+
|
| 199 |
+
This method is automatically called every `every` minutes. You can also call it manually to trigger a commit
|
| 200 |
+
immediately, without waiting for the next scheduled commit.
|
| 201 |
+
"""
|
| 202 |
+
return self.api.run_as_future(self._push_to_hub)
|
| 203 |
+
|
| 204 |
+
def _push_to_hub(self) -> CommitInfo | None:
|
| 205 |
+
if self.__stopped: # If stopped, already scheduled commits are ignored
|
| 206 |
+
return None
|
| 207 |
+
|
| 208 |
+
logger.info("(Background) scheduled commit triggered.")
|
| 209 |
+
try:
|
| 210 |
+
value = self.push_to_hub()
|
| 211 |
+
if self.squash_history:
|
| 212 |
+
logger.info("(Background) squashing repo history.")
|
| 213 |
+
self.api.super_squash_history(
|
| 214 |
+
repo_id=self.repo_id, repo_type=self.repo_type, branch=self.revision
|
| 215 |
+
)
|
| 216 |
+
return value
|
| 217 |
+
except Exception as e:
|
| 218 |
+
logger.error(
|
| 219 |
+
f"Error while pushing to Hub: {e}"
|
| 220 |
+
) # Depending on the setup, error might be silenced
|
| 221 |
+
raise
|
| 222 |
+
|
| 223 |
+
def push_to_hub(self) -> CommitInfo | None:
|
| 224 |
+
"""
|
| 225 |
+
Push folder to the Hub and return the commit info.
|
| 226 |
+
|
| 227 |
+
<Tip warning={true}>
|
| 228 |
+
|
| 229 |
+
This method is not meant to be called directly. It is run in the background by the scheduler, respecting a
|
| 230 |
+
queue mechanism to avoid concurrent commits. Making a direct call to the method might lead to concurrency
|
| 231 |
+
issues.
|
| 232 |
+
|
| 233 |
+
</Tip>
|
| 234 |
+
|
| 235 |
+
The default behavior of `push_to_hub` is to assume an append-only folder. It lists all files in the folder and
|
| 236 |
+
uploads only changed files. If no changes are found, the method returns without committing anything. If you want
|
| 237 |
+
to change this behavior, you can inherit from [`CommitScheduler`] and override this method. This can be useful
|
| 238 |
+
for example to compress data together in a single file before committing. For more details and examples, check
|
| 239 |
+
out our [integration guide](https://huggingface.co/docs/huggingface_hub/main/en/guides/upload#scheduled-uploads).
|
| 240 |
+
"""
|
| 241 |
+
# Check files to upload (with lock)
|
| 242 |
+
with self.lock:
|
| 243 |
+
if self.on_before_commit is not None:
|
| 244 |
+
self.on_before_commit()
|
| 245 |
+
|
| 246 |
+
logger.debug("Listing files to upload for scheduled commit.")
|
| 247 |
+
|
| 248 |
+
# List files from folder (taken from `_prepare_upload_folder_additions`)
|
| 249 |
+
relpath_to_abspath = {
|
| 250 |
+
path.relative_to(self.folder_path).as_posix(): path
|
| 251 |
+
for path in sorted(
|
| 252 |
+
self.folder_path.glob("**/*")
|
| 253 |
+
) # sorted to be deterministic
|
| 254 |
+
if path.is_file()
|
| 255 |
+
}
|
| 256 |
+
prefix = f"{self.path_in_repo.strip('/')}/" if self.path_in_repo else ""
|
| 257 |
+
|
| 258 |
+
# Filter with pattern + filter out unchanged files + retrieve current file size
|
| 259 |
+
files_to_upload: List[_FileToUpload] = []
|
| 260 |
+
for relpath in filter_repo_objects(
|
| 261 |
+
relpath_to_abspath.keys(),
|
| 262 |
+
allow_patterns=self.allow_patterns,
|
| 263 |
+
ignore_patterns=self.ignore_patterns,
|
| 264 |
+
):
|
| 265 |
+
local_path = relpath_to_abspath[relpath]
|
| 266 |
+
stat = local_path.stat()
|
| 267 |
+
if (
|
| 268 |
+
self.last_uploaded.get(local_path) is None
|
| 269 |
+
or self.last_uploaded[local_path] != stat.st_mtime
|
| 270 |
+
):
|
| 271 |
+
files_to_upload.append(
|
| 272 |
+
_FileToUpload(
|
| 273 |
+
local_path=local_path,
|
| 274 |
+
path_in_repo=prefix + relpath,
|
| 275 |
+
size_limit=stat.st_size,
|
| 276 |
+
last_modified=stat.st_mtime,
|
| 277 |
+
)
|
| 278 |
+
)
|
| 279 |
+
|
| 280 |
+
# Return if nothing to upload
|
| 281 |
+
if len(files_to_upload) == 0:
|
| 282 |
+
logger.debug("Dropping schedule commit: no changed file to upload.")
|
| 283 |
+
return None
|
| 284 |
+
|
| 285 |
+
# Convert `_FileToUpload` as `CommitOperationAdd` (=> compute file shas + limit to file size)
|
| 286 |
+
logger.debug("Removing unchanged files since previous scheduled commit.")
|
| 287 |
+
add_operations = [
|
| 288 |
+
CommitOperationAdd(
|
| 289 |
+
# TODO: Cap the file to its current size, even if the user append data to it while a scheduled commit is happening
|
| 290 |
+
# (requires an upstream fix for XET-535: `hf_xet` should support `BinaryIO` for upload)
|
| 291 |
+
path_or_fileobj=file_to_upload.local_path,
|
| 292 |
+
path_in_repo=file_to_upload.path_in_repo,
|
| 293 |
+
)
|
| 294 |
+
for file_to_upload in files_to_upload
|
| 295 |
+
]
|
| 296 |
+
|
| 297 |
+
# Upload files (append mode expected - no need for lock)
|
| 298 |
+
logger.debug("Uploading files for scheduled commit.")
|
| 299 |
+
commit_info = self.api.create_commit(
|
| 300 |
+
repo_id=self.repo_id,
|
| 301 |
+
repo_type=self.repo_type,
|
| 302 |
+
operations=add_operations,
|
| 303 |
+
commit_message="Scheduled Commit",
|
| 304 |
+
revision=self.revision,
|
| 305 |
+
)
|
| 306 |
+
|
| 307 |
+
for file in files_to_upload:
|
| 308 |
+
self.last_uploaded[file.local_path] = file.last_modified
|
| 309 |
+
|
| 310 |
+
self.last_push_time = time.time()
|
| 311 |
+
|
| 312 |
+
return commit_info
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
class PartialFileIO(BytesIO):
|
| 316 |
+
"""A file-like object that reads only the first part of a file.
|
| 317 |
+
|
| 318 |
+
Useful to upload a file to the Hub when the user might still be appending data to it. Only the first part of the
|
| 319 |
+
file is uploaded (i.e. the part that was available when the filesystem was first scanned).
|
| 320 |
+
|
| 321 |
+
In practice, only used internally by the CommitScheduler to regularly push a folder to the Hub with minimal
|
| 322 |
+
disturbance for the user. The object is passed to `CommitOperationAdd`.
|
| 323 |
+
|
| 324 |
+
Only supports `read`, `tell` and `seek` methods.
|
| 325 |
+
|
| 326 |
+
Args:
|
| 327 |
+
file_path (`str` or `Path`):
|
| 328 |
+
Path to the file to read.
|
| 329 |
+
size_limit (`int`):
|
| 330 |
+
The maximum number of bytes to read from the file. If the file is larger than this, only the first part
|
| 331 |
+
will be read (and uploaded).
|
| 332 |
+
"""
|
| 333 |
+
|
| 334 |
+
def __init__(self, file_path: Union[str, Path], size_limit: int) -> None:
|
| 335 |
+
self._file_path = Path(file_path)
|
| 336 |
+
self._file = self._file_path.open("rb")
|
| 337 |
+
self._size_limit = min(size_limit, os.fstat(self._file.fileno()).st_size)
|
| 338 |
+
|
| 339 |
+
def __del__(self) -> None:
|
| 340 |
+
self._file.close()
|
| 341 |
+
return super().__del__()
|
| 342 |
+
|
| 343 |
+
def __repr__(self) -> str:
|
| 344 |
+
return (
|
| 345 |
+
f"<PartialFileIO file_path={self._file_path} size_limit={self._size_limit}>"
|
| 346 |
+
)
|
| 347 |
+
|
| 348 |
+
def __len__(self) -> int:
|
| 349 |
+
return self._size_limit
|
| 350 |
+
|
| 351 |
+
def __getattribute__(self, name: str):
|
| 352 |
+
if name.startswith("_") or name in (
|
| 353 |
+
"read",
|
| 354 |
+
"tell",
|
| 355 |
+
"seek",
|
| 356 |
+
): # only 3 public methods supported
|
| 357 |
+
return super().__getattribute__(name)
|
| 358 |
+
raise NotImplementedError(f"PartialFileIO does not support '{name}'.")
|
| 359 |
+
|
| 360 |
+
def tell(self) -> int:
|
| 361 |
+
"""Return the current file position."""
|
| 362 |
+
return self._file.tell()
|
| 363 |
+
|
| 364 |
+
def seek(self, __offset: int, __whence: int = SEEK_SET) -> int:
|
| 365 |
+
"""Change the stream position to the given offset.
|
| 366 |
+
|
| 367 |
+
Behavior is the same as a regular file, except that the position is capped to the size limit.
|
| 368 |
+
"""
|
| 369 |
+
if __whence == SEEK_END:
|
| 370 |
+
# SEEK_END => set from the truncated end
|
| 371 |
+
__offset = len(self) + __offset
|
| 372 |
+
__whence = SEEK_SET
|
| 373 |
+
|
| 374 |
+
pos = self._file.seek(__offset, __whence)
|
| 375 |
+
if pos > self._size_limit:
|
| 376 |
+
return self._file.seek(self._size_limit)
|
| 377 |
+
return pos
|
| 378 |
+
|
| 379 |
+
def read(self, __size: int | None = -1) -> bytes:
|
| 380 |
+
"""Read at most `__size` bytes from the file.
|
| 381 |
+
|
| 382 |
+
Behavior is the same as a regular file, except that it is capped to the size limit.
|
| 383 |
+
"""
|
| 384 |
+
current = self._file.tell()
|
| 385 |
+
if __size is None or __size < 0:
|
| 386 |
+
# Read until file limit
|
| 387 |
+
truncated_size = self._size_limit - current
|
| 388 |
+
else:
|
| 389 |
+
# Read until file limit or __size
|
| 390 |
+
truncated_size = min(__size, self._size_limit - current)
|
| 391 |
+
return self._file.read(truncated_size)
|
context_vars.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import contextvars
|
| 2 |
+
from typing import TYPE_CHECKING
|
| 3 |
+
|
| 4 |
+
if TYPE_CHECKING:
|
| 5 |
+
from trackio.run import Run
|
| 6 |
+
|
| 7 |
+
current_run: contextvars.ContextVar["Run | None"] = contextvars.ContextVar(
|
| 8 |
+
"current_run", default=None
|
| 9 |
+
)
|
| 10 |
+
current_project: contextvars.ContextVar[str | None] = contextvars.ContextVar(
|
| 11 |
+
"current_project", default=None
|
| 12 |
+
)
|
| 13 |
+
current_server: contextvars.ContextVar[str | None] = contextvars.ContextVar(
|
| 14 |
+
"current_server", default=None
|
| 15 |
+
)
|
| 16 |
+
current_share_server: contextvars.ContextVar[str | None] = contextvars.ContextVar(
|
| 17 |
+
"current_share_server", default=None
|
| 18 |
+
)
|
deploy.py
ADDED
|
@@ -0,0 +1,258 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import importlib.metadata
|
| 2 |
+
import io
|
| 3 |
+
import os
|
| 4 |
+
import time
|
| 5 |
+
from importlib.resources import files
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
import gradio
|
| 9 |
+
import huggingface_hub
|
| 10 |
+
from gradio_client import Client, handle_file
|
| 11 |
+
from httpx import ReadTimeout
|
| 12 |
+
from huggingface_hub.errors import RepositoryNotFoundError
|
| 13 |
+
from requests import HTTPError
|
| 14 |
+
|
| 15 |
+
import trackio
|
| 16 |
+
from trackio.sqlite_storage import SQLiteStorage
|
| 17 |
+
|
| 18 |
+
SPACE_HOST_URL = "https://{user_name}-{space_name}.hf.space/"
|
| 19 |
+
SPACE_URL = "https://huggingface.co/spaces/{space_id}"
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def _is_trackio_installed_from_source() -> bool:
|
| 23 |
+
"""Check if trackio is installed from source/editable install vs PyPI."""
|
| 24 |
+
try:
|
| 25 |
+
trackio_file = trackio.__file__
|
| 26 |
+
if "site-packages" not in trackio_file:
|
| 27 |
+
return True
|
| 28 |
+
|
| 29 |
+
dist = importlib.metadata.distribution("trackio")
|
| 30 |
+
if dist.files:
|
| 31 |
+
files = list(dist.files)
|
| 32 |
+
has_pth = any(".pth" in str(f) for f in files)
|
| 33 |
+
if has_pth:
|
| 34 |
+
return True
|
| 35 |
+
|
| 36 |
+
return False
|
| 37 |
+
except (
|
| 38 |
+
AttributeError,
|
| 39 |
+
importlib.metadata.PackageNotFoundError,
|
| 40 |
+
importlib.metadata.MetadataError,
|
| 41 |
+
ValueError,
|
| 42 |
+
TypeError,
|
| 43 |
+
):
|
| 44 |
+
return True
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def deploy_as_space(
|
| 48 |
+
space_id: str,
|
| 49 |
+
space_storage: huggingface_hub.SpaceStorage | None = None,
|
| 50 |
+
dataset_id: str | None = None,
|
| 51 |
+
private: bool | None = None,
|
| 52 |
+
):
|
| 53 |
+
if (
|
| 54 |
+
os.getenv("SYSTEM") == "spaces"
|
| 55 |
+
): # in case a repo with this function is uploaded to spaces
|
| 56 |
+
return
|
| 57 |
+
|
| 58 |
+
trackio_path = files("trackio")
|
| 59 |
+
|
| 60 |
+
hf_api = huggingface_hub.HfApi()
|
| 61 |
+
|
| 62 |
+
try:
|
| 63 |
+
huggingface_hub.create_repo(
|
| 64 |
+
space_id,
|
| 65 |
+
private=private,
|
| 66 |
+
space_sdk="gradio",
|
| 67 |
+
space_storage=space_storage,
|
| 68 |
+
repo_type="space",
|
| 69 |
+
exist_ok=True,
|
| 70 |
+
)
|
| 71 |
+
except HTTPError as e:
|
| 72 |
+
if e.response.status_code in [401, 403]: # unauthorized or forbidden
|
| 73 |
+
print("Need 'write' access token to create a Spaces repo.")
|
| 74 |
+
huggingface_hub.login(add_to_git_credential=False)
|
| 75 |
+
huggingface_hub.create_repo(
|
| 76 |
+
space_id,
|
| 77 |
+
private=private,
|
| 78 |
+
space_sdk="gradio",
|
| 79 |
+
space_storage=space_storage,
|
| 80 |
+
repo_type="space",
|
| 81 |
+
exist_ok=True,
|
| 82 |
+
)
|
| 83 |
+
else:
|
| 84 |
+
raise ValueError(f"Failed to create Space: {e}")
|
| 85 |
+
|
| 86 |
+
with open(Path(trackio_path, "README.md"), "r") as f:
|
| 87 |
+
readme_content = f.read()
|
| 88 |
+
readme_content = readme_content.replace("{GRADIO_VERSION}", gradio.__version__)
|
| 89 |
+
readme_buffer = io.BytesIO(readme_content.encode("utf-8"))
|
| 90 |
+
hf_api.upload_file(
|
| 91 |
+
path_or_fileobj=readme_buffer,
|
| 92 |
+
path_in_repo="README.md",
|
| 93 |
+
repo_id=space_id,
|
| 94 |
+
repo_type="space",
|
| 95 |
+
)
|
| 96 |
+
|
| 97 |
+
# We can assume pandas, gradio, and huggingface-hub are already installed in a Gradio Space.
|
| 98 |
+
# Make sure necessary dependencies are installed by creating a requirements.txt.
|
| 99 |
+
is_source_install = _is_trackio_installed_from_source()
|
| 100 |
+
|
| 101 |
+
if is_source_install:
|
| 102 |
+
requirements_content = """pyarrow>=21.0
|
| 103 |
+
plotly>=6.0.0,<7.0.0"""
|
| 104 |
+
else:
|
| 105 |
+
requirements_content = f"""pyarrow>=21.0
|
| 106 |
+
trackio=={trackio.__version__}
|
| 107 |
+
plotly>=6.0.0,<7.0.0"""
|
| 108 |
+
|
| 109 |
+
requirements_buffer = io.BytesIO(requirements_content.encode("utf-8"))
|
| 110 |
+
hf_api.upload_file(
|
| 111 |
+
path_or_fileobj=requirements_buffer,
|
| 112 |
+
path_in_repo="requirements.txt",
|
| 113 |
+
repo_id=space_id,
|
| 114 |
+
repo_type="space",
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
huggingface_hub.utils.disable_progress_bars()
|
| 118 |
+
|
| 119 |
+
if is_source_install:
|
| 120 |
+
hf_api.upload_folder(
|
| 121 |
+
repo_id=space_id,
|
| 122 |
+
repo_type="space",
|
| 123 |
+
folder_path=trackio_path,
|
| 124 |
+
ignore_patterns=["README.md"],
|
| 125 |
+
)
|
| 126 |
+
else:
|
| 127 |
+
app_file_content = """import trackio
|
| 128 |
+
trackio.show()"""
|
| 129 |
+
app_file_buffer = io.BytesIO(app_file_content.encode("utf-8"))
|
| 130 |
+
hf_api.upload_file(
|
| 131 |
+
path_or_fileobj=app_file_buffer,
|
| 132 |
+
path_in_repo="ui/main.py",
|
| 133 |
+
repo_id=space_id,
|
| 134 |
+
repo_type="space",
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
if hf_token := huggingface_hub.utils.get_token():
|
| 138 |
+
huggingface_hub.add_space_secret(space_id, "HF_TOKEN", hf_token)
|
| 139 |
+
if dataset_id is not None:
|
| 140 |
+
huggingface_hub.add_space_variable(space_id, "TRACKIO_DATASET_ID", dataset_id)
|
| 141 |
+
|
| 142 |
+
if logo_light_url := os.environ.get("TRACKIO_LOGO_LIGHT_URL"):
|
| 143 |
+
huggingface_hub.add_space_variable(
|
| 144 |
+
space_id, "TRACKIO_LOGO_LIGHT_URL", logo_light_url
|
| 145 |
+
)
|
| 146 |
+
if logo_dark_url := os.environ.get("TRACKIO_LOGO_DARK_URL"):
|
| 147 |
+
huggingface_hub.add_space_variable(
|
| 148 |
+
space_id, "TRACKIO_LOGO_DARK_URL", logo_dark_url
|
| 149 |
+
)
|
| 150 |
+
|
| 151 |
+
if plot_order := os.environ.get("TRACKIO_PLOT_ORDER"):
|
| 152 |
+
huggingface_hub.add_space_variable(space_id, "TRACKIO_PLOT_ORDER", plot_order)
|
| 153 |
+
|
| 154 |
+
if theme := os.environ.get("TRACKIO_THEME"):
|
| 155 |
+
huggingface_hub.add_space_variable(space_id, "TRACKIO_THEME", theme)
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
def create_space_if_not_exists(
|
| 159 |
+
space_id: str,
|
| 160 |
+
space_storage: huggingface_hub.SpaceStorage | None = None,
|
| 161 |
+
dataset_id: str | None = None,
|
| 162 |
+
private: bool | None = None,
|
| 163 |
+
) -> None:
|
| 164 |
+
"""
|
| 165 |
+
Creates a new Hugging Face Space if it does not exist. If a dataset_id is provided, it will be added as a space variable.
|
| 166 |
+
|
| 167 |
+
Args:
|
| 168 |
+
space_id: The ID of the Space to create.
|
| 169 |
+
dataset_id: The ID of the Dataset to add to the Space.
|
| 170 |
+
private: Whether to make the Space private. If None (default), the repo will be
|
| 171 |
+
public unless the organization's default is private. This value is ignored if
|
| 172 |
+
the repo already exists.
|
| 173 |
+
"""
|
| 174 |
+
if "/" not in space_id:
|
| 175 |
+
raise ValueError(
|
| 176 |
+
f"Invalid space ID: {space_id}. Must be in the format: username/reponame or orgname/reponame."
|
| 177 |
+
)
|
| 178 |
+
if dataset_id is not None and "/" not in dataset_id:
|
| 179 |
+
raise ValueError(
|
| 180 |
+
f"Invalid dataset ID: {dataset_id}. Must be in the format: username/datasetname or orgname/datasetname."
|
| 181 |
+
)
|
| 182 |
+
try:
|
| 183 |
+
huggingface_hub.repo_info(space_id, repo_type="space")
|
| 184 |
+
print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")
|
| 185 |
+
if dataset_id is not None:
|
| 186 |
+
huggingface_hub.add_space_variable(
|
| 187 |
+
space_id, "TRACKIO_DATASET_ID", dataset_id
|
| 188 |
+
)
|
| 189 |
+
if logo_light_url := os.environ.get("TRACKIO_LOGO_LIGHT_URL"):
|
| 190 |
+
huggingface_hub.add_space_variable(
|
| 191 |
+
space_id, "TRACKIO_LOGO_LIGHT_URL", logo_light_url
|
| 192 |
+
)
|
| 193 |
+
if logo_dark_url := os.environ.get("TRACKIO_LOGO_DARK_URL"):
|
| 194 |
+
huggingface_hub.add_space_variable(
|
| 195 |
+
space_id, "TRACKIO_LOGO_DARK_URL", logo_dark_url
|
| 196 |
+
)
|
| 197 |
+
|
| 198 |
+
if plot_order := os.environ.get("TRACKIO_PLOT_ORDER"):
|
| 199 |
+
huggingface_hub.add_space_variable(
|
| 200 |
+
space_id, "TRACKIO_PLOT_ORDER", plot_order
|
| 201 |
+
)
|
| 202 |
+
|
| 203 |
+
if theme := os.environ.get("TRACKIO_THEME"):
|
| 204 |
+
huggingface_hub.add_space_variable(space_id, "TRACKIO_THEME", theme)
|
| 205 |
+
return
|
| 206 |
+
except RepositoryNotFoundError:
|
| 207 |
+
pass
|
| 208 |
+
except HTTPError as e:
|
| 209 |
+
if e.response.status_code in [401, 403]: # unauthorized or forbidden
|
| 210 |
+
print("Need 'write' access token to create a Spaces repo.")
|
| 211 |
+
huggingface_hub.login(add_to_git_credential=False)
|
| 212 |
+
huggingface_hub.add_space_variable(
|
| 213 |
+
space_id, "TRACKIO_DATASET_ID", dataset_id
|
| 214 |
+
)
|
| 215 |
+
else:
|
| 216 |
+
raise ValueError(f"Failed to create Space: {e}")
|
| 217 |
+
|
| 218 |
+
print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")
|
| 219 |
+
deploy_as_space(space_id, space_storage, dataset_id, private)
|
| 220 |
+
|
| 221 |
+
|
| 222 |
+
def wait_until_space_exists(
|
| 223 |
+
space_id: str,
|
| 224 |
+
) -> None:
|
| 225 |
+
"""
|
| 226 |
+
Blocks the current thread until the space exists.
|
| 227 |
+
May raise a TimeoutError if this takes quite a while.
|
| 228 |
+
|
| 229 |
+
Args:
|
| 230 |
+
space_id: The ID of the Space to wait for.
|
| 231 |
+
"""
|
| 232 |
+
delay = 1
|
| 233 |
+
for _ in range(10):
|
| 234 |
+
try:
|
| 235 |
+
Client(space_id, verbose=False)
|
| 236 |
+
return
|
| 237 |
+
except (ReadTimeout, ValueError):
|
| 238 |
+
time.sleep(delay)
|
| 239 |
+
delay = min(delay * 2, 30)
|
| 240 |
+
raise TimeoutError("Waiting for space to exist took longer than expected")
|
| 241 |
+
|
| 242 |
+
|
| 243 |
+
def upload_db_to_space(project: str, space_id: str) -> None:
|
| 244 |
+
"""
|
| 245 |
+
Uploads the database of a local Trackio project to a Hugging Face Space.
|
| 246 |
+
|
| 247 |
+
Args:
|
| 248 |
+
project: The name of the project to upload.
|
| 249 |
+
space_id: The ID of the Space to upload to.
|
| 250 |
+
"""
|
| 251 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 252 |
+
client = Client(space_id, verbose=False)
|
| 253 |
+
client.predict(
|
| 254 |
+
api_name="/upload_db_to_space",
|
| 255 |
+
project=project,
|
| 256 |
+
uploaded_db=handle_file(db_path),
|
| 257 |
+
hf_token=huggingface_hub.utils.get_token(),
|
| 258 |
+
)
|
dummy_commit_scheduler.py
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# A dummy object to fit the interface of huggingface_hub's CommitScheduler
|
| 2 |
+
class DummyCommitSchedulerLock:
|
| 3 |
+
def __enter__(self):
|
| 4 |
+
return None
|
| 5 |
+
|
| 6 |
+
def __exit__(self, exception_type, exception_value, exception_traceback):
|
| 7 |
+
pass
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
class DummyCommitScheduler:
|
| 11 |
+
def __init__(self):
|
| 12 |
+
self.lock = DummyCommitSchedulerLock()
|
file_storage.py
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pathlib import Path
|
| 2 |
+
|
| 3 |
+
try: # absolute imports when installed
|
| 4 |
+
from trackio.utils import MEDIA_DIR
|
| 5 |
+
except ImportError: # relative imports for local execution on Spaces
|
| 6 |
+
from utils import MEDIA_DIR
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class FileStorage:
|
| 10 |
+
@staticmethod
|
| 11 |
+
def get_project_media_path(
|
| 12 |
+
project: str,
|
| 13 |
+
run: str | None = None,
|
| 14 |
+
step: int | None = None,
|
| 15 |
+
filename: str | None = None,
|
| 16 |
+
) -> Path:
|
| 17 |
+
if filename is not None and step is None:
|
| 18 |
+
raise ValueError("filename requires step")
|
| 19 |
+
if step is not None and run is None:
|
| 20 |
+
raise ValueError("step requires run")
|
| 21 |
+
|
| 22 |
+
path = MEDIA_DIR / project
|
| 23 |
+
if run:
|
| 24 |
+
path /= run
|
| 25 |
+
if step is not None:
|
| 26 |
+
path /= str(step)
|
| 27 |
+
if filename:
|
| 28 |
+
path /= filename
|
| 29 |
+
return path
|
| 30 |
+
|
| 31 |
+
@staticmethod
|
| 32 |
+
def init_project_media_path(
|
| 33 |
+
project: str, run: str | None = None, step: int | None = None
|
| 34 |
+
) -> Path:
|
| 35 |
+
path = FileStorage.get_project_media_path(project, run, step)
|
| 36 |
+
path.mkdir(parents=True, exist_ok=True)
|
| 37 |
+
return path
|
histogram.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from typing import Any
|
| 2 |
+
|
| 3 |
+
import numpy as np
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class Histogram:
|
| 7 |
+
"""
|
| 8 |
+
Histogram data type for Trackio, compatible with wandb.Histogram.
|
| 9 |
+
|
| 10 |
+
Example:
|
| 11 |
+
```python
|
| 12 |
+
import trackio
|
| 13 |
+
import numpy as np
|
| 14 |
+
|
| 15 |
+
# Create histogram from sequence
|
| 16 |
+
data = np.random.randn(1000)
|
| 17 |
+
trackio.log({"distribution": trackio.Histogram(data)})
|
| 18 |
+
|
| 19 |
+
# Create histogram from numpy histogram
|
| 20 |
+
hist, bins = np.histogram(data, bins=30)
|
| 21 |
+
trackio.log({"distribution": trackio.Histogram(np_histogram=(hist, bins))})
|
| 22 |
+
|
| 23 |
+
# Specify custom number of bins
|
| 24 |
+
trackio.log({"distribution": trackio.Histogram(data, num_bins=50)})
|
| 25 |
+
```
|
| 26 |
+
|
| 27 |
+
Args:
|
| 28 |
+
sequence: Optional sequence of values to create histogram from
|
| 29 |
+
np_histogram: Optional pre-computed numpy histogram (hist, bins) tuple
|
| 30 |
+
num_bins: Number of bins for the histogram (default 64, max 512)
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
TYPE = "trackio.histogram"
|
| 34 |
+
|
| 35 |
+
def __init__(
|
| 36 |
+
self,
|
| 37 |
+
sequence: Any = None,
|
| 38 |
+
np_histogram: tuple | None = None,
|
| 39 |
+
num_bins: int = 64,
|
| 40 |
+
):
|
| 41 |
+
if sequence is None and np_histogram is None:
|
| 42 |
+
raise ValueError("Must provide either sequence or np_histogram")
|
| 43 |
+
|
| 44 |
+
if sequence is not None and np_histogram is not None:
|
| 45 |
+
raise ValueError("Cannot provide both sequence and np_histogram")
|
| 46 |
+
|
| 47 |
+
num_bins = min(num_bins, 512)
|
| 48 |
+
|
| 49 |
+
if np_histogram is not None:
|
| 50 |
+
self.histogram, self.bins = np_histogram
|
| 51 |
+
self.histogram = np.asarray(self.histogram)
|
| 52 |
+
self.bins = np.asarray(self.bins)
|
| 53 |
+
else:
|
| 54 |
+
data = np.asarray(sequence).flatten()
|
| 55 |
+
data = data[np.isfinite(data)]
|
| 56 |
+
if len(data) == 0:
|
| 57 |
+
self.histogram = np.array([])
|
| 58 |
+
self.bins = np.array([])
|
| 59 |
+
else:
|
| 60 |
+
self.histogram, self.bins = np.histogram(data, bins=num_bins)
|
| 61 |
+
|
| 62 |
+
def _to_dict(self) -> dict:
|
| 63 |
+
"""Convert histogram to dictionary for storage."""
|
| 64 |
+
return {
|
| 65 |
+
"_type": self.TYPE,
|
| 66 |
+
"bins": self.bins.tolist(),
|
| 67 |
+
"values": self.histogram.tolist(),
|
| 68 |
+
}
|
imports.py
ADDED
|
@@ -0,0 +1,302 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
from pathlib import Path
|
| 3 |
+
|
| 4 |
+
import pandas as pd
|
| 5 |
+
|
| 6 |
+
from trackio import deploy, utils
|
| 7 |
+
from trackio.sqlite_storage import SQLiteStorage
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
def import_csv(
|
| 11 |
+
csv_path: str | Path,
|
| 12 |
+
project: str,
|
| 13 |
+
name: str | None = None,
|
| 14 |
+
space_id: str | None = None,
|
| 15 |
+
dataset_id: str | None = None,
|
| 16 |
+
private: bool | None = None,
|
| 17 |
+
) -> None:
|
| 18 |
+
"""
|
| 19 |
+
Imports a CSV file into a Trackio project. The CSV file must contain a `"step"`
|
| 20 |
+
column, may optionally contain a `"timestamp"` column, and any other columns will be
|
| 21 |
+
treated as metrics. It should also include a header row with the column names.
|
| 22 |
+
|
| 23 |
+
TODO: call init() and return a Run object so that the user can continue to log metrics to it.
|
| 24 |
+
|
| 25 |
+
Args:
|
| 26 |
+
csv_path (`str` or `Path`):
|
| 27 |
+
The str or Path to the CSV file to import.
|
| 28 |
+
project (`str`):
|
| 29 |
+
The name of the project to import the CSV file into. Must not be an existing
|
| 30 |
+
project.
|
| 31 |
+
name (`str`, *optional*):
|
| 32 |
+
The name of the Run to import the CSV file into. If not provided, a default
|
| 33 |
+
name will be generated.
|
| 34 |
+
name (`str`, *optional*):
|
| 35 |
+
The name of the run (if not provided, a default name will be generated).
|
| 36 |
+
space_id (`str`, *optional*):
|
| 37 |
+
If provided, the project will be logged to a Hugging Face Space instead of a
|
| 38 |
+
local directory. Should be a complete Space name like `"username/reponame"`
|
| 39 |
+
or `"orgname/reponame"`, or just `"reponame"` in which case the Space will
|
| 40 |
+
be created in the currently-logged-in Hugging Face user's namespace. If the
|
| 41 |
+
Space does not exist, it will be created. If the Space already exists, the
|
| 42 |
+
project will be logged to it.
|
| 43 |
+
dataset_id (`str`, *optional*):
|
| 44 |
+
If provided, a persistent Hugging Face Dataset will be created and the
|
| 45 |
+
metrics will be synced to it every 5 minutes. Should be a complete Dataset
|
| 46 |
+
name like `"username/datasetname"` or `"orgname/datasetname"`, or just
|
| 47 |
+
`"datasetname"` in which case the Dataset will be created in the
|
| 48 |
+
currently-logged-in Hugging Face user's namespace. If the Dataset does not
|
| 49 |
+
exist, it will be created. If the Dataset already exists, the project will
|
| 50 |
+
be appended to it. If not provided, the metrics will be logged to a local
|
| 51 |
+
SQLite database, unless a `space_id` is provided, in which case a Dataset
|
| 52 |
+
will be automatically created with the same name as the Space but with the
|
| 53 |
+
`"_dataset"` suffix.
|
| 54 |
+
private (`bool`, *optional*):
|
| 55 |
+
Whether to make the Space private. If None (default), the repo will be
|
| 56 |
+
public unless the organization's default is private. This value is ignored
|
| 57 |
+
if the repo already exists.
|
| 58 |
+
"""
|
| 59 |
+
if SQLiteStorage.get_runs(project):
|
| 60 |
+
raise ValueError(
|
| 61 |
+
f"Project '{project}' already exists. Cannot import CSV into existing project."
|
| 62 |
+
)
|
| 63 |
+
|
| 64 |
+
csv_path = Path(csv_path)
|
| 65 |
+
if not csv_path.exists():
|
| 66 |
+
raise FileNotFoundError(f"CSV file not found: {csv_path}")
|
| 67 |
+
|
| 68 |
+
df = pd.read_csv(csv_path)
|
| 69 |
+
if df.empty:
|
| 70 |
+
raise ValueError("CSV file is empty")
|
| 71 |
+
|
| 72 |
+
column_mapping = utils.simplify_column_names(df.columns.tolist())
|
| 73 |
+
df = df.rename(columns=column_mapping)
|
| 74 |
+
|
| 75 |
+
step_column = None
|
| 76 |
+
for col in df.columns:
|
| 77 |
+
if col.lower() == "step":
|
| 78 |
+
step_column = col
|
| 79 |
+
break
|
| 80 |
+
|
| 81 |
+
if step_column is None:
|
| 82 |
+
raise ValueError("CSV file must contain a 'step' or 'Step' column")
|
| 83 |
+
|
| 84 |
+
if name is None:
|
| 85 |
+
name = csv_path.stem
|
| 86 |
+
|
| 87 |
+
metrics_list = []
|
| 88 |
+
steps = []
|
| 89 |
+
timestamps = []
|
| 90 |
+
|
| 91 |
+
numeric_columns = []
|
| 92 |
+
for column in df.columns:
|
| 93 |
+
if column == step_column:
|
| 94 |
+
continue
|
| 95 |
+
if column == "timestamp":
|
| 96 |
+
continue
|
| 97 |
+
|
| 98 |
+
try:
|
| 99 |
+
pd.to_numeric(df[column], errors="raise")
|
| 100 |
+
numeric_columns.append(column)
|
| 101 |
+
except (ValueError, TypeError):
|
| 102 |
+
continue
|
| 103 |
+
|
| 104 |
+
for _, row in df.iterrows():
|
| 105 |
+
metrics = {}
|
| 106 |
+
for column in numeric_columns:
|
| 107 |
+
value = row[column]
|
| 108 |
+
if bool(pd.notna(value)):
|
| 109 |
+
metrics[column] = float(value)
|
| 110 |
+
|
| 111 |
+
if metrics:
|
| 112 |
+
metrics_list.append(metrics)
|
| 113 |
+
steps.append(int(row[step_column]))
|
| 114 |
+
|
| 115 |
+
if "timestamp" in df.columns and bool(pd.notna(row["timestamp"])):
|
| 116 |
+
timestamps.append(str(row["timestamp"]))
|
| 117 |
+
else:
|
| 118 |
+
timestamps.append("")
|
| 119 |
+
|
| 120 |
+
if metrics_list:
|
| 121 |
+
SQLiteStorage.bulk_log(
|
| 122 |
+
project=project,
|
| 123 |
+
run=name,
|
| 124 |
+
metrics_list=metrics_list,
|
| 125 |
+
steps=steps,
|
| 126 |
+
timestamps=timestamps,
|
| 127 |
+
)
|
| 128 |
+
|
| 129 |
+
print(
|
| 130 |
+
f"* Imported {len(metrics_list)} rows from {csv_path} into project '{project}' as run '{name}'"
|
| 131 |
+
)
|
| 132 |
+
print(f"* Metrics found: {', '.join(metrics_list[0].keys())}")
|
| 133 |
+
|
| 134 |
+
space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)
|
| 135 |
+
if dataset_id is not None:
|
| 136 |
+
os.environ["TRACKIO_DATASET_ID"] = dataset_id
|
| 137 |
+
print(f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}")
|
| 138 |
+
|
| 139 |
+
if space_id is None:
|
| 140 |
+
utils.print_dashboard_instructions(project)
|
| 141 |
+
else:
|
| 142 |
+
deploy.create_space_if_not_exists(
|
| 143 |
+
space_id=space_id, dataset_id=dataset_id, private=private
|
| 144 |
+
)
|
| 145 |
+
deploy.wait_until_space_exists(space_id=space_id)
|
| 146 |
+
deploy.upload_db_to_space(project=project, space_id=space_id)
|
| 147 |
+
print(
|
| 148 |
+
f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"
|
| 149 |
+
)
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
def import_tf_events(
|
| 153 |
+
log_dir: str | Path,
|
| 154 |
+
project: str,
|
| 155 |
+
name: str | None = None,
|
| 156 |
+
space_id: str | None = None,
|
| 157 |
+
dataset_id: str | None = None,
|
| 158 |
+
private: bool | None = None,
|
| 159 |
+
) -> None:
|
| 160 |
+
"""
|
| 161 |
+
Imports TensorFlow Events files from a directory into a Trackio project. Each
|
| 162 |
+
subdirectory in the log directory will be imported as a separate run.
|
| 163 |
+
|
| 164 |
+
Args:
|
| 165 |
+
log_dir (`str` or `Path`):
|
| 166 |
+
The str or Path to the directory containing TensorFlow Events files.
|
| 167 |
+
project (`str`):
|
| 168 |
+
The name of the project to import the TensorFlow Events files into. Must not
|
| 169 |
+
be an existing project.
|
| 170 |
+
name (`str`, *optional*):
|
| 171 |
+
The name prefix for runs (if not provided, will use directory names). Each
|
| 172 |
+
subdirectory will create a separate run.
|
| 173 |
+
space_id (`str`, *optional*):
|
| 174 |
+
If provided, the project will be logged to a Hugging Face Space instead of a
|
| 175 |
+
local directory. Should be a complete Space name like `"username/reponame"`
|
| 176 |
+
or `"orgname/reponame"`, or just `"reponame"` in which case the Space will
|
| 177 |
+
be created in the currently-logged-in Hugging Face user's namespace. If the
|
| 178 |
+
Space does not exist, it will be created. If the Space already exists, the
|
| 179 |
+
project will be logged to it.
|
| 180 |
+
dataset_id (`str`, *optional*):
|
| 181 |
+
If provided, a persistent Hugging Face Dataset will be created and the
|
| 182 |
+
metrics will be synced to it every 5 minutes. Should be a complete Dataset
|
| 183 |
+
name like `"username/datasetname"` or `"orgname/datasetname"`, or just
|
| 184 |
+
`"datasetname"` in which case the Dataset will be created in the
|
| 185 |
+
currently-logged-in Hugging Face user's namespace. If the Dataset does not
|
| 186 |
+
exist, it will be created. If the Dataset already exists, the project will
|
| 187 |
+
be appended to it. If not provided, the metrics will be logged to a local
|
| 188 |
+
SQLite database, unless a `space_id` is provided, in which case a Dataset
|
| 189 |
+
will be automatically created with the same name as the Space but with the
|
| 190 |
+
`"_dataset"` suffix.
|
| 191 |
+
private (`bool`, *optional*):
|
| 192 |
+
Whether to make the Space private. If None (default), the repo will be
|
| 193 |
+
public unless the organization's default is private. This value is ignored
|
| 194 |
+
if the repo already exists.
|
| 195 |
+
"""
|
| 196 |
+
try:
|
| 197 |
+
from tbparse import SummaryReader
|
| 198 |
+
except ImportError:
|
| 199 |
+
raise ImportError(
|
| 200 |
+
"The `tbparse` package is not installed but is required for `import_tf_events`. Please install trackio with the `tensorboard` extra: `pip install trackio[tensorboard]`."
|
| 201 |
+
)
|
| 202 |
+
|
| 203 |
+
if SQLiteStorage.get_runs(project):
|
| 204 |
+
raise ValueError(
|
| 205 |
+
f"Project '{project}' already exists. Cannot import TF events into existing project."
|
| 206 |
+
)
|
| 207 |
+
|
| 208 |
+
path = Path(log_dir)
|
| 209 |
+
if not path.exists():
|
| 210 |
+
raise FileNotFoundError(f"TF events directory not found: {path}")
|
| 211 |
+
|
| 212 |
+
# Use tbparse to read all tfevents files in the directory structure
|
| 213 |
+
reader = SummaryReader(str(path), extra_columns={"dir_name"})
|
| 214 |
+
df = reader.scalars
|
| 215 |
+
|
| 216 |
+
if df.empty:
|
| 217 |
+
raise ValueError(f"No TensorFlow events data found in {path}")
|
| 218 |
+
|
| 219 |
+
total_imported = 0
|
| 220 |
+
imported_runs = []
|
| 221 |
+
|
| 222 |
+
# Group by dir_name to create separate runs
|
| 223 |
+
for dir_name, group_df in df.groupby("dir_name"):
|
| 224 |
+
try:
|
| 225 |
+
# Determine run name based on directory name
|
| 226 |
+
if dir_name == "":
|
| 227 |
+
run_name = "main" # For files in the root directory
|
| 228 |
+
else:
|
| 229 |
+
run_name = dir_name # Use directory name
|
| 230 |
+
|
| 231 |
+
if name:
|
| 232 |
+
run_name = f"{name}_{run_name}"
|
| 233 |
+
|
| 234 |
+
if group_df.empty:
|
| 235 |
+
print(f"* Skipping directory {dir_name}: no scalar data found")
|
| 236 |
+
continue
|
| 237 |
+
|
| 238 |
+
metrics_list = []
|
| 239 |
+
steps = []
|
| 240 |
+
timestamps = []
|
| 241 |
+
|
| 242 |
+
for _, row in group_df.iterrows():
|
| 243 |
+
# Convert row values to appropriate types
|
| 244 |
+
tag = str(row["tag"])
|
| 245 |
+
value = float(row["value"])
|
| 246 |
+
step = int(row["step"])
|
| 247 |
+
|
| 248 |
+
metrics = {tag: value}
|
| 249 |
+
metrics_list.append(metrics)
|
| 250 |
+
steps.append(step)
|
| 251 |
+
|
| 252 |
+
# Use wall_time if present, else fallback
|
| 253 |
+
if "wall_time" in group_df.columns and not bool(
|
| 254 |
+
pd.isna(row["wall_time"])
|
| 255 |
+
):
|
| 256 |
+
timestamps.append(str(row["wall_time"]))
|
| 257 |
+
else:
|
| 258 |
+
timestamps.append("")
|
| 259 |
+
|
| 260 |
+
if metrics_list:
|
| 261 |
+
SQLiteStorage.bulk_log(
|
| 262 |
+
project=project,
|
| 263 |
+
run=str(run_name),
|
| 264 |
+
metrics_list=metrics_list,
|
| 265 |
+
steps=steps,
|
| 266 |
+
timestamps=timestamps,
|
| 267 |
+
)
|
| 268 |
+
|
| 269 |
+
total_imported += len(metrics_list)
|
| 270 |
+
imported_runs.append(run_name)
|
| 271 |
+
|
| 272 |
+
print(
|
| 273 |
+
f"* Imported {len(metrics_list)} scalar events from directory '{dir_name}' as run '{run_name}'"
|
| 274 |
+
)
|
| 275 |
+
print(f"* Metrics in this run: {', '.join(set(group_df['tag']))}")
|
| 276 |
+
|
| 277 |
+
except Exception as e:
|
| 278 |
+
print(f"* Error processing directory {dir_name}: {e}")
|
| 279 |
+
continue
|
| 280 |
+
|
| 281 |
+
if not imported_runs:
|
| 282 |
+
raise ValueError("No valid TensorFlow events data could be imported")
|
| 283 |
+
|
| 284 |
+
print(f"* Total imported events: {total_imported}")
|
| 285 |
+
print(f"* Created runs: {', '.join(imported_runs)}")
|
| 286 |
+
|
| 287 |
+
space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)
|
| 288 |
+
if dataset_id is not None:
|
| 289 |
+
os.environ["TRACKIO_DATASET_ID"] = dataset_id
|
| 290 |
+
print(f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}")
|
| 291 |
+
|
| 292 |
+
if space_id is None:
|
| 293 |
+
utils.print_dashboard_instructions(project)
|
| 294 |
+
else:
|
| 295 |
+
deploy.create_space_if_not_exists(
|
| 296 |
+
space_id, dataset_id=dataset_id, private=private
|
| 297 |
+
)
|
| 298 |
+
deploy.wait_until_space_exists(space_id)
|
| 299 |
+
deploy.upload_db_to_space(project, space_id)
|
| 300 |
+
print(
|
| 301 |
+
f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"
|
| 302 |
+
)
|
media.py
ADDED
|
@@ -0,0 +1,299 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import shutil
|
| 3 |
+
import uuid
|
| 4 |
+
from abc import ABC, abstractmethod
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
from typing import Literal
|
| 7 |
+
|
| 8 |
+
import numpy as np
|
| 9 |
+
from PIL import Image as PILImage
|
| 10 |
+
|
| 11 |
+
try: # absolute imports when installed
|
| 12 |
+
from trackio.file_storage import FileStorage
|
| 13 |
+
from trackio.utils import MEDIA_DIR
|
| 14 |
+
from trackio.video_writer import write_video
|
| 15 |
+
except ImportError: # relative imports for local execution on Spaces
|
| 16 |
+
from file_storage import FileStorage
|
| 17 |
+
from utils import MEDIA_DIR
|
| 18 |
+
from video_writer import write_video
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class TrackioMedia(ABC):
|
| 22 |
+
"""
|
| 23 |
+
Abstract base class for Trackio media objects
|
| 24 |
+
Provides shared functionality for file handling and serialization.
|
| 25 |
+
"""
|
| 26 |
+
|
| 27 |
+
TYPE: str
|
| 28 |
+
|
| 29 |
+
def __init_subclass__(cls, **kwargs):
|
| 30 |
+
"""Ensure subclasses define the TYPE attribute."""
|
| 31 |
+
super().__init_subclass__(**kwargs)
|
| 32 |
+
if not hasattr(cls, "TYPE") or cls.TYPE is None:
|
| 33 |
+
raise TypeError(f"Class {cls.__name__} must define TYPE attribute")
|
| 34 |
+
|
| 35 |
+
def __init__(self, value, caption: str | None = None):
|
| 36 |
+
self.caption = caption
|
| 37 |
+
self._value = value
|
| 38 |
+
self._file_path: Path | None = None
|
| 39 |
+
|
| 40 |
+
if isinstance(self._value, str | Path):
|
| 41 |
+
if not os.path.isfile(self._value):
|
| 42 |
+
raise ValueError(f"File not found: {self._value}")
|
| 43 |
+
elif isinstance(self._value, np.ndarray) and self._value.dtype != np.uint8:
|
| 44 |
+
raise ValueError(
|
| 45 |
+
f"Invalid value dtype, expected np.uint8, got {self._value.dtype}"
|
| 46 |
+
)
|
| 47 |
+
|
| 48 |
+
def _file_extension(self) -> str:
|
| 49 |
+
if self._file_path:
|
| 50 |
+
return self._file_path.suffix[1:].lower()
|
| 51 |
+
if isinstance(self._value, str | Path):
|
| 52 |
+
path = Path(self._value)
|
| 53 |
+
return path.suffix[1:].lower()
|
| 54 |
+
if hasattr(self, "_format") and self._format:
|
| 55 |
+
return self._format
|
| 56 |
+
return "unknown"
|
| 57 |
+
|
| 58 |
+
def _get_relative_file_path(self) -> Path | None:
|
| 59 |
+
return self._file_path
|
| 60 |
+
|
| 61 |
+
def _get_absolute_file_path(self) -> Path | None:
|
| 62 |
+
if self._file_path:
|
| 63 |
+
return MEDIA_DIR / self._file_path
|
| 64 |
+
return None
|
| 65 |
+
|
| 66 |
+
def _save(self, project: str, run: str, step: int = 0):
|
| 67 |
+
if self._file_path:
|
| 68 |
+
return
|
| 69 |
+
|
| 70 |
+
media_dir = FileStorage.init_project_media_path(project, run, step)
|
| 71 |
+
filename = f"{uuid.uuid4()}.{self._file_extension()}"
|
| 72 |
+
file_path = media_dir / filename
|
| 73 |
+
|
| 74 |
+
self._save_media(file_path)
|
| 75 |
+
self._file_path = file_path.relative_to(MEDIA_DIR)
|
| 76 |
+
|
| 77 |
+
@abstractmethod
|
| 78 |
+
def _save_media(self, file_path: Path):
|
| 79 |
+
"""
|
| 80 |
+
Performs the actual media saving logic.
|
| 81 |
+
"""
|
| 82 |
+
pass
|
| 83 |
+
|
| 84 |
+
def _to_dict(self) -> dict:
|
| 85 |
+
if not self._file_path:
|
| 86 |
+
raise ValueError("Media must be saved to file before serialization")
|
| 87 |
+
return {
|
| 88 |
+
"_type": self.TYPE,
|
| 89 |
+
"file_path": str(self._get_relative_file_path()),
|
| 90 |
+
"caption": self.caption,
|
| 91 |
+
}
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
TrackioImageSourceType = str | Path | np.ndarray | PILImage.Image
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
class TrackioImage(TrackioMedia):
|
| 98 |
+
"""
|
| 99 |
+
Initializes an Image object.
|
| 100 |
+
|
| 101 |
+
Example:
|
| 102 |
+
```python
|
| 103 |
+
import trackio
|
| 104 |
+
import numpy as np
|
| 105 |
+
from PIL import Image
|
| 106 |
+
|
| 107 |
+
# Create an image from numpy array
|
| 108 |
+
image_data = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
|
| 109 |
+
image = trackio.Image(image_data, caption="Random image")
|
| 110 |
+
trackio.log({"my_image": image})
|
| 111 |
+
|
| 112 |
+
# Create an image from PIL Image
|
| 113 |
+
pil_image = Image.new('RGB', (100, 100), color='red')
|
| 114 |
+
image = trackio.Image(pil_image, caption="Red square")
|
| 115 |
+
trackio.log({"red_image": image})
|
| 116 |
+
|
| 117 |
+
# Create an image from file path
|
| 118 |
+
image = trackio.Image("path/to/image.jpg", caption="Photo from file")
|
| 119 |
+
trackio.log({"file_image": image})
|
| 120 |
+
```
|
| 121 |
+
|
| 122 |
+
Args:
|
| 123 |
+
value (`str`, `Path`, `numpy.ndarray`, or `PIL.Image`, *optional*):
|
| 124 |
+
A path to an image, a PIL Image, or a numpy array of shape (height, width, channels).
|
| 125 |
+
If numpy array, should be of type `np.uint8` with RGB values in the range `[0, 255]`.
|
| 126 |
+
caption (`str`, *optional*):
|
| 127 |
+
A string caption for the image.
|
| 128 |
+
"""
|
| 129 |
+
|
| 130 |
+
TYPE = "trackio.image"
|
| 131 |
+
|
| 132 |
+
def __init__(self, value: TrackioImageSourceType, caption: str | None = None):
|
| 133 |
+
super().__init__(value, caption)
|
| 134 |
+
self._format: str | None = None
|
| 135 |
+
|
| 136 |
+
if not isinstance(self._value, TrackioImageSourceType):
|
| 137 |
+
raise ValueError(
|
| 138 |
+
f"Invalid value type, expected {TrackioImageSourceType}, got {type(self._value)}"
|
| 139 |
+
)
|
| 140 |
+
|
| 141 |
+
if (
|
| 142 |
+
isinstance(self._value, np.ndarray | PILImage.Image)
|
| 143 |
+
and self._format is None
|
| 144 |
+
):
|
| 145 |
+
self._format = "png"
|
| 146 |
+
|
| 147 |
+
def _as_pil(self) -> PILImage.Image | None:
|
| 148 |
+
try:
|
| 149 |
+
if isinstance(self._value, np.ndarray):
|
| 150 |
+
arr = np.asarray(self._value).astype("uint8")
|
| 151 |
+
return PILImage.fromarray(arr).convert("RGBA")
|
| 152 |
+
if isinstance(self._value, PILImage.Image):
|
| 153 |
+
return self._value.convert("RGBA")
|
| 154 |
+
except Exception as e:
|
| 155 |
+
raise ValueError(f"Failed to process image data: {self._value}") from e
|
| 156 |
+
return None
|
| 157 |
+
|
| 158 |
+
def _save_media(self, file_path: Path):
|
| 159 |
+
if pil := self._as_pil():
|
| 160 |
+
pil.save(file_path, format=self._format)
|
| 161 |
+
elif isinstance(self._value, str | Path):
|
| 162 |
+
if os.path.isfile(self._value):
|
| 163 |
+
shutil.copy(self._value, file_path)
|
| 164 |
+
else:
|
| 165 |
+
raise ValueError(f"File not found: {self._value}")
|
| 166 |
+
|
| 167 |
+
|
| 168 |
+
TrackioVideoSourceType = str | Path | np.ndarray
|
| 169 |
+
TrackioVideoFormatType = Literal["gif", "mp4", "webm"]
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
class TrackioVideo(TrackioMedia):
|
| 173 |
+
"""
|
| 174 |
+
Initializes a Video object.
|
| 175 |
+
|
| 176 |
+
Example:
|
| 177 |
+
```python
|
| 178 |
+
import trackio
|
| 179 |
+
import numpy as np
|
| 180 |
+
|
| 181 |
+
# Create a simple video from numpy array
|
| 182 |
+
frames = np.random.randint(0, 255, (10, 3, 64, 64), dtype=np.uint8)
|
| 183 |
+
video = trackio.Video(frames, caption="Random video", fps=30)
|
| 184 |
+
|
| 185 |
+
# Create a batch of videos
|
| 186 |
+
batch_frames = np.random.randint(0, 255, (3, 10, 3, 64, 64), dtype=np.uint8)
|
| 187 |
+
batch_video = trackio.Video(batch_frames, caption="Batch of videos", fps=15)
|
| 188 |
+
|
| 189 |
+
# Create video from file path
|
| 190 |
+
video = trackio.Video("path/to/video.mp4", caption="Video from file")
|
| 191 |
+
```
|
| 192 |
+
|
| 193 |
+
Args:
|
| 194 |
+
value (`str`, `Path`, or `numpy.ndarray`, *optional*):
|
| 195 |
+
A path to a video file, or a numpy array.
|
| 196 |
+
If numpy array, should be of type `np.uint8` with RGB values in the range `[0, 255]`.
|
| 197 |
+
It is expected to have shape of either (frames, channels, height, width) or (batch, frames, channels, height, width).
|
| 198 |
+
For the latter, the videos will be tiled into a grid.
|
| 199 |
+
caption (`str`, *optional*):
|
| 200 |
+
A string caption for the video.
|
| 201 |
+
fps (`int`, *optional*):
|
| 202 |
+
Frames per second for the video. Only used when value is an ndarray. Default is `24`.
|
| 203 |
+
format (`Literal["gif", "mp4", "webm"]`, *optional*):
|
| 204 |
+
Video format ("gif", "mp4", or "webm"). Only used when value is an ndarray. Default is "gif".
|
| 205 |
+
"""
|
| 206 |
+
|
| 207 |
+
TYPE = "trackio.video"
|
| 208 |
+
|
| 209 |
+
def __init__(
|
| 210 |
+
self,
|
| 211 |
+
value: TrackioVideoSourceType,
|
| 212 |
+
caption: str | None = None,
|
| 213 |
+
fps: int | None = None,
|
| 214 |
+
format: TrackioVideoFormatType | None = None,
|
| 215 |
+
):
|
| 216 |
+
super().__init__(value, caption)
|
| 217 |
+
|
| 218 |
+
if not isinstance(self._value, TrackioVideoSourceType):
|
| 219 |
+
raise ValueError(
|
| 220 |
+
f"Invalid value type, expected {TrackioVideoSourceType}, got {type(self._value)}"
|
| 221 |
+
)
|
| 222 |
+
|
| 223 |
+
if isinstance(value, np.ndarray):
|
| 224 |
+
if format is None:
|
| 225 |
+
format = "gif"
|
| 226 |
+
if fps is None:
|
| 227 |
+
fps = 24
|
| 228 |
+
self._fps = fps
|
| 229 |
+
self._format = format
|
| 230 |
+
|
| 231 |
+
@property
|
| 232 |
+
def _codec(self) -> str:
|
| 233 |
+
match self._format:
|
| 234 |
+
case "gif":
|
| 235 |
+
return "gif"
|
| 236 |
+
case "mp4":
|
| 237 |
+
return "h264"
|
| 238 |
+
case "webm":
|
| 239 |
+
return "vp9"
|
| 240 |
+
case _:
|
| 241 |
+
raise ValueError(f"Unsupported format: {self._format}")
|
| 242 |
+
|
| 243 |
+
def _save_media(self, file_path: Path):
|
| 244 |
+
if isinstance(self._value, np.ndarray):
|
| 245 |
+
video = TrackioVideo._process_ndarray(self._value)
|
| 246 |
+
write_video(file_path, video, fps=self._fps, codec=self._codec)
|
| 247 |
+
elif isinstance(self._value, str | Path):
|
| 248 |
+
if os.path.isfile(self._value):
|
| 249 |
+
shutil.copy(self._value, file_path)
|
| 250 |
+
else:
|
| 251 |
+
raise ValueError(f"File not found: {self._value}")
|
| 252 |
+
|
| 253 |
+
@staticmethod
|
| 254 |
+
def _process_ndarray(value: np.ndarray) -> np.ndarray:
|
| 255 |
+
# Verify value is either 4D (single video) or 5D array (batched videos).
|
| 256 |
+
# Expected format: (frames, channels, height, width) or (batch, frames, channels, height, width)
|
| 257 |
+
if value.ndim < 4:
|
| 258 |
+
raise ValueError(
|
| 259 |
+
"Video requires at least 4 dimensions (frames, channels, height, width)"
|
| 260 |
+
)
|
| 261 |
+
if value.ndim > 5:
|
| 262 |
+
raise ValueError(
|
| 263 |
+
"Videos can have at most 5 dimensions (batch, frames, channels, height, width)"
|
| 264 |
+
)
|
| 265 |
+
if value.ndim == 4:
|
| 266 |
+
# Reshape to 5D with single batch: (1, frames, channels, height, width)
|
| 267 |
+
value = value[np.newaxis, ...]
|
| 268 |
+
|
| 269 |
+
value = TrackioVideo._tile_batched_videos(value)
|
| 270 |
+
return value
|
| 271 |
+
|
| 272 |
+
@staticmethod
|
| 273 |
+
def _tile_batched_videos(video: np.ndarray) -> np.ndarray:
|
| 274 |
+
"""
|
| 275 |
+
Tiles a batch of videos into a grid of videos.
|
| 276 |
+
|
| 277 |
+
Input format: (batch, frames, channels, height, width) - original FCHW format
|
| 278 |
+
Output format: (frames, total_height, total_width, channels)
|
| 279 |
+
"""
|
| 280 |
+
batch_size, frames, channels, height, width = video.shape
|
| 281 |
+
|
| 282 |
+
next_pow2 = 1 << (batch_size - 1).bit_length()
|
| 283 |
+
if batch_size != next_pow2:
|
| 284 |
+
pad_len = next_pow2 - batch_size
|
| 285 |
+
pad_shape = (pad_len, frames, channels, height, width)
|
| 286 |
+
padding = np.zeros(pad_shape, dtype=video.dtype)
|
| 287 |
+
video = np.concatenate((video, padding), axis=0)
|
| 288 |
+
batch_size = next_pow2
|
| 289 |
+
|
| 290 |
+
n_rows = 1 << ((batch_size.bit_length() - 1) // 2)
|
| 291 |
+
n_cols = batch_size // n_rows
|
| 292 |
+
|
| 293 |
+
# Reshape to grid layout: (n_rows, n_cols, frames, channels, height, width)
|
| 294 |
+
video = video.reshape(n_rows, n_cols, frames, channels, height, width)
|
| 295 |
+
|
| 296 |
+
# Rearrange dimensions to (frames, total_height, total_width, channels)
|
| 297 |
+
video = video.transpose(2, 0, 4, 1, 5, 3)
|
| 298 |
+
video = video.reshape(frames, n_rows * height, n_cols * width, channels)
|
| 299 |
+
return video
|
package.json
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"name": "trackio",
|
| 3 |
+
"version": "0.6.0",
|
| 4 |
+
"description": "",
|
| 5 |
+
"python": "true"
|
| 6 |
+
}
|
py.typed
ADDED
|
File without changes
|
run.py
ADDED
|
@@ -0,0 +1,192 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import threading
|
| 2 |
+
import time
|
| 3 |
+
import warnings
|
| 4 |
+
from datetime import datetime, timezone
|
| 5 |
+
|
| 6 |
+
import huggingface_hub
|
| 7 |
+
from gradio_client import Client, handle_file
|
| 8 |
+
|
| 9 |
+
from trackio import utils
|
| 10 |
+
from trackio.histogram import Histogram
|
| 11 |
+
from trackio.media import TrackioMedia
|
| 12 |
+
from trackio.sqlite_storage import SQLiteStorage
|
| 13 |
+
from trackio.table import Table
|
| 14 |
+
from trackio.typehints import LogEntry, UploadEntry
|
| 15 |
+
|
| 16 |
+
BATCH_SEND_INTERVAL = 0.5
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
class Run:
|
| 20 |
+
def __init__(
|
| 21 |
+
self,
|
| 22 |
+
url: str,
|
| 23 |
+
project: str,
|
| 24 |
+
client: Client | None,
|
| 25 |
+
name: str | None = None,
|
| 26 |
+
group: str | None = None,
|
| 27 |
+
config: dict | None = None,
|
| 28 |
+
space_id: str | None = None,
|
| 29 |
+
):
|
| 30 |
+
self.url = url
|
| 31 |
+
self.project = project
|
| 32 |
+
self._client_lock = threading.Lock()
|
| 33 |
+
self._client_thread = None
|
| 34 |
+
self._client = client
|
| 35 |
+
self._space_id = space_id
|
| 36 |
+
self.name = name or utils.generate_readable_name(
|
| 37 |
+
SQLiteStorage.get_runs(project), space_id
|
| 38 |
+
)
|
| 39 |
+
self.group = group
|
| 40 |
+
self.config = utils.to_json_safe(config or {})
|
| 41 |
+
|
| 42 |
+
if isinstance(self.config, dict):
|
| 43 |
+
for key in self.config:
|
| 44 |
+
if key.startswith("_"):
|
| 45 |
+
raise ValueError(
|
| 46 |
+
f"Config key '{key}' is reserved (keys starting with '_' are reserved for internal use)"
|
| 47 |
+
)
|
| 48 |
+
|
| 49 |
+
self.config["_Username"] = self._get_username()
|
| 50 |
+
self.config["_Created"] = datetime.now(timezone.utc).isoformat()
|
| 51 |
+
self.config["_Group"] = self.group
|
| 52 |
+
|
| 53 |
+
self._queued_logs: list[LogEntry] = []
|
| 54 |
+
self._queued_uploads: list[UploadEntry] = []
|
| 55 |
+
self._stop_flag = threading.Event()
|
| 56 |
+
self._config_logged = False
|
| 57 |
+
|
| 58 |
+
self._client_thread = threading.Thread(target=self._init_client_background)
|
| 59 |
+
self._client_thread.daemon = True
|
| 60 |
+
self._client_thread.start()
|
| 61 |
+
|
| 62 |
+
def _get_username(self) -> str | None:
|
| 63 |
+
"""Get the current HuggingFace username if logged in, otherwise None."""
|
| 64 |
+
try:
|
| 65 |
+
who = huggingface_hub.whoami()
|
| 66 |
+
return who["name"] if who else None
|
| 67 |
+
except Exception:
|
| 68 |
+
return None
|
| 69 |
+
|
| 70 |
+
def _batch_sender(self):
|
| 71 |
+
"""Send batched logs every BATCH_SEND_INTERVAL."""
|
| 72 |
+
while not self._stop_flag.is_set() or len(self._queued_logs) > 0:
|
| 73 |
+
# If the stop flag has been set, then just quickly send all
|
| 74 |
+
# the logs and exit.
|
| 75 |
+
if not self._stop_flag.is_set():
|
| 76 |
+
time.sleep(BATCH_SEND_INTERVAL)
|
| 77 |
+
|
| 78 |
+
with self._client_lock:
|
| 79 |
+
if self._client is None:
|
| 80 |
+
return
|
| 81 |
+
if self._queued_logs:
|
| 82 |
+
logs_to_send = self._queued_logs.copy()
|
| 83 |
+
self._queued_logs.clear()
|
| 84 |
+
self._client.predict(
|
| 85 |
+
api_name="/bulk_log",
|
| 86 |
+
logs=logs_to_send,
|
| 87 |
+
hf_token=huggingface_hub.utils.get_token(),
|
| 88 |
+
)
|
| 89 |
+
if self._queued_uploads:
|
| 90 |
+
uploads_to_send = self._queued_uploads.copy()
|
| 91 |
+
self._queued_uploads.clear()
|
| 92 |
+
self._client.predict(
|
| 93 |
+
api_name="/bulk_upload_media",
|
| 94 |
+
uploads=uploads_to_send,
|
| 95 |
+
hf_token=huggingface_hub.utils.get_token(),
|
| 96 |
+
)
|
| 97 |
+
|
| 98 |
+
def _init_client_background(self):
|
| 99 |
+
if self._client is None:
|
| 100 |
+
fib = utils.fibo()
|
| 101 |
+
for sleep_coefficient in fib:
|
| 102 |
+
try:
|
| 103 |
+
client = Client(self.url, verbose=False)
|
| 104 |
+
|
| 105 |
+
with self._client_lock:
|
| 106 |
+
self._client = client
|
| 107 |
+
break
|
| 108 |
+
except Exception:
|
| 109 |
+
pass
|
| 110 |
+
if sleep_coefficient is not None:
|
| 111 |
+
time.sleep(0.1 * sleep_coefficient)
|
| 112 |
+
|
| 113 |
+
self._batch_sender()
|
| 114 |
+
|
| 115 |
+
def _process_media(self, metrics, step: int | None) -> dict:
|
| 116 |
+
"""
|
| 117 |
+
Serialize media in metrics and upload to space if needed.
|
| 118 |
+
"""
|
| 119 |
+
serializable_metrics = {}
|
| 120 |
+
if not step:
|
| 121 |
+
step = 0
|
| 122 |
+
for key, value in metrics.items():
|
| 123 |
+
if isinstance(value, TrackioMedia):
|
| 124 |
+
value._save(self.project, self.name, step)
|
| 125 |
+
serializable_metrics[key] = value._to_dict()
|
| 126 |
+
if self._space_id:
|
| 127 |
+
# Upload local media when deploying to space
|
| 128 |
+
upload_entry: UploadEntry = {
|
| 129 |
+
"project": self.project,
|
| 130 |
+
"run": self.name,
|
| 131 |
+
"step": step,
|
| 132 |
+
"uploaded_file": handle_file(value._get_absolute_file_path()),
|
| 133 |
+
}
|
| 134 |
+
with self._client_lock:
|
| 135 |
+
self._queued_uploads.append(upload_entry)
|
| 136 |
+
else:
|
| 137 |
+
serializable_metrics[key] = value
|
| 138 |
+
return serializable_metrics
|
| 139 |
+
|
| 140 |
+
@staticmethod
|
| 141 |
+
def _replace_tables(metrics):
|
| 142 |
+
for k, v in metrics.items():
|
| 143 |
+
if isinstance(v, (Table, Histogram)):
|
| 144 |
+
metrics[k] = v._to_dict()
|
| 145 |
+
|
| 146 |
+
def log(self, metrics: dict, step: int | None = None):
|
| 147 |
+
renamed_keys = []
|
| 148 |
+
new_metrics = {}
|
| 149 |
+
|
| 150 |
+
for k, v in metrics.items():
|
| 151 |
+
if k in utils.RESERVED_KEYS or k.startswith("__"):
|
| 152 |
+
new_key = f"__{k}"
|
| 153 |
+
renamed_keys.append(k)
|
| 154 |
+
new_metrics[new_key] = v
|
| 155 |
+
else:
|
| 156 |
+
new_metrics[k] = v
|
| 157 |
+
|
| 158 |
+
if renamed_keys:
|
| 159 |
+
warnings.warn(f"Reserved keys renamed: {renamed_keys} → '__{{key}}'")
|
| 160 |
+
|
| 161 |
+
metrics = new_metrics
|
| 162 |
+
Run._replace_tables(metrics)
|
| 163 |
+
|
| 164 |
+
metrics = self._process_media(metrics, step)
|
| 165 |
+
metrics = utils.serialize_values(metrics)
|
| 166 |
+
|
| 167 |
+
config_to_log = None
|
| 168 |
+
if not self._config_logged and self.config:
|
| 169 |
+
config_to_log = utils.to_json_safe(self.config)
|
| 170 |
+
self._config_logged = True
|
| 171 |
+
|
| 172 |
+
log_entry: LogEntry = {
|
| 173 |
+
"project": self.project,
|
| 174 |
+
"run": self.name,
|
| 175 |
+
"metrics": metrics,
|
| 176 |
+
"step": step,
|
| 177 |
+
"config": config_to_log,
|
| 178 |
+
}
|
| 179 |
+
|
| 180 |
+
with self._client_lock:
|
| 181 |
+
self._queued_logs.append(log_entry)
|
| 182 |
+
|
| 183 |
+
def finish(self):
|
| 184 |
+
"""Cleanup when run is finished."""
|
| 185 |
+
self._stop_flag.set()
|
| 186 |
+
|
| 187 |
+
# Wait for the batch sender to finish before joining the client thread.
|
| 188 |
+
time.sleep(2 * BATCH_SEND_INTERVAL)
|
| 189 |
+
|
| 190 |
+
if self._client_thread is not None:
|
| 191 |
+
print("* Run finished. Uploading logs to Trackio (please wait...)")
|
| 192 |
+
self._client_thread.join()
|
sqlite_storage.py
ADDED
|
@@ -0,0 +1,677 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import platform
|
| 3 |
+
import sqlite3
|
| 4 |
+
import time
|
| 5 |
+
from datetime import datetime
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
from threading import Lock
|
| 8 |
+
|
| 9 |
+
try:
|
| 10 |
+
import fcntl
|
| 11 |
+
except ImportError: # fcntl is not available on Windows
|
| 12 |
+
fcntl = None
|
| 13 |
+
|
| 14 |
+
import huggingface_hub as hf
|
| 15 |
+
import orjson
|
| 16 |
+
import pandas as pd
|
| 17 |
+
|
| 18 |
+
try: # absolute imports when installed from PyPI
|
| 19 |
+
from trackio.commit_scheduler import CommitScheduler
|
| 20 |
+
from trackio.dummy_commit_scheduler import DummyCommitScheduler
|
| 21 |
+
from trackio.utils import (
|
| 22 |
+
TRACKIO_DIR,
|
| 23 |
+
deserialize_values,
|
| 24 |
+
serialize_values,
|
| 25 |
+
)
|
| 26 |
+
except ImportError: # relative imports when installed from source on Spaces
|
| 27 |
+
from commit_scheduler import CommitScheduler
|
| 28 |
+
from dummy_commit_scheduler import DummyCommitScheduler
|
| 29 |
+
from utils import TRACKIO_DIR, deserialize_values, serialize_values
|
| 30 |
+
|
| 31 |
+
DB_EXT = ".db"
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class ProcessLock:
|
| 35 |
+
"""A file-based lock that works across processes. Is a no-op on Windows."""
|
| 36 |
+
|
| 37 |
+
def __init__(self, lockfile_path: Path):
|
| 38 |
+
self.lockfile_path = lockfile_path
|
| 39 |
+
self.lockfile = None
|
| 40 |
+
self.is_windows = platform.system() == "Windows"
|
| 41 |
+
|
| 42 |
+
def __enter__(self):
|
| 43 |
+
"""Acquire the lock with retry logic."""
|
| 44 |
+
if self.is_windows:
|
| 45 |
+
return self
|
| 46 |
+
self.lockfile_path.parent.mkdir(parents=True, exist_ok=True)
|
| 47 |
+
self.lockfile = open(self.lockfile_path, "w")
|
| 48 |
+
|
| 49 |
+
max_retries = 100
|
| 50 |
+
for attempt in range(max_retries):
|
| 51 |
+
try:
|
| 52 |
+
fcntl.flock(self.lockfile.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
| 53 |
+
return self
|
| 54 |
+
except IOError:
|
| 55 |
+
if attempt < max_retries - 1:
|
| 56 |
+
time.sleep(0.1)
|
| 57 |
+
else:
|
| 58 |
+
raise IOError("Could not acquire database lock after 10 seconds")
|
| 59 |
+
|
| 60 |
+
def __exit__(self, exc_type, exc_val, exc_tb):
|
| 61 |
+
"""Release the lock."""
|
| 62 |
+
if self.is_windows:
|
| 63 |
+
return
|
| 64 |
+
|
| 65 |
+
if self.lockfile:
|
| 66 |
+
fcntl.flock(self.lockfile.fileno(), fcntl.LOCK_UN)
|
| 67 |
+
self.lockfile.close()
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
class SQLiteStorage:
|
| 71 |
+
_dataset_import_attempted = False
|
| 72 |
+
_current_scheduler: CommitScheduler | DummyCommitScheduler | None = None
|
| 73 |
+
_scheduler_lock = Lock()
|
| 74 |
+
|
| 75 |
+
@staticmethod
|
| 76 |
+
def _get_connection(db_path: Path) -> sqlite3.Connection:
|
| 77 |
+
conn = sqlite3.connect(str(db_path), timeout=30.0)
|
| 78 |
+
# Keep WAL for concurrency + performance on many small writes
|
| 79 |
+
conn.execute("PRAGMA journal_mode = WAL")
|
| 80 |
+
# ---- Minimal perf tweaks for many tiny transactions ----
|
| 81 |
+
# NORMAL = fsync at critical points only (safer than OFF, much faster than FULL)
|
| 82 |
+
conn.execute("PRAGMA synchronous = NORMAL")
|
| 83 |
+
# Keep temp data in memory to avoid disk hits during small writes
|
| 84 |
+
conn.execute("PRAGMA temp_store = MEMORY")
|
| 85 |
+
# Give SQLite a bit more room for cache (negative = KB, engine-managed)
|
| 86 |
+
conn.execute("PRAGMA cache_size = -20000")
|
| 87 |
+
# --------------------------------------------------------
|
| 88 |
+
conn.row_factory = sqlite3.Row
|
| 89 |
+
return conn
|
| 90 |
+
|
| 91 |
+
@staticmethod
|
| 92 |
+
def _get_process_lock(project: str) -> ProcessLock:
|
| 93 |
+
lockfile_path = TRACKIO_DIR / f"{project}.lock"
|
| 94 |
+
return ProcessLock(lockfile_path)
|
| 95 |
+
|
| 96 |
+
@staticmethod
|
| 97 |
+
def get_project_db_filename(project: str) -> str:
|
| 98 |
+
"""Get the database filename for a specific project."""
|
| 99 |
+
safe_project_name = "".join(
|
| 100 |
+
c for c in project if c.isalnum() or c in ("-", "_")
|
| 101 |
+
).rstrip()
|
| 102 |
+
if not safe_project_name:
|
| 103 |
+
safe_project_name = "default"
|
| 104 |
+
return f"{safe_project_name}{DB_EXT}"
|
| 105 |
+
|
| 106 |
+
@staticmethod
|
| 107 |
+
def get_project_db_path(project: str) -> Path:
|
| 108 |
+
"""Get the database path for a specific project."""
|
| 109 |
+
filename = SQLiteStorage.get_project_db_filename(project)
|
| 110 |
+
return TRACKIO_DIR / filename
|
| 111 |
+
|
| 112 |
+
@staticmethod
|
| 113 |
+
def init_db(project: str) -> Path:
|
| 114 |
+
"""
|
| 115 |
+
Initialize the SQLite database with required tables.
|
| 116 |
+
Returns the database path.
|
| 117 |
+
"""
|
| 118 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 119 |
+
db_path.parent.mkdir(parents=True, exist_ok=True)
|
| 120 |
+
with SQLiteStorage._get_process_lock(project):
|
| 121 |
+
with sqlite3.connect(str(db_path), timeout=30.0) as conn:
|
| 122 |
+
conn.execute("PRAGMA journal_mode = WAL")
|
| 123 |
+
conn.execute("PRAGMA synchronous = NORMAL")
|
| 124 |
+
conn.execute("PRAGMA temp_store = MEMORY")
|
| 125 |
+
conn.execute("PRAGMA cache_size = -20000")
|
| 126 |
+
cursor = conn.cursor()
|
| 127 |
+
cursor.execute(
|
| 128 |
+
"""
|
| 129 |
+
CREATE TABLE IF NOT EXISTS metrics (
|
| 130 |
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
| 131 |
+
timestamp TEXT NOT NULL,
|
| 132 |
+
run_name TEXT NOT NULL,
|
| 133 |
+
step INTEGER NOT NULL,
|
| 134 |
+
metrics TEXT NOT NULL
|
| 135 |
+
)
|
| 136 |
+
"""
|
| 137 |
+
)
|
| 138 |
+
cursor.execute(
|
| 139 |
+
"""
|
| 140 |
+
CREATE TABLE IF NOT EXISTS configs (
|
| 141 |
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
| 142 |
+
run_name TEXT NOT NULL,
|
| 143 |
+
config TEXT NOT NULL,
|
| 144 |
+
created_at TEXT NOT NULL,
|
| 145 |
+
UNIQUE(run_name)
|
| 146 |
+
)
|
| 147 |
+
"""
|
| 148 |
+
)
|
| 149 |
+
cursor.execute(
|
| 150 |
+
"""
|
| 151 |
+
CREATE INDEX IF NOT EXISTS idx_metrics_run_step
|
| 152 |
+
ON metrics(run_name, step)
|
| 153 |
+
"""
|
| 154 |
+
)
|
| 155 |
+
cursor.execute(
|
| 156 |
+
"""
|
| 157 |
+
CREATE INDEX IF NOT EXISTS idx_configs_run_name
|
| 158 |
+
ON configs(run_name)
|
| 159 |
+
"""
|
| 160 |
+
)
|
| 161 |
+
cursor.execute(
|
| 162 |
+
"""
|
| 163 |
+
CREATE INDEX IF NOT EXISTS idx_metrics_run_timestamp
|
| 164 |
+
ON metrics(run_name, timestamp)
|
| 165 |
+
"""
|
| 166 |
+
)
|
| 167 |
+
conn.commit()
|
| 168 |
+
return db_path
|
| 169 |
+
|
| 170 |
+
@staticmethod
|
| 171 |
+
def export_to_parquet():
|
| 172 |
+
"""
|
| 173 |
+
Exports all projects' DB files as Parquet under the same path but with extension ".parquet".
|
| 174 |
+
"""
|
| 175 |
+
# don't attempt to export (potentially wrong/blank) data before importing for the first time
|
| 176 |
+
if not SQLiteStorage._dataset_import_attempted:
|
| 177 |
+
return
|
| 178 |
+
if not TRACKIO_DIR.exists():
|
| 179 |
+
return
|
| 180 |
+
|
| 181 |
+
all_paths = os.listdir(TRACKIO_DIR)
|
| 182 |
+
db_names = [f for f in all_paths if f.endswith(DB_EXT)]
|
| 183 |
+
for db_name in db_names:
|
| 184 |
+
db_path = TRACKIO_DIR / db_name
|
| 185 |
+
parquet_path = db_path.with_suffix(".parquet")
|
| 186 |
+
if (not parquet_path.exists()) or (
|
| 187 |
+
db_path.stat().st_mtime > parquet_path.stat().st_mtime
|
| 188 |
+
):
|
| 189 |
+
with sqlite3.connect(str(db_path)) as conn:
|
| 190 |
+
df = pd.read_sql("SELECT * FROM metrics", conn)
|
| 191 |
+
# break out the single JSON metrics column into individual columns
|
| 192 |
+
metrics = df["metrics"].copy()
|
| 193 |
+
metrics = pd.DataFrame(
|
| 194 |
+
metrics.apply(
|
| 195 |
+
lambda x: deserialize_values(orjson.loads(x))
|
| 196 |
+
).values.tolist(),
|
| 197 |
+
index=df.index,
|
| 198 |
+
)
|
| 199 |
+
del df["metrics"]
|
| 200 |
+
for col in metrics.columns:
|
| 201 |
+
df[col] = metrics[col]
|
| 202 |
+
|
| 203 |
+
df.to_parquet(parquet_path)
|
| 204 |
+
|
| 205 |
+
@staticmethod
|
| 206 |
+
def _cleanup_wal_sidecars(db_path: Path) -> None:
|
| 207 |
+
"""Remove leftover -wal/-shm files for a DB basename (prevents disk I/O errors)."""
|
| 208 |
+
for suffix in ("-wal", "-shm"):
|
| 209 |
+
sidecar = Path(str(db_path) + suffix)
|
| 210 |
+
try:
|
| 211 |
+
if sidecar.exists():
|
| 212 |
+
sidecar.unlink()
|
| 213 |
+
except Exception:
|
| 214 |
+
pass
|
| 215 |
+
|
| 216 |
+
@staticmethod
|
| 217 |
+
def import_from_parquet():
|
| 218 |
+
"""
|
| 219 |
+
Imports to all DB files that have matching files under the same path but with extension ".parquet".
|
| 220 |
+
"""
|
| 221 |
+
if not TRACKIO_DIR.exists():
|
| 222 |
+
return
|
| 223 |
+
|
| 224 |
+
all_paths = os.listdir(TRACKIO_DIR)
|
| 225 |
+
parquet_names = [f for f in all_paths if f.endswith(".parquet")]
|
| 226 |
+
for pq_name in parquet_names:
|
| 227 |
+
parquet_path = TRACKIO_DIR / pq_name
|
| 228 |
+
db_path = parquet_path.with_suffix(DB_EXT)
|
| 229 |
+
|
| 230 |
+
SQLiteStorage._cleanup_wal_sidecars(db_path)
|
| 231 |
+
|
| 232 |
+
df = pd.read_parquet(parquet_path)
|
| 233 |
+
# fix up df to have a single JSON metrics column
|
| 234 |
+
if "metrics" not in df.columns:
|
| 235 |
+
# separate other columns from metrics
|
| 236 |
+
metrics = df.copy()
|
| 237 |
+
other_cols = ["id", "timestamp", "run_name", "step"]
|
| 238 |
+
df = df[other_cols]
|
| 239 |
+
for col in other_cols:
|
| 240 |
+
del metrics[col]
|
| 241 |
+
# combine them all into a single metrics col
|
| 242 |
+
metrics = orjson.loads(metrics.to_json(orient="records"))
|
| 243 |
+
df["metrics"] = [orjson.dumps(serialize_values(row)) for row in metrics]
|
| 244 |
+
|
| 245 |
+
with sqlite3.connect(str(db_path), timeout=30.0) as conn:
|
| 246 |
+
df.to_sql("metrics", conn, if_exists="replace", index=False)
|
| 247 |
+
conn.commit()
|
| 248 |
+
|
| 249 |
+
@staticmethod
|
| 250 |
+
def get_scheduler():
|
| 251 |
+
"""
|
| 252 |
+
Get the scheduler for the database based on the environment variables.
|
| 253 |
+
This applies to both local and Spaces.
|
| 254 |
+
"""
|
| 255 |
+
with SQLiteStorage._scheduler_lock:
|
| 256 |
+
if SQLiteStorage._current_scheduler is not None:
|
| 257 |
+
return SQLiteStorage._current_scheduler
|
| 258 |
+
hf_token = os.environ.get("HF_TOKEN")
|
| 259 |
+
dataset_id = os.environ.get("TRACKIO_DATASET_ID")
|
| 260 |
+
space_repo_name = os.environ.get("SPACE_REPO_NAME")
|
| 261 |
+
if dataset_id is None or space_repo_name is None:
|
| 262 |
+
scheduler = DummyCommitScheduler()
|
| 263 |
+
else:
|
| 264 |
+
scheduler = CommitScheduler(
|
| 265 |
+
repo_id=dataset_id,
|
| 266 |
+
repo_type="dataset",
|
| 267 |
+
folder_path=TRACKIO_DIR,
|
| 268 |
+
private=True,
|
| 269 |
+
allow_patterns=["*.parquet", "media/**/*"],
|
| 270 |
+
squash_history=True,
|
| 271 |
+
token=hf_token,
|
| 272 |
+
on_before_commit=SQLiteStorage.export_to_parquet,
|
| 273 |
+
)
|
| 274 |
+
SQLiteStorage._current_scheduler = scheduler
|
| 275 |
+
return scheduler
|
| 276 |
+
|
| 277 |
+
@staticmethod
|
| 278 |
+
def log(project: str, run: str, metrics: dict, step: int | None = None):
|
| 279 |
+
"""
|
| 280 |
+
Safely log metrics to the database. Before logging, this method will ensure the database exists
|
| 281 |
+
and is set up with the correct tables. It also uses a cross-process lock to prevent
|
| 282 |
+
database locking errors when multiple processes access the same database.
|
| 283 |
+
|
| 284 |
+
This method is not used in the latest versions of Trackio (replaced by bulk_log) but
|
| 285 |
+
is kept for backwards compatibility for users who are connecting to a newer version of
|
| 286 |
+
a Trackio Spaces dashboard with an older version of Trackio installed locally.
|
| 287 |
+
"""
|
| 288 |
+
db_path = SQLiteStorage.init_db(project)
|
| 289 |
+
with SQLiteStorage._get_process_lock(project):
|
| 290 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 291 |
+
cursor = conn.cursor()
|
| 292 |
+
cursor.execute(
|
| 293 |
+
"""
|
| 294 |
+
SELECT MAX(step)
|
| 295 |
+
FROM metrics
|
| 296 |
+
WHERE run_name = ?
|
| 297 |
+
""",
|
| 298 |
+
(run,),
|
| 299 |
+
)
|
| 300 |
+
last_step = cursor.fetchone()[0]
|
| 301 |
+
current_step = (
|
| 302 |
+
0
|
| 303 |
+
if step is None and last_step is None
|
| 304 |
+
else (step if step is not None else last_step + 1)
|
| 305 |
+
)
|
| 306 |
+
current_timestamp = datetime.now().isoformat()
|
| 307 |
+
cursor.execute(
|
| 308 |
+
"""
|
| 309 |
+
INSERT INTO metrics
|
| 310 |
+
(timestamp, run_name, step, metrics)
|
| 311 |
+
VALUES (?, ?, ?, ?)
|
| 312 |
+
""",
|
| 313 |
+
(
|
| 314 |
+
current_timestamp,
|
| 315 |
+
run,
|
| 316 |
+
current_step,
|
| 317 |
+
orjson.dumps(serialize_values(metrics)),
|
| 318 |
+
),
|
| 319 |
+
)
|
| 320 |
+
conn.commit()
|
| 321 |
+
|
| 322 |
+
@staticmethod
|
| 323 |
+
def bulk_log(
|
| 324 |
+
project: str,
|
| 325 |
+
run: str,
|
| 326 |
+
metrics_list: list[dict],
|
| 327 |
+
steps: list[int] | None = None,
|
| 328 |
+
timestamps: list[str] | None = None,
|
| 329 |
+
config: dict | None = None,
|
| 330 |
+
):
|
| 331 |
+
"""
|
| 332 |
+
Safely log bulk metrics to the database. Before logging, this method will ensure the database exists
|
| 333 |
+
and is set up with the correct tables. It also uses a cross-process lock to prevent
|
| 334 |
+
database locking errors when multiple processes access the same database.
|
| 335 |
+
"""
|
| 336 |
+
if not metrics_list:
|
| 337 |
+
return
|
| 338 |
+
|
| 339 |
+
if timestamps is None:
|
| 340 |
+
timestamps = [datetime.now().isoformat()] * len(metrics_list)
|
| 341 |
+
|
| 342 |
+
db_path = SQLiteStorage.init_db(project)
|
| 343 |
+
with SQLiteStorage._get_process_lock(project):
|
| 344 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 345 |
+
cursor = conn.cursor()
|
| 346 |
+
|
| 347 |
+
if steps is None:
|
| 348 |
+
steps = list(range(len(metrics_list)))
|
| 349 |
+
elif any(s is None for s in steps):
|
| 350 |
+
cursor.execute(
|
| 351 |
+
"SELECT MAX(step) FROM metrics WHERE run_name = ?", (run,)
|
| 352 |
+
)
|
| 353 |
+
last_step = cursor.fetchone()[0]
|
| 354 |
+
current_step = 0 if last_step is None else last_step + 1
|
| 355 |
+
processed_steps = []
|
| 356 |
+
for step in steps:
|
| 357 |
+
if step is None:
|
| 358 |
+
processed_steps.append(current_step)
|
| 359 |
+
current_step += 1
|
| 360 |
+
else:
|
| 361 |
+
processed_steps.append(step)
|
| 362 |
+
steps = processed_steps
|
| 363 |
+
|
| 364 |
+
if len(metrics_list) != len(steps) or len(metrics_list) != len(
|
| 365 |
+
timestamps
|
| 366 |
+
):
|
| 367 |
+
raise ValueError(
|
| 368 |
+
"metrics_list, steps, and timestamps must have the same length"
|
| 369 |
+
)
|
| 370 |
+
|
| 371 |
+
data = []
|
| 372 |
+
for i, metrics in enumerate(metrics_list):
|
| 373 |
+
data.append(
|
| 374 |
+
(
|
| 375 |
+
timestamps[i],
|
| 376 |
+
run,
|
| 377 |
+
steps[i],
|
| 378 |
+
orjson.dumps(serialize_values(metrics)),
|
| 379 |
+
)
|
| 380 |
+
)
|
| 381 |
+
|
| 382 |
+
cursor.executemany(
|
| 383 |
+
"""
|
| 384 |
+
INSERT INTO metrics
|
| 385 |
+
(timestamp, run_name, step, metrics)
|
| 386 |
+
VALUES (?, ?, ?, ?)
|
| 387 |
+
""",
|
| 388 |
+
data,
|
| 389 |
+
)
|
| 390 |
+
|
| 391 |
+
if config:
|
| 392 |
+
current_timestamp = datetime.now().isoformat()
|
| 393 |
+
cursor.execute(
|
| 394 |
+
"""
|
| 395 |
+
INSERT OR REPLACE INTO configs
|
| 396 |
+
(run_name, config, created_at)
|
| 397 |
+
VALUES (?, ?, ?)
|
| 398 |
+
""",
|
| 399 |
+
(
|
| 400 |
+
run,
|
| 401 |
+
orjson.dumps(serialize_values(config)),
|
| 402 |
+
current_timestamp,
|
| 403 |
+
),
|
| 404 |
+
)
|
| 405 |
+
|
| 406 |
+
conn.commit()
|
| 407 |
+
|
| 408 |
+
@staticmethod
|
| 409 |
+
def get_logs(project: str, run: str) -> list[dict]:
|
| 410 |
+
"""Retrieve logs for a specific run. Logs include the step count (int) and the timestamp (datetime object)."""
|
| 411 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 412 |
+
if not db_path.exists():
|
| 413 |
+
return []
|
| 414 |
+
|
| 415 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 416 |
+
cursor = conn.cursor()
|
| 417 |
+
cursor.execute(
|
| 418 |
+
"""
|
| 419 |
+
SELECT timestamp, step, metrics
|
| 420 |
+
FROM metrics
|
| 421 |
+
WHERE run_name = ?
|
| 422 |
+
ORDER BY timestamp
|
| 423 |
+
""",
|
| 424 |
+
(run,),
|
| 425 |
+
)
|
| 426 |
+
|
| 427 |
+
rows = cursor.fetchall()
|
| 428 |
+
results = []
|
| 429 |
+
for row in rows:
|
| 430 |
+
metrics = orjson.loads(row["metrics"])
|
| 431 |
+
metrics = deserialize_values(metrics)
|
| 432 |
+
metrics["timestamp"] = row["timestamp"]
|
| 433 |
+
metrics["step"] = row["step"]
|
| 434 |
+
results.append(metrics)
|
| 435 |
+
return results
|
| 436 |
+
|
| 437 |
+
@staticmethod
|
| 438 |
+
def load_from_dataset():
|
| 439 |
+
dataset_id = os.environ.get("TRACKIO_DATASET_ID")
|
| 440 |
+
space_repo_name = os.environ.get("SPACE_REPO_NAME")
|
| 441 |
+
if dataset_id is not None and space_repo_name is not None:
|
| 442 |
+
hfapi = hf.HfApi()
|
| 443 |
+
updated = False
|
| 444 |
+
if not TRACKIO_DIR.exists():
|
| 445 |
+
TRACKIO_DIR.mkdir(parents=True, exist_ok=True)
|
| 446 |
+
with SQLiteStorage.get_scheduler().lock:
|
| 447 |
+
try:
|
| 448 |
+
files = hfapi.list_repo_files(dataset_id, repo_type="dataset")
|
| 449 |
+
for file in files:
|
| 450 |
+
# Download parquet and media assets
|
| 451 |
+
if not (file.endswith(".parquet") or file.startswith("media/")):
|
| 452 |
+
continue
|
| 453 |
+
if (TRACKIO_DIR / file).exists():
|
| 454 |
+
continue
|
| 455 |
+
hf.hf_hub_download(
|
| 456 |
+
dataset_id, file, repo_type="dataset", local_dir=TRACKIO_DIR
|
| 457 |
+
)
|
| 458 |
+
updated = True
|
| 459 |
+
except hf.errors.EntryNotFoundError:
|
| 460 |
+
pass
|
| 461 |
+
except hf.errors.RepositoryNotFoundError:
|
| 462 |
+
pass
|
| 463 |
+
if updated:
|
| 464 |
+
SQLiteStorage.import_from_parquet()
|
| 465 |
+
SQLiteStorage._dataset_import_attempted = True
|
| 466 |
+
|
| 467 |
+
@staticmethod
|
| 468 |
+
def get_projects() -> list[str]:
|
| 469 |
+
"""
|
| 470 |
+
Get list of all projects by scanning the database files in the trackio directory.
|
| 471 |
+
"""
|
| 472 |
+
if not SQLiteStorage._dataset_import_attempted:
|
| 473 |
+
SQLiteStorage.load_from_dataset()
|
| 474 |
+
|
| 475 |
+
projects: set[str] = set()
|
| 476 |
+
if not TRACKIO_DIR.exists():
|
| 477 |
+
return []
|
| 478 |
+
|
| 479 |
+
for db_file in TRACKIO_DIR.glob(f"*{DB_EXT}"):
|
| 480 |
+
project_name = db_file.stem
|
| 481 |
+
projects.add(project_name)
|
| 482 |
+
return sorted(projects)
|
| 483 |
+
|
| 484 |
+
@staticmethod
|
| 485 |
+
def get_runs(project: str) -> list[str]:
|
| 486 |
+
"""Get list of all runs for a project."""
|
| 487 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 488 |
+
if not db_path.exists():
|
| 489 |
+
return []
|
| 490 |
+
|
| 491 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 492 |
+
cursor = conn.cursor()
|
| 493 |
+
cursor.execute(
|
| 494 |
+
"SELECT DISTINCT run_name FROM metrics",
|
| 495 |
+
)
|
| 496 |
+
return [row[0] for row in cursor.fetchall()]
|
| 497 |
+
|
| 498 |
+
@staticmethod
|
| 499 |
+
def get_max_steps_for_runs(project: str) -> dict[str, int]:
|
| 500 |
+
"""Get the maximum step for each run in a project."""
|
| 501 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 502 |
+
if not db_path.exists():
|
| 503 |
+
return {}
|
| 504 |
+
|
| 505 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 506 |
+
cursor = conn.cursor()
|
| 507 |
+
cursor.execute(
|
| 508 |
+
"""
|
| 509 |
+
SELECT run_name, MAX(step) as max_step
|
| 510 |
+
FROM metrics
|
| 511 |
+
GROUP BY run_name
|
| 512 |
+
"""
|
| 513 |
+
)
|
| 514 |
+
|
| 515 |
+
results = {}
|
| 516 |
+
for row in cursor.fetchall():
|
| 517 |
+
results[row["run_name"]] = row["max_step"]
|
| 518 |
+
|
| 519 |
+
return results
|
| 520 |
+
|
| 521 |
+
@staticmethod
|
| 522 |
+
def store_config(project: str, run: str, config: dict) -> None:
|
| 523 |
+
"""Store configuration for a run."""
|
| 524 |
+
db_path = SQLiteStorage.init_db(project)
|
| 525 |
+
|
| 526 |
+
with SQLiteStorage._get_process_lock(project):
|
| 527 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 528 |
+
cursor = conn.cursor()
|
| 529 |
+
current_timestamp = datetime.now().isoformat()
|
| 530 |
+
|
| 531 |
+
cursor.execute(
|
| 532 |
+
"""
|
| 533 |
+
INSERT OR REPLACE INTO configs
|
| 534 |
+
(run_name, config, created_at)
|
| 535 |
+
VALUES (?, ?, ?)
|
| 536 |
+
""",
|
| 537 |
+
(run, orjson.dumps(serialize_values(config)), current_timestamp),
|
| 538 |
+
)
|
| 539 |
+
conn.commit()
|
| 540 |
+
|
| 541 |
+
@staticmethod
|
| 542 |
+
def get_run_config(project: str, run: str) -> dict | None:
|
| 543 |
+
"""Get configuration for a specific run."""
|
| 544 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 545 |
+
if not db_path.exists():
|
| 546 |
+
return None
|
| 547 |
+
|
| 548 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 549 |
+
cursor = conn.cursor()
|
| 550 |
+
try:
|
| 551 |
+
cursor.execute(
|
| 552 |
+
"""
|
| 553 |
+
SELECT config FROM configs WHERE run_name = ?
|
| 554 |
+
""",
|
| 555 |
+
(run,),
|
| 556 |
+
)
|
| 557 |
+
|
| 558 |
+
row = cursor.fetchone()
|
| 559 |
+
if row:
|
| 560 |
+
config = orjson.loads(row["config"])
|
| 561 |
+
return deserialize_values(config)
|
| 562 |
+
return None
|
| 563 |
+
except sqlite3.OperationalError as e:
|
| 564 |
+
if "no such table: configs" in str(e):
|
| 565 |
+
return None
|
| 566 |
+
raise
|
| 567 |
+
|
| 568 |
+
@staticmethod
|
| 569 |
+
def delete_run(project: str, run: str) -> bool:
|
| 570 |
+
"""Delete a run from the database (both metrics and config)."""
|
| 571 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 572 |
+
if not db_path.exists():
|
| 573 |
+
return False
|
| 574 |
+
|
| 575 |
+
with SQLiteStorage._get_process_lock(project):
|
| 576 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 577 |
+
cursor = conn.cursor()
|
| 578 |
+
try:
|
| 579 |
+
cursor.execute("DELETE FROM metrics WHERE run_name = ?", (run,))
|
| 580 |
+
cursor.execute("DELETE FROM configs WHERE run_name = ?", (run,))
|
| 581 |
+
conn.commit()
|
| 582 |
+
return True
|
| 583 |
+
except sqlite3.Error:
|
| 584 |
+
return False
|
| 585 |
+
|
| 586 |
+
@staticmethod
|
| 587 |
+
def get_all_run_configs(project: str) -> dict[str, dict]:
|
| 588 |
+
"""Get configurations for all runs in a project."""
|
| 589 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 590 |
+
if not db_path.exists():
|
| 591 |
+
return {}
|
| 592 |
+
|
| 593 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 594 |
+
cursor = conn.cursor()
|
| 595 |
+
try:
|
| 596 |
+
cursor.execute(
|
| 597 |
+
"""
|
| 598 |
+
SELECT run_name, config FROM configs
|
| 599 |
+
"""
|
| 600 |
+
)
|
| 601 |
+
|
| 602 |
+
results = {}
|
| 603 |
+
for row in cursor.fetchall():
|
| 604 |
+
config = orjson.loads(row["config"])
|
| 605 |
+
results[row["run_name"]] = deserialize_values(config)
|
| 606 |
+
return results
|
| 607 |
+
except sqlite3.OperationalError as e:
|
| 608 |
+
if "no such table: configs" in str(e):
|
| 609 |
+
return {}
|
| 610 |
+
raise
|
| 611 |
+
|
| 612 |
+
@staticmethod
|
| 613 |
+
def get_metric_values(project: str, run: str, metric_name: str) -> list[dict]:
|
| 614 |
+
"""Get all values for a specific metric in a project/run."""
|
| 615 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 616 |
+
if not db_path.exists():
|
| 617 |
+
return []
|
| 618 |
+
|
| 619 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 620 |
+
cursor = conn.cursor()
|
| 621 |
+
cursor.execute(
|
| 622 |
+
"""
|
| 623 |
+
SELECT timestamp, step, metrics
|
| 624 |
+
FROM metrics
|
| 625 |
+
WHERE run_name = ?
|
| 626 |
+
ORDER BY timestamp
|
| 627 |
+
""",
|
| 628 |
+
(run,),
|
| 629 |
+
)
|
| 630 |
+
|
| 631 |
+
rows = cursor.fetchall()
|
| 632 |
+
results = []
|
| 633 |
+
for row in rows:
|
| 634 |
+
metrics = orjson.loads(row["metrics"])
|
| 635 |
+
metrics = deserialize_values(metrics)
|
| 636 |
+
if metric_name in metrics:
|
| 637 |
+
results.append(
|
| 638 |
+
{
|
| 639 |
+
"timestamp": row["timestamp"],
|
| 640 |
+
"step": row["step"],
|
| 641 |
+
"value": metrics[metric_name],
|
| 642 |
+
}
|
| 643 |
+
)
|
| 644 |
+
return results
|
| 645 |
+
|
| 646 |
+
@staticmethod
|
| 647 |
+
def get_all_metrics_for_run(project: str, run: str) -> list[str]:
|
| 648 |
+
"""Get all metric names for a specific project/run."""
|
| 649 |
+
db_path = SQLiteStorage.get_project_db_path(project)
|
| 650 |
+
if not db_path.exists():
|
| 651 |
+
return []
|
| 652 |
+
|
| 653 |
+
with SQLiteStorage._get_connection(db_path) as conn:
|
| 654 |
+
cursor = conn.cursor()
|
| 655 |
+
cursor.execute(
|
| 656 |
+
"""
|
| 657 |
+
SELECT metrics
|
| 658 |
+
FROM metrics
|
| 659 |
+
WHERE run_name = ?
|
| 660 |
+
ORDER BY timestamp
|
| 661 |
+
""",
|
| 662 |
+
(run,),
|
| 663 |
+
)
|
| 664 |
+
|
| 665 |
+
rows = cursor.fetchall()
|
| 666 |
+
all_metrics = set()
|
| 667 |
+
for row in rows:
|
| 668 |
+
metrics = orjson.loads(row["metrics"])
|
| 669 |
+
metrics = deserialize_values(metrics)
|
| 670 |
+
for key in metrics.keys():
|
| 671 |
+
if key not in ["timestamp", "step"]:
|
| 672 |
+
all_metrics.add(key)
|
| 673 |
+
return sorted(list(all_metrics))
|
| 674 |
+
|
| 675 |
+
def finish(self):
|
| 676 |
+
"""Cleanup when run is finished."""
|
| 677 |
+
pass
|
table.py
ADDED
|
@@ -0,0 +1,53 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from typing import Any, Literal
|
| 2 |
+
|
| 3 |
+
from pandas import DataFrame
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class Table:
|
| 7 |
+
"""
|
| 8 |
+
Initializes a Table object.
|
| 9 |
+
|
| 10 |
+
Args:
|
| 11 |
+
columns (`list[str]`, *optional*):
|
| 12 |
+
Names of the columns in the table. Optional if `data` is provided. Not
|
| 13 |
+
expected if `dataframe` is provided. Currently ignored.
|
| 14 |
+
data (`list[list[Any]]`, *optional*):
|
| 15 |
+
2D row-oriented array of values.
|
| 16 |
+
dataframe (`pandas.`DataFrame``, *optional*):
|
| 17 |
+
DataFrame object used to create the table. When set, `data` and `columns`
|
| 18 |
+
arguments are ignored.
|
| 19 |
+
rows (`list[list[any]]`, *optional*):
|
| 20 |
+
Currently ignored.
|
| 21 |
+
optional (`bool` or `list[bool]`, *optional*, defaults to `True`):
|
| 22 |
+
Currently ignored.
|
| 23 |
+
allow_mixed_types (`bool`, *optional*, defaults to `False`):
|
| 24 |
+
Currently ignored.
|
| 25 |
+
log_mode: (`Literal["IMMUTABLE", "MUTABLE", "INCREMENTAL"]` or `None`, *optional*, defaults to `"IMMUTABLE"`):
|
| 26 |
+
Currently ignored.
|
| 27 |
+
"""
|
| 28 |
+
|
| 29 |
+
TYPE = "trackio.table"
|
| 30 |
+
|
| 31 |
+
def __init__(
|
| 32 |
+
self,
|
| 33 |
+
columns: list[str] | None = None,
|
| 34 |
+
data: list[list[Any]] | None = None,
|
| 35 |
+
dataframe: DataFrame | None = None,
|
| 36 |
+
rows: list[list[Any]] | None = None,
|
| 37 |
+
optional: bool | list[bool] = True,
|
| 38 |
+
allow_mixed_types: bool = False,
|
| 39 |
+
log_mode: Literal["IMMUTABLE", "MUTABLE", "INCREMENTAL"] | None = "IMMUTABLE",
|
| 40 |
+
):
|
| 41 |
+
# TODO: implement support for columns, dtype, optional, allow_mixed_types, and log_mode.
|
| 42 |
+
# for now (like `rows`) they are included for API compat but don't do anything.
|
| 43 |
+
|
| 44 |
+
if dataframe is None:
|
| 45 |
+
self.data = data
|
| 46 |
+
else:
|
| 47 |
+
self.data = dataframe.to_dict(orient="records")
|
| 48 |
+
|
| 49 |
+
def _to_dict(self):
|
| 50 |
+
return {
|
| 51 |
+
"_type": self.TYPE,
|
| 52 |
+
"_value": self.data,
|
| 53 |
+
}
|
typehints.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from typing import Any, TypedDict
|
| 2 |
+
|
| 3 |
+
from gradio import FileData
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class LogEntry(TypedDict):
|
| 7 |
+
project: str
|
| 8 |
+
run: str
|
| 9 |
+
metrics: dict[str, Any]
|
| 10 |
+
step: int | None
|
| 11 |
+
config: dict[str, Any] | None
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
class UploadEntry(TypedDict):
|
| 15 |
+
project: str
|
| 16 |
+
run: str
|
| 17 |
+
step: int | None
|
| 18 |
+
uploaded_file: FileData
|
ui/__init__.py
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
try:
|
| 2 |
+
from trackio.ui.main import demo
|
| 3 |
+
from trackio.ui.run_detail import run_detail_page
|
| 4 |
+
from trackio.ui.runs import run_page
|
| 5 |
+
except ImportError:
|
| 6 |
+
from ui.main import demo
|
| 7 |
+
from ui.run_detail import run_detail_page
|
| 8 |
+
from ui.runs import run_page
|
| 9 |
+
|
| 10 |
+
__all__ = ["demo", "run_page", "run_detail_page"]
|
ui/__pycache__/__init__.cpython-312.pyc
ADDED
|
Binary file (507 Bytes). View file
|
|
|
ui/__pycache__/fns.cpython-312.pyc
ADDED
|
Binary file (10.8 kB). View file
|
|
|
ui/__pycache__/main.cpython-312.pyc
ADDED
|
Binary file (46.9 kB). View file
|
|
|
ui/__pycache__/run_detail.cpython-312.pyc
ADDED
|
Binary file (4.07 kB). View file
|
|
|
ui/__pycache__/runs.cpython-312.pyc
ADDED
|
Binary file (12.1 kB). View file
|
|
|
ui/fns.py
ADDED
|
@@ -0,0 +1,241 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Shared functions for the Trackio UI."""
|
| 2 |
+
|
| 3 |
+
import os
|
| 4 |
+
|
| 5 |
+
import gradio as gr
|
| 6 |
+
import huggingface_hub as hf
|
| 7 |
+
|
| 8 |
+
try:
|
| 9 |
+
import trackio.utils as utils
|
| 10 |
+
from trackio.sqlite_storage import SQLiteStorage
|
| 11 |
+
from trackio.ui.helpers.run_selection import RunSelection
|
| 12 |
+
except ImportError:
|
| 13 |
+
import utils
|
| 14 |
+
from sqlite_storage import SQLiteStorage
|
| 15 |
+
from ui.helpers.run_selection import RunSelection
|
| 16 |
+
|
| 17 |
+
CONFIG_COLUMN_MAPPINGS = {
|
| 18 |
+
"_Username": "Username",
|
| 19 |
+
"_Created": "Created",
|
| 20 |
+
"_Group": "Group",
|
| 21 |
+
}
|
| 22 |
+
CONFIG_COLUMN_MAPPINGS_REVERSE = {v: k for k, v in CONFIG_COLUMN_MAPPINGS.items()}
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
HfApi = hf.HfApi()
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def get_project_info() -> str | None:
|
| 29 |
+
dataset_id = os.environ.get("TRACKIO_DATASET_ID")
|
| 30 |
+
space_id = utils.get_space()
|
| 31 |
+
if utils.persistent_storage_enabled():
|
| 32 |
+
return "✨ Persistent Storage is enabled, logs are stored directly in this Space."
|
| 33 |
+
if dataset_id:
|
| 34 |
+
sync_status = utils.get_sync_status(SQLiteStorage.get_scheduler())
|
| 35 |
+
upgrade_message = f"New changes are synced every 5 min <span class='info-container'><input type='checkbox' class='info-checkbox' id='upgrade-info'><label for='upgrade-info' class='info-icon'>ⓘ</label><span class='info-expandable'> To avoid losing data between syncs, <a href='https://huggingface.co/spaces/{space_id}/settings' class='accent-link'>click here</a> to open this Space's settings and add Persistent Storage. Make sure data is synced prior to enabling.</span></span>"
|
| 36 |
+
if sync_status is not None:
|
| 37 |
+
info = f"↻ Backed up {sync_status} min ago to <a href='https://huggingface.co/datasets/{dataset_id}' target='_blank' class='accent-link'>{dataset_id}</a> | {upgrade_message}"
|
| 38 |
+
else:
|
| 39 |
+
info = f"↻ Not backed up yet to <a href='https://huggingface.co/datasets/{dataset_id}' target='_blank' class='accent-link'>{dataset_id}</a> | {upgrade_message}"
|
| 40 |
+
return info
|
| 41 |
+
return None
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def get_projects(request: gr.Request):
|
| 45 |
+
projects = SQLiteStorage.get_projects()
|
| 46 |
+
if project := request.query_params.get("project"):
|
| 47 |
+
interactive = False
|
| 48 |
+
else:
|
| 49 |
+
interactive = True
|
| 50 |
+
if selected_project := request.query_params.get("selected_project"):
|
| 51 |
+
project = selected_project
|
| 52 |
+
else:
|
| 53 |
+
project = projects[0] if projects else None
|
| 54 |
+
|
| 55 |
+
return gr.Dropdown(
|
| 56 |
+
label="Project",
|
| 57 |
+
choices=projects,
|
| 58 |
+
value=project,
|
| 59 |
+
allow_custom_value=True,
|
| 60 |
+
interactive=interactive,
|
| 61 |
+
info=get_project_info(),
|
| 62 |
+
)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def update_navbar_value(project_dd, request: gr.Request):
|
| 66 |
+
write_token = None
|
| 67 |
+
if hasattr(request, "query_params") and request.query_params:
|
| 68 |
+
write_token = request.query_params.get("write_token")
|
| 69 |
+
|
| 70 |
+
metrics_url = f"?selected_project={project_dd}"
|
| 71 |
+
runs_url = f"runs?selected_project={project_dd}"
|
| 72 |
+
|
| 73 |
+
if write_token:
|
| 74 |
+
metrics_url += f"&write_token={write_token}"
|
| 75 |
+
runs_url += f"&write_token={write_token}"
|
| 76 |
+
|
| 77 |
+
return gr.Navbar(
|
| 78 |
+
value=[
|
| 79 |
+
("Metrics", metrics_url),
|
| 80 |
+
("Runs", runs_url),
|
| 81 |
+
]
|
| 82 |
+
)
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def check_hf_token_has_write_access(hf_token: str | None) -> None:
|
| 86 |
+
"""
|
| 87 |
+
Checks to see if the provided hf_token is valid and has write access to the Space
|
| 88 |
+
that Trackio is running in. If the hf_token is valid or if Trackio is not running
|
| 89 |
+
on a Space, this function does nothing. Otherwise, it raises a PermissionError.
|
| 90 |
+
"""
|
| 91 |
+
if os.getenv("SYSTEM") == "spaces": # if we are running in Spaces
|
| 92 |
+
# check auth token passed in
|
| 93 |
+
if hf_token is None:
|
| 94 |
+
raise PermissionError(
|
| 95 |
+
"Expected a HF_TOKEN to be provided when logging to a Space"
|
| 96 |
+
)
|
| 97 |
+
who = HfApi.whoami(hf_token)
|
| 98 |
+
owner_name = os.getenv("SPACE_AUTHOR_NAME")
|
| 99 |
+
repo_name = os.getenv("SPACE_REPO_NAME")
|
| 100 |
+
# make sure the token user is either the author of the space,
|
| 101 |
+
# or is a member of an org that is the author.
|
| 102 |
+
orgs = [o["name"] for o in who["orgs"]]
|
| 103 |
+
if owner_name != who["name"] and owner_name not in orgs:
|
| 104 |
+
raise PermissionError(
|
| 105 |
+
"Expected the provided hf_token to be the user owner of the space, or be a member of the org owner of the space"
|
| 106 |
+
)
|
| 107 |
+
# reject fine-grained tokens without specific repo access
|
| 108 |
+
access_token = who["auth"]["accessToken"]
|
| 109 |
+
if access_token["role"] == "fineGrained":
|
| 110 |
+
matched = False
|
| 111 |
+
for item in access_token["fineGrained"]["scoped"]:
|
| 112 |
+
if (
|
| 113 |
+
item["entity"]["type"] == "space"
|
| 114 |
+
and item["entity"]["name"] == f"{owner_name}/{repo_name}"
|
| 115 |
+
and "repo.write" in item["permissions"]
|
| 116 |
+
):
|
| 117 |
+
matched = True
|
| 118 |
+
break
|
| 119 |
+
if (
|
| 120 |
+
(
|
| 121 |
+
item["entity"]["type"] == "user"
|
| 122 |
+
or item["entity"]["type"] == "org"
|
| 123 |
+
)
|
| 124 |
+
and item["entity"]["name"] == owner_name
|
| 125 |
+
and "repo.write" in item["permissions"]
|
| 126 |
+
):
|
| 127 |
+
matched = True
|
| 128 |
+
break
|
| 129 |
+
if not matched:
|
| 130 |
+
raise PermissionError(
|
| 131 |
+
"Expected the provided hf_token with fine grained permissions to provide write access to the space"
|
| 132 |
+
)
|
| 133 |
+
# reject read-only tokens
|
| 134 |
+
elif access_token["role"] != "write":
|
| 135 |
+
raise PermissionError(
|
| 136 |
+
"Expected the provided hf_token to provide write permissions"
|
| 137 |
+
)
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
def check_oauth_token_has_write_access(oauth_token: str | None) -> None:
|
| 141 |
+
"""
|
| 142 |
+
Checks to see if the oauth token provided via Gradio's OAuth is valid and has write access
|
| 143 |
+
to the Space that Trackio is running in. If the oauth token is valid or if Trackio is not running
|
| 144 |
+
on a Space, this function does nothing. Otherwise, it raises a PermissionError.
|
| 145 |
+
"""
|
| 146 |
+
if not os.getenv("SYSTEM") == "spaces":
|
| 147 |
+
return
|
| 148 |
+
if oauth_token is None:
|
| 149 |
+
raise PermissionError(
|
| 150 |
+
"Expected an oauth to be provided when logging to a Space"
|
| 151 |
+
)
|
| 152 |
+
who = HfApi.whoami(oauth_token)
|
| 153 |
+
user_name = who["name"]
|
| 154 |
+
owner_name = os.getenv("SPACE_AUTHOR_NAME")
|
| 155 |
+
if user_name == owner_name:
|
| 156 |
+
return
|
| 157 |
+
# check if user is a member of an org that owns the space with write permissions
|
| 158 |
+
for org in who["orgs"]:
|
| 159 |
+
if org["name"] == owner_name and org["roleInOrg"] == "write":
|
| 160 |
+
return
|
| 161 |
+
raise PermissionError(
|
| 162 |
+
"Expected the oauth token to be the user owner of the space, or be a member of the org owner of the space"
|
| 163 |
+
)
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def get_group_by_fields(project: str):
|
| 167 |
+
configs = SQLiteStorage.get_all_run_configs(project) if project else {}
|
| 168 |
+
keys = set()
|
| 169 |
+
for config in configs.values():
|
| 170 |
+
keys.update(config.keys())
|
| 171 |
+
keys.discard("_Created")
|
| 172 |
+
keys = [CONFIG_COLUMN_MAPPINGS.get(key, key) for key in keys]
|
| 173 |
+
choices = [None] + sorted(keys)
|
| 174 |
+
return gr.Dropdown(
|
| 175 |
+
choices=choices,
|
| 176 |
+
value=None,
|
| 177 |
+
interactive=True,
|
| 178 |
+
)
|
| 179 |
+
|
| 180 |
+
|
| 181 |
+
def group_runs_by_config(
|
| 182 |
+
project: str, config_key: str, filter_text: str | None = None
|
| 183 |
+
) -> dict[str, list[str]]:
|
| 184 |
+
if not project or not config_key:
|
| 185 |
+
return {}
|
| 186 |
+
display_key = config_key
|
| 187 |
+
config_key = CONFIG_COLUMN_MAPPINGS_REVERSE.get(config_key, config_key)
|
| 188 |
+
configs = SQLiteStorage.get_all_run_configs(project)
|
| 189 |
+
groups: dict[str, list[str]] = {}
|
| 190 |
+
for run_name, config in configs.items():
|
| 191 |
+
if filter_text and filter_text not in run_name:
|
| 192 |
+
continue
|
| 193 |
+
group_name = config.get(config_key, "None")
|
| 194 |
+
label = f"{display_key}: {group_name}"
|
| 195 |
+
groups.setdefault(label, []).append(run_name)
|
| 196 |
+
for label in groups:
|
| 197 |
+
groups[label].sort()
|
| 198 |
+
sorted_groups = dict(sorted(groups.items(), key=lambda kv: kv[0].lower()))
|
| 199 |
+
return sorted_groups
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
def run_checkbox_update(selection: RunSelection, **kwargs) -> gr.CheckboxGroup:
|
| 203 |
+
return gr.CheckboxGroup(
|
| 204 |
+
choices=selection.choices,
|
| 205 |
+
value=selection.selected,
|
| 206 |
+
**kwargs,
|
| 207 |
+
)
|
| 208 |
+
|
| 209 |
+
|
| 210 |
+
def handle_run_checkbox_change(
|
| 211 |
+
selected_runs: list[str] | None, selection: RunSelection
|
| 212 |
+
) -> RunSelection:
|
| 213 |
+
selection.select(selected_runs or [])
|
| 214 |
+
return selection
|
| 215 |
+
|
| 216 |
+
|
| 217 |
+
def handle_group_checkbox_change(
|
| 218 |
+
group_selected: list[str] | None,
|
| 219 |
+
selection: RunSelection,
|
| 220 |
+
group_runs: list[str] | None,
|
| 221 |
+
):
|
| 222 |
+
subset, _ = selection.replace_group(group_runs or [], group_selected or [])
|
| 223 |
+
return (
|
| 224 |
+
selection,
|
| 225 |
+
gr.CheckboxGroup(value=subset),
|
| 226 |
+
run_checkbox_update(selection),
|
| 227 |
+
)
|
| 228 |
+
|
| 229 |
+
|
| 230 |
+
def handle_group_toggle(
|
| 231 |
+
select_all: bool,
|
| 232 |
+
selection: RunSelection,
|
| 233 |
+
group_runs: list[str] | None,
|
| 234 |
+
):
|
| 235 |
+
target = list(group_runs or []) if select_all else []
|
| 236 |
+
subset, _ = selection.replace_group(group_runs or [], target)
|
| 237 |
+
return (
|
| 238 |
+
selection,
|
| 239 |
+
gr.CheckboxGroup(value=subset),
|
| 240 |
+
run_checkbox_update(selection),
|
| 241 |
+
)
|
ui/helpers/__pycache__/run_selection.cpython-312.pyc
ADDED
|
Binary file (2.78 kB). View file
|
|
|
ui/helpers/run_selection.py
ADDED
|
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from dataclasses import dataclass, field
|
| 2 |
+
|
| 3 |
+
try:
|
| 4 |
+
import trackio.utils as utils
|
| 5 |
+
except ImportError:
|
| 6 |
+
import utils
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
@dataclass
|
| 10 |
+
class RunSelection:
|
| 11 |
+
choices: list[str] = field(default_factory=list)
|
| 12 |
+
selected: list[str] = field(default_factory=list)
|
| 13 |
+
locked: bool = False
|
| 14 |
+
|
| 15 |
+
def update_choices(
|
| 16 |
+
self, runs: list[str], preferred: list[str] | None = None
|
| 17 |
+
) -> bool:
|
| 18 |
+
if self.choices == runs:
|
| 19 |
+
return False
|
| 20 |
+
new_choices = set(runs) - set(self.choices)
|
| 21 |
+
self.choices = list(runs)
|
| 22 |
+
if self.locked:
|
| 23 |
+
base = set(self.selected) | new_choices
|
| 24 |
+
elif preferred:
|
| 25 |
+
base = set(preferred)
|
| 26 |
+
else:
|
| 27 |
+
base = set(runs)
|
| 28 |
+
self.selected = [run for run in self.choices if run in base]
|
| 29 |
+
return True
|
| 30 |
+
|
| 31 |
+
def select(self, runs: list[str]) -> list[str]:
|
| 32 |
+
choice_set = set(self.choices)
|
| 33 |
+
self.selected = [run for run in runs if run in choice_set]
|
| 34 |
+
self.locked = True
|
| 35 |
+
return self.selected
|
| 36 |
+
|
| 37 |
+
def replace_group(
|
| 38 |
+
self, group_runs: list[str], new_subset: list[str] | None
|
| 39 |
+
) -> tuple[list[str], list[str]]:
|
| 40 |
+
new_subset = utils.ordered_subset(group_runs, new_subset)
|
| 41 |
+
selection_set = set(self.selected)
|
| 42 |
+
selection_set.difference_update(group_runs)
|
| 43 |
+
selection_set.update(new_subset)
|
| 44 |
+
self.selected = [run for run in self.choices if run in selection_set]
|
| 45 |
+
self.locked = True
|
| 46 |
+
return new_subset, self.selected
|