abidlabs HF Staff commited on
Commit
51d7dfa
·
verified ·
1 Parent(s): a3d2a61

Upload folder using huggingface_hub

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +1 -0
  2. trackio/CHANGELOG.md +152 -0
  3. trackio/__init__.py +601 -0
  4. trackio/__pycache__/__init__.cpython-310.pyc +0 -0
  5. trackio/__pycache__/api.cpython-310.pyc +0 -0
  6. trackio/__pycache__/commit_scheduler.cpython-310.pyc +0 -0
  7. trackio/__pycache__/context_vars.cpython-310.pyc +0 -0
  8. trackio/__pycache__/deploy.cpython-310.pyc +0 -0
  9. trackio/__pycache__/dummy_commit_scheduler.cpython-310.pyc +0 -0
  10. trackio/__pycache__/gpu.cpython-310.pyc +0 -0
  11. trackio/__pycache__/histogram.cpython-310.pyc +0 -0
  12. trackio/__pycache__/imports.cpython-310.pyc +0 -0
  13. trackio/__pycache__/run.cpython-310.pyc +0 -0
  14. trackio/__pycache__/sqlite_storage.cpython-310.pyc +0 -0
  15. trackio/__pycache__/table.cpython-310.pyc +0 -0
  16. trackio/__pycache__/typehints.cpython-310.pyc +0 -0
  17. trackio/__pycache__/utils.cpython-310.pyc +0 -0
  18. trackio/api.py +66 -0
  19. trackio/assets/badge.png +0 -0
  20. trackio/assets/trackio_logo_dark.png +0 -0
  21. trackio/assets/trackio_logo_light.png +0 -0
  22. trackio/assets/trackio_logo_old.png +3 -0
  23. trackio/assets/trackio_logo_type_dark.png +0 -0
  24. trackio/assets/trackio_logo_type_dark_transparent.png +0 -0
  25. trackio/assets/trackio_logo_type_light.png +0 -0
  26. trackio/assets/trackio_logo_type_light_transparent.png +0 -0
  27. trackio/cli.py +514 -0
  28. trackio/cli_helpers.py +118 -0
  29. trackio/commit_scheduler.py +310 -0
  30. trackio/context_vars.py +18 -0
  31. trackio/deploy.py +433 -0
  32. trackio/dummy_commit_scheduler.py +12 -0
  33. trackio/gpu.py +357 -0
  34. trackio/histogram.py +71 -0
  35. trackio/imports.py +304 -0
  36. trackio/media/__init__.py +27 -0
  37. trackio/media/__pycache__/__init__.cpython-310.pyc +0 -0
  38. trackio/media/__pycache__/audio.cpython-310.pyc +0 -0
  39. trackio/media/__pycache__/image.cpython-310.pyc +0 -0
  40. trackio/media/__pycache__/media.cpython-310.pyc +0 -0
  41. trackio/media/__pycache__/utils.cpython-310.pyc +0 -0
  42. trackio/media/__pycache__/video.cpython-310.pyc +0 -0
  43. trackio/media/audio.py +167 -0
  44. trackio/media/image.py +84 -0
  45. trackio/media/media.py +79 -0
  46. trackio/media/utils.py +60 -0
  47. trackio/media/video.py +246 -0
  48. trackio/package.json +6 -0
  49. trackio/py.typed +0 -0
  50. trackio/run.py +586 -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
+ trackio/assets/trackio_logo_old.png filter=lfs diff=lfs merge=lfs -text
trackio/CHANGELOG.md ADDED
@@ -0,0 +1,152 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # trackio
2
+
3
+ ## 0.16.1
4
+
5
+ ### Features
6
+
7
+ - [#431](https://github.com/gradio-app/trackio/pull/431) [`c7ce55b`](https://github.com/gradio-app/trackio/commit/c7ce55b14dd5eb0c2165fb15df17dd60721c9325) - Lazy load the UI when trackio is imported. Thanks @abidlabs!
8
+
9
+ ## 0.16.0
10
+
11
+ ### Features
12
+
13
+ - [#426](https://github.com/gradio-app/trackio/pull/426) [`ead4dc8`](https://github.com/gradio-app/trackio/commit/ead4dc8e74ee2d8e47d61bca0a7668456acf49be) - Fix redundant double rendering of group checkboxes. Thanks @abidlabs!
14
+ - [#413](https://github.com/gradio-app/trackio/pull/413) [`39c4750`](https://github.com/gradio-app/trackio/commit/39c4750951d554ba6eb4d58847c6bb444b2891a8) - Check `dist-packages` when checking for source installation. Thanks @sergiopaniego!
15
+ - [#423](https://github.com/gradio-app/trackio/pull/423) [`2e52ab3`](https://github.com/gradio-app/trackio/commit/2e52ab303e3041718a6a56fbf84d0848aca9ad67) - Fix legend outline visibility issue. Thanks @Raghunath-Balaji!
16
+ - [#407](https://github.com/gradio-app/trackio/pull/407) [`c8a384d`](https://github.com/gradio-app/trackio/commit/c8a384ddfe5a295cecf862a26178d40e48acb424) - Fix pytests that were failling locally on MacOS. Thanks @abidlabs!
17
+ - [#405](https://github.com/gradio-app/trackio/pull/405) [`35aae4e`](https://github.com/gradio-app/trackio/commit/35aae4e3aa3e2b2888887528478b9dc6a9808bda) - Add conditional padding for HF Space dashboard when not in iframe. Thanks @znation!
18
+
19
+ ## 0.15.0
20
+
21
+ ### Features
22
+
23
+ - [#397](https://github.com/gradio-app/trackio/pull/397) [`6b38ad0`](https://github.com/gradio-app/trackio/commit/6b38ad02e5d73a0df49c4eede7e91331282ece04) - Adds `--host` cli option support. Thanks @abidlabs!
24
+ - [#396](https://github.com/gradio-app/trackio/pull/396) [`4a4d1ab`](https://github.com/gradio-app/trackio/commit/4a4d1ab85e63d923132a3fa7afa5d90e16431bec) - Fix run selection issue. Thanks @abidlabs!
25
+ - [#394](https://github.com/gradio-app/trackio/pull/394) [`c47a3a3`](https://github.com/gradio-app/trackio/commit/c47a3a31f8c4b83bce1aa7fc22eeba3d9021ad3d) - Add wandb-compatible API for trackio. Thanks @abidlabs!
26
+ - [#378](https://github.com/gradio-app/trackio/pull/378) [`b02046a`](https://github.com/gradio-app/trackio/commit/b02046a5b0dad7c9854e099a87f884afba4aecb2) - Add JSON export button for line plots and upgrade gradio dependency. Thanks @JamshedAli18!
27
+
28
+ ## 0.14.2
29
+
30
+ ### Features
31
+
32
+ - [#386](https://github.com/gradio-app/trackio/pull/386) [`f9452cd`](https://github.com/gradio-app/trackio/commit/f9452cdb8f0819368f3610f7ac0ed08957305275) - Fixing some issues related to deployed Trackio Spaces. Thanks @abidlabs!
33
+
34
+ ## 0.14.1
35
+
36
+ ### Features
37
+
38
+ - [#382](https://github.com/gradio-app/trackio/pull/382) [`44fe9bb`](https://github.com/gradio-app/trackio/commit/44fe9bb264fb2aafb0ec302ff15227c045819a2c) - Fix app file path when Trackio is not installed from source. Thanks @abidlabs!
39
+ - [#380](https://github.com/gradio-app/trackio/pull/380) [`c3f4cff`](https://github.com/gradio-app/trackio/commit/c3f4cff74bc5676e812773d8571454894fcdc7cc) - Add CLI commands for querying projects, runs, and metrics. Thanks @abidlabs!
40
+
41
+ ## 0.14.0
42
+
43
+ ### Features
44
+
45
+ - [#377](https://github.com/gradio-app/trackio/pull/377) [`5c5015b`](https://github.com/gradio-app/trackio/commit/5c5015b68c85c5de51111dad983f735c27b9a05f) - fixed wrapping issue in Runs table. Thanks @gaganchapa!
46
+ - [#374](https://github.com/gradio-app/trackio/pull/374) [`388e26b`](https://github.com/gradio-app/trackio/commit/388e26b9e9f24cd7ad203affe9b709be885b3d24) - Save Optimized Parquet files. Thanks @lhoestq!
47
+ - [#371](https://github.com/gradio-app/trackio/pull/371) [`fbace9c`](https://github.com/gradio-app/trackio/commit/fbace9cd7732c166f34d268f54b05bb06846cc5d) - Add GPU metrics logging. Thanks @kashif!
48
+ - [#367](https://github.com/gradio-app/trackio/pull/367) [`862840c`](https://github.com/gradio-app/trackio/commit/862840c13e30fc960cbee5b9eac4d3c25beba9de) - Add option to only show latest run, and fix the double logo issue. Thanks @abidlabs!
49
+
50
+ ## 0.13.1
51
+
52
+ ### Features
53
+
54
+ - [#369](https://github.com/gradio-app/trackio/pull/369) [`767e9fe`](https://github.com/gradio-app/trackio/commit/767e9fe095d7c6ed102016caf927c1517fb8618c) - tiny pr removing unnecessary code. Thanks @abidlabs!
55
+
56
+ ## 0.13.0
57
+
58
+ ### Features
59
+
60
+ - [#358](https://github.com/gradio-app/trackio/pull/358) [`073715d`](https://github.com/gradio-app/trackio/commit/073715d1caf8282f68890117f09c3ac301205312) - Improvements to `trackio.sync()`. Thanks @abidlabs!
61
+
62
+ ## 0.12.0
63
+
64
+ ### Features
65
+
66
+ - [#357](https://github.com/gradio-app/trackio/pull/357) [`02ba815`](https://github.com/gradio-app/trackio/commit/02ba815358060f1966052de051a5bdb09702920e) - Redesign media and tables to show up on separate page. Thanks @abidlabs!
67
+ - [#359](https://github.com/gradio-app/trackio/pull/359) [`08fe9c9`](https://github.com/gradio-app/trackio/commit/08fe9c9ddd7fe99ee811555fdfb62df9ab88e939) - docs: Improve docstrings. Thanks @qgallouedec!
68
+
69
+ ## 0.11.0
70
+
71
+ ### Features
72
+
73
+ - [#355](https://github.com/gradio-app/trackio/pull/355) [`ea51f49`](https://github.com/gradio-app/trackio/commit/ea51f4954922f21be76ef828700420fe9a912c4b) - Color code run checkboxes and match with plot lines. Thanks @abidlabs!
74
+ - [#353](https://github.com/gradio-app/trackio/pull/353) [`8abe691`](https://github.com/gradio-app/trackio/commit/8abe6919aeefe21fc7a23af814883efbb037c21f) - Remove show_api from demo.launch. Thanks @sergiopaniego!
75
+ - [#351](https://github.com/gradio-app/trackio/pull/351) [`8a8957e`](https://github.com/gradio-app/trackio/commit/8a8957e530dd7908d1fef7f2df030303f808101f) - Add `trackio.save()`. Thanks @abidlabs!
76
+
77
+ ## 0.10.0
78
+
79
+ ### Features
80
+
81
+ - [#305](https://github.com/gradio-app/trackio/pull/305) [`e64883a`](https://github.com/gradio-app/trackio/commit/e64883a51f7b8b93f7d48b8afe55acdb62238b71) - bump to gradio 6.0, make `trackio` compatible, and fix related issues. Thanks @abidlabs!
82
+
83
+ ## 0.9.1
84
+
85
+ ### Features
86
+
87
+ - [#344](https://github.com/gradio-app/trackio/pull/344) [`7e01024`](https://github.com/gradio-app/trackio/commit/7e010241d9a34794e0ce0dc19c1a6f0cf94ba856) - Avoid redundant calls to /whoami-v2. Thanks @Wauplin!
88
+
89
+ ## 0.9.0
90
+
91
+ ### Features
92
+
93
+ - [#343](https://github.com/gradio-app/trackio/pull/343) [`51bea30`](https://github.com/gradio-app/trackio/commit/51bea30f2877adff8e6497466d3a799400a0a049) - Sync offline projects to Hugging Face spaces. Thanks @candemircan!
94
+ - [#341](https://github.com/gradio-app/trackio/pull/341) [`4fd841f`](https://github.com/gradio-app/trackio/commit/4fd841fa190e15071b02f6fba7683ef4f393a654) - Adds a basic UI test to `trackio`. Thanks @abidlabs!
95
+ - [#339](https://github.com/gradio-app/trackio/pull/339) [`011d91b`](https://github.com/gradio-app/trackio/commit/011d91bb6ae266516fd250a349285670a8049d05) - Allow customzing the trackio color palette. Thanks @abidlabs!
96
+
97
+ ## 0.8.1
98
+
99
+ ### Features
100
+
101
+ - [#336](https://github.com/gradio-app/trackio/pull/336) [`5f9f51d`](https://github.com/gradio-app/trackio/commit/5f9f51dac8677f240d7c42c3e3b2660a22aee138) - Support a list of `Trackio.Image` in a `trackio.Table` cell. Thanks @abidlabs!
102
+
103
+ ## 0.8.0
104
+
105
+ ### Features
106
+
107
+ - [#331](https://github.com/gradio-app/trackio/pull/331) [`2c02d0f`](https://github.com/gradio-app/trackio/commit/2c02d0fd0a5824160528782402bb0dd4083396d5) - Truncate table string values that are greater than 250 characters (configuirable via env variable). Thanks @abidlabs!
108
+ - [#324](https://github.com/gradio-app/trackio/pull/324) [`50b2122`](https://github.com/gradio-app/trackio/commit/50b2122e7965ac82a72e6cb3b7d048bc10a2a6b1) - Add log y-axis functionality to UI. Thanks @abidlabs!
109
+ - [#326](https://github.com/gradio-app/trackio/pull/326) [`61dc1f4`](https://github.com/gradio-app/trackio/commit/61dc1f40af2f545f8e70395ddf0dbb8aee6b60d5) - Fix: improve table rendering for metrics in Trackio Dashboard. Thanks @vigneshwaran!
110
+ - [#328](https://github.com/gradio-app/trackio/pull/328) [`6857cbb`](https://github.com/gradio-app/trackio/commit/6857cbbe557a59a4642f210ec42566d108294e63) - Support trackio.Table with trackio.Image columns. Thanks @abidlabs!
111
+ - [#323](https://github.com/gradio-app/trackio/pull/323) [`6857cbb`](https://github.com/gradio-app/trackio/commit/6857cbbe557a59a4642f210ec42566d108294e63) - add Trackio client implementations in Go, Rust, and JS. Thanks @vaibhav-research!
112
+
113
+ ## 0.7.0
114
+
115
+ ### Features
116
+
117
+ - [#277](https://github.com/gradio-app/trackio/pull/277) [`db35601`](https://github.com/gradio-app/trackio/commit/db35601b9c023423c4654c9909b8ab73e58737de) - fix: make grouped runs view reflect live updates. Thanks @Saba9!
118
+ - [#320](https://github.com/gradio-app/trackio/pull/320) [`24ae739`](https://github.com/gradio-app/trackio/commit/24ae73969b09fb3126acd2f91647cdfbf8cf72a1) - Add additional query parms for xmin, xmax, and smoothing. Thanks @abidlabs!
119
+ - [#270](https://github.com/gradio-app/trackio/pull/270) [`cd1dfc3`](https://github.com/gradio-app/trackio/commit/cd1dfc3dc641b4499ac6d4a1b066fa8e2b52c57b) - feature: add support for logging audio. Thanks @Saba9!
120
+
121
+ ## 0.6.0
122
+
123
+ ### Features
124
+
125
+ - [#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!
126
+ - [#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!
127
+ - [#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!
128
+ - [#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!
129
+ - [#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!
130
+
131
+ ## 0.5.3
132
+
133
+ ### Features
134
+
135
+ - [#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!
136
+
137
+ ## 0.5.2
138
+
139
+ ### Features
140
+
141
+ - [#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!
142
+ - [#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!
143
+
144
+ ### Fixes
145
+
146
+ - [#291](https://github.com/gradio-app/trackio/pull/291) [`3b5adc3`](https://github.com/gradio-app/trackio/commit/3b5adc3d1f452dbab7a714d235f4974782f93730) - Fix the wheel build. Thanks @pngwn!
147
+
148
+ ## 0.5.1
149
+
150
+ ### Fixes
151
+
152
+ - [#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!
trackio/__init__.py ADDED
@@ -0,0 +1,601 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import atexit
2
+ import glob
3
+ import json
4
+ import logging
5
+ import os
6
+ import shutil
7
+ import warnings
8
+ import webbrowser
9
+ from pathlib import Path
10
+ from typing import Any
11
+
12
+ import huggingface_hub
13
+ from gradio.themes import ThemeClass
14
+ from gradio.utils import TupleNoPrint
15
+ from gradio_client import Client, handle_file
16
+ from huggingface_hub import SpaceStorage
17
+ from huggingface_hub.errors import LocalTokenNotFoundError
18
+
19
+ from trackio import context_vars, deploy, utils
20
+ from trackio.api import Api
21
+ from trackio.deploy import sync
22
+ from trackio.gpu import gpu_available, log_gpu
23
+ from trackio.histogram import Histogram
24
+ from trackio.imports import import_csv, import_tf_events
25
+ from trackio.media import (
26
+ TrackioAudio,
27
+ TrackioImage,
28
+ TrackioVideo,
29
+ get_project_media_path,
30
+ )
31
+ from trackio.run import Run
32
+ from trackio.sqlite_storage import SQLiteStorage
33
+ from trackio.table import Table
34
+ from trackio.typehints import UploadEntry
35
+ from trackio.utils import TRACKIO_DIR, TRACKIO_LOGO_DIR
36
+
37
+ logging.getLogger("httpx").setLevel(logging.WARNING)
38
+
39
+ warnings.filterwarnings(
40
+ "ignore",
41
+ message="Empty session being created. Install gradio\\[oauth\\]",
42
+ category=UserWarning,
43
+ module="gradio.helpers",
44
+ )
45
+
46
+ __version__ = json.loads(Path(__file__).parent.joinpath("package.json").read_text())[
47
+ "version"
48
+ ]
49
+
50
+ __all__ = [
51
+ "init",
52
+ "log",
53
+ "log_system",
54
+ "log_gpu",
55
+ "finish",
56
+ "show",
57
+ "sync",
58
+ "delete_project",
59
+ "import_csv",
60
+ "import_tf_events",
61
+ "save",
62
+ "Image",
63
+ "Video",
64
+ "Audio",
65
+ "Table",
66
+ "Histogram",
67
+ "Api",
68
+ ]
69
+
70
+ Image = TrackioImage
71
+ Video = TrackioVideo
72
+ Audio = TrackioAudio
73
+
74
+
75
+ config = {}
76
+
77
+ _atexit_registered = False
78
+
79
+
80
+ def _cleanup_current_run():
81
+ run = context_vars.current_run.get()
82
+ if run is not None:
83
+ try:
84
+ run.finish()
85
+ except Exception:
86
+ pass
87
+
88
+
89
+ def _get_demo():
90
+ # Lazy import to avoid initializing Gradio Blocks (and FastAPI) at import time,
91
+ # which causes import lock errors for libraries that just `import trackio`.
92
+ from trackio.ui.main import CSS, HEAD, demo
93
+
94
+ return demo, CSS, HEAD
95
+
96
+
97
+ def init(
98
+ project: str,
99
+ name: str | None = None,
100
+ group: str | None = None,
101
+ space_id: str | None = None,
102
+ space_storage: SpaceStorage | None = None,
103
+ dataset_id: str | None = None,
104
+ config: dict | None = None,
105
+ resume: str = "never",
106
+ settings: Any = None,
107
+ private: bool | None = None,
108
+ embed: bool = True,
109
+ auto_log_gpu: bool | None = None,
110
+ gpu_log_interval: float = 10.0,
111
+ ) -> Run:
112
+ """
113
+ Creates a new Trackio project and returns a [`Run`] object.
114
+
115
+ Args:
116
+ project (`str`):
117
+ The name of the project (can be an existing project to continue tracking or
118
+ a new project to start tracking from scratch).
119
+ name (`str`, *optional*):
120
+ The name of the run (if not provided, a default name will be generated).
121
+ group (`str`, *optional*):
122
+ The name of the group which this run belongs to in order to help organize
123
+ related runs together. You can toggle the entire group's visibilitiy in the
124
+ dashboard.
125
+ space_id (`str`, *optional*):
126
+ If provided, the project will be logged to a Hugging Face Space instead of
127
+ a local directory. Should be a complete Space name like
128
+ `"username/reponame"` or `"orgname/reponame"`, or just `"reponame"` in which
129
+ case the Space will be created in the currently-logged-in Hugging Face
130
+ user's namespace. If the Space does not exist, it will be created. If the
131
+ Space already exists, the project will be logged to it.
132
+ space_storage ([`~huggingface_hub.SpaceStorage`], *optional*):
133
+ Choice of persistent storage tier.
134
+ dataset_id (`str`, *optional*):
135
+ If a `space_id` is provided, a persistent Hugging Face Dataset will be
136
+ created and the metrics will be synced to it every 5 minutes. Specify a
137
+ Dataset with name like `"username/datasetname"` or `"orgname/datasetname"`,
138
+ or `"datasetname"` (uses currently-logged-in Hugging Face user's namespace),
139
+ or `None` (uses the same name as the Space but with the `"_dataset"`
140
+ suffix). If the Dataset does not exist, it will be created. If the Dataset
141
+ already exists, the project will be appended to it.
142
+ config (`dict`, *optional*):
143
+ A dictionary of configuration options. Provided for compatibility with
144
+ `wandb.init()`.
145
+ resume (`str`, *optional*, defaults to `"never"`):
146
+ Controls how to handle resuming a run. Can be one of:
147
+
148
+ - `"must"`: Must resume the run with the given name, raises error if run
149
+ doesn't exist
150
+ - `"allow"`: Resume the run if it exists, otherwise create a new run
151
+ - `"never"`: Never resume a run, always create a new one
152
+ private (`bool`, *optional*):
153
+ Whether to make the Space private. If None (default), the repo will be
154
+ public unless the organization's default is private. This value is ignored
155
+ if the repo already exists.
156
+ settings (`Any`, *optional*):
157
+ Not used. Provided for compatibility with `wandb.init()`.
158
+ embed (`bool`, *optional*, defaults to `True`):
159
+ If running inside a jupyter/Colab notebook, whether the dashboard should
160
+ automatically be embedded in the cell when trackio.init() is called.
161
+ auto_log_gpu (`bool` or `None`, *optional*, defaults to `None`):
162
+ Controls automatic GPU metrics logging. If `None` (default), GPU logging
163
+ is automatically enabled when `nvidia-ml-py` is installed and an NVIDIA
164
+ GPU is detected. Set to `True` to force enable or `False` to disable.
165
+ gpu_log_interval (`float`, *optional*, defaults to `10.0`):
166
+ The interval in seconds between automatic GPU metric logs.
167
+ Only used when `auto_log_gpu=True`.
168
+
169
+ Returns:
170
+ `Run`: A [`Run`] object that can be used to log metrics and finish the run.
171
+ """
172
+ if settings is not None:
173
+ warnings.warn(
174
+ "* 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."
175
+ )
176
+
177
+ if space_id is None and dataset_id is not None:
178
+ raise ValueError("Must provide a `space_id` when `dataset_id` is provided.")
179
+ try:
180
+ space_id, dataset_id = utils.preprocess_space_and_dataset_ids(
181
+ space_id, dataset_id
182
+ )
183
+ except LocalTokenNotFoundError as e:
184
+ raise LocalTokenNotFoundError(
185
+ f"You must be logged in to Hugging Face locally when `space_id` is provided to deploy to a Space. {e}"
186
+ ) from e
187
+
188
+ url = context_vars.current_server.get()
189
+
190
+ if space_id is not None:
191
+ if url is None:
192
+ url = space_id
193
+ context_vars.current_server.set(url)
194
+ context_vars.current_space_id.set(space_id)
195
+
196
+ if (
197
+ context_vars.current_project.get() is None
198
+ or context_vars.current_project.get() != project
199
+ ):
200
+ print(f"* Trackio project initialized: {project}")
201
+
202
+ if dataset_id is not None:
203
+ os.environ["TRACKIO_DATASET_ID"] = dataset_id
204
+ print(
205
+ f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}"
206
+ )
207
+ if space_id is None:
208
+ print(f"* Trackio metrics logged to: {TRACKIO_DIR}")
209
+ utils.print_dashboard_instructions(project)
210
+ else:
211
+ deploy.create_space_if_not_exists(
212
+ space_id, space_storage, dataset_id, private
213
+ )
214
+ user_name, space_name = space_id.split("/")
215
+ space_url = deploy.SPACE_HOST_URL.format(
216
+ user_name=user_name, space_name=space_name
217
+ )
218
+ print(f"* View dashboard by going to: {space_url}")
219
+ if utils.is_in_notebook() and embed:
220
+ utils.embed_url_in_notebook(space_url)
221
+ context_vars.current_project.set(project)
222
+
223
+ if resume == "must":
224
+ if name is None:
225
+ raise ValueError("Must provide a run name when resume='must'")
226
+ if name not in SQLiteStorage.get_runs(project):
227
+ raise ValueError(f"Run '{name}' does not exist in project '{project}'")
228
+ resumed = True
229
+ elif resume == "allow":
230
+ resumed = name is not None and name in SQLiteStorage.get_runs(project)
231
+ elif resume == "never":
232
+ if name is not None and name in SQLiteStorage.get_runs(project):
233
+ warnings.warn(
234
+ f"* Warning: resume='never' but a run '{name}' already exists in "
235
+ f"project '{project}'. Generating a new name and instead. If you want "
236
+ "to resume this run, call init() with resume='must' or resume='allow'."
237
+ )
238
+ name = None
239
+ resumed = False
240
+ else:
241
+ raise ValueError("resume must be one of: 'must', 'allow', or 'never'")
242
+
243
+ if auto_log_gpu is None:
244
+ auto_log_gpu = gpu_available()
245
+ if auto_log_gpu:
246
+ print("* GPU detected, enabling automatic GPU metrics logging")
247
+
248
+ run = Run(
249
+ url=url,
250
+ project=project,
251
+ client=None,
252
+ name=name,
253
+ group=group,
254
+ config=config,
255
+ space_id=space_id,
256
+ auto_log_gpu=auto_log_gpu,
257
+ gpu_log_interval=gpu_log_interval,
258
+ )
259
+
260
+ if space_id is not None:
261
+ SQLiteStorage.set_project_metadata(project, "space_id", space_id)
262
+ if SQLiteStorage.has_pending_data(project):
263
+ run._has_local_buffer = True
264
+
265
+ global _atexit_registered
266
+ if not _atexit_registered:
267
+ atexit.register(_cleanup_current_run)
268
+ _atexit_registered = True
269
+
270
+ if resumed:
271
+ print(f"* Resumed existing run: {run.name}")
272
+ else:
273
+ print(f"* Created new run: {run.name}")
274
+
275
+ context_vars.current_run.set(run)
276
+ globals()["config"] = run.config
277
+ return run
278
+
279
+
280
+ def log(metrics: dict, step: int | None = None) -> None:
281
+ """
282
+ Logs metrics to the current run.
283
+
284
+ Args:
285
+ metrics (`dict`):
286
+ A dictionary of metrics to log.
287
+ step (`int`, *optional*):
288
+ The step number. If not provided, the step will be incremented
289
+ automatically.
290
+ """
291
+ run = context_vars.current_run.get()
292
+ if run is None:
293
+ raise RuntimeError("Call trackio.init() before trackio.log().")
294
+ run.log(
295
+ metrics=metrics,
296
+ step=step,
297
+ )
298
+
299
+
300
+ def log_system(metrics: dict) -> None:
301
+ """
302
+ Logs system metrics (GPU, etc.) to the current run using timestamps instead of steps.
303
+
304
+ Args:
305
+ metrics (`dict`):
306
+ A dictionary of system metrics to log.
307
+ """
308
+ run = context_vars.current_run.get()
309
+ if run is None:
310
+ raise RuntimeError("Call trackio.init() before trackio.log_system().")
311
+ run.log_system(metrics=metrics)
312
+
313
+
314
+ def finish():
315
+ """
316
+ Finishes the current run.
317
+ """
318
+ run = context_vars.current_run.get()
319
+ if run is None:
320
+ raise RuntimeError("Call trackio.init() before trackio.finish().")
321
+ run.finish()
322
+
323
+
324
+ def delete_project(project: str, force: bool = False) -> bool:
325
+ """
326
+ Deletes a project by removing its local SQLite database.
327
+
328
+ Args:
329
+ project (`str`):
330
+ The name of the project to delete.
331
+ force (`bool`, *optional*, defaults to `False`):
332
+ If `True`, deletes the project without prompting for confirmation.
333
+ If `False`, prompts the user to confirm before deleting.
334
+
335
+ Returns:
336
+ `bool`: `True` if the project was deleted, `False` otherwise.
337
+ """
338
+ db_path = SQLiteStorage.get_project_db_path(project)
339
+
340
+ if not db_path.exists():
341
+ print(f"* Project '{project}' does not exist.")
342
+ return False
343
+
344
+ if not force:
345
+ response = input(
346
+ f"Are you sure you want to delete project '{project}'? "
347
+ f"This will permanently delete all runs and metrics. (y/N): "
348
+ )
349
+ if response.lower() not in ["y", "yes"]:
350
+ print("* Deletion cancelled.")
351
+ return False
352
+
353
+ try:
354
+ db_path.unlink()
355
+
356
+ for suffix in ("-wal", "-shm"):
357
+ sidecar = Path(str(db_path) + suffix)
358
+ if sidecar.exists():
359
+ sidecar.unlink()
360
+
361
+ print(f"* Project '{project}' has been deleted.")
362
+ return True
363
+ except Exception as e:
364
+ print(f"* Error deleting project '{project}': {e}")
365
+ return False
366
+
367
+
368
+ def save(
369
+ glob_str: str | Path,
370
+ project: str | None = None,
371
+ ) -> str:
372
+ """
373
+ Saves files to a project (not linked to a specific run). If Trackio is running
374
+ locally, the file(s) will be copied to the project's files directory. If Trackio is
375
+ running in a Space, the file(s) will be uploaded to the Space's files directory.
376
+
377
+ Args:
378
+ glob_str (`str` or `Path`):
379
+ The file path or glob pattern to save. Can be a single file or a pattern
380
+ matching multiple files (e.g., `"*.py"`, `"models/**/*.pth"`).
381
+ project (`str`, *optional*):
382
+ The name of the project to save files to. If not provided, uses the current
383
+ project from `trackio.init()`. If no project is initialized, raises an
384
+ error.
385
+
386
+ Returns:
387
+ `str`: The path where the file(s) were saved (project's files directory).
388
+
389
+ Example:
390
+ ```python
391
+ import trackio
392
+
393
+ trackio.init(project="my-project")
394
+ trackio.save("config.yaml")
395
+ trackio.save("models/*.pth")
396
+ ```
397
+ """
398
+ if project is None:
399
+ project = context_vars.current_project.get()
400
+ if project is None:
401
+ raise RuntimeError(
402
+ "No project specified. Either call trackio.init() first or provide a "
403
+ "project parameter to trackio.save()."
404
+ )
405
+
406
+ glob_str = Path(glob_str)
407
+ base_path = Path.cwd().resolve()
408
+
409
+ matched_files = []
410
+ if glob_str.is_file():
411
+ matched_files = [glob_str.resolve()]
412
+ else:
413
+ pattern = str(glob_str)
414
+ if not glob_str.is_absolute():
415
+ pattern = str((Path.cwd() / glob_str).resolve())
416
+ matched_files = [
417
+ Path(f).resolve()
418
+ for f in glob.glob(pattern, recursive=True)
419
+ if Path(f).is_file()
420
+ ]
421
+
422
+ if not matched_files:
423
+ raise ValueError(f"No files found matching pattern: {glob_str}")
424
+
425
+ current_run = context_vars.current_run.get()
426
+ is_local = (
427
+ current_run._is_local
428
+ if current_run is not None
429
+ else (context_vars.current_space_id.get() is None)
430
+ )
431
+
432
+ if is_local:
433
+ for file_path in matched_files:
434
+ try:
435
+ relative_to_base = file_path.relative_to(base_path)
436
+ except ValueError:
437
+ relative_to_base = Path(file_path.name)
438
+
439
+ if current_run is not None:
440
+ current_run._queue_upload(
441
+ file_path,
442
+ step=None,
443
+ relative_path=str(relative_to_base.parent),
444
+ use_run_name=False,
445
+ )
446
+ else:
447
+ media_path = get_project_media_path(
448
+ project=project,
449
+ run=None,
450
+ step=None,
451
+ relative_path=str(relative_to_base),
452
+ )
453
+ shutil.copy(str(file_path), str(media_path))
454
+ else:
455
+ url = context_vars.current_server.get()
456
+
457
+ upload_entries = []
458
+ for file_path in matched_files:
459
+ try:
460
+ relative_to_base = file_path.relative_to(base_path)
461
+ except ValueError:
462
+ relative_to_base = Path(file_path.name)
463
+
464
+ if current_run is not None:
465
+ current_run._queue_upload(
466
+ file_path,
467
+ step=None,
468
+ relative_path=str(relative_to_base.parent),
469
+ use_run_name=False,
470
+ )
471
+ else:
472
+ upload_entry: UploadEntry = {
473
+ "project": project,
474
+ "run": None,
475
+ "step": None,
476
+ "relative_path": str(relative_to_base),
477
+ "uploaded_file": handle_file(file_path),
478
+ }
479
+ upload_entries.append(upload_entry)
480
+
481
+ if upload_entries:
482
+ if url is None:
483
+ raise RuntimeError(
484
+ "No server available. Call trackio.init() before trackio.save() to start the server."
485
+ )
486
+
487
+ try:
488
+ client = Client(url, verbose=False, httpx_kwargs={"timeout": 90})
489
+ client.predict(
490
+ api_name="/bulk_upload_media",
491
+ uploads=upload_entries,
492
+ hf_token=huggingface_hub.utils.get_token(),
493
+ )
494
+ except Exception as e:
495
+ warnings.warn(
496
+ f"Failed to upload files: {e}. "
497
+ "Files may not be available in the dashboard."
498
+ )
499
+
500
+ return str(utils.MEDIA_DIR / project / "files")
501
+
502
+
503
+ def show(
504
+ project: str | None = None,
505
+ *,
506
+ theme: str | ThemeClass | None = None,
507
+ mcp_server: bool | None = None,
508
+ footer: bool = True,
509
+ color_palette: list[str] | None = None,
510
+ open_browser: bool = True,
511
+ block_thread: bool | None = None,
512
+ host: str | None = None,
513
+ ):
514
+ """
515
+ Launches the Trackio dashboard.
516
+
517
+ Args:
518
+ project (`str`, *optional*):
519
+ The name of the project whose runs to show. If not provided, all projects
520
+ will be shown and the user can select one.
521
+ theme (`str` or `ThemeClass`, *optional*):
522
+ A Gradio Theme to use for the dashboard instead of the default Gradio theme,
523
+ can be a built-in theme (e.g. `'soft'`, `'citrus'`), a theme from the Hub
524
+ (e.g. `"gstaff/xkcd"`), or a custom Theme class. If not provided, the
525
+ `TRACKIO_THEME` environment variable will be used, or if that is not set,
526
+ the default Gradio theme will be used.
527
+ mcp_server (`bool`, *optional*):
528
+ If `True`, the Trackio dashboard will be set up as an MCP server and certain
529
+ functions will be added as MCP tools. If `None` (default behavior), then the
530
+ `GRADIO_MCP_SERVER` environment variable will be used to determine if the
531
+ MCP server should be enabled (which is `"True"` on Hugging Face Spaces).
532
+ footer (`bool`, *optional*, defaults to `True`):
533
+ Whether to show the Gradio footer. When `False`, the footer will be hidden.
534
+ This can also be controlled via the `footer` query parameter in the URL.
535
+ color_palette (`list[str]`, *optional*):
536
+ A list of hex color codes to use for plot lines. If not provided, the
537
+ `TRACKIO_COLOR_PALETTE` environment variable will be used (comma-separated
538
+ hex codes), or if that is not set, the default color palette will be used.
539
+ Example: `['#FF0000', '#00FF00', '#0000FF']`
540
+ open_browser (`bool`, *optional*, defaults to `True`):
541
+ If `True` and not in a notebook, a new browser tab will be opened with the
542
+ dashboard. If `False`, the browser will not be opened.
543
+ block_thread (`bool`, *optional*):
544
+ If `True`, the main thread will be blocked until the dashboard is closed.
545
+ If `None` (default behavior), then the main thread will not be blocked if the
546
+ dashboard is launched in a notebook, otherwise the main thread will be blocked.
547
+ host (`str`, *optional*):
548
+ The host to bind the server to. If not provided, defaults to `'127.0.0.1'`
549
+ (localhost only). Set to `'0.0.0.0'` to allow remote access.
550
+
551
+ Returns:
552
+ `app`: The Gradio app object corresponding to the dashboard launched by Trackio.
553
+ `url`: The local URL of the dashboard.
554
+ `share_url`: The public share URL of the dashboard.
555
+ `full_url`: The full URL of the dashboard including the write token (will use the public share URL if launched publicly, otherwise the local URL).
556
+ """
557
+ demo, CSS, HEAD = _get_demo()
558
+
559
+ if color_palette is not None:
560
+ os.environ["TRACKIO_COLOR_PALETTE"] = ",".join(color_palette)
561
+
562
+ theme = theme or os.environ.get("TRACKIO_THEME")
563
+
564
+ _mcp_server = (
565
+ mcp_server
566
+ if mcp_server is not None
567
+ else os.environ.get("GRADIO_MCP_SERVER", "False") == "True"
568
+ )
569
+
570
+ app, url, share_url = demo.launch(
571
+ css=CSS,
572
+ head=HEAD,
573
+ footer_links=["gradio", "settings"] + (["api"] if _mcp_server else []),
574
+ quiet=True,
575
+ inline=False,
576
+ prevent_thread_lock=True,
577
+ favicon_path=TRACKIO_LOGO_DIR / "trackio_logo_light.png",
578
+ allowed_paths=[TRACKIO_LOGO_DIR, TRACKIO_DIR],
579
+ mcp_server=_mcp_server,
580
+ theme=theme,
581
+ ssr_mode=False,
582
+ server_name=host,
583
+ )
584
+
585
+ base_url = share_url + "/" if share_url else url
586
+ full_url = utils.get_full_url(
587
+ base_url, project=project, write_token=demo.write_token, footer=footer
588
+ )
589
+
590
+ if not utils.is_in_notebook():
591
+ print(f"* Trackio UI launched at: {full_url}")
592
+ if open_browser:
593
+ webbrowser.open(full_url)
594
+ block_thread = block_thread if block_thread is not None else True
595
+ else:
596
+ utils.embed_url_in_notebook(full_url)
597
+ block_thread = block_thread if block_thread is not None else False
598
+
599
+ if block_thread:
600
+ utils.block_main_thread_until_keyboard_interrupt()
601
+ return TupleNoPrint((demo, url, share_url, full_url))
trackio/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (18.3 kB). View file
 
trackio/__pycache__/api.cpython-310.pyc ADDED
Binary file (3.13 kB). View file
 
trackio/__pycache__/commit_scheduler.cpython-310.pyc ADDED
Binary file (10.8 kB). View file
 
trackio/__pycache__/context_vars.cpython-310.pyc ADDED
Binary file (553 Bytes). View file
 
trackio/__pycache__/deploy.cpython-310.pyc ADDED
Binary file (12.1 kB). View file
 
trackio/__pycache__/dummy_commit_scheduler.cpython-310.pyc ADDED
Binary file (942 Bytes). View file
 
trackio/__pycache__/gpu.cpython-310.pyc ADDED
Binary file (9.13 kB). View file
 
trackio/__pycache__/histogram.cpython-310.pyc ADDED
Binary file (2.37 kB). View file
 
trackio/__pycache__/imports.cpython-310.pyc ADDED
Binary file (9.41 kB). View file
 
trackio/__pycache__/run.cpython-310.pyc ADDED
Binary file (14.3 kB). View file
 
trackio/__pycache__/sqlite_storage.cpython-310.pyc ADDED
Binary file (37.7 kB). View file
 
trackio/__pycache__/table.cpython-310.pyc ADDED
Binary file (6.57 kB). View file
 
trackio/__pycache__/typehints.cpython-310.pyc ADDED
Binary file (1.11 kB). View file
 
trackio/__pycache__/utils.cpython-310.pyc ADDED
Binary file (20 kB). View file
 
trackio/api.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Iterator
2
+
3
+ from trackio.sqlite_storage import SQLiteStorage
4
+
5
+
6
+ class Run:
7
+ def __init__(self, project: str, name: str):
8
+ self.project = project
9
+ self.name = name
10
+ self._config = None
11
+
12
+ @property
13
+ def id(self) -> str:
14
+ return self.name
15
+
16
+ @property
17
+ def config(self) -> dict | None:
18
+ if self._config is None:
19
+ self._config = SQLiteStorage.get_run_config(self.project, self.name)
20
+ return self._config
21
+
22
+ def delete(self) -> bool:
23
+ return SQLiteStorage.delete_run(self.project, self.name)
24
+
25
+ def move(self, new_project: str) -> bool:
26
+ success = SQLiteStorage.move_run(self.project, self.name, new_project)
27
+ if success:
28
+ self.project = new_project
29
+ return success
30
+
31
+ def __repr__(self) -> str:
32
+ return f"<Run {self.name} in project {self.project}>"
33
+
34
+
35
+ class Runs:
36
+ def __init__(self, project: str):
37
+ self.project = project
38
+ self._runs = None
39
+
40
+ def _load_runs(self):
41
+ if self._runs is None:
42
+ run_names = SQLiteStorage.get_runs(self.project)
43
+ self._runs = [Run(self.project, name) for name in run_names]
44
+
45
+ def __iter__(self) -> Iterator[Run]:
46
+ self._load_runs()
47
+ return iter(self._runs)
48
+
49
+ def __getitem__(self, index: int) -> Run:
50
+ self._load_runs()
51
+ return self._runs[index]
52
+
53
+ def __len__(self) -> int:
54
+ self._load_runs()
55
+ return len(self._runs)
56
+
57
+ def __repr__(self) -> str:
58
+ self._load_runs()
59
+ return f"<Runs project={self.project} count={len(self._runs)}>"
60
+
61
+
62
+ class Api:
63
+ def runs(self, project: str) -> Runs:
64
+ if not SQLiteStorage.get_project_db_path(project).exists():
65
+ raise ValueError(f"Project '{project}' does not exist")
66
+ return Runs(project)
trackio/assets/badge.png ADDED
trackio/assets/trackio_logo_dark.png ADDED
trackio/assets/trackio_logo_light.png ADDED
trackio/assets/trackio_logo_old.png ADDED

Git LFS Details

  • SHA256: 3922c4d1e465270ad4d8abb12023f3beed5d9f7f338528a4c0ac21dcf358a1c8
  • Pointer size: 131 Bytes
  • Size of remote file: 487 kB
trackio/assets/trackio_logo_type_dark.png ADDED
trackio/assets/trackio_logo_type_dark_transparent.png ADDED
trackio/assets/trackio_logo_type_light.png ADDED
trackio/assets/trackio_logo_type_light_transparent.png ADDED
trackio/cli.py ADDED
@@ -0,0 +1,514 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+
3
+ from trackio import show, sync
4
+ from trackio.cli_helpers import (
5
+ error_exit,
6
+ format_json,
7
+ format_list,
8
+ format_metric_values,
9
+ format_project_summary,
10
+ format_run_summary,
11
+ format_system_metric_names,
12
+ format_system_metrics,
13
+ )
14
+ from trackio.sqlite_storage import SQLiteStorage
15
+ from trackio.ui.main import get_project_summary, get_run_summary
16
+
17
+
18
+ def _handle_status():
19
+ print("Reading local Trackio projects...\n")
20
+ projects = SQLiteStorage.get_projects()
21
+ if not projects:
22
+ print("No Trackio projects found.")
23
+ return
24
+
25
+ local_projects = []
26
+ synced_projects = []
27
+ unsynced_projects = []
28
+
29
+ for project in projects:
30
+ space_id = SQLiteStorage.get_space_id(project)
31
+ if space_id is None:
32
+ local_projects.append(project)
33
+ elif SQLiteStorage.has_pending_data(project):
34
+ unsynced_projects.append(project)
35
+ else:
36
+ synced_projects.append(project)
37
+
38
+ print("Finished reading Trackio projects")
39
+ if local_projects:
40
+ print(f" * {len(local_projects)} local trackio project(s) [OK]")
41
+ if synced_projects:
42
+ print(f" * {len(synced_projects)} trackio project(s) synced to Spaces [OK]")
43
+ if unsynced_projects:
44
+ print(
45
+ f" * {len(unsynced_projects)} trackio project(s) with unsynced changes [WARNING]:"
46
+ )
47
+ for p in unsynced_projects:
48
+ print(f" - {p}")
49
+
50
+ if unsynced_projects:
51
+ print(
52
+ f"\nRun `trackio sync --project {unsynced_projects[0]}` to sync. "
53
+ "Or run `trackio sync --all` to sync all unsynced changes."
54
+ )
55
+
56
+
57
+ def _handle_sync(args):
58
+ from trackio.deploy import sync_incremental
59
+
60
+ if args.sync_all and args.project:
61
+ error_exit("Cannot use --all and --project together.")
62
+ if not args.sync_all and not args.project:
63
+ error_exit("Must provide either --project or --all.")
64
+
65
+ if args.sync_all:
66
+ projects = SQLiteStorage.get_projects()
67
+ synced_any = False
68
+ for project in projects:
69
+ space_id = SQLiteStorage.get_space_id(project)
70
+ if space_id and SQLiteStorage.has_pending_data(project):
71
+ sync_incremental(
72
+ project, space_id, private=args.private, pending_only=True
73
+ )
74
+ synced_any = True
75
+ if not synced_any:
76
+ print("No projects with unsynced data found.")
77
+ else:
78
+ space_id = args.space_id
79
+ if space_id is None:
80
+ space_id = SQLiteStorage.get_space_id(args.project)
81
+ sync(
82
+ project=args.project,
83
+ space_id=space_id,
84
+ private=args.private,
85
+ force=args.force,
86
+ )
87
+
88
+
89
+ def main():
90
+ parser = argparse.ArgumentParser(description="Trackio CLI")
91
+ subparsers = parser.add_subparsers(dest="command")
92
+
93
+ ui_parser = subparsers.add_parser(
94
+ "show", help="Show the Trackio dashboard UI for a project"
95
+ )
96
+ ui_parser.add_argument(
97
+ "--project", required=False, help="Project name to show in the dashboard"
98
+ )
99
+ ui_parser.add_argument(
100
+ "--theme",
101
+ required=False,
102
+ default="default",
103
+ 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').",
104
+ )
105
+ ui_parser.add_argument(
106
+ "--mcp-server",
107
+ action="store_true",
108
+ 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.",
109
+ )
110
+ ui_parser.add_argument(
111
+ "--footer",
112
+ action="store_true",
113
+ default=True,
114
+ help="Show the Gradio footer. Use --no-footer to hide it.",
115
+ )
116
+ ui_parser.add_argument(
117
+ "--no-footer",
118
+ dest="footer",
119
+ action="store_false",
120
+ help="Hide the Gradio footer.",
121
+ )
122
+ ui_parser.add_argument(
123
+ "--color-palette",
124
+ required=False,
125
+ help="Comma-separated list of hex color codes for plot lines (e.g. '#FF0000,#00FF00,#0000FF'). If not provided, the TRACKIO_COLOR_PALETTE environment variable will be used, or the default palette if not set.",
126
+ )
127
+ ui_parser.add_argument(
128
+ "--host",
129
+ required=False,
130
+ help="Host to bind the server to (e.g. '0.0.0.0' for remote access). If not provided, defaults to '127.0.0.1' (localhost only).",
131
+ )
132
+
133
+ subparsers.add_parser(
134
+ "status",
135
+ help="Show the status of all local Trackio projects, including sync status.",
136
+ )
137
+
138
+ sync_parser = subparsers.add_parser(
139
+ "sync",
140
+ help="Sync a local project's database to a Hugging Face Space. If the Space does not exist, it will be created.",
141
+ )
142
+ sync_parser.add_argument(
143
+ "--project",
144
+ required=False,
145
+ help="The name of the local project.",
146
+ )
147
+ sync_parser.add_argument(
148
+ "--space-id",
149
+ required=False,
150
+ help="The Hugging Face Space ID where the project will be synced (e.g. username/space_id). If not provided, uses the previously-configured Space.",
151
+ )
152
+ sync_parser.add_argument(
153
+ "--all",
154
+ action="store_true",
155
+ dest="sync_all",
156
+ help="Sync all projects that have unsynced data to their configured Spaces.",
157
+ )
158
+ sync_parser.add_argument(
159
+ "--private",
160
+ action="store_true",
161
+ help="Make the Hugging Face Space private if creating a new Space. By default, the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.",
162
+ )
163
+ sync_parser.add_argument(
164
+ "--force",
165
+ action="store_true",
166
+ help="Overwrite the existing database without prompting for confirmation.",
167
+ )
168
+
169
+ list_parser = subparsers.add_parser(
170
+ "list",
171
+ help="List projects, runs, or metrics",
172
+ )
173
+ list_subparsers = list_parser.add_subparsers(dest="list_type", required=True)
174
+
175
+ list_projects_parser = list_subparsers.add_parser(
176
+ "projects",
177
+ help="List all projects",
178
+ )
179
+ list_projects_parser.add_argument(
180
+ "--json",
181
+ action="store_true",
182
+ help="Output in JSON format",
183
+ )
184
+
185
+ list_runs_parser = list_subparsers.add_parser(
186
+ "runs",
187
+ help="List runs for a project",
188
+ )
189
+ list_runs_parser.add_argument(
190
+ "--project",
191
+ required=True,
192
+ help="Project name",
193
+ )
194
+ list_runs_parser.add_argument(
195
+ "--json",
196
+ action="store_true",
197
+ help="Output in JSON format",
198
+ )
199
+
200
+ list_metrics_parser = list_subparsers.add_parser(
201
+ "metrics",
202
+ help="List metrics for a run",
203
+ )
204
+ list_metrics_parser.add_argument(
205
+ "--project",
206
+ required=True,
207
+ help="Project name",
208
+ )
209
+ list_metrics_parser.add_argument(
210
+ "--run",
211
+ required=True,
212
+ help="Run name",
213
+ )
214
+ list_metrics_parser.add_argument(
215
+ "--json",
216
+ action="store_true",
217
+ help="Output in JSON format",
218
+ )
219
+
220
+ list_system_metrics_parser = list_subparsers.add_parser(
221
+ "system-metrics",
222
+ help="List system metrics for a run",
223
+ )
224
+ list_system_metrics_parser.add_argument(
225
+ "--project",
226
+ required=True,
227
+ help="Project name",
228
+ )
229
+ list_system_metrics_parser.add_argument(
230
+ "--run",
231
+ required=True,
232
+ help="Run name",
233
+ )
234
+ list_system_metrics_parser.add_argument(
235
+ "--json",
236
+ action="store_true",
237
+ help="Output in JSON format",
238
+ )
239
+
240
+ get_parser = subparsers.add_parser(
241
+ "get",
242
+ help="Get project, run, or metric information",
243
+ )
244
+ get_subparsers = get_parser.add_subparsers(dest="get_type", required=True)
245
+
246
+ get_project_parser = get_subparsers.add_parser(
247
+ "project",
248
+ help="Get project summary",
249
+ )
250
+ get_project_parser.add_argument(
251
+ "--project",
252
+ required=True,
253
+ help="Project name",
254
+ )
255
+ get_project_parser.add_argument(
256
+ "--json",
257
+ action="store_true",
258
+ help="Output in JSON format",
259
+ )
260
+
261
+ get_run_parser = get_subparsers.add_parser(
262
+ "run",
263
+ help="Get run summary",
264
+ )
265
+ get_run_parser.add_argument(
266
+ "--project",
267
+ required=True,
268
+ help="Project name",
269
+ )
270
+ get_run_parser.add_argument(
271
+ "--run",
272
+ required=True,
273
+ help="Run name",
274
+ )
275
+ get_run_parser.add_argument(
276
+ "--json",
277
+ action="store_true",
278
+ help="Output in JSON format",
279
+ )
280
+
281
+ get_metric_parser = get_subparsers.add_parser(
282
+ "metric",
283
+ help="Get metric values for a run",
284
+ )
285
+ get_metric_parser.add_argument(
286
+ "--project",
287
+ required=True,
288
+ help="Project name",
289
+ )
290
+ get_metric_parser.add_argument(
291
+ "--run",
292
+ required=True,
293
+ help="Run name",
294
+ )
295
+ get_metric_parser.add_argument(
296
+ "--metric",
297
+ required=True,
298
+ help="Metric name",
299
+ )
300
+ get_metric_parser.add_argument(
301
+ "--json",
302
+ action="store_true",
303
+ help="Output in JSON format",
304
+ )
305
+
306
+ get_system_metric_parser = get_subparsers.add_parser(
307
+ "system-metric",
308
+ help="Get system metric values for a run",
309
+ )
310
+ get_system_metric_parser.add_argument(
311
+ "--project",
312
+ required=True,
313
+ help="Project name",
314
+ )
315
+ get_system_metric_parser.add_argument(
316
+ "--run",
317
+ required=True,
318
+ help="Run name",
319
+ )
320
+ get_system_metric_parser.add_argument(
321
+ "--metric",
322
+ required=False,
323
+ help="System metric name (optional, if not provided returns all system metrics)",
324
+ )
325
+ get_system_metric_parser.add_argument(
326
+ "--json",
327
+ action="store_true",
328
+ help="Output in JSON format",
329
+ )
330
+
331
+ args = parser.parse_args()
332
+
333
+ if args.command == "show":
334
+ color_palette = None
335
+ if args.color_palette:
336
+ color_palette = [color.strip() for color in args.color_palette.split(",")]
337
+ show(
338
+ project=args.project,
339
+ theme=args.theme,
340
+ mcp_server=args.mcp_server,
341
+ footer=args.footer,
342
+ color_palette=color_palette,
343
+ host=args.host,
344
+ )
345
+ elif args.command == "status":
346
+ _handle_status()
347
+ elif args.command == "sync":
348
+ _handle_sync(args)
349
+ elif args.command == "list":
350
+ if args.list_type == "projects":
351
+ projects = SQLiteStorage.get_projects()
352
+ if args.json:
353
+ print(format_json({"projects": projects}))
354
+ else:
355
+ print(format_list(projects, "Projects"))
356
+ elif args.list_type == "runs":
357
+ db_path = SQLiteStorage.get_project_db_path(args.project)
358
+ if not db_path.exists():
359
+ error_exit(f"Project '{args.project}' not found.")
360
+ runs = SQLiteStorage.get_runs(args.project)
361
+ if args.json:
362
+ print(format_json({"project": args.project, "runs": runs}))
363
+ else:
364
+ print(format_list(runs, f"Runs in '{args.project}'"))
365
+ elif args.list_type == "metrics":
366
+ db_path = SQLiteStorage.get_project_db_path(args.project)
367
+ if not db_path.exists():
368
+ error_exit(f"Project '{args.project}' not found.")
369
+ runs = SQLiteStorage.get_runs(args.project)
370
+ if args.run not in runs:
371
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
372
+ metrics = SQLiteStorage.get_all_metrics_for_run(args.project, args.run)
373
+ if args.json:
374
+ print(
375
+ format_json(
376
+ {"project": args.project, "run": args.run, "metrics": metrics}
377
+ )
378
+ )
379
+ else:
380
+ print(
381
+ format_list(
382
+ metrics, f"Metrics for '{args.run}' in '{args.project}'"
383
+ )
384
+ )
385
+ elif args.list_type == "system-metrics":
386
+ db_path = SQLiteStorage.get_project_db_path(args.project)
387
+ if not db_path.exists():
388
+ error_exit(f"Project '{args.project}' not found.")
389
+ runs = SQLiteStorage.get_runs(args.project)
390
+ if args.run not in runs:
391
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
392
+ system_metrics = SQLiteStorage.get_all_system_metrics_for_run(
393
+ args.project, args.run
394
+ )
395
+ if args.json:
396
+ print(
397
+ format_json(
398
+ {
399
+ "project": args.project,
400
+ "run": args.run,
401
+ "system_metrics": system_metrics,
402
+ }
403
+ )
404
+ )
405
+ else:
406
+ print(format_system_metric_names(system_metrics))
407
+ elif args.command == "get":
408
+ if args.get_type == "project":
409
+ db_path = SQLiteStorage.get_project_db_path(args.project)
410
+ if not db_path.exists():
411
+ error_exit(f"Project '{args.project}' not found.")
412
+ summary = get_project_summary(args.project)
413
+ if args.json:
414
+ print(format_json(summary))
415
+ else:
416
+ print(format_project_summary(summary))
417
+ elif args.get_type == "run":
418
+ db_path = SQLiteStorage.get_project_db_path(args.project)
419
+ if not db_path.exists():
420
+ error_exit(f"Project '{args.project}' not found.")
421
+ runs = SQLiteStorage.get_runs(args.project)
422
+ if args.run not in runs:
423
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
424
+ summary = get_run_summary(args.project, args.run)
425
+ if args.json:
426
+ print(format_json(summary))
427
+ else:
428
+ print(format_run_summary(summary))
429
+ elif args.get_type == "metric":
430
+ db_path = SQLiteStorage.get_project_db_path(args.project)
431
+ if not db_path.exists():
432
+ error_exit(f"Project '{args.project}' not found.")
433
+ runs = SQLiteStorage.get_runs(args.project)
434
+ if args.run not in runs:
435
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
436
+ metrics = SQLiteStorage.get_all_metrics_for_run(args.project, args.run)
437
+ if args.metric not in metrics:
438
+ error_exit(
439
+ f"Metric '{args.metric}' not found in run '{args.run}' of project '{args.project}'."
440
+ )
441
+ values = SQLiteStorage.get_metric_values(
442
+ args.project, args.run, args.metric
443
+ )
444
+ if args.json:
445
+ print(
446
+ format_json(
447
+ {
448
+ "project": args.project,
449
+ "run": args.run,
450
+ "metric": args.metric,
451
+ "values": values,
452
+ }
453
+ )
454
+ )
455
+ else:
456
+ print(format_metric_values(values))
457
+ elif args.get_type == "system-metric":
458
+ db_path = SQLiteStorage.get_project_db_path(args.project)
459
+ if not db_path.exists():
460
+ error_exit(f"Project '{args.project}' not found.")
461
+ runs = SQLiteStorage.get_runs(args.project)
462
+ if args.run not in runs:
463
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
464
+ if args.metric:
465
+ system_metrics = SQLiteStorage.get_system_logs(args.project, args.run)
466
+ all_system_metric_names = SQLiteStorage.get_all_system_metrics_for_run(
467
+ args.project, args.run
468
+ )
469
+ if args.metric not in all_system_metric_names:
470
+ error_exit(
471
+ f"System metric '{args.metric}' not found in run '{args.run}' of project '{args.project}'."
472
+ )
473
+ filtered_metrics = [
474
+ {
475
+ k: v
476
+ for k, v in entry.items()
477
+ if k == "timestamp" or k == args.metric
478
+ }
479
+ for entry in system_metrics
480
+ if args.metric in entry
481
+ ]
482
+ if args.json:
483
+ print(
484
+ format_json(
485
+ {
486
+ "project": args.project,
487
+ "run": args.run,
488
+ "metric": args.metric,
489
+ "values": filtered_metrics,
490
+ }
491
+ )
492
+ )
493
+ else:
494
+ print(format_system_metrics(filtered_metrics))
495
+ else:
496
+ system_metrics = SQLiteStorage.get_system_logs(args.project, args.run)
497
+ if args.json:
498
+ print(
499
+ format_json(
500
+ {
501
+ "project": args.project,
502
+ "run": args.run,
503
+ "system_metrics": system_metrics,
504
+ }
505
+ )
506
+ )
507
+ else:
508
+ print(format_system_metrics(system_metrics))
509
+ else:
510
+ parser.print_help()
511
+
512
+
513
+ if __name__ == "__main__":
514
+ main()
trackio/cli_helpers.py ADDED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import sys
3
+ from typing import Any
4
+
5
+
6
+ def format_json(data: Any) -> str:
7
+ """Format data as JSON."""
8
+ return json.dumps(data, indent=2)
9
+
10
+
11
+ def format_list(items: list[str], title: str | None = None) -> str:
12
+ """Format a list of items in human-readable format."""
13
+ if not items:
14
+ return f"No {title.lower() if title else 'items'} found."
15
+
16
+ output = []
17
+ if title:
18
+ output.append(f"{title}:")
19
+
20
+ for item in items:
21
+ output.append(f" - {item}")
22
+
23
+ return "\n".join(output)
24
+
25
+
26
+ def format_project_summary(summary: dict) -> str:
27
+ """Format project summary in human-readable format."""
28
+ output = [f"Project: {summary['project']}"]
29
+ output.append(f"Number of runs: {summary['num_runs']}")
30
+
31
+ if summary["runs"]:
32
+ output.append("\nRuns:")
33
+ for run in summary["runs"]:
34
+ output.append(f" - {run}")
35
+ else:
36
+ output.append("\nNo runs found.")
37
+
38
+ if summary.get("last_activity"):
39
+ output.append(f"\nLast activity (max step): {summary['last_activity']}")
40
+
41
+ return "\n".join(output)
42
+
43
+
44
+ def format_run_summary(summary: dict) -> str:
45
+ """Format run summary in human-readable format."""
46
+ output = [f"Project: {summary['project']}"]
47
+ output.append(f"Run: {summary['run']}")
48
+ output.append(f"Number of logs: {summary['num_logs']}")
49
+
50
+ if summary.get("last_step") is not None:
51
+ output.append(f"Last step: {summary['last_step']}")
52
+
53
+ if summary.get("metrics"):
54
+ output.append("\nMetrics:")
55
+ for metric in summary["metrics"]:
56
+ output.append(f" - {metric}")
57
+ else:
58
+ output.append("\nNo metrics found.")
59
+
60
+ config = summary.get("config")
61
+ if config:
62
+ output.append("\nConfig:")
63
+ config_display = {k: v for k, v in config.items() if not k.startswith("_")}
64
+ if config_display:
65
+ for key, value in config_display.items():
66
+ output.append(f" {key}: {value}")
67
+ else:
68
+ output.append(" (no config)")
69
+ else:
70
+ output.append("\nConfig: (no config)")
71
+
72
+ return "\n".join(output)
73
+
74
+
75
+ def format_metric_values(values: list[dict]) -> str:
76
+ """Format metric values in human-readable format."""
77
+ if not values:
78
+ return "No metric values found."
79
+
80
+ output = [f"Found {len(values)} value(s):\n"]
81
+ output.append("Step | Timestamp | Value")
82
+ output.append("-" * 50)
83
+
84
+ for value in values:
85
+ step = value.get("step", "N/A")
86
+ timestamp = value.get("timestamp", "N/A")
87
+ val = value.get("value", "N/A")
88
+ output.append(f"{step} | {timestamp} | {val}")
89
+
90
+ return "\n".join(output)
91
+
92
+
93
+ def format_system_metrics(metrics: list[dict]) -> str:
94
+ """Format system metrics in human-readable format."""
95
+ if not metrics:
96
+ return "No system metrics found."
97
+
98
+ output = [f"Found {len(metrics)} system metric entry/entries:\n"]
99
+
100
+ for i, entry in enumerate(metrics):
101
+ timestamp = entry.get("timestamp", "N/A")
102
+ output.append(f"\nEntry {i + 1} (Timestamp: {timestamp}):")
103
+ for key, value in entry.items():
104
+ if key != "timestamp":
105
+ output.append(f" {key}: {value}")
106
+
107
+ return "\n".join(output)
108
+
109
+
110
+ def format_system_metric_names(names: list[str]) -> str:
111
+ """Format system metric names in human-readable format."""
112
+ return format_list(names, "System Metrics")
113
+
114
+
115
+ def error_exit(message: str, code: int = 1) -> None:
116
+ """Print error message and exit."""
117
+ print(f"Error: {message}", file=sys.stderr)
118
+ sys.exit(code)
trackio/commit_scheduler.py ADDED
@@ -0,0 +1,310 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 time
6
+ from concurrent.futures import Future
7
+ from dataclasses import dataclass
8
+ from pathlib import Path
9
+ from threading import Lock, Thread
10
+ from typing import Callable, Dict, List, Union
11
+
12
+ from huggingface_hub.hf_api import (
13
+ DEFAULT_IGNORE_PATTERNS,
14
+ CommitInfo,
15
+ CommitOperationAdd,
16
+ HfApi,
17
+ )
18
+ from huggingface_hub.utils import filter_repo_objects
19
+
20
+ logger = logging.getLogger(__name__)
21
+
22
+
23
+ @dataclass(frozen=True)
24
+ class _FileToUpload:
25
+ """Temporary dataclass to store info about files to upload. Not meant to be used directly."""
26
+
27
+ local_path: Path
28
+ path_in_repo: str
29
+ size_limit: int
30
+ last_modified: float
31
+
32
+
33
+ class CommitScheduler:
34
+ """
35
+ Scheduler to upload a local folder to the Hub at regular intervals (e.g. push to hub every 5 minutes).
36
+
37
+ The recommended way to use the scheduler is to use it as a context manager. This ensures that the scheduler is
38
+ properly stopped and the last commit is triggered when the script ends. The scheduler can also be stopped manually
39
+ with the `stop` method. Checkout the [upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#scheduled-uploads)
40
+ to learn more about how to use it.
41
+
42
+ Args:
43
+ repo_id (`str`):
44
+ The id of the repo to commit to.
45
+ folder_path (`str` or `Path`):
46
+ Path to the local folder to upload regularly.
47
+ every (`int` or `float`, *optional*):
48
+ The number of minutes between each commit. Defaults to 5 minutes.
49
+ path_in_repo (`str`, *optional*):
50
+ Relative path of the directory in the repo, for example: `"checkpoints/"`. Defaults to the root folder
51
+ of the repository.
52
+ repo_type (`str`, *optional*):
53
+ The type of the repo to commit to. Defaults to `model`.
54
+ revision (`str`, *optional*):
55
+ The revision of the repo to commit to. Defaults to `main`.
56
+ private (`bool`, *optional*):
57
+ 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.
58
+ token (`str`, *optional*):
59
+ The token to use to commit to the repo. Defaults to the token saved on the machine.
60
+ allow_patterns (`List[str]` or `str`, *optional*):
61
+ If provided, only files matching at least one pattern are uploaded.
62
+ ignore_patterns (`List[str]` or `str`, *optional*):
63
+ If provided, files matching any of the patterns are not uploaded.
64
+ squash_history (`bool`, *optional*):
65
+ Whether to squash the history of the repo after each commit. Defaults to `False`. Squashing commits is
66
+ useful to avoid degraded performances on the repo when it grows too large.
67
+ hf_api (`HfApi`, *optional*):
68
+ The [`HfApi`] client to use to commit to the Hub. Can be set with custom settings (user agent, token,...).
69
+ on_before_commit (`Callable[[], None]`, *optional*):
70
+ If specified, a function that will be called before the CommitScheduler lists files to create a commit.
71
+
72
+ Example:
73
+ ```py
74
+ >>> from pathlib import Path
75
+ >>> from huggingface_hub import CommitScheduler
76
+
77
+ # Scheduler uploads every 10 minutes
78
+ >>> csv_path = Path("watched_folder/data.csv")
79
+ >>> CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path=csv_path.parent, every=10)
80
+
81
+ >>> with csv_path.open("a") as f:
82
+ ... f.write("first line")
83
+
84
+ # Some time later (...)
85
+ >>> with csv_path.open("a") as f:
86
+ ... f.write("second line")
87
+ ```
88
+
89
+ Example using a context manager:
90
+ ```py
91
+ >>> from pathlib import Path
92
+ >>> from huggingface_hub import CommitScheduler
93
+
94
+ >>> with CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path="watched_folder", every=10) as scheduler:
95
+ ... csv_path = Path("watched_folder/data.csv")
96
+ ... with csv_path.open("a") as f:
97
+ ... f.write("first line")
98
+ ... (...)
99
+ ... with csv_path.open("a") as f:
100
+ ... f.write("second line")
101
+
102
+ # Scheduler is now stopped and last commit have been triggered
103
+ ```
104
+ """
105
+
106
+ def __init__(
107
+ self,
108
+ *,
109
+ repo_id: str,
110
+ folder_path: Union[str, Path],
111
+ every: Union[int, float] = 5,
112
+ path_in_repo: str | None = None,
113
+ repo_type: str | None = None,
114
+ revision: str | None = None,
115
+ private: bool | None = None,
116
+ token: str | None = None,
117
+ allow_patterns: list[str] | str | None = None,
118
+ ignore_patterns: list[str] | str | None = None,
119
+ squash_history: bool = False,
120
+ hf_api: HfApi | None = None,
121
+ on_before_commit: Callable[[], None] | None = None,
122
+ ) -> None:
123
+ self.api = hf_api or HfApi(token=token)
124
+ self.on_before_commit = on_before_commit
125
+
126
+ # Folder
127
+ self.folder_path = Path(folder_path).expanduser().resolve()
128
+ self.path_in_repo = path_in_repo or ""
129
+ self.allow_patterns = allow_patterns
130
+
131
+ if ignore_patterns is None:
132
+ ignore_patterns = []
133
+ elif isinstance(ignore_patterns, str):
134
+ ignore_patterns = [ignore_patterns]
135
+ self.ignore_patterns = ignore_patterns + DEFAULT_IGNORE_PATTERNS
136
+
137
+ if self.folder_path.is_file():
138
+ raise ValueError(
139
+ f"'folder_path' must be a directory, not a file: '{self.folder_path}'."
140
+ )
141
+ self.folder_path.mkdir(parents=True, exist_ok=True)
142
+
143
+ # Repository
144
+ repo_url = self.api.create_repo(
145
+ repo_id=repo_id, private=private, repo_type=repo_type, exist_ok=True
146
+ )
147
+ self.repo_id = repo_url.repo_id
148
+ self.repo_type = repo_type
149
+ self.revision = revision
150
+ self.token = token
151
+
152
+ self.last_uploaded: Dict[Path, float] = {}
153
+ self.last_push_time: float | None = None
154
+
155
+ if not every > 0:
156
+ raise ValueError(f"'every' must be a positive integer, not '{every}'.")
157
+ self.lock = Lock()
158
+ self.every = every
159
+ self.squash_history = squash_history
160
+
161
+ logger.info(
162
+ f"Scheduled job to push '{self.folder_path}' to '{self.repo_id}' every {self.every} minutes."
163
+ )
164
+ self._scheduler_thread = Thread(target=self._run_scheduler, daemon=True)
165
+ self._scheduler_thread.start()
166
+ atexit.register(self._push_to_hub)
167
+
168
+ self.__stopped = False
169
+
170
+ def stop(self) -> None:
171
+ """Stop the scheduler.
172
+
173
+ A stopped scheduler cannot be restarted. Mostly for tests purposes.
174
+ """
175
+ self.__stopped = True
176
+
177
+ def __enter__(self) -> "CommitScheduler":
178
+ return self
179
+
180
+ def __exit__(self, exc_type, exc_value, traceback) -> None:
181
+ # Upload last changes before exiting
182
+ self.trigger().result()
183
+ self.stop()
184
+ return
185
+
186
+ def _run_scheduler(self) -> None:
187
+ """Dumb thread waiting between each scheduled push to Hub."""
188
+ while True:
189
+ self.last_future = self.trigger()
190
+ time.sleep(self.every * 60)
191
+ if self.__stopped:
192
+ break
193
+
194
+ def trigger(self) -> Future:
195
+ """Trigger a `push_to_hub` and return a future.
196
+
197
+ This method is automatically called every `every` minutes. You can also call it manually to trigger a commit
198
+ immediately, without waiting for the next scheduled commit.
199
+ """
200
+ return self.api.run_as_future(self._push_to_hub)
201
+
202
+ def _push_to_hub(self) -> CommitInfo | None:
203
+ if self.__stopped: # If stopped, already scheduled commits are ignored
204
+ return None
205
+
206
+ logger.info("(Background) scheduled commit triggered.")
207
+ try:
208
+ value = self.push_to_hub()
209
+ if self.squash_history:
210
+ logger.info("(Background) squashing repo history.")
211
+ self.api.super_squash_history(
212
+ repo_id=self.repo_id, repo_type=self.repo_type, branch=self.revision
213
+ )
214
+ return value
215
+ except Exception as e:
216
+ logger.error(
217
+ f"Error while pushing to Hub: {e}"
218
+ ) # Depending on the setup, error might be silenced
219
+ raise
220
+
221
+ def push_to_hub(self) -> CommitInfo | None:
222
+ """
223
+ Push folder to the Hub and return the commit info.
224
+
225
+ <Tip warning={true}>
226
+
227
+ This method is not meant to be called directly. It is run in the background by the scheduler, respecting a
228
+ queue mechanism to avoid concurrent commits. Making a direct call to the method might lead to concurrency
229
+ issues.
230
+
231
+ </Tip>
232
+
233
+ The default behavior of `push_to_hub` is to assume an append-only folder. It lists all files in the folder and
234
+ uploads only changed files. If no changes are found, the method returns without committing anything. If you want
235
+ to change this behavior, you can inherit from [`CommitScheduler`] and override this method. This can be useful
236
+ for example to compress data together in a single file before committing. For more details and examples, check
237
+ out our [integration guide](https://huggingface.co/docs/huggingface_hub/main/en/guides/upload#scheduled-uploads).
238
+ """
239
+ # Check files to upload (with lock)
240
+ with self.lock:
241
+ if self.on_before_commit is not None:
242
+ self.on_before_commit()
243
+
244
+ logger.debug("Listing files to upload for scheduled commit.")
245
+
246
+ # List files from folder (taken from `_prepare_upload_folder_additions`)
247
+ relpath_to_abspath = {
248
+ path.relative_to(self.folder_path).as_posix(): path
249
+ for path in sorted(
250
+ self.folder_path.glob("**/*")
251
+ ) # sorted to be deterministic
252
+ if path.is_file()
253
+ }
254
+ prefix = f"{self.path_in_repo.strip('/')}/" if self.path_in_repo else ""
255
+
256
+ # Filter with pattern + filter out unchanged files + retrieve current file size
257
+ files_to_upload: List[_FileToUpload] = []
258
+ for relpath in filter_repo_objects(
259
+ relpath_to_abspath.keys(),
260
+ allow_patterns=self.allow_patterns,
261
+ ignore_patterns=self.ignore_patterns,
262
+ ):
263
+ local_path = relpath_to_abspath[relpath]
264
+ stat = local_path.stat()
265
+ if (
266
+ self.last_uploaded.get(local_path) is None
267
+ or self.last_uploaded[local_path] != stat.st_mtime
268
+ ):
269
+ files_to_upload.append(
270
+ _FileToUpload(
271
+ local_path=local_path,
272
+ path_in_repo=prefix + relpath,
273
+ size_limit=stat.st_size,
274
+ last_modified=stat.st_mtime,
275
+ )
276
+ )
277
+
278
+ # Return if nothing to upload
279
+ if len(files_to_upload) == 0:
280
+ logger.debug("Dropping schedule commit: no changed file to upload.")
281
+ return None
282
+
283
+ # Convert `_FileToUpload` as `CommitOperationAdd` (=> compute file shas + limit to file size)
284
+ logger.debug("Removing unchanged files since previous scheduled commit.")
285
+ add_operations = [
286
+ CommitOperationAdd(
287
+ # TODO: Cap the file to its current size, even if the user append data to it while a scheduled commit is happening
288
+ # (requires an upstream fix for XET-535: `hf_xet` should support `BinaryIO` for upload)
289
+ path_or_fileobj=file_to_upload.local_path,
290
+ path_in_repo=file_to_upload.path_in_repo,
291
+ )
292
+ for file_to_upload in files_to_upload
293
+ ]
294
+
295
+ # Upload files (append mode expected - no need for lock)
296
+ logger.debug("Uploading files for scheduled commit.")
297
+ commit_info = self.api.create_commit(
298
+ repo_id=self.repo_id,
299
+ repo_type=self.repo_type,
300
+ operations=add_operations,
301
+ commit_message="Scheduled Commit",
302
+ revision=self.revision,
303
+ )
304
+
305
+ for file in files_to_upload:
306
+ self.last_uploaded[file.local_path] = file.last_modified
307
+
308
+ self.last_push_time = time.time()
309
+
310
+ return commit_info
trackio/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_space_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
17
+ "current_space_id", default=None
18
+ )
trackio/deploy.py ADDED
@@ -0,0 +1,433 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import importlib.metadata
2
+ import io
3
+ import os
4
+ import sys
5
+ import threading
6
+ import time
7
+ from importlib.resources import files
8
+ from pathlib import Path
9
+
10
+ if sys.version_info >= (3, 11):
11
+ import tomllib
12
+ else:
13
+ import tomli as tomllib
14
+
15
+ import gradio
16
+ import huggingface_hub
17
+ from gradio_client import Client, handle_file
18
+ from httpx import ReadTimeout
19
+ from huggingface_hub.errors import HfHubHTTPError, RepositoryNotFoundError
20
+
21
+ import trackio
22
+ from trackio.sqlite_storage import SQLiteStorage
23
+ from trackio.utils import get_or_create_project_hash, preprocess_space_and_dataset_ids
24
+
25
+ SPACE_HOST_URL = "https://{user_name}-{space_name}.hf.space/"
26
+ SPACE_URL = "https://huggingface.co/spaces/{space_id}"
27
+
28
+
29
+ def _get_source_install_dependencies() -> str:
30
+ """Get trackio dependencies from pyproject.toml for source installs."""
31
+ trackio_path = files("trackio")
32
+ pyproject_path = Path(trackio_path).parent / "pyproject.toml"
33
+ with open(pyproject_path, "rb") as f:
34
+ pyproject = tomllib.load(f)
35
+ deps = pyproject["project"]["dependencies"]
36
+ spaces_deps = (
37
+ pyproject["project"].get("optional-dependencies", {}).get("spaces", [])
38
+ )
39
+ return "\n".join(deps + spaces_deps)
40
+
41
+
42
+ def _is_trackio_installed_from_source() -> bool:
43
+ """Check if trackio is installed from source/editable install vs PyPI."""
44
+ try:
45
+ trackio_file = trackio.__file__
46
+ if "site-packages" not in trackio_file and "dist-packages" not in trackio_file:
47
+ return True
48
+
49
+ dist = importlib.metadata.distribution("trackio")
50
+ if dist.files:
51
+ files = list(dist.files)
52
+ has_pth = any(".pth" in str(f) for f in files)
53
+ if has_pth:
54
+ return True
55
+
56
+ return False
57
+ except (
58
+ AttributeError,
59
+ importlib.metadata.PackageNotFoundError,
60
+ importlib.metadata.MetadataError,
61
+ ValueError,
62
+ TypeError,
63
+ ):
64
+ return True
65
+
66
+
67
+ def deploy_as_space(
68
+ space_id: str,
69
+ space_storage: huggingface_hub.SpaceStorage | None = None,
70
+ dataset_id: str | None = None,
71
+ private: bool | None = None,
72
+ ):
73
+ if (
74
+ os.getenv("SYSTEM") == "spaces"
75
+ ): # in case a repo with this function is uploaded to spaces
76
+ return
77
+
78
+ trackio_path = files("trackio")
79
+
80
+ hf_api = huggingface_hub.HfApi()
81
+
82
+ try:
83
+ huggingface_hub.create_repo(
84
+ space_id,
85
+ private=private,
86
+ space_sdk="gradio",
87
+ space_storage=space_storage,
88
+ repo_type="space",
89
+ exist_ok=True,
90
+ )
91
+ except HfHubHTTPError as e:
92
+ if e.response.status_code in [401, 403]: # unauthorized or forbidden
93
+ print("Need 'write' access token to create a Spaces repo.")
94
+ huggingface_hub.login(add_to_git_credential=False)
95
+ huggingface_hub.create_repo(
96
+ space_id,
97
+ private=private,
98
+ space_sdk="gradio",
99
+ space_storage=space_storage,
100
+ repo_type="space",
101
+ exist_ok=True,
102
+ )
103
+ else:
104
+ raise ValueError(f"Failed to create Space: {e}")
105
+
106
+ # We can assume pandas, gradio, and huggingface-hub are already installed in a Gradio Space.
107
+ # Make sure necessary dependencies are installed by creating a requirements.txt.
108
+ is_source_install = _is_trackio_installed_from_source()
109
+
110
+ with open(Path(trackio_path, "README.md"), "r") as f:
111
+ readme_content = f.read()
112
+ readme_content = readme_content.replace("{GRADIO_VERSION}", gradio.__version__)
113
+ if is_source_install:
114
+ readme_content = readme_content.replace("{APP_FILE}", "trackio/ui/main.py")
115
+ else:
116
+ readme_content = readme_content.replace("{APP_FILE}", "app.py")
117
+ readme_buffer = io.BytesIO(readme_content.encode("utf-8"))
118
+ hf_api.upload_file(
119
+ path_or_fileobj=readme_buffer,
120
+ path_in_repo="README.md",
121
+ repo_id=space_id,
122
+ repo_type="space",
123
+ )
124
+
125
+ if is_source_install:
126
+ requirements_content = _get_source_install_dependencies()
127
+ else:
128
+ requirements_content = f"trackio[spaces]=={trackio.__version__}"
129
+
130
+ requirements_buffer = io.BytesIO(requirements_content.encode("utf-8"))
131
+ hf_api.upload_file(
132
+ path_or_fileobj=requirements_buffer,
133
+ path_in_repo="requirements.txt",
134
+ repo_id=space_id,
135
+ repo_type="space",
136
+ )
137
+
138
+ huggingface_hub.utils.disable_progress_bars()
139
+
140
+ if is_source_install:
141
+ hf_api.upload_folder(
142
+ repo_id=space_id,
143
+ repo_type="space",
144
+ folder_path=trackio_path,
145
+ path_in_repo="trackio",
146
+ ignore_patterns=["README.md"],
147
+ )
148
+
149
+ app_file_content = """import trackio
150
+ trackio.show()"""
151
+ app_file_buffer = io.BytesIO(app_file_content.encode("utf-8"))
152
+ hf_api.upload_file(
153
+ path_or_fileobj=app_file_buffer,
154
+ path_in_repo="app.py",
155
+ repo_id=space_id,
156
+ repo_type="space",
157
+ )
158
+
159
+ if hf_token := huggingface_hub.utils.get_token():
160
+ huggingface_hub.add_space_secret(space_id, "HF_TOKEN", hf_token)
161
+ if dataset_id is not None:
162
+ huggingface_hub.add_space_variable(space_id, "TRACKIO_DATASET_ID", dataset_id)
163
+ if logo_light_url := os.environ.get("TRACKIO_LOGO_LIGHT_URL"):
164
+ huggingface_hub.add_space_variable(
165
+ space_id, "TRACKIO_LOGO_LIGHT_URL", logo_light_url
166
+ )
167
+ if logo_dark_url := os.environ.get("TRACKIO_LOGO_DARK_URL"):
168
+ huggingface_hub.add_space_variable(
169
+ space_id, "TRACKIO_LOGO_DARK_URL", logo_dark_url
170
+ )
171
+ if plot_order := os.environ.get("TRACKIO_PLOT_ORDER"):
172
+ huggingface_hub.add_space_variable(space_id, "TRACKIO_PLOT_ORDER", plot_order)
173
+ if theme := os.environ.get("TRACKIO_THEME"):
174
+ huggingface_hub.add_space_variable(space_id, "TRACKIO_THEME", theme)
175
+ huggingface_hub.add_space_variable(space_id, "GRADIO_MCP_SERVER", "True")
176
+
177
+
178
+ def create_space_if_not_exists(
179
+ space_id: str,
180
+ space_storage: huggingface_hub.SpaceStorage | None = None,
181
+ dataset_id: str | None = None,
182
+ private: bool | None = None,
183
+ ) -> None:
184
+ """
185
+ Creates a new Hugging Face Space if it does not exist.
186
+
187
+ Args:
188
+ space_id (`str`):
189
+ The ID of the Space to create.
190
+ space_storage ([`~huggingface_hub.SpaceStorage`], *optional*):
191
+ Choice of persistent storage tier for the Space.
192
+ dataset_id (`str`, *optional*):
193
+ The ID of the Dataset to add to the Space as a space variable.
194
+ private (`bool`, *optional*):
195
+ Whether to make the Space private. If `None` (default), the repo will be
196
+ public unless the organization's default is private. This value is ignored
197
+ if the repo already exists.
198
+ """
199
+ if "/" not in space_id:
200
+ raise ValueError(
201
+ f"Invalid space ID: {space_id}. Must be in the format: username/reponame or orgname/reponame."
202
+ )
203
+ if dataset_id is not None and "/" not in dataset_id:
204
+ raise ValueError(
205
+ f"Invalid dataset ID: {dataset_id}. Must be in the format: username/datasetname or orgname/datasetname."
206
+ )
207
+ try:
208
+ huggingface_hub.repo_info(space_id, repo_type="space")
209
+ print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")
210
+ return
211
+ except RepositoryNotFoundError:
212
+ pass
213
+ except HfHubHTTPError as e:
214
+ if e.response.status_code in [401, 403]: # unauthorized or forbidden
215
+ print("Need 'write' access token to create a Spaces repo.")
216
+ huggingface_hub.login(add_to_git_credential=False)
217
+ else:
218
+ raise ValueError(f"Failed to create Space: {e}")
219
+
220
+ print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")
221
+ deploy_as_space(space_id, space_storage, dataset_id, private)
222
+
223
+
224
+ def wait_until_space_exists(
225
+ space_id: str,
226
+ ) -> None:
227
+ """
228
+ Blocks the current thread until the Space exists.
229
+
230
+ Args:
231
+ space_id (`str`):
232
+ The ID of the Space to wait for.
233
+
234
+ Raises:
235
+ `TimeoutError`: If waiting for the Space takes longer than expected.
236
+ """
237
+ hf_api = huggingface_hub.HfApi()
238
+ delay = 1
239
+ for _ in range(30):
240
+ try:
241
+ hf_api.space_info(space_id)
242
+ return
243
+ except (huggingface_hub.utils.HfHubHTTPError, ReadTimeout):
244
+ time.sleep(delay)
245
+ delay = min(delay * 2, 60)
246
+ raise TimeoutError("Waiting for space to exist took longer than expected")
247
+
248
+
249
+ def upload_db_to_space(project: str, space_id: str, force: bool = False) -> None:
250
+ """
251
+ Uploads the database of a local Trackio project to a Hugging Face Space.
252
+
253
+ This uses the Gradio Client to upload since we do not want to trigger a new build of
254
+ the Space, which would happen if we used `huggingface_hub.upload_file`.
255
+
256
+ Args:
257
+ project (`str`):
258
+ The name of the project to upload.
259
+ space_id (`str`):
260
+ The ID of the Space to upload to.
261
+ force (`bool`, *optional*, defaults to `False`):
262
+ If `True`, overwrites the existing database without prompting. If `False`,
263
+ prompts for confirmation.
264
+ """
265
+ db_path = SQLiteStorage.get_project_db_path(project)
266
+ client = Client(space_id, verbose=False, httpx_kwargs={"timeout": 90})
267
+
268
+ if not force:
269
+ try:
270
+ existing_projects = client.predict(api_name="/get_all_projects")
271
+ if project in existing_projects:
272
+ response = input(
273
+ f"Database for project '{project}' already exists on Space '{space_id}'. "
274
+ f"Overwrite it? (y/N): "
275
+ )
276
+ if response.lower() not in ["y", "yes"]:
277
+ print("* Upload cancelled.")
278
+ return
279
+ except Exception as e:
280
+ print(f"* Warning: Could not check if project exists on Space: {e}")
281
+ print("* Proceeding with upload...")
282
+
283
+ client.predict(
284
+ api_name="/upload_db_to_space",
285
+ project=project,
286
+ uploaded_db=handle_file(db_path),
287
+ hf_token=huggingface_hub.utils.get_token(),
288
+ )
289
+
290
+
291
+ SYNC_BATCH_SIZE = 500
292
+
293
+
294
+ def sync_incremental(
295
+ project: str,
296
+ space_id: str,
297
+ private: bool | None = None,
298
+ pending_only: bool = False,
299
+ ) -> None:
300
+ """
301
+ Syncs a local Trackio project to a Space via the bulk_log API endpoints
302
+ instead of uploading the entire DB file. Supports incremental sync.
303
+
304
+ Args:
305
+ project: The name of the project to sync.
306
+ space_id: The HF Space ID to sync to.
307
+ private: Whether to make the Space private if creating.
308
+ pending_only: If True, only sync rows tagged with space_id (pending data).
309
+ """
310
+ print(
311
+ f"* Syncing project '{project}' to: {SPACE_URL.format(space_id=space_id)} (please wait...)"
312
+ )
313
+ create_space_if_not_exists(space_id, private=private)
314
+ wait_until_space_exists(space_id)
315
+
316
+ client = Client(space_id, verbose=False, httpx_kwargs={"timeout": 90})
317
+ hf_token = huggingface_hub.utils.get_token()
318
+
319
+ if pending_only:
320
+ pending_logs = SQLiteStorage.get_pending_logs(project)
321
+ if pending_logs:
322
+ logs = pending_logs["logs"]
323
+ for i in range(0, len(logs), SYNC_BATCH_SIZE):
324
+ batch = logs[i : i + SYNC_BATCH_SIZE]
325
+ print(
326
+ f" Syncing metrics: {min(i + SYNC_BATCH_SIZE, len(logs))}/{len(logs)}..."
327
+ )
328
+ client.predict(api_name="/bulk_log", logs=batch, hf_token=hf_token)
329
+ SQLiteStorage.clear_pending_logs(project, pending_logs["ids"])
330
+
331
+ pending_sys = SQLiteStorage.get_pending_system_logs(project)
332
+ if pending_sys:
333
+ logs = pending_sys["logs"]
334
+ for i in range(0, len(logs), SYNC_BATCH_SIZE):
335
+ batch = logs[i : i + SYNC_BATCH_SIZE]
336
+ print(
337
+ f" Syncing system metrics: {min(i + SYNC_BATCH_SIZE, len(logs))}/{len(logs)}..."
338
+ )
339
+ client.predict(
340
+ api_name="/bulk_log_system", logs=batch, hf_token=hf_token
341
+ )
342
+ SQLiteStorage.clear_pending_system_logs(project, pending_sys["ids"])
343
+
344
+ pending_uploads = SQLiteStorage.get_pending_uploads(project)
345
+ if pending_uploads:
346
+ upload_entries = []
347
+ for u in pending_uploads["uploads"]:
348
+ fp = u["file_path"]
349
+ if os.path.exists(fp):
350
+ upload_entries.append(
351
+ {
352
+ "project": u["project"],
353
+ "run": u["run"],
354
+ "step": u["step"],
355
+ "relative_path": u["relative_path"],
356
+ "uploaded_file": handle_file(fp),
357
+ }
358
+ )
359
+ if upload_entries:
360
+ print(f" Syncing {len(upload_entries)} media files...")
361
+ client.predict(
362
+ api_name="/bulk_upload_media",
363
+ uploads=upload_entries,
364
+ hf_token=hf_token,
365
+ )
366
+ SQLiteStorage.clear_pending_uploads(project, pending_uploads["ids"])
367
+ else:
368
+ all_logs = SQLiteStorage.get_all_logs_for_sync(project)
369
+ if all_logs:
370
+ for i in range(0, len(all_logs), SYNC_BATCH_SIZE):
371
+ batch = all_logs[i : i + SYNC_BATCH_SIZE]
372
+ print(
373
+ f" Syncing metrics: {min(i + SYNC_BATCH_SIZE, len(all_logs))}/{len(all_logs)}..."
374
+ )
375
+ client.predict(api_name="/bulk_log", logs=batch, hf_token=hf_token)
376
+
377
+ all_sys_logs = SQLiteStorage.get_all_system_logs_for_sync(project)
378
+ if all_sys_logs:
379
+ for i in range(0, len(all_sys_logs), SYNC_BATCH_SIZE):
380
+ batch = all_sys_logs[i : i + SYNC_BATCH_SIZE]
381
+ print(
382
+ f" Syncing system metrics: {min(i + SYNC_BATCH_SIZE, len(all_sys_logs))}/{len(all_sys_logs)}..."
383
+ )
384
+ client.predict(
385
+ api_name="/bulk_log_system", logs=batch, hf_token=hf_token
386
+ )
387
+
388
+ SQLiteStorage.set_project_metadata(project, "space_id", space_id)
389
+ print(f"* Synced successfully to space: {SPACE_URL.format(space_id=space_id)}")
390
+
391
+
392
+ def sync(
393
+ project: str,
394
+ space_id: str | None = None,
395
+ private: bool | None = None,
396
+ force: bool = False,
397
+ run_in_background: bool = False,
398
+ ) -> str:
399
+ """
400
+ Syncs a local Trackio project's database to a Hugging Face Space.
401
+ If the Space does not exist, it will be created.
402
+
403
+ Args:
404
+ project (`str`): The name of the project to upload.
405
+ space_id (`str`, *optional*): The ID of the Space to upload to (e.g., `"username/space_id"`).
406
+ If not provided, checks project metadata first, then generates a random space_id.
407
+ private (`bool`, *optional*):
408
+ Whether to make the Space private. If None (default), the repo will be
409
+ public unless the organization's default is private. This value is ignored
410
+ if the repo already exists.
411
+ force (`bool`, *optional*, defaults to `False`):
412
+ If `True`, overwrite the existing database without prompting for confirmation.
413
+ If `False`, prompt the user before overwriting an existing database.
414
+ run_in_background (`bool`, *optional*, defaults to `False`):
415
+ If `True`, the Space creation and database upload will be run in a background thread.
416
+ If `False`, all the steps will be run synchronously.
417
+ Returns:
418
+ `str`: The Space ID of the synced project.
419
+ """
420
+ if space_id is None:
421
+ space_id = SQLiteStorage.get_space_id(project)
422
+ if space_id is None:
423
+ space_id = f"{project}-{get_or_create_project_hash(project)}"
424
+ space_id, _ = preprocess_space_and_dataset_ids(space_id, None)
425
+
426
+ def _do_sync(space_id: str, private: bool | None = None):
427
+ sync_incremental(project, space_id, private=private, pending_only=False)
428
+
429
+ if run_in_background:
430
+ threading.Thread(target=_do_sync, args=(space_id, private)).start()
431
+ else:
432
+ _do_sync(space_id, private)
433
+ return space_id
trackio/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()
trackio/gpu.py ADDED
@@ -0,0 +1,357 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import threading
3
+ import warnings
4
+ from typing import TYPE_CHECKING, Any
5
+
6
+ if TYPE_CHECKING:
7
+ from trackio.run import Run
8
+
9
+ pynvml: Any = None
10
+ PYNVML_AVAILABLE = False
11
+ _nvml_initialized = False
12
+ _nvml_lock = threading.Lock()
13
+ _energy_baseline: dict[int, float] = {}
14
+
15
+
16
+ def _ensure_pynvml():
17
+ global PYNVML_AVAILABLE, pynvml
18
+ if PYNVML_AVAILABLE:
19
+ return pynvml
20
+ try:
21
+ import pynvml as _pynvml
22
+
23
+ pynvml = _pynvml
24
+ PYNVML_AVAILABLE = True
25
+ return pynvml
26
+ except ImportError:
27
+ raise ImportError(
28
+ "nvidia-ml-py is required for GPU monitoring. "
29
+ "Install it with: pip install nvidia-ml-py"
30
+ )
31
+
32
+
33
+ def _init_nvml() -> bool:
34
+ global _nvml_initialized
35
+ with _nvml_lock:
36
+ if _nvml_initialized:
37
+ return True
38
+ try:
39
+ nvml = _ensure_pynvml()
40
+ nvml.nvmlInit()
41
+ _nvml_initialized = True
42
+ return True
43
+ except Exception:
44
+ return False
45
+
46
+
47
+ def get_gpu_count() -> tuple[int, list[int]]:
48
+ """
49
+ Get the number of GPUs visible to this process and their physical indices.
50
+ Respects CUDA_VISIBLE_DEVICES environment variable.
51
+
52
+ Returns:
53
+ Tuple of (count, physical_indices) where:
54
+ - count: Number of visible GPUs
55
+ - physical_indices: List mapping logical index to physical GPU index.
56
+ e.g., if CUDA_VISIBLE_DEVICES=2,3 returns (2, [2, 3])
57
+ meaning logical GPU 0 = physical GPU 2, logical GPU 1 = physical GPU 3
58
+ """
59
+ if not _init_nvml():
60
+ return 0, []
61
+
62
+ cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES")
63
+ if cuda_visible is not None and cuda_visible.strip():
64
+ try:
65
+ indices = [int(x.strip()) for x in cuda_visible.split(",") if x.strip()]
66
+ return len(indices), indices
67
+ except ValueError:
68
+ pass
69
+
70
+ try:
71
+ total = pynvml.nvmlDeviceGetCount()
72
+ return total, list(range(total))
73
+ except Exception:
74
+ return 0, []
75
+
76
+
77
+ def gpu_available() -> bool:
78
+ """
79
+ Check if GPU monitoring is available.
80
+
81
+ Returns True if nvidia-ml-py is installed and at least one NVIDIA GPU is detected.
82
+ This is used for auto-detection of GPU logging.
83
+ """
84
+ try:
85
+ _ensure_pynvml()
86
+ count, _ = get_gpu_count()
87
+ return count > 0
88
+ except ImportError:
89
+ return False
90
+ except Exception:
91
+ return False
92
+
93
+
94
+ def reset_energy_baseline():
95
+ """Reset the energy baseline for all GPUs. Called when a new run starts."""
96
+ global _energy_baseline
97
+ _energy_baseline = {}
98
+
99
+
100
+ def collect_gpu_metrics(device: int | None = None) -> dict:
101
+ """
102
+ Collect GPU metrics for visible GPUs.
103
+
104
+ Args:
105
+ device: CUDA device index to collect metrics from. If None, collects
106
+ from all GPUs visible to this process (respects CUDA_VISIBLE_DEVICES).
107
+ The device index is the logical CUDA index (0, 1, 2...), not the
108
+ physical GPU index.
109
+
110
+ Returns:
111
+ Dictionary of GPU metrics. Keys use logical device indices (gpu/0/, gpu/1/, etc.)
112
+ which correspond to CUDA device indices, not physical GPU indices.
113
+ """
114
+ if not _init_nvml():
115
+ return {}
116
+
117
+ gpu_count, visible_gpus = get_gpu_count()
118
+ if gpu_count == 0:
119
+ return {}
120
+
121
+ if device is not None:
122
+ if device < 0 or device >= gpu_count:
123
+ return {}
124
+ gpu_indices = [(device, visible_gpus[device])]
125
+ else:
126
+ gpu_indices = list(enumerate(visible_gpus))
127
+
128
+ metrics = {}
129
+ total_util = 0.0
130
+ total_mem_used_gib = 0.0
131
+ total_power = 0.0
132
+ max_temp = 0.0
133
+ valid_util_count = 0
134
+
135
+ for logical_idx, physical_idx in gpu_indices:
136
+ prefix = f"gpu/{logical_idx}"
137
+ try:
138
+ handle = pynvml.nvmlDeviceGetHandleByIndex(physical_idx)
139
+
140
+ try:
141
+ util = pynvml.nvmlDeviceGetUtilizationRates(handle)
142
+ metrics[f"{prefix}/utilization"] = util.gpu
143
+ metrics[f"{prefix}/memory_utilization"] = util.memory
144
+ total_util += util.gpu
145
+ valid_util_count += 1
146
+ except Exception:
147
+ pass
148
+
149
+ try:
150
+ mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
151
+ mem_used_gib = mem.used / (1024**3)
152
+ mem_total_gib = mem.total / (1024**3)
153
+ metrics[f"{prefix}/allocated_memory"] = mem_used_gib
154
+ metrics[f"{prefix}/total_memory"] = mem_total_gib
155
+ if mem.total > 0:
156
+ metrics[f"{prefix}/memory_usage"] = mem.used / mem.total
157
+ total_mem_used_gib += mem_used_gib
158
+ except Exception:
159
+ pass
160
+
161
+ try:
162
+ power_mw = pynvml.nvmlDeviceGetPowerUsage(handle)
163
+ power_w = power_mw / 1000.0
164
+ metrics[f"{prefix}/power"] = power_w
165
+ total_power += power_w
166
+ except Exception:
167
+ pass
168
+
169
+ try:
170
+ power_limit_mw = pynvml.nvmlDeviceGetPowerManagementLimit(handle)
171
+ power_limit_w = power_limit_mw / 1000.0
172
+ metrics[f"{prefix}/power_limit"] = power_limit_w
173
+ if power_limit_w > 0 and f"{prefix}/power" in metrics:
174
+ metrics[f"{prefix}/power_percent"] = (
175
+ metrics[f"{prefix}/power"] / power_limit_w
176
+ ) * 100
177
+ except Exception:
178
+ pass
179
+
180
+ try:
181
+ temp = pynvml.nvmlDeviceGetTemperature(
182
+ handle, pynvml.NVML_TEMPERATURE_GPU
183
+ )
184
+ metrics[f"{prefix}/temp"] = temp
185
+ max_temp = max(max_temp, temp)
186
+ except Exception:
187
+ pass
188
+
189
+ try:
190
+ sm_clock = pynvml.nvmlDeviceGetClockInfo(handle, pynvml.NVML_CLOCK_SM)
191
+ metrics[f"{prefix}/sm_clock"] = sm_clock
192
+ except Exception:
193
+ pass
194
+
195
+ try:
196
+ mem_clock = pynvml.nvmlDeviceGetClockInfo(handle, pynvml.NVML_CLOCK_MEM)
197
+ metrics[f"{prefix}/memory_clock"] = mem_clock
198
+ except Exception:
199
+ pass
200
+
201
+ try:
202
+ fan_speed = pynvml.nvmlDeviceGetFanSpeed(handle)
203
+ metrics[f"{prefix}/fan_speed"] = fan_speed
204
+ except Exception:
205
+ pass
206
+
207
+ try:
208
+ pstate = pynvml.nvmlDeviceGetPerformanceState(handle)
209
+ metrics[f"{prefix}/performance_state"] = pstate
210
+ except Exception:
211
+ pass
212
+
213
+ try:
214
+ energy_mj = pynvml.nvmlDeviceGetTotalEnergyConsumption(handle)
215
+ if logical_idx not in _energy_baseline:
216
+ _energy_baseline[logical_idx] = energy_mj
217
+ energy_consumed_mj = energy_mj - _energy_baseline[logical_idx]
218
+ metrics[f"{prefix}/energy_consumed"] = energy_consumed_mj / 1000.0
219
+ except Exception:
220
+ pass
221
+
222
+ try:
223
+ pcie_tx = pynvml.nvmlDeviceGetPcieThroughput(
224
+ handle, pynvml.NVML_PCIE_UTIL_TX_BYTES
225
+ )
226
+ pcie_rx = pynvml.nvmlDeviceGetPcieThroughput(
227
+ handle, pynvml.NVML_PCIE_UTIL_RX_BYTES
228
+ )
229
+ metrics[f"{prefix}/pcie_tx"] = pcie_tx / 1024.0
230
+ metrics[f"{prefix}/pcie_rx"] = pcie_rx / 1024.0
231
+ except Exception:
232
+ pass
233
+
234
+ try:
235
+ throttle = pynvml.nvmlDeviceGetCurrentClocksThrottleReasons(handle)
236
+ metrics[f"{prefix}/throttle_thermal"] = int(
237
+ bool(throttle & pynvml.nvmlClocksThrottleReasonSwThermalSlowdown)
238
+ )
239
+ metrics[f"{prefix}/throttle_power"] = int(
240
+ bool(throttle & pynvml.nvmlClocksThrottleReasonSwPowerCap)
241
+ )
242
+ metrics[f"{prefix}/throttle_hw_slowdown"] = int(
243
+ bool(throttle & pynvml.nvmlClocksThrottleReasonHwSlowdown)
244
+ )
245
+ metrics[f"{prefix}/throttle_apps"] = int(
246
+ bool(
247
+ throttle
248
+ & pynvml.nvmlClocksThrottleReasonApplicationsClocksSetting
249
+ )
250
+ )
251
+ except Exception:
252
+ pass
253
+
254
+ try:
255
+ ecc_corrected = pynvml.nvmlDeviceGetTotalEccErrors(
256
+ handle,
257
+ pynvml.NVML_MEMORY_ERROR_TYPE_CORRECTED,
258
+ pynvml.NVML_VOLATILE_ECC,
259
+ )
260
+ metrics[f"{prefix}/corrected_memory_errors"] = ecc_corrected
261
+ except Exception:
262
+ pass
263
+
264
+ try:
265
+ ecc_uncorrected = pynvml.nvmlDeviceGetTotalEccErrors(
266
+ handle,
267
+ pynvml.NVML_MEMORY_ERROR_TYPE_UNCORRECTED,
268
+ pynvml.NVML_VOLATILE_ECC,
269
+ )
270
+ metrics[f"{prefix}/uncorrected_memory_errors"] = ecc_uncorrected
271
+ except Exception:
272
+ pass
273
+
274
+ except Exception:
275
+ continue
276
+
277
+ if valid_util_count > 0:
278
+ metrics["gpu/mean_utilization"] = total_util / valid_util_count
279
+ if total_mem_used_gib > 0:
280
+ metrics["gpu/total_allocated_memory"] = total_mem_used_gib
281
+ if total_power > 0:
282
+ metrics["gpu/total_power"] = total_power
283
+ if max_temp > 0:
284
+ metrics["gpu/max_temp"] = max_temp
285
+
286
+ return metrics
287
+
288
+
289
+ class GpuMonitor:
290
+ def __init__(self, run: "Run", interval: float = 10.0):
291
+ self._run = run
292
+ self._interval = interval
293
+ self._stop_flag = threading.Event()
294
+ self._thread: "threading.Thread | None" = None
295
+
296
+ def start(self):
297
+ count, _ = get_gpu_count()
298
+ if count == 0:
299
+ warnings.warn(
300
+ "auto_log_gpu=True but no NVIDIA GPUs detected. GPU logging disabled."
301
+ )
302
+ return
303
+
304
+ reset_energy_baseline()
305
+ self._thread = threading.Thread(target=self._monitor_loop, daemon=True)
306
+ self._thread.start()
307
+
308
+ def stop(self):
309
+ self._stop_flag.set()
310
+ if self._thread is not None:
311
+ self._thread.join(timeout=2.0)
312
+
313
+ def _monitor_loop(self):
314
+ while not self._stop_flag.is_set():
315
+ try:
316
+ metrics = collect_gpu_metrics()
317
+ if metrics:
318
+ self._run.log_system(metrics)
319
+ except Exception:
320
+ pass
321
+
322
+ self._stop_flag.wait(timeout=self._interval)
323
+
324
+
325
+ def log_gpu(run: "Run | None" = None, device: int | None = None) -> dict:
326
+ """
327
+ Log GPU metrics to the current or specified run as system metrics.
328
+
329
+ Args:
330
+ run: Optional Run instance. If None, uses current run from context.
331
+ device: CUDA device index to collect metrics from. If None, collects
332
+ from all GPUs visible to this process (respects CUDA_VISIBLE_DEVICES).
333
+
334
+ Returns:
335
+ dict: The GPU metrics that were logged.
336
+
337
+ Example:
338
+ ```python
339
+ import trackio
340
+
341
+ run = trackio.init(project="my-project")
342
+ trackio.log({"loss": 0.5})
343
+ trackio.log_gpu() # logs all visible GPUs
344
+ trackio.log_gpu(device=0) # logs only CUDA device 0
345
+ ```
346
+ """
347
+ from trackio import context_vars
348
+
349
+ if run is None:
350
+ run = context_vars.current_run.get()
351
+ if run is None:
352
+ raise RuntimeError("Call trackio.init() before trackio.log_gpu().")
353
+
354
+ metrics = collect_gpu_metrics(device=device)
355
+ if metrics:
356
+ run.log_system(metrics)
357
+ return metrics
trackio/histogram.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Sequence
2
+
3
+ import numpy as np
4
+
5
+
6
+ class Histogram:
7
+ """
8
+ Histogram data type for Trackio, compatible with wandb.Histogram.
9
+
10
+ Args:
11
+ sequence (`np.ndarray` or `Sequence[float]` or `Sequence[int]`, *optional*):
12
+ Sequence of values to create the histogram from.
13
+ np_histogram (`tuple`, *optional*):
14
+ Pre-computed NumPy histogram as a `(hist, bins)` tuple.
15
+ num_bins (`int`, *optional*, defaults to `64`):
16
+ Number of bins for the histogram (maximum `512`).
17
+
18
+ Example:
19
+ ```python
20
+ import trackio
21
+ import numpy as np
22
+
23
+ # Create histogram from sequence
24
+ data = np.random.randn(1000)
25
+ trackio.log({"distribution": trackio.Histogram(data)})
26
+
27
+ # Create histogram from numpy histogram
28
+ hist, bins = np.histogram(data, bins=30)
29
+ trackio.log({"distribution": trackio.Histogram(np_histogram=(hist, bins))})
30
+
31
+ # Specify custom number of bins
32
+ trackio.log({"distribution": trackio.Histogram(data, num_bins=50)})
33
+ ```
34
+ """
35
+
36
+ TYPE = "trackio.histogram"
37
+
38
+ def __init__(
39
+ self,
40
+ sequence: np.ndarray | Sequence[float] | Sequence[int] | None = None,
41
+ np_histogram: tuple | None = None,
42
+ num_bins: int = 64,
43
+ ):
44
+ if sequence is None and np_histogram is None:
45
+ raise ValueError("Must provide either sequence or np_histogram")
46
+
47
+ if sequence is not None and np_histogram is not None:
48
+ raise ValueError("Cannot provide both sequence and np_histogram")
49
+
50
+ num_bins = min(num_bins, 512)
51
+
52
+ if np_histogram is not None:
53
+ self.histogram, self.bins = np_histogram
54
+ self.histogram = np.asarray(self.histogram)
55
+ self.bins = np.asarray(self.bins)
56
+ else:
57
+ data = np.asarray(sequence).flatten()
58
+ data = data[np.isfinite(data)]
59
+ if len(data) == 0:
60
+ self.histogram = np.array([])
61
+ self.bins = np.array([])
62
+ else:
63
+ self.histogram, self.bins = np.histogram(data, bins=num_bins)
64
+
65
+ def _to_dict(self) -> dict:
66
+ """Convert histogram to dictionary for storage."""
67
+ return {
68
+ "_type": self.TYPE,
69
+ "bins": self.bins.tolist(),
70
+ "values": self.histogram.tolist(),
71
+ }
trackio/imports.py ADDED
@@ -0,0 +1,304 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ force: bool = False,
18
+ ) -> None:
19
+ """
20
+ Imports a CSV file into a Trackio project. The CSV file must contain a `"step"`
21
+ column, may optionally contain a `"timestamp"` column, and any other columns will be
22
+ treated as metrics. It should also include a header row with the column names.
23
+
24
+ TODO: call init() and return a Run object so that the user can continue to log metrics to it.
25
+
26
+ Args:
27
+ csv_path (`str` or `Path`):
28
+ The str or Path to the CSV file to import.
29
+ project (`str`):
30
+ The name of the project to import the CSV file into. Must not be an existing
31
+ project.
32
+ name (`str`, *optional*):
33
+ The name of the Run to import the CSV file into. If not provided, a default
34
+ name will be generated.
35
+ name (`str`, *optional*):
36
+ The name of the run (if not provided, a default name will be generated).
37
+ space_id (`str`, *optional*):
38
+ If provided, the project will be logged to a Hugging Face Space instead of a
39
+ local directory. Should be a complete Space name like `"username/reponame"`
40
+ or `"orgname/reponame"`, or just `"reponame"` in which case the Space will
41
+ be created in the currently-logged-in Hugging Face user's namespace. If the
42
+ Space does not exist, it will be created. If the Space already exists, the
43
+ project will be logged to it.
44
+ dataset_id (`str`, *optional*):
45
+ If provided, a persistent Hugging Face Dataset will be created and the
46
+ metrics will be synced to it every 5 minutes. Should be a complete Dataset
47
+ name like `"username/datasetname"` or `"orgname/datasetname"`, or just
48
+ `"datasetname"` in which case the Dataset will be created in the
49
+ currently-logged-in Hugging Face user's namespace. If the Dataset does not
50
+ exist, it will be created. If the Dataset already exists, the project will
51
+ be appended to it. If not provided, the metrics will be logged to a local
52
+ SQLite database, unless a `space_id` is provided, in which case a Dataset
53
+ will be automatically created with the same name as the Space but with the
54
+ `"_dataset"` suffix.
55
+ private (`bool`, *optional*):
56
+ Whether to make the Space private. If None (default), the repo will be
57
+ public unless the organization's default is private. This value is ignored
58
+ if the repo already exists.
59
+ """
60
+ if SQLiteStorage.get_runs(project):
61
+ raise ValueError(
62
+ f"Project '{project}' already exists. Cannot import CSV into existing project."
63
+ )
64
+
65
+ csv_path = Path(csv_path)
66
+ if not csv_path.exists():
67
+ raise FileNotFoundError(f"CSV file not found: {csv_path}")
68
+
69
+ df = pd.read_csv(csv_path)
70
+ if df.empty:
71
+ raise ValueError("CSV file is empty")
72
+
73
+ column_mapping = utils.simplify_column_names(df.columns.tolist())
74
+ df = df.rename(columns=column_mapping)
75
+
76
+ step_column = None
77
+ for col in df.columns:
78
+ if col.lower() == "step":
79
+ step_column = col
80
+ break
81
+
82
+ if step_column is None:
83
+ raise ValueError("CSV file must contain a 'step' or 'Step' column")
84
+
85
+ if name is None:
86
+ name = csv_path.stem
87
+
88
+ metrics_list = []
89
+ steps = []
90
+ timestamps = []
91
+
92
+ numeric_columns = []
93
+ for column in df.columns:
94
+ if column == step_column:
95
+ continue
96
+ if column == "timestamp":
97
+ continue
98
+
99
+ try:
100
+ pd.to_numeric(df[column], errors="raise")
101
+ numeric_columns.append(column)
102
+ except (ValueError, TypeError):
103
+ continue
104
+
105
+ for _, row in df.iterrows():
106
+ metrics = {}
107
+ for column in numeric_columns:
108
+ value = row[column]
109
+ if bool(pd.notna(value)):
110
+ metrics[column] = float(value)
111
+
112
+ if metrics:
113
+ metrics_list.append(metrics)
114
+ steps.append(int(row[step_column]))
115
+
116
+ if "timestamp" in df.columns and bool(pd.notna(row["timestamp"])):
117
+ timestamps.append(str(row["timestamp"]))
118
+ else:
119
+ timestamps.append("")
120
+
121
+ if metrics_list:
122
+ SQLiteStorage.bulk_log(
123
+ project=project,
124
+ run=name,
125
+ metrics_list=metrics_list,
126
+ steps=steps,
127
+ timestamps=timestamps,
128
+ )
129
+
130
+ print(
131
+ f"* Imported {len(metrics_list)} rows from {csv_path} into project '{project}' as run '{name}'"
132
+ )
133
+ print(f"* Metrics found: {', '.join(metrics_list[0].keys())}")
134
+
135
+ space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)
136
+ if dataset_id is not None:
137
+ os.environ["TRACKIO_DATASET_ID"] = dataset_id
138
+ print(f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}")
139
+
140
+ if space_id is None:
141
+ utils.print_dashboard_instructions(project)
142
+ else:
143
+ deploy.create_space_if_not_exists(
144
+ space_id=space_id, dataset_id=dataset_id, private=private
145
+ )
146
+ deploy.wait_until_space_exists(space_id=space_id)
147
+ deploy.upload_db_to_space(project=project, space_id=space_id, force=force)
148
+ print(
149
+ f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"
150
+ )
151
+
152
+
153
+ def import_tf_events(
154
+ log_dir: str | Path,
155
+ project: str,
156
+ name: str | None = None,
157
+ space_id: str | None = None,
158
+ dataset_id: str | None = None,
159
+ private: bool | None = None,
160
+ force: bool = False,
161
+ ) -> None:
162
+ """
163
+ Imports TensorFlow Events files from a directory into a Trackio project. Each
164
+ subdirectory in the log directory will be imported as a separate run.
165
+
166
+ Args:
167
+ log_dir (`str` or `Path`):
168
+ The str or Path to the directory containing TensorFlow Events files.
169
+ project (`str`):
170
+ The name of the project to import the TensorFlow Events files into. Must not
171
+ be an existing project.
172
+ name (`str`, *optional*):
173
+ The name prefix for runs (if not provided, will use directory names). Each
174
+ subdirectory will create a separate run.
175
+ space_id (`str`, *optional*):
176
+ If provided, the project will be logged to a Hugging Face Space instead of a
177
+ local directory. Should be a complete Space name like `"username/reponame"`
178
+ or `"orgname/reponame"`, or just `"reponame"` in which case the Space will
179
+ be created in the currently-logged-in Hugging Face user's namespace. If the
180
+ Space does not exist, it will be created. If the Space already exists, the
181
+ project will be logged to it.
182
+ dataset_id (`str`, *optional*):
183
+ If provided, a persistent Hugging Face Dataset will be created and the
184
+ metrics will be synced to it every 5 minutes. Should be a complete Dataset
185
+ name like `"username/datasetname"` or `"orgname/datasetname"`, or just
186
+ `"datasetname"` in which case the Dataset will be created in the
187
+ currently-logged-in Hugging Face user's namespace. If the Dataset does not
188
+ exist, it will be created. If the Dataset already exists, the project will
189
+ be appended to it. If not provided, the metrics will be logged to a local
190
+ SQLite database, unless a `space_id` is provided, in which case a Dataset
191
+ will be automatically created with the same name as the Space but with the
192
+ `"_dataset"` suffix.
193
+ private (`bool`, *optional*):
194
+ Whether to make the Space private. If None (default), the repo will be
195
+ public unless the organization's default is private. This value is ignored
196
+ if the repo already exists.
197
+ """
198
+ try:
199
+ from tbparse import SummaryReader
200
+ except ImportError:
201
+ raise ImportError(
202
+ "The `tbparse` package is not installed but is required for `import_tf_events`. Please install trackio with the `tensorboard` extra: `pip install trackio[tensorboard]`."
203
+ )
204
+
205
+ if SQLiteStorage.get_runs(project):
206
+ raise ValueError(
207
+ f"Project '{project}' already exists. Cannot import TF events into existing project."
208
+ )
209
+
210
+ path = Path(log_dir)
211
+ if not path.exists():
212
+ raise FileNotFoundError(f"TF events directory not found: {path}")
213
+
214
+ # Use tbparse to read all tfevents files in the directory structure
215
+ reader = SummaryReader(str(path), extra_columns={"dir_name"})
216
+ df = reader.scalars
217
+
218
+ if df.empty:
219
+ raise ValueError(f"No TensorFlow events data found in {path}")
220
+
221
+ total_imported = 0
222
+ imported_runs = []
223
+
224
+ # Group by dir_name to create separate runs
225
+ for dir_name, group_df in df.groupby("dir_name"):
226
+ try:
227
+ # Determine run name based on directory name
228
+ if dir_name == "":
229
+ run_name = "main" # For files in the root directory
230
+ else:
231
+ run_name = dir_name # Use directory name
232
+
233
+ if name:
234
+ run_name = f"{name}_{run_name}"
235
+
236
+ if group_df.empty:
237
+ print(f"* Skipping directory {dir_name}: no scalar data found")
238
+ continue
239
+
240
+ metrics_list = []
241
+ steps = []
242
+ timestamps = []
243
+
244
+ for _, row in group_df.iterrows():
245
+ # Convert row values to appropriate types
246
+ tag = str(row["tag"])
247
+ value = float(row["value"])
248
+ step = int(row["step"])
249
+
250
+ metrics = {tag: value}
251
+ metrics_list.append(metrics)
252
+ steps.append(step)
253
+
254
+ # Use wall_time if present, else fallback
255
+ if "wall_time" in group_df.columns and not bool(
256
+ pd.isna(row["wall_time"])
257
+ ):
258
+ timestamps.append(str(row["wall_time"]))
259
+ else:
260
+ timestamps.append("")
261
+
262
+ if metrics_list:
263
+ SQLiteStorage.bulk_log(
264
+ project=project,
265
+ run=str(run_name),
266
+ metrics_list=metrics_list,
267
+ steps=steps,
268
+ timestamps=timestamps,
269
+ )
270
+
271
+ total_imported += len(metrics_list)
272
+ imported_runs.append(run_name)
273
+
274
+ print(
275
+ f"* Imported {len(metrics_list)} scalar events from directory '{dir_name}' as run '{run_name}'"
276
+ )
277
+ print(f"* Metrics in this run: {', '.join(set(group_df['tag']))}")
278
+
279
+ except Exception as e:
280
+ print(f"* Error processing directory {dir_name}: {e}")
281
+ continue
282
+
283
+ if not imported_runs:
284
+ raise ValueError("No valid TensorFlow events data could be imported")
285
+
286
+ print(f"* Total imported events: {total_imported}")
287
+ print(f"* Created runs: {', '.join(imported_runs)}")
288
+
289
+ space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)
290
+ if dataset_id is not None:
291
+ os.environ["TRACKIO_DATASET_ID"] = dataset_id
292
+ print(f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}")
293
+
294
+ if space_id is None:
295
+ utils.print_dashboard_instructions(project)
296
+ else:
297
+ deploy.create_space_if_not_exists(
298
+ space_id, dataset_id=dataset_id, private=private
299
+ )
300
+ deploy.wait_until_space_exists(space_id)
301
+ deploy.upload_db_to_space(project, space_id, force=force)
302
+ print(
303
+ f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"
304
+ )
trackio/media/__init__.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Media module for Trackio.
3
+
4
+ This module contains all media-related functionality including:
5
+ - TrackioImage, TrackioVideo, TrackioAudio classes
6
+ - Video writing utilities
7
+ - Audio conversion utilities
8
+ """
9
+
10
+ from trackio.media.audio import TrackioAudio
11
+ from trackio.media.image import TrackioImage
12
+ from trackio.media.media import TrackioMedia
13
+ from trackio.media.utils import get_project_media_path
14
+ from trackio.media.video import TrackioVideo
15
+
16
+ write_audio = TrackioAudio.write_audio
17
+ write_video = TrackioVideo.write_video
18
+
19
+ __all__ = [
20
+ "TrackioMedia",
21
+ "TrackioImage",
22
+ "TrackioVideo",
23
+ "TrackioAudio",
24
+ "get_project_media_path",
25
+ "write_video",
26
+ "write_audio",
27
+ ]
trackio/media/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (755 Bytes). View file
 
trackio/media/__pycache__/audio.cpython-310.pyc ADDED
Binary file (5.61 kB). View file
 
trackio/media/__pycache__/image.cpython-310.pyc ADDED
Binary file (3.1 kB). View file
 
trackio/media/__pycache__/media.cpython-310.pyc ADDED
Binary file (3.1 kB). View file
 
trackio/media/__pycache__/utils.cpython-310.pyc ADDED
Binary file (2.01 kB). View file
 
trackio/media/__pycache__/video.cpython-310.pyc ADDED
Binary file (7 kB). View file
 
trackio/media/audio.py ADDED
@@ -0,0 +1,167 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import shutil
3
+ import warnings
4
+ from pathlib import Path
5
+ from typing import Literal
6
+
7
+ import numpy as np
8
+ from pydub import AudioSegment
9
+
10
+ from trackio.media.media import TrackioMedia
11
+ from trackio.media.utils import check_ffmpeg_installed, check_path
12
+
13
+ SUPPORTED_FORMATS = ["wav", "mp3"]
14
+ AudioFormatType = Literal["wav", "mp3"]
15
+ TrackioAudioSourceType = str | Path | np.ndarray
16
+
17
+
18
+ class TrackioAudio(TrackioMedia):
19
+ """
20
+ Initializes an Audio object.
21
+
22
+ Example:
23
+ ```python
24
+ import trackio
25
+ import numpy as np
26
+
27
+ # Generate a 1-second 440 Hz sine wave (mono)
28
+ sr = 16000
29
+ t = np.linspace(0, 1, sr, endpoint=False)
30
+ wave = 0.2 * np.sin(2 * np.pi * 440 * t)
31
+ audio = trackio.Audio(wave, caption="A4 sine", sample_rate=sr, format="wav")
32
+ trackio.log({"tone": audio})
33
+
34
+ # Stereo from numpy array (shape: samples, 2)
35
+ stereo = np.stack([wave, wave], axis=1)
36
+ audio = trackio.Audio(stereo, caption="Stereo", sample_rate=sr, format="mp3")
37
+ trackio.log({"stereo": audio})
38
+
39
+ # From an existing file
40
+ audio = trackio.Audio("path/to/audio.wav", caption="From file")
41
+ trackio.log({"file_audio": audio})
42
+ ```
43
+
44
+ Args:
45
+ value (`str`, `Path`, or `numpy.ndarray`, *optional*):
46
+ A path to an audio file, or a numpy array.
47
+ The array should be shaped `(samples,)` for mono or `(samples, 2)` for stereo.
48
+ Float arrays will be peak-normalized and converted to 16-bit PCM; integer arrays will be converted to 16-bit PCM as needed.
49
+ caption (`str`, *optional*):
50
+ A string caption for the audio.
51
+ sample_rate (`int`, *optional*):
52
+ Sample rate in Hz. Required when `value` is a numpy array.
53
+ format (`Literal["wav", "mp3"]`, *optional*):
54
+ Audio format used when `value` is a numpy array. Default is "wav".
55
+ """
56
+
57
+ TYPE = "trackio.audio"
58
+
59
+ def __init__(
60
+ self,
61
+ value: TrackioAudioSourceType,
62
+ caption: str | None = None,
63
+ sample_rate: int | None = None,
64
+ format: AudioFormatType | None = None,
65
+ ):
66
+ super().__init__(value, caption)
67
+ if isinstance(value, np.ndarray):
68
+ if sample_rate is None:
69
+ raise ValueError("Sample rate is required when value is an ndarray")
70
+ if format is None:
71
+ format = "wav"
72
+ self._format = format
73
+ self._sample_rate = sample_rate
74
+
75
+ def _save_media(self, file_path: Path):
76
+ if isinstance(self._value, np.ndarray):
77
+ TrackioAudio.write_audio(
78
+ data=self._value,
79
+ sample_rate=self._sample_rate,
80
+ filename=file_path,
81
+ format=self._format,
82
+ )
83
+ elif isinstance(self._value, str | Path):
84
+ if os.path.isfile(self._value):
85
+ shutil.copy(self._value, file_path)
86
+ else:
87
+ raise ValueError(f"File not found: {self._value}")
88
+
89
+ @staticmethod
90
+ def ensure_int16_pcm(data: np.ndarray) -> np.ndarray:
91
+ """
92
+ Convert input audio array to contiguous int16 PCM.
93
+ Peak normalization is applied to floating inputs.
94
+ """
95
+ arr = np.asarray(data)
96
+ if arr.ndim not in (1, 2):
97
+ raise ValueError("Audio data must be 1D (mono) or 2D ([samples, channels])")
98
+
99
+ if arr.dtype != np.int16:
100
+ warnings.warn(
101
+ f"Converting {arr.dtype} audio to int16 PCM; pass int16 to avoid conversion.",
102
+ stacklevel=2,
103
+ )
104
+
105
+ arr = np.nan_to_num(arr, copy=False)
106
+
107
+ # Floating types: normalize to peak 1.0, then scale to int16
108
+ if np.issubdtype(arr.dtype, np.floating):
109
+ max_abs = float(np.max(np.abs(arr))) if arr.size else 0.0
110
+ if max_abs > 0.0:
111
+ arr = arr / max_abs
112
+ out = (arr * 32767.0).clip(-32768, 32767).astype(np.int16, copy=False)
113
+ return np.ascontiguousarray(out)
114
+
115
+ converters: dict[np.dtype, callable] = {
116
+ np.dtype(np.int16): lambda a: a,
117
+ np.dtype(np.int32): lambda a: (a.astype(np.int32) // 65536).astype(
118
+ np.int16, copy=False
119
+ ),
120
+ np.dtype(np.uint16): lambda a: (a.astype(np.int32) - 32768).astype(
121
+ np.int16, copy=False
122
+ ),
123
+ np.dtype(np.uint8): lambda a: (a.astype(np.int32) * 257 - 32768).astype(
124
+ np.int16, copy=False
125
+ ),
126
+ np.dtype(np.int8): lambda a: (a.astype(np.int32) * 256).astype(
127
+ np.int16, copy=False
128
+ ),
129
+ }
130
+
131
+ conv = converters.get(arr.dtype)
132
+ if conv is not None:
133
+ out = conv(arr)
134
+ return np.ascontiguousarray(out)
135
+ raise TypeError(f"Unsupported audio dtype: {arr.dtype}")
136
+
137
+ @staticmethod
138
+ def write_audio(
139
+ data: np.ndarray,
140
+ sample_rate: int,
141
+ filename: str | Path,
142
+ format: AudioFormatType = "wav",
143
+ ) -> None:
144
+ if not isinstance(sample_rate, int) or sample_rate <= 0:
145
+ raise ValueError(f"Invalid sample_rate: {sample_rate}")
146
+ if format not in SUPPORTED_FORMATS:
147
+ raise ValueError(
148
+ f"Unsupported format: {format}. Supported: {SUPPORTED_FORMATS}"
149
+ )
150
+
151
+ check_path(filename)
152
+
153
+ pcm = TrackioAudio.ensure_int16_pcm(data)
154
+
155
+ if format != "wav":
156
+ check_ffmpeg_installed()
157
+
158
+ channels = 1 if pcm.ndim == 1 else pcm.shape[1]
159
+ audio = AudioSegment(
160
+ pcm.tobytes(),
161
+ frame_rate=sample_rate,
162
+ sample_width=2, # int16
163
+ channels=channels,
164
+ )
165
+
166
+ file = audio.export(str(filename), format=format)
167
+ file.close()
trackio/media/image.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import shutil
3
+ from pathlib import Path
4
+
5
+ import numpy as np
6
+ from PIL import Image as PILImage
7
+
8
+ from trackio.media.media import TrackioMedia
9
+
10
+ TrackioImageSourceType = str | Path | np.ndarray | PILImage.Image
11
+
12
+
13
+ class TrackioImage(TrackioMedia):
14
+ """
15
+ Initializes an Image object.
16
+
17
+ Example:
18
+ ```python
19
+ import trackio
20
+ import numpy as np
21
+ from PIL import Image
22
+
23
+ # Create an image from numpy array
24
+ image_data = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
25
+ image = trackio.Image(image_data, caption="Random image")
26
+ trackio.log({"my_image": image})
27
+
28
+ # Create an image from PIL Image
29
+ pil_image = Image.new('RGB', (100, 100), color='red')
30
+ image = trackio.Image(pil_image, caption="Red square")
31
+ trackio.log({"red_image": image})
32
+
33
+ # Create an image from file path
34
+ image = trackio.Image("path/to/image.jpg", caption="Photo from file")
35
+ trackio.log({"file_image": image})
36
+ ```
37
+
38
+ Args:
39
+ value (`str`, `Path`, `numpy.ndarray`, or `PIL.Image`, *optional*):
40
+ A path to an image, a PIL Image, or a numpy array of shape (height, width, channels).
41
+ If numpy array, should be of type `np.uint8` with RGB values in the range `[0, 255]`.
42
+ caption (`str`, *optional*):
43
+ A string caption for the image.
44
+ """
45
+
46
+ TYPE = "trackio.image"
47
+
48
+ def __init__(self, value: TrackioImageSourceType, caption: str | None = None):
49
+ super().__init__(value, caption)
50
+ self._format: str | None = None
51
+
52
+ if not isinstance(self._value, TrackioImageSourceType):
53
+ raise ValueError(
54
+ f"Invalid value type, expected {TrackioImageSourceType}, got {type(self._value)}"
55
+ )
56
+ if isinstance(self._value, np.ndarray) and self._value.dtype != np.uint8:
57
+ raise ValueError(
58
+ f"Invalid value dtype, expected np.uint8, got {self._value.dtype}"
59
+ )
60
+ if (
61
+ isinstance(self._value, np.ndarray | PILImage.Image)
62
+ and self._format is None
63
+ ):
64
+ self._format = "png"
65
+
66
+ def _as_pil(self) -> PILImage.Image | None:
67
+ try:
68
+ if isinstance(self._value, np.ndarray):
69
+ arr = np.asarray(self._value).astype("uint8")
70
+ return PILImage.fromarray(arr).convert("RGBA")
71
+ if isinstance(self._value, PILImage.Image):
72
+ return self._value.convert("RGBA")
73
+ except Exception as e:
74
+ raise ValueError(f"Failed to process image data: {self._value}") from e
75
+ return None
76
+
77
+ def _save_media(self, file_path: Path):
78
+ if pil := self._as_pil():
79
+ pil.save(file_path, format=self._format)
80
+ elif isinstance(self._value, str | Path):
81
+ if os.path.isfile(self._value):
82
+ shutil.copy(self._value, file_path)
83
+ else:
84
+ raise ValueError(f"File not found: {self._value}")
trackio/media/media.py ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import uuid
3
+ from abc import ABC, abstractmethod
4
+ from pathlib import Path
5
+
6
+ from trackio.media.utils import get_project_media_path
7
+ from trackio.utils import MEDIA_DIR
8
+
9
+
10
+ class TrackioMedia(ABC):
11
+ """
12
+ Abstract base class for Trackio media objects
13
+ Provides shared functionality for file handling and serialization.
14
+ """
15
+
16
+ TYPE: str
17
+
18
+ def __init_subclass__(cls, **kwargs):
19
+ """Ensure subclasses define the TYPE attribute."""
20
+ super().__init_subclass__(**kwargs)
21
+ if not hasattr(cls, "TYPE") or cls.TYPE is None:
22
+ raise TypeError(f"Class {cls.__name__} must define TYPE attribute")
23
+
24
+ def __init__(self, value, caption: str | None = None):
25
+ """
26
+ Saves the value and caption, and if the value is a file path, checks if the file exists.
27
+ """
28
+ self.caption = caption
29
+ self._value = value
30
+ self._file_path: Path | None = None
31
+
32
+ if isinstance(self._value, str | Path):
33
+ if not os.path.isfile(self._value):
34
+ raise ValueError(f"File not found: {self._value}")
35
+
36
+ def _file_extension(self) -> str:
37
+ if self._file_path:
38
+ return self._file_path.suffix[1:].lower()
39
+ if isinstance(self._value, str | Path):
40
+ path = Path(self._value)
41
+ return path.suffix[1:].lower()
42
+ if hasattr(self, "_format") and self._format:
43
+ return self._format
44
+ return "unknown"
45
+
46
+ def _get_relative_file_path(self) -> Path | None:
47
+ return self._file_path
48
+
49
+ def _get_absolute_file_path(self) -> Path | None:
50
+ if self._file_path:
51
+ return MEDIA_DIR / self._file_path
52
+ return None
53
+
54
+ def _save(self, project: str, run: str, step: int = 0):
55
+ if self._file_path:
56
+ return
57
+
58
+ media_dir = get_project_media_path(project=project, run=run, step=step)
59
+ filename = f"{uuid.uuid4()}.{self._file_extension()}"
60
+ file_path = media_dir / filename
61
+
62
+ self._save_media(file_path)
63
+ self._file_path = file_path.relative_to(MEDIA_DIR)
64
+
65
+ @abstractmethod
66
+ def _save_media(self, file_path: Path):
67
+ """
68
+ Performs the actual media saving logic.
69
+ """
70
+ pass
71
+
72
+ def _to_dict(self) -> dict:
73
+ if not self._file_path:
74
+ raise ValueError("Media must be saved to file before serialization")
75
+ return {
76
+ "_type": self.TYPE,
77
+ "file_path": str(self._get_relative_file_path()),
78
+ "caption": self.caption,
79
+ }
trackio/media/utils.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import shutil
2
+ from pathlib import Path
3
+
4
+ from trackio.utils import MEDIA_DIR
5
+
6
+
7
+ def check_path(file_path: str | Path) -> None:
8
+ """Raise an error if the parent directory does not exist."""
9
+ file_path = Path(file_path)
10
+ if not file_path.parent.exists():
11
+ try:
12
+ file_path.parent.mkdir(parents=True, exist_ok=True)
13
+ except OSError as e:
14
+ raise ValueError(
15
+ f"Failed to create parent directory {file_path.parent}: {e}"
16
+ )
17
+
18
+
19
+ def check_ffmpeg_installed() -> None:
20
+ """Raise an error if ffmpeg is not available on the system PATH."""
21
+ if shutil.which("ffmpeg") is None:
22
+ raise RuntimeError(
23
+ "ffmpeg is required to write video but was not found on your system. "
24
+ "Please install ffmpeg and ensure it is available on your PATH."
25
+ )
26
+
27
+
28
+ def get_project_media_path(
29
+ project: str,
30
+ run: str | None = None,
31
+ step: int | None = None,
32
+ relative_path: str | Path | None = None,
33
+ ) -> Path:
34
+ """
35
+ Get the full path where uploaded files are stored for a Trackio project (and create the directory if it doesn't exist).
36
+ If a run is not provided, the files are stored in a project-level directory with the given relative path.
37
+
38
+ Args:
39
+ project: The project name
40
+ run: The run name
41
+ step: The step number
42
+ relative_path: The relative path within the directory (only used if run is not provided)
43
+
44
+ Returns:
45
+ The full path to the media file
46
+ """
47
+ if step is not None and run is None:
48
+ raise ValueError("Uploading files at a specific step requires a run")
49
+
50
+ path = MEDIA_DIR / project
51
+ if run:
52
+ path /= run
53
+ if step is not None:
54
+ path /= str(step)
55
+ else:
56
+ path /= "files"
57
+ if relative_path:
58
+ path /= relative_path
59
+ path.mkdir(parents=True, exist_ok=True)
60
+ return path
trackio/media/video.py ADDED
@@ -0,0 +1,246 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import shutil
3
+ import subprocess
4
+ from pathlib import Path
5
+ from typing import Literal
6
+
7
+ import numpy as np
8
+
9
+ from trackio.media.media import TrackioMedia
10
+ from trackio.media.utils import check_ffmpeg_installed, check_path
11
+
12
+ TrackioVideoSourceType = str | Path | np.ndarray
13
+ TrackioVideoFormatType = Literal["gif", "mp4", "webm"]
14
+ VideoCodec = Literal["h264", "vp9", "gif"]
15
+
16
+
17
+ class TrackioVideo(TrackioMedia):
18
+ """
19
+ Initializes a Video object.
20
+
21
+ Example:
22
+ ```python
23
+ import trackio
24
+ import numpy as np
25
+
26
+ # Create a simple video from numpy array
27
+ frames = np.random.randint(0, 255, (10, 3, 64, 64), dtype=np.uint8)
28
+ video = trackio.Video(frames, caption="Random video", fps=30)
29
+
30
+ # Create a batch of videos
31
+ batch_frames = np.random.randint(0, 255, (3, 10, 3, 64, 64), dtype=np.uint8)
32
+ batch_video = trackio.Video(batch_frames, caption="Batch of videos", fps=15)
33
+
34
+ # Create video from file path
35
+ video = trackio.Video("path/to/video.mp4", caption="Video from file")
36
+ ```
37
+
38
+ Args:
39
+ value (`str`, `Path`, or `numpy.ndarray`, *optional*):
40
+ A path to a video file, or a numpy array.
41
+ If numpy array, should be of type `np.uint8` with RGB values in the range `[0, 255]`.
42
+ It is expected to have shape of either (frames, channels, height, width) or (batch, frames, channels, height, width).
43
+ For the latter, the videos will be tiled into a grid.
44
+ caption (`str`, *optional*):
45
+ A string caption for the video.
46
+ fps (`int`, *optional*):
47
+ Frames per second for the video. Only used when value is an ndarray. Default is `24`.
48
+ format (`Literal["gif", "mp4", "webm"]`, *optional*):
49
+ Video format ("gif", "mp4", or "webm"). Only used when value is an ndarray. Default is "gif".
50
+ """
51
+
52
+ TYPE = "trackio.video"
53
+
54
+ def __init__(
55
+ self,
56
+ value: TrackioVideoSourceType,
57
+ caption: str | None = None,
58
+ fps: int | None = None,
59
+ format: TrackioVideoFormatType | None = None,
60
+ ):
61
+ super().__init__(value, caption)
62
+
63
+ if not isinstance(self._value, TrackioVideoSourceType):
64
+ raise ValueError(
65
+ f"Invalid value type, expected {TrackioVideoSourceType}, got {type(self._value)}"
66
+ )
67
+ if isinstance(self._value, np.ndarray):
68
+ if self._value.dtype != np.uint8:
69
+ raise ValueError(
70
+ f"Invalid value dtype, expected np.uint8, got {self._value.dtype}"
71
+ )
72
+ if format is None:
73
+ format = "gif"
74
+ if fps is None:
75
+ fps = 24
76
+ self._fps = fps
77
+ self._format = format
78
+
79
+ @staticmethod
80
+ def _check_array_format(video: np.ndarray) -> None:
81
+ """Raise an error if the array is not in the expected format."""
82
+ if not (video.ndim == 4 and video.shape[-1] == 3):
83
+ raise ValueError(
84
+ f"Expected RGB input shaped (F, H, W, 3), got {video.shape}. "
85
+ f"Input has {video.ndim} dimensions, expected 4."
86
+ )
87
+ if video.dtype != np.uint8:
88
+ raise TypeError(
89
+ f"Expected dtype=uint8, got {video.dtype}. "
90
+ "Please convert your video data to uint8 format."
91
+ )
92
+
93
+ @staticmethod
94
+ def write_video(
95
+ file_path: str | Path, video: np.ndarray, fps: float, codec: VideoCodec
96
+ ) -> None:
97
+ """RGB uint8 only, shape (F, H, W, 3)."""
98
+ check_ffmpeg_installed()
99
+ check_path(file_path)
100
+
101
+ if codec not in {"h264", "vp9", "gif"}:
102
+ raise ValueError("Unsupported codec. Use h264, vp9, or gif.")
103
+
104
+ arr = np.asarray(video)
105
+ TrackioVideo._check_array_format(arr)
106
+
107
+ frames = np.ascontiguousarray(arr)
108
+ _, height, width, _ = frames.shape
109
+ out_path = str(file_path)
110
+
111
+ cmd = [
112
+ "ffmpeg",
113
+ "-y",
114
+ "-f",
115
+ "rawvideo",
116
+ "-s",
117
+ f"{width}x{height}",
118
+ "-pix_fmt",
119
+ "rgb24",
120
+ "-r",
121
+ str(fps),
122
+ "-i",
123
+ "-",
124
+ "-an",
125
+ ]
126
+
127
+ if codec == "gif":
128
+ video_filter = "split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse"
129
+ cmd += [
130
+ "-vf",
131
+ video_filter,
132
+ "-loop",
133
+ "0",
134
+ ]
135
+ elif codec == "h264":
136
+ cmd += [
137
+ "-vcodec",
138
+ "libx264",
139
+ "-pix_fmt",
140
+ "yuv420p",
141
+ "-movflags",
142
+ "+faststart",
143
+ ]
144
+ elif codec == "vp9":
145
+ bpp = 0.08
146
+ bps = int(width * height * fps * bpp)
147
+ if bps >= 1_000_000:
148
+ bitrate = f"{round(bps / 1_000_000)}M"
149
+ elif bps >= 1_000:
150
+ bitrate = f"{round(bps / 1_000)}k"
151
+ else:
152
+ bitrate = str(max(bps, 1))
153
+ cmd += [
154
+ "-vcodec",
155
+ "libvpx-vp9",
156
+ "-b:v",
157
+ bitrate,
158
+ "-pix_fmt",
159
+ "yuv420p",
160
+ ]
161
+ cmd += [out_path]
162
+ proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)
163
+ try:
164
+ for frame in frames:
165
+ proc.stdin.write(frame.tobytes())
166
+ finally:
167
+ if proc.stdin:
168
+ proc.stdin.close()
169
+ stderr = (
170
+ proc.stderr.read().decode("utf-8", errors="ignore")
171
+ if proc.stderr
172
+ else ""
173
+ )
174
+ ret = proc.wait()
175
+ if ret != 0:
176
+ raise RuntimeError(f"ffmpeg failed with code {ret}\n{stderr}")
177
+
178
+ @property
179
+ def _codec(self) -> str:
180
+ match self._format:
181
+ case "gif":
182
+ return "gif"
183
+ case "mp4":
184
+ return "h264"
185
+ case "webm":
186
+ return "vp9"
187
+ case _:
188
+ raise ValueError(f"Unsupported format: {self._format}")
189
+
190
+ def _save_media(self, file_path: Path):
191
+ if isinstance(self._value, np.ndarray):
192
+ video = TrackioVideo._process_ndarray(self._value)
193
+ TrackioVideo.write_video(file_path, video, fps=self._fps, codec=self._codec)
194
+ elif isinstance(self._value, str | Path):
195
+ if os.path.isfile(self._value):
196
+ shutil.copy(self._value, file_path)
197
+ else:
198
+ raise ValueError(f"File not found: {self._value}")
199
+
200
+ @staticmethod
201
+ def _process_ndarray(value: np.ndarray) -> np.ndarray:
202
+ # Verify value is either 4D (single video) or 5D array (batched videos).
203
+ # Expected format: (frames, channels, height, width) or (batch, frames, channels, height, width)
204
+ if value.ndim < 4:
205
+ raise ValueError(
206
+ "Video requires at least 4 dimensions (frames, channels, height, width)"
207
+ )
208
+ if value.ndim > 5:
209
+ raise ValueError(
210
+ "Videos can have at most 5 dimensions (batch, frames, channels, height, width)"
211
+ )
212
+ if value.ndim == 4:
213
+ # Reshape to 5D with single batch: (1, frames, channels, height, width)
214
+ value = value[np.newaxis, ...]
215
+
216
+ value = TrackioVideo._tile_batched_videos(value)
217
+ return value
218
+
219
+ @staticmethod
220
+ def _tile_batched_videos(video: np.ndarray) -> np.ndarray:
221
+ """
222
+ Tiles a batch of videos into a grid of videos.
223
+
224
+ Input format: (batch, frames, channels, height, width) - original FCHW format
225
+ Output format: (frames, total_height, total_width, channels)
226
+ """
227
+ batch_size, frames, channels, height, width = video.shape
228
+
229
+ next_pow2 = 1 << (batch_size - 1).bit_length()
230
+ if batch_size != next_pow2:
231
+ pad_len = next_pow2 - batch_size
232
+ pad_shape = (pad_len, frames, channels, height, width)
233
+ padding = np.zeros(pad_shape, dtype=video.dtype)
234
+ video = np.concatenate((video, padding), axis=0)
235
+ batch_size = next_pow2
236
+
237
+ n_rows = 1 << ((batch_size.bit_length() - 1) // 2)
238
+ n_cols = batch_size // n_rows
239
+
240
+ # Reshape to grid layout: (n_rows, n_cols, frames, channels, height, width)
241
+ video = video.reshape(n_rows, n_cols, frames, channels, height, width)
242
+
243
+ # Rearrange dimensions to (frames, total_height, total_width, channels)
244
+ video = video.transpose(2, 0, 4, 1, 5, 3)
245
+ video = video.reshape(frames, n_rows * height, n_cols * width, channels)
246
+ return video
trackio/package.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "name": "trackio",
3
+ "version": "0.16.1",
4
+ "description": "",
5
+ "python": "true"
6
+ }
trackio/py.typed ADDED
File without changes
trackio/run.py ADDED
@@ -0,0 +1,586 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import shutil
2
+ import threading
3
+ import time
4
+ import uuid
5
+ import warnings
6
+ from datetime import datetime, timezone
7
+ from pathlib import Path
8
+
9
+ import huggingface_hub
10
+ from gradio_client import Client, handle_file
11
+
12
+ from trackio import utils
13
+ from trackio.gpu import GpuMonitor
14
+ from trackio.histogram import Histogram
15
+ from trackio.media import TrackioMedia, get_project_media_path
16
+ from trackio.sqlite_storage import SQLiteStorage
17
+ from trackio.table import Table
18
+ from trackio.typehints import LogEntry, SystemLogEntry, UploadEntry
19
+ from trackio.utils import _get_default_namespace
20
+
21
+ BATCH_SEND_INTERVAL = 0.5
22
+ MAX_BACKOFF = 30
23
+
24
+
25
+ class Run:
26
+ def __init__(
27
+ self,
28
+ url: str | None,
29
+ project: str,
30
+ client: Client | None,
31
+ name: str | None = None,
32
+ group: str | None = None,
33
+ config: dict | None = None,
34
+ space_id: str | None = None,
35
+ auto_log_gpu: bool = False,
36
+ gpu_log_interval: float = 10.0,
37
+ ):
38
+ """
39
+ Initialize a Run for logging metrics to Trackio.
40
+
41
+ Args:
42
+ url: The URL of the Trackio server (local Gradio app or HF Space).
43
+ project: The name of the project to log metrics to.
44
+ client: A pre-configured gradio_client.Client instance, or None to
45
+ create one automatically in a background thread with retry logic.
46
+ Passing None is recommended for normal usage. Passing a client
47
+ is useful for testing (e.g., injecting a mock client).
48
+ name: The name of this run. If None, a readable name like
49
+ "brave-sunset-0" is auto-generated. If space_id is provided,
50
+ generates a "username-timestamp" format instead.
51
+ group: Optional group name to organize related runs together.
52
+ config: A dictionary of configuration/hyperparameters for this run.
53
+ Keys starting with '_' are reserved for internal use.
54
+ space_id: The HF Space ID if logging to a Space (e.g., "user/space").
55
+ If provided, media files will be uploaded to the Space.
56
+ auto_log_gpu: Whether to automatically log GPU metrics (utilization,
57
+ memory, temperature) at regular intervals.
58
+ gpu_log_interval: The interval in seconds between GPU metric logs.
59
+ Only used when auto_log_gpu is True.
60
+ """
61
+ self.url = url
62
+ self.project = project
63
+ self._client_lock = threading.Lock()
64
+ self._client_thread = None
65
+ self._client = client
66
+ self._space_id = space_id
67
+ self.name = name or utils.generate_readable_name(
68
+ SQLiteStorage.get_runs(project), space_id
69
+ )
70
+ self.group = group
71
+ self.config = utils.to_json_safe(config or {})
72
+
73
+ if isinstance(self.config, dict):
74
+ for key in self.config:
75
+ if key.startswith("_"):
76
+ raise ValueError(
77
+ f"Config key '{key}' is reserved (keys starting with '_' are reserved for internal use)"
78
+ )
79
+
80
+ self.config["_Username"] = self._get_username()
81
+ self.config["_Created"] = datetime.now(timezone.utc).isoformat()
82
+ self.config["_Group"] = self.group
83
+
84
+ self._queued_logs: list[LogEntry] = []
85
+ self._queued_system_logs: list[SystemLogEntry] = []
86
+ self._queued_uploads: list[UploadEntry] = []
87
+ self._stop_flag = threading.Event()
88
+ self._config_logged = False
89
+ max_step = SQLiteStorage.get_max_step_for_run(self.project, self.name)
90
+ self._next_step = 0 if max_step is None else max_step + 1
91
+ self._has_local_buffer = False
92
+
93
+ self._is_local = space_id is None
94
+
95
+ if self._is_local:
96
+ self._local_sender_thread = threading.Thread(
97
+ target=self._local_batch_sender
98
+ )
99
+ self._local_sender_thread.daemon = True
100
+ self._local_sender_thread.start()
101
+ else:
102
+ self._client_thread = threading.Thread(target=self._init_client_background)
103
+ self._client_thread.daemon = True
104
+ self._client_thread.start()
105
+
106
+ self._gpu_monitor: "GpuMonitor | None" = None
107
+ if auto_log_gpu:
108
+ self._gpu_monitor = GpuMonitor(self, interval=gpu_log_interval)
109
+ self._gpu_monitor.start()
110
+
111
+ def _get_username(self) -> str | None:
112
+ try:
113
+ return _get_default_namespace()
114
+ except Exception:
115
+ return None
116
+
117
+ def _local_batch_sender(self):
118
+ while (
119
+ not self._stop_flag.is_set()
120
+ or len(self._queued_logs) > 0
121
+ or len(self._queued_system_logs) > 0
122
+ ):
123
+ if not self._stop_flag.is_set():
124
+ time.sleep(BATCH_SEND_INTERVAL)
125
+
126
+ with self._client_lock:
127
+ if self._queued_logs:
128
+ logs_to_send = self._queued_logs.copy()
129
+ self._queued_logs.clear()
130
+ self._write_logs_to_sqlite(logs_to_send)
131
+
132
+ if self._queued_system_logs:
133
+ system_logs_to_send = self._queued_system_logs.copy()
134
+ self._queued_system_logs.clear()
135
+ self._write_system_logs_to_sqlite(system_logs_to_send)
136
+
137
+ def _write_logs_to_sqlite(self, logs: list[LogEntry]):
138
+ logs_by_run: dict[tuple, dict] = {}
139
+ for entry in logs:
140
+ key = (entry["project"], entry["run"])
141
+ if key not in logs_by_run:
142
+ logs_by_run[key] = {
143
+ "metrics": [],
144
+ "steps": [],
145
+ "log_ids": [],
146
+ "config": None,
147
+ }
148
+ logs_by_run[key]["metrics"].append(entry["metrics"])
149
+ logs_by_run[key]["steps"].append(entry.get("step"))
150
+ logs_by_run[key]["log_ids"].append(entry.get("log_id"))
151
+ if entry.get("config") and logs_by_run[key]["config"] is None:
152
+ logs_by_run[key]["config"] = entry["config"]
153
+
154
+ for (project, run), data in logs_by_run.items():
155
+ has_log_ids = any(lid is not None for lid in data["log_ids"])
156
+ SQLiteStorage.bulk_log(
157
+ project=project,
158
+ run=run,
159
+ metrics_list=data["metrics"],
160
+ steps=data["steps"],
161
+ config=data["config"],
162
+ log_ids=data["log_ids"] if has_log_ids else None,
163
+ )
164
+
165
+ def _write_system_logs_to_sqlite(self, logs: list[SystemLogEntry]):
166
+ logs_by_run: dict[tuple, dict] = {}
167
+ for entry in logs:
168
+ key = (entry["project"], entry["run"])
169
+ if key not in logs_by_run:
170
+ logs_by_run[key] = {"metrics": [], "timestamps": [], "log_ids": []}
171
+ logs_by_run[key]["metrics"].append(entry["metrics"])
172
+ logs_by_run[key]["timestamps"].append(entry.get("timestamp"))
173
+ logs_by_run[key]["log_ids"].append(entry.get("log_id"))
174
+
175
+ for (project, run), data in logs_by_run.items():
176
+ has_log_ids = any(lid is not None for lid in data["log_ids"])
177
+ SQLiteStorage.bulk_log_system(
178
+ project=project,
179
+ run=run,
180
+ metrics_list=data["metrics"],
181
+ timestamps=data["timestamps"],
182
+ log_ids=data["log_ids"] if has_log_ids else None,
183
+ )
184
+
185
+ def _batch_sender(self):
186
+ consecutive_failures = 0
187
+ while (
188
+ not self._stop_flag.is_set()
189
+ or len(self._queued_logs) > 0
190
+ or len(self._queued_system_logs) > 0
191
+ or len(self._queued_uploads) > 0
192
+ ):
193
+ if not self._stop_flag.is_set():
194
+ if consecutive_failures:
195
+ sleep_time = min(
196
+ BATCH_SEND_INTERVAL * (2**consecutive_failures), MAX_BACKOFF
197
+ )
198
+ else:
199
+ sleep_time = BATCH_SEND_INTERVAL
200
+ time.sleep(sleep_time)
201
+
202
+ with self._client_lock:
203
+ if self._client is None:
204
+ return
205
+
206
+ failed = False
207
+
208
+ if self._queued_logs:
209
+ logs_to_send = self._queued_logs.copy()
210
+ self._queued_logs.clear()
211
+ try:
212
+ self._client.predict(
213
+ api_name="/bulk_log",
214
+ logs=logs_to_send,
215
+ hf_token=huggingface_hub.utils.get_token(),
216
+ )
217
+ except Exception:
218
+ self._persist_logs_locally(logs_to_send)
219
+ failed = True
220
+
221
+ if self._queued_system_logs:
222
+ system_logs_to_send = self._queued_system_logs.copy()
223
+ self._queued_system_logs.clear()
224
+ try:
225
+ self._client.predict(
226
+ api_name="/bulk_log_system",
227
+ logs=system_logs_to_send,
228
+ hf_token=huggingface_hub.utils.get_token(),
229
+ )
230
+ except Exception:
231
+ self._persist_system_logs_locally(system_logs_to_send)
232
+ failed = True
233
+
234
+ if self._queued_uploads:
235
+ uploads_to_send = self._queued_uploads.copy()
236
+ self._queued_uploads.clear()
237
+ try:
238
+ self._client.predict(
239
+ api_name="/bulk_upload_media",
240
+ uploads=uploads_to_send,
241
+ hf_token=huggingface_hub.utils.get_token(),
242
+ )
243
+ except Exception:
244
+ self._persist_uploads_locally(uploads_to_send)
245
+ failed = True
246
+
247
+ if failed:
248
+ consecutive_failures += 1
249
+ else:
250
+ consecutive_failures = 0
251
+ if self._has_local_buffer:
252
+ self._flush_local_buffer()
253
+
254
+ def _persist_logs_locally(self, logs: list[LogEntry]):
255
+ if not self._space_id:
256
+ return
257
+ logs_by_run: dict[tuple, dict] = {}
258
+ for entry in logs:
259
+ key = (entry["project"], entry["run"])
260
+ if key not in logs_by_run:
261
+ logs_by_run[key] = {
262
+ "metrics": [],
263
+ "steps": [],
264
+ "log_ids": [],
265
+ "config": None,
266
+ }
267
+ logs_by_run[key]["metrics"].append(entry["metrics"])
268
+ logs_by_run[key]["steps"].append(entry.get("step"))
269
+ logs_by_run[key]["log_ids"].append(entry.get("log_id"))
270
+ if entry.get("config") and logs_by_run[key]["config"] is None:
271
+ logs_by_run[key]["config"] = entry["config"]
272
+
273
+ for (project, run), data in logs_by_run.items():
274
+ SQLiteStorage.bulk_log(
275
+ project=project,
276
+ run=run,
277
+ metrics_list=data["metrics"],
278
+ steps=data["steps"],
279
+ log_ids=data["log_ids"],
280
+ config=data["config"],
281
+ space_id=self._space_id,
282
+ )
283
+ self._has_local_buffer = True
284
+
285
+ def _persist_system_logs_locally(self, logs: list[SystemLogEntry]):
286
+ if not self._space_id:
287
+ return
288
+ logs_by_run: dict[tuple, dict] = {}
289
+ for entry in logs:
290
+ key = (entry["project"], entry["run"])
291
+ if key not in logs_by_run:
292
+ logs_by_run[key] = {"metrics": [], "timestamps": [], "log_ids": []}
293
+ logs_by_run[key]["metrics"].append(entry["metrics"])
294
+ logs_by_run[key]["timestamps"].append(entry.get("timestamp"))
295
+ logs_by_run[key]["log_ids"].append(entry.get("log_id"))
296
+
297
+ for (project, run), data in logs_by_run.items():
298
+ SQLiteStorage.bulk_log_system(
299
+ project=project,
300
+ run=run,
301
+ metrics_list=data["metrics"],
302
+ timestamps=data["timestamps"],
303
+ log_ids=data["log_ids"],
304
+ space_id=self._space_id,
305
+ )
306
+ self._has_local_buffer = True
307
+
308
+ def _persist_uploads_locally(self, uploads: list[UploadEntry]):
309
+ if not self._space_id:
310
+ return
311
+ for entry in uploads:
312
+ file_data = entry.get("uploaded_file")
313
+ file_path = ""
314
+ if isinstance(file_data, dict):
315
+ file_path = file_data.get("path", "")
316
+ elif hasattr(file_data, "path"):
317
+ file_path = str(file_data.path)
318
+ else:
319
+ file_path = str(file_data)
320
+ SQLiteStorage.add_pending_upload(
321
+ project=entry["project"],
322
+ space_id=self._space_id,
323
+ run_name=entry.get("run"),
324
+ step=entry.get("step"),
325
+ file_path=file_path,
326
+ relative_path=entry.get("relative_path"),
327
+ )
328
+ self._has_local_buffer = True
329
+
330
+ def _flush_local_buffer(self):
331
+ try:
332
+ buffered_logs = SQLiteStorage.get_pending_logs(self.project)
333
+ if buffered_logs:
334
+ self._client.predict(
335
+ api_name="/bulk_log",
336
+ logs=buffered_logs["logs"],
337
+ hf_token=huggingface_hub.utils.get_token(),
338
+ )
339
+ SQLiteStorage.clear_pending_logs(self.project, buffered_logs["ids"])
340
+
341
+ buffered_sys = SQLiteStorage.get_pending_system_logs(self.project)
342
+ if buffered_sys:
343
+ self._client.predict(
344
+ api_name="/bulk_log_system",
345
+ logs=buffered_sys["logs"],
346
+ hf_token=huggingface_hub.utils.get_token(),
347
+ )
348
+ SQLiteStorage.clear_pending_system_logs(
349
+ self.project, buffered_sys["ids"]
350
+ )
351
+
352
+ buffered_uploads = SQLiteStorage.get_pending_uploads(self.project)
353
+ if buffered_uploads:
354
+ upload_entries = []
355
+ for u in buffered_uploads["uploads"]:
356
+ fp = u["file_path"]
357
+ if Path(fp).exists():
358
+ upload_entries.append(
359
+ {
360
+ "project": u["project"],
361
+ "run": u["run"],
362
+ "step": u["step"],
363
+ "relative_path": u["relative_path"],
364
+ "uploaded_file": handle_file(fp),
365
+ }
366
+ )
367
+ if upload_entries:
368
+ self._client.predict(
369
+ api_name="/bulk_upload_media",
370
+ uploads=upload_entries,
371
+ hf_token=huggingface_hub.utils.get_token(),
372
+ )
373
+ SQLiteStorage.clear_pending_uploads(
374
+ self.project, buffered_uploads["ids"]
375
+ )
376
+
377
+ self._has_local_buffer = False
378
+ except Exception:
379
+ pass
380
+
381
+ def _init_client_background(self):
382
+ if self._client is None:
383
+ fib = utils.fibo()
384
+ for sleep_coefficient in fib:
385
+ try:
386
+ client = Client(self.url, verbose=False)
387
+
388
+ with self._client_lock:
389
+ self._client = client
390
+ break
391
+ except Exception:
392
+ pass
393
+ if sleep_coefficient is not None:
394
+ time.sleep(0.1 * sleep_coefficient)
395
+
396
+ self._batch_sender()
397
+
398
+ def _queue_upload(
399
+ self,
400
+ file_path,
401
+ step: int | None,
402
+ relative_path: str | None = None,
403
+ use_run_name: bool = True,
404
+ ):
405
+ if self._is_local:
406
+ self._save_upload_locally(file_path, step, relative_path, use_run_name)
407
+ else:
408
+ upload_entry: UploadEntry = {
409
+ "project": self.project,
410
+ "run": self.name if use_run_name else None,
411
+ "step": step,
412
+ "relative_path": relative_path,
413
+ "uploaded_file": handle_file(file_path),
414
+ }
415
+ with self._client_lock:
416
+ self._queued_uploads.append(upload_entry)
417
+
418
+ def _save_upload_locally(
419
+ self,
420
+ file_path,
421
+ step: int | None,
422
+ relative_path: str | None = None,
423
+ use_run_name: bool = True,
424
+ ):
425
+ media_path = get_project_media_path(
426
+ project=self.project,
427
+ run=self.name if use_run_name else None,
428
+ step=step,
429
+ relative_path=relative_path,
430
+ )
431
+ src = Path(file_path)
432
+ if src.exists() and str(src.resolve()) != str(Path(media_path).resolve()):
433
+ shutil.copy(str(src), str(media_path))
434
+
435
+ def _process_media(self, value: TrackioMedia, step: int | None) -> dict:
436
+ value._save(self.project, self.name, step if step is not None else 0)
437
+ if self._space_id:
438
+ self._queue_upload(value._get_absolute_file_path(), step)
439
+ return value._to_dict()
440
+
441
+ def _scan_and_queue_media_uploads(self, table_dict: dict, step: int | None):
442
+ if not self._space_id:
443
+ return
444
+
445
+ table_data = table_dict.get("_value", [])
446
+ for row in table_data:
447
+ for value in row.values():
448
+ if isinstance(value, dict) and value.get("_type") in [
449
+ "trackio.image",
450
+ "trackio.video",
451
+ "trackio.audio",
452
+ ]:
453
+ file_path = value.get("file_path")
454
+ if file_path:
455
+ from trackio.utils import MEDIA_DIR
456
+
457
+ absolute_path = MEDIA_DIR / file_path
458
+ self._queue_upload(absolute_path, step)
459
+ elif isinstance(value, list):
460
+ for item in value:
461
+ if isinstance(item, dict) and item.get("_type") in [
462
+ "trackio.image",
463
+ "trackio.video",
464
+ "trackio.audio",
465
+ ]:
466
+ file_path = item.get("file_path")
467
+ if file_path:
468
+ from trackio.utils import MEDIA_DIR
469
+
470
+ absolute_path = MEDIA_DIR / file_path
471
+ self._queue_upload(absolute_path, step)
472
+
473
+ def _ensure_sender_alive(self):
474
+ if self._is_local:
475
+ if (
476
+ hasattr(self, "_local_sender_thread")
477
+ and not self._local_sender_thread.is_alive()
478
+ and not self._stop_flag.is_set()
479
+ ):
480
+ self._local_sender_thread = threading.Thread(
481
+ target=self._local_batch_sender
482
+ )
483
+ self._local_sender_thread.daemon = True
484
+ self._local_sender_thread.start()
485
+ else:
486
+ if (
487
+ self._client_thread is not None
488
+ and not self._client_thread.is_alive()
489
+ and not self._stop_flag.is_set()
490
+ ):
491
+ self._client_thread = threading.Thread(
492
+ target=self._init_client_background
493
+ )
494
+ self._client_thread.daemon = True
495
+ self._client_thread.start()
496
+
497
+ def log(self, metrics: dict, step: int | None = None):
498
+ renamed_keys = []
499
+ new_metrics = {}
500
+
501
+ for k, v in metrics.items():
502
+ if k in utils.RESERVED_KEYS or k.startswith("__"):
503
+ new_key = f"__{k}"
504
+ renamed_keys.append(k)
505
+ new_metrics[new_key] = v
506
+ else:
507
+ new_metrics[k] = v
508
+
509
+ if renamed_keys:
510
+ warnings.warn(f"Reserved keys renamed: {renamed_keys} → '__{{key}}'")
511
+
512
+ metrics = new_metrics
513
+ for key, value in metrics.items():
514
+ if isinstance(value, Table):
515
+ metrics[key] = value._to_dict(
516
+ project=self.project, run=self.name, step=step
517
+ )
518
+ self._scan_and_queue_media_uploads(metrics[key], step)
519
+ elif isinstance(value, Histogram):
520
+ metrics[key] = value._to_dict()
521
+ elif isinstance(value, TrackioMedia):
522
+ metrics[key] = self._process_media(value, step)
523
+ metrics = utils.serialize_values(metrics)
524
+
525
+ if step is None:
526
+ step = self._next_step
527
+ self._next_step = max(self._next_step, step + 1)
528
+
529
+ config_to_log = None
530
+ if not self._config_logged and self.config:
531
+ config_to_log = utils.to_json_safe(self.config)
532
+ self._config_logged = True
533
+
534
+ log_entry: LogEntry = {
535
+ "project": self.project,
536
+ "run": self.name,
537
+ "metrics": metrics,
538
+ "step": step,
539
+ "config": config_to_log,
540
+ "log_id": uuid.uuid4().hex,
541
+ }
542
+
543
+ with self._client_lock:
544
+ self._queued_logs.append(log_entry)
545
+ self._ensure_sender_alive()
546
+
547
+ def log_system(self, metrics: dict):
548
+ metrics = utils.serialize_values(metrics)
549
+ timestamp = datetime.now(timezone.utc).isoformat()
550
+
551
+ system_log_entry: SystemLogEntry = {
552
+ "project": self.project,
553
+ "run": self.name,
554
+ "metrics": metrics,
555
+ "timestamp": timestamp,
556
+ "log_id": uuid.uuid4().hex,
557
+ }
558
+
559
+ with self._client_lock:
560
+ self._queued_system_logs.append(system_log_entry)
561
+ self._ensure_sender_alive()
562
+
563
+ def finish(self):
564
+ if self._gpu_monitor is not None:
565
+ self._gpu_monitor.stop()
566
+
567
+ self._stop_flag.set()
568
+
569
+ if self._is_local:
570
+ if hasattr(self, "_local_sender_thread"):
571
+ print("* Run finished. Uploading logs to Trackio (please wait...)")
572
+ self._local_sender_thread.join(timeout=30)
573
+ if self._local_sender_thread.is_alive():
574
+ warnings.warn(
575
+ "Could not flush all logs within 30s. Some data may be buffered locally."
576
+ )
577
+ else:
578
+ if self._client_thread is not None:
579
+ print(
580
+ "* Run finished. Uploading logs to Trackio Space (please wait...)"
581
+ )
582
+ self._client_thread.join(timeout=30)
583
+ if self._client_thread.is_alive():
584
+ warnings.warn(
585
+ "Could not flush all logs within 30s. Some data may be buffered locally."
586
+ )