abidlabs HF Staff commited on
Commit
579d97a
·
verified ·
1 Parent(s): 8de9556

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 +160 -0
  3. trackio/__init__.py +613 -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__/markdown.cpython-310.pyc +0 -0
  14. trackio/__pycache__/run.cpython-310.pyc +0 -0
  15. trackio/__pycache__/sqlite_storage.cpython-310.pyc +0 -0
  16. trackio/__pycache__/table.cpython-310.pyc +0 -0
  17. trackio/__pycache__/typehints.cpython-310.pyc +0 -0
  18. trackio/__pycache__/utils.cpython-310.pyc +0 -0
  19. trackio/api.py +71 -0
  20. trackio/assets/badge.png +0 -0
  21. trackio/assets/trackio_logo_dark.png +0 -0
  22. trackio/assets/trackio_logo_light.png +0 -0
  23. trackio/assets/trackio_logo_old.png +3 -0
  24. trackio/assets/trackio_logo_type_dark.png +0 -0
  25. trackio/assets/trackio_logo_type_dark_transparent.png +0 -0
  26. trackio/assets/trackio_logo_type_light.png +0 -0
  27. trackio/assets/trackio_logo_type_light_transparent.png +0 -0
  28. trackio/cli.py +659 -0
  29. trackio/cli_helpers.py +118 -0
  30. trackio/commit_scheduler.py +310 -0
  31. trackio/context_vars.py +18 -0
  32. trackio/deploy.py +433 -0
  33. trackio/dummy_commit_scheduler.py +12 -0
  34. trackio/gpu.py +357 -0
  35. trackio/histogram.py +71 -0
  36. trackio/imports.py +304 -0
  37. trackio/markdown.py +21 -0
  38. trackio/media/__init__.py +27 -0
  39. trackio/media/__pycache__/__init__.cpython-310.pyc +0 -0
  40. trackio/media/__pycache__/audio.cpython-310.pyc +0 -0
  41. trackio/media/__pycache__/image.cpython-310.pyc +0 -0
  42. trackio/media/__pycache__/media.cpython-310.pyc +0 -0
  43. trackio/media/__pycache__/utils.cpython-310.pyc +0 -0
  44. trackio/media/__pycache__/video.cpython-310.pyc +0 -0
  45. trackio/media/audio.py +167 -0
  46. trackio/media/image.py +84 -0
  47. trackio/media/media.py +79 -0
  48. trackio/media/utils.py +60 -0
  49. trackio/media/video.py +246 -0
  50. trackio/package.json +6 -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,160 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # trackio
2
+
3
+ ## 0.17.0
4
+
5
+ ### Features
6
+
7
+ - [#428](https://github.com/gradio-app/trackio/pull/428) [`f7dd1ce`](https://github.com/gradio-app/trackio/commit/f7dd1ce2dc8a1936f9983467fcbcf93bfef01e09) - feat: add ability to rename runs. Thanks @Saba9!
8
+ - [#437](https://github.com/gradio-app/trackio/pull/437) [`2727c0b`](https://github.com/gradio-app/trackio/commit/2727c0b0755f48f7f186162ea45185c98f6b5516) - Add markdown reports across Trackio. Thanks @abidlabs!
9
+ - [#427](https://github.com/gradio-app/trackio/pull/427) [`5aeb9ed`](https://github.com/gradio-app/trackio/commit/5aeb9edcfd2068d309d9d64f172dcbcc327be1ab) - Make Trackio logging much more robust. Thanks @abidlabs!
10
+
11
+ ## 0.16.1
12
+
13
+ ### Features
14
+
15
+ - [#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!
16
+
17
+ ## 0.16.0
18
+
19
+ ### Features
20
+
21
+ - [#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!
22
+ - [#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!
23
+ - [#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!
24
+ - [#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!
25
+ - [#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!
26
+
27
+ ## 0.15.0
28
+
29
+ ### Features
30
+
31
+ - [#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!
32
+ - [#396](https://github.com/gradio-app/trackio/pull/396) [`4a4d1ab`](https://github.com/gradio-app/trackio/commit/4a4d1ab85e63d923132a3fa7afa5d90e16431bec) - Fix run selection issue. Thanks @abidlabs!
33
+ - [#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!
34
+ - [#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!
35
+
36
+ ## 0.14.2
37
+
38
+ ### Features
39
+
40
+ - [#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!
41
+
42
+ ## 0.14.1
43
+
44
+ ### Features
45
+
46
+ - [#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!
47
+ - [#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!
48
+
49
+ ## 0.14.0
50
+
51
+ ### Features
52
+
53
+ - [#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!
54
+ - [#374](https://github.com/gradio-app/trackio/pull/374) [`388e26b`](https://github.com/gradio-app/trackio/commit/388e26b9e9f24cd7ad203affe9b709be885b3d24) - Save Optimized Parquet files. Thanks @lhoestq!
55
+ - [#371](https://github.com/gradio-app/trackio/pull/371) [`fbace9c`](https://github.com/gradio-app/trackio/commit/fbace9cd7732c166f34d268f54b05bb06846cc5d) - Add GPU metrics logging. Thanks @kashif!
56
+ - [#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!
57
+
58
+ ## 0.13.1
59
+
60
+ ### Features
61
+
62
+ - [#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!
63
+
64
+ ## 0.13.0
65
+
66
+ ### Features
67
+
68
+ - [#358](https://github.com/gradio-app/trackio/pull/358) [`073715d`](https://github.com/gradio-app/trackio/commit/073715d1caf8282f68890117f09c3ac301205312) - Improvements to `trackio.sync()`. Thanks @abidlabs!
69
+
70
+ ## 0.12.0
71
+
72
+ ### Features
73
+
74
+ - [#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!
75
+ - [#359](https://github.com/gradio-app/trackio/pull/359) [`08fe9c9`](https://github.com/gradio-app/trackio/commit/08fe9c9ddd7fe99ee811555fdfb62df9ab88e939) - docs: Improve docstrings. Thanks @qgallouedec!
76
+
77
+ ## 0.11.0
78
+
79
+ ### Features
80
+
81
+ - [#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!
82
+ - [#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!
83
+ - [#351](https://github.com/gradio-app/trackio/pull/351) [`8a8957e`](https://github.com/gradio-app/trackio/commit/8a8957e530dd7908d1fef7f2df030303f808101f) - Add `trackio.save()`. Thanks @abidlabs!
84
+
85
+ ## 0.10.0
86
+
87
+ ### Features
88
+
89
+ - [#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!
90
+
91
+ ## 0.9.1
92
+
93
+ ### Features
94
+
95
+ - [#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!
96
+
97
+ ## 0.9.0
98
+
99
+ ### Features
100
+
101
+ - [#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!
102
+ - [#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!
103
+ - [#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!
104
+
105
+ ## 0.8.1
106
+
107
+ ### Features
108
+
109
+ - [#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!
110
+
111
+ ## 0.8.0
112
+
113
+ ### Features
114
+
115
+ - [#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!
116
+ - [#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!
117
+ - [#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!
118
+ - [#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!
119
+ - [#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!
120
+
121
+ ## 0.7.0
122
+
123
+ ### Features
124
+
125
+ - [#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!
126
+ - [#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!
127
+ - [#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!
128
+
129
+ ## 0.6.0
130
+
131
+ ### Features
132
+
133
+ - [#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!
134
+ - [#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!
135
+ - [#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!
136
+ - [#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!
137
+ - [#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!
138
+
139
+ ## 0.5.3
140
+
141
+ ### Features
142
+
143
+ - [#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!
144
+
145
+ ## 0.5.2
146
+
147
+ ### Features
148
+
149
+ - [#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!
150
+ - [#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!
151
+
152
+ ### Fixes
153
+
154
+ - [#291](https://github.com/gradio-app/trackio/pull/291) [`3b5adc3`](https://github.com/gradio-app/trackio/commit/3b5adc3d1f452dbab7a714d235f4974782f93730) - Fix the wheel build. Thanks @pngwn!
155
+
156
+ ## 0.5.1
157
+
158
+ ### Fixes
159
+
160
+ - [#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,613 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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.markdown import Markdown
26
+ from trackio.media import (
27
+ TrackioAudio,
28
+ TrackioImage,
29
+ TrackioVideo,
30
+ get_project_media_path,
31
+ )
32
+ from trackio.run import Run
33
+ from trackio.sqlite_storage import SQLiteStorage
34
+ from trackio.table import Table
35
+ from trackio.typehints import UploadEntry
36
+ from trackio.utils import TRACKIO_DIR, TRACKIO_LOGO_DIR
37
+
38
+ logging.getLogger("httpx").setLevel(logging.WARNING)
39
+
40
+ warnings.filterwarnings(
41
+ "ignore",
42
+ message="Empty session being created. Install gradio\\[oauth\\]",
43
+ category=UserWarning,
44
+ module="gradio.helpers",
45
+ )
46
+
47
+ __version__ = json.loads(Path(__file__).parent.joinpath("package.json").read_text())[
48
+ "version"
49
+ ]
50
+
51
+ __all__ = [
52
+ "init",
53
+ "log",
54
+ "log_system",
55
+ "log_gpu",
56
+ "finish",
57
+ "show",
58
+ "sync",
59
+ "delete_project",
60
+ "import_csv",
61
+ "import_tf_events",
62
+ "save",
63
+ "Image",
64
+ "Video",
65
+ "Audio",
66
+ "Table",
67
+ "Histogram",
68
+ "Markdown",
69
+ "Api",
70
+ ]
71
+
72
+ Image = TrackioImage
73
+ Video = TrackioVideo
74
+ Audio = TrackioAudio
75
+
76
+
77
+ config = {}
78
+
79
+ _atexit_registered = False
80
+
81
+
82
+ def _cleanup_current_run():
83
+ run = context_vars.current_run.get()
84
+ if run is not None:
85
+ try:
86
+ run.finish()
87
+ except Exception:
88
+ pass
89
+
90
+
91
+ def _get_demo():
92
+ # Lazy import to avoid initializing Gradio Blocks (and FastAPI) at import time,
93
+ # which causes import lock errors for libraries that just `import trackio`.
94
+ from trackio.ui.main import CSS, HEAD, demo
95
+
96
+ return demo, CSS, HEAD
97
+
98
+
99
+ def init(
100
+ project: str,
101
+ name: str | None = None,
102
+ group: str | None = None,
103
+ space_id: str | None = None,
104
+ space_storage: SpaceStorage | None = None,
105
+ dataset_id: str | None = None,
106
+ config: dict | None = None,
107
+ resume: str = "never",
108
+ settings: Any = None,
109
+ private: bool | None = None,
110
+ embed: bool = True,
111
+ auto_log_gpu: bool | None = None,
112
+ gpu_log_interval: float = 10.0,
113
+ ) -> Run:
114
+ """
115
+ Creates a new Trackio project and returns a [`Run`] object.
116
+
117
+ Args:
118
+ project (`str`):
119
+ The name of the project (can be an existing project to continue tracking or
120
+ a new project to start tracking from scratch).
121
+ name (`str`, *optional*):
122
+ The name of the run (if not provided, a default name will be generated).
123
+ group (`str`, *optional*):
124
+ The name of the group which this run belongs to in order to help organize
125
+ related runs together. You can toggle the entire group's visibilitiy in the
126
+ dashboard.
127
+ space_id (`str`, *optional*):
128
+ If provided, the project will be logged to a Hugging Face Space instead of
129
+ a local directory. Should be a complete Space name like
130
+ `"username/reponame"` or `"orgname/reponame"`, or just `"reponame"` in which
131
+ case the Space will be created in the currently-logged-in Hugging Face
132
+ user's namespace. If the Space does not exist, it will be created. If the
133
+ Space already exists, the project will be logged to it.
134
+ space_storage ([`~huggingface_hub.SpaceStorage`], *optional*):
135
+ Choice of persistent storage tier.
136
+ dataset_id (`str`, *optional*):
137
+ If a `space_id` is provided, a persistent Hugging Face Dataset will be
138
+ created and the metrics will be synced to it every 5 minutes. Specify a
139
+ Dataset with name like `"username/datasetname"` or `"orgname/datasetname"`,
140
+ or `"datasetname"` (uses currently-logged-in Hugging Face user's namespace),
141
+ or `None` (uses the same name as the Space but with the `"_dataset"`
142
+ suffix). If the Dataset does not exist, it will be created. If the Dataset
143
+ already exists, the project will be appended to it.
144
+ config (`dict`, *optional*):
145
+ A dictionary of configuration options. Provided for compatibility with
146
+ `wandb.init()`.
147
+ resume (`str`, *optional*, defaults to `"never"`):
148
+ Controls how to handle resuming a run. Can be one of:
149
+
150
+ - `"must"`: Must resume the run with the given name, raises error if run
151
+ doesn't exist
152
+ - `"allow"`: Resume the run if it exists, otherwise create a new run
153
+ - `"never"`: Never resume a run, always create a new one
154
+ private (`bool`, *optional*):
155
+ Whether to make the Space private. If None (default), the repo will be
156
+ public unless the organization's default is private. This value is ignored
157
+ if the repo already exists.
158
+ settings (`Any`, *optional*):
159
+ Not used. Provided for compatibility with `wandb.init()`.
160
+ embed (`bool`, *optional*, defaults to `True`):
161
+ If running inside a Jupyter/Colab notebook, whether the dashboard should
162
+ automatically be embedded in the cell when trackio.init() is called. For
163
+ local runs, this launches a local Gradio app and embeds it. For Space runs,
164
+ this embeds the Space URL. In Colab, the local dashboard will be accessible
165
+ via a public share URL (default Gradio behavior).
166
+ auto_log_gpu (`bool` or `None`, *optional*, defaults to `None`):
167
+ Controls automatic GPU metrics logging. If `None` (default), GPU logging
168
+ is automatically enabled when `nvidia-ml-py` is installed and an NVIDIA
169
+ GPU is detected. Set to `True` to force enable or `False` to disable.
170
+ gpu_log_interval (`float`, *optional*, defaults to `10.0`):
171
+ The interval in seconds between automatic GPU metric logs.
172
+ Only used when `auto_log_gpu=True`.
173
+ Returns:
174
+ `Run`: A [`Run`] object that can be used to log metrics and finish the run.
175
+ """
176
+ if settings is not None:
177
+ warnings.warn(
178
+ "* 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."
179
+ )
180
+
181
+ if space_id is None and dataset_id is not None:
182
+ raise ValueError("Must provide a `space_id` when `dataset_id` is provided.")
183
+ try:
184
+ space_id, dataset_id = utils.preprocess_space_and_dataset_ids(
185
+ space_id, dataset_id
186
+ )
187
+ except LocalTokenNotFoundError as e:
188
+ raise LocalTokenNotFoundError(
189
+ f"You must be logged in to Hugging Face locally when `space_id` is provided to deploy to a Space. {e}"
190
+ ) from e
191
+
192
+ url = context_vars.current_server.get()
193
+
194
+ if space_id is not None:
195
+ if url is None:
196
+ url = space_id
197
+ context_vars.current_server.set(url)
198
+ context_vars.current_space_id.set(space_id)
199
+
200
+ _should_embed_local = False
201
+
202
+ if (
203
+ context_vars.current_project.get() is None
204
+ or context_vars.current_project.get() != project
205
+ ):
206
+ print(f"* Trackio project initialized: {project}")
207
+
208
+ if dataset_id is not None:
209
+ os.environ["TRACKIO_DATASET_ID"] = dataset_id
210
+ print(
211
+ f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}"
212
+ )
213
+ if space_id is None:
214
+ print(f"* Trackio metrics logged to: {TRACKIO_DIR}")
215
+ _should_embed_local = embed and utils.is_in_notebook()
216
+ if not _should_embed_local:
217
+ utils.print_dashboard_instructions(project)
218
+ else:
219
+ deploy.create_space_if_not_exists(
220
+ space_id, space_storage, dataset_id, private
221
+ )
222
+ user_name, space_name = space_id.split("/")
223
+ space_url = deploy.SPACE_HOST_URL.format(
224
+ user_name=user_name, space_name=space_name
225
+ )
226
+ print(f"* View dashboard by going to: {space_url}")
227
+ if utils.is_in_notebook() and embed:
228
+ utils.embed_url_in_notebook(space_url)
229
+ context_vars.current_project.set(project)
230
+
231
+ if resume == "must":
232
+ if name is None:
233
+ raise ValueError("Must provide a run name when resume='must'")
234
+ if name not in SQLiteStorage.get_runs(project):
235
+ raise ValueError(f"Run '{name}' does not exist in project '{project}'")
236
+ resumed = True
237
+ elif resume == "allow":
238
+ resumed = name is not None and name in SQLiteStorage.get_runs(project)
239
+ elif resume == "never":
240
+ if name is not None and name in SQLiteStorage.get_runs(project):
241
+ warnings.warn(
242
+ f"* Warning: resume='never' but a run '{name}' already exists in "
243
+ f"project '{project}'. Generating a new name and instead. If you want "
244
+ "to resume this run, call init() with resume='must' or resume='allow'."
245
+ )
246
+ name = None
247
+ resumed = False
248
+ else:
249
+ raise ValueError("resume must be one of: 'must', 'allow', or 'never'")
250
+
251
+ if auto_log_gpu is None:
252
+ auto_log_gpu = gpu_available()
253
+ if auto_log_gpu:
254
+ print("* GPU detected, enabling automatic GPU metrics logging")
255
+
256
+ run = Run(
257
+ url=url,
258
+ project=project,
259
+ client=None,
260
+ name=name,
261
+ group=group,
262
+ config=config,
263
+ space_id=space_id,
264
+ auto_log_gpu=auto_log_gpu,
265
+ gpu_log_interval=gpu_log_interval,
266
+ )
267
+
268
+ if space_id is not None:
269
+ SQLiteStorage.set_project_metadata(project, "space_id", space_id)
270
+ if SQLiteStorage.has_pending_data(project):
271
+ run._has_local_buffer = True
272
+
273
+ global _atexit_registered
274
+ if not _atexit_registered:
275
+ atexit.register(_cleanup_current_run)
276
+ _atexit_registered = True
277
+
278
+ if resumed:
279
+ print(f"* Resumed existing run: {run.name}")
280
+ else:
281
+ print(f"* Created new run: {run.name}")
282
+
283
+ context_vars.current_run.set(run)
284
+ globals()["config"] = run.config
285
+
286
+ if _should_embed_local:
287
+ show(project=project, open_browser=False, block_thread=False)
288
+
289
+ return run
290
+
291
+
292
+ def log(metrics: dict, step: int | None = None) -> None:
293
+ """
294
+ Logs metrics to the current run.
295
+
296
+ Args:
297
+ metrics (`dict`):
298
+ A dictionary of metrics to log.
299
+ step (`int`, *optional*):
300
+ The step number. If not provided, the step will be incremented
301
+ automatically.
302
+ """
303
+ run = context_vars.current_run.get()
304
+ if run is None:
305
+ raise RuntimeError("Call trackio.init() before trackio.log().")
306
+ run.log(
307
+ metrics=metrics,
308
+ step=step,
309
+ )
310
+
311
+
312
+ def log_system(metrics: dict) -> None:
313
+ """
314
+ Logs system metrics (GPU, etc.) to the current run using timestamps instead of steps.
315
+
316
+ Args:
317
+ metrics (`dict`):
318
+ A dictionary of system metrics to log.
319
+ """
320
+ run = context_vars.current_run.get()
321
+ if run is None:
322
+ raise RuntimeError("Call trackio.init() before trackio.log_system().")
323
+ run.log_system(metrics=metrics)
324
+
325
+
326
+ def finish():
327
+ """
328
+ Finishes the current run.
329
+ """
330
+ run = context_vars.current_run.get()
331
+ if run is None:
332
+ raise RuntimeError("Call trackio.init() before trackio.finish().")
333
+ run.finish()
334
+
335
+
336
+ def delete_project(project: str, force: bool = False) -> bool:
337
+ """
338
+ Deletes a project by removing its local SQLite database.
339
+
340
+ Args:
341
+ project (`str`):
342
+ The name of the project to delete.
343
+ force (`bool`, *optional*, defaults to `False`):
344
+ If `True`, deletes the project without prompting for confirmation.
345
+ If `False`, prompts the user to confirm before deleting.
346
+
347
+ Returns:
348
+ `bool`: `True` if the project was deleted, `False` otherwise.
349
+ """
350
+ db_path = SQLiteStorage.get_project_db_path(project)
351
+
352
+ if not db_path.exists():
353
+ print(f"* Project '{project}' does not exist.")
354
+ return False
355
+
356
+ if not force:
357
+ response = input(
358
+ f"Are you sure you want to delete project '{project}'? "
359
+ f"This will permanently delete all runs and metrics. (y/N): "
360
+ )
361
+ if response.lower() not in ["y", "yes"]:
362
+ print("* Deletion cancelled.")
363
+ return False
364
+
365
+ try:
366
+ db_path.unlink()
367
+
368
+ for suffix in ("-wal", "-shm"):
369
+ sidecar = Path(str(db_path) + suffix)
370
+ if sidecar.exists():
371
+ sidecar.unlink()
372
+
373
+ print(f"* Project '{project}' has been deleted.")
374
+ return True
375
+ except Exception as e:
376
+ print(f"* Error deleting project '{project}': {e}")
377
+ return False
378
+
379
+
380
+ def save(
381
+ glob_str: str | Path,
382
+ project: str | None = None,
383
+ ) -> str:
384
+ """
385
+ Saves files to a project (not linked to a specific run). If Trackio is running
386
+ locally, the file(s) will be copied to the project's files directory. If Trackio is
387
+ running in a Space, the file(s) will be uploaded to the Space's files directory.
388
+
389
+ Args:
390
+ glob_str (`str` or `Path`):
391
+ The file path or glob pattern to save. Can be a single file or a pattern
392
+ matching multiple files (e.g., `"*.py"`, `"models/**/*.pth"`).
393
+ project (`str`, *optional*):
394
+ The name of the project to save files to. If not provided, uses the current
395
+ project from `trackio.init()`. If no project is initialized, raises an
396
+ error.
397
+
398
+ Returns:
399
+ `str`: The path where the file(s) were saved (project's files directory).
400
+
401
+ Example:
402
+ ```python
403
+ import trackio
404
+
405
+ trackio.init(project="my-project")
406
+ trackio.save("config.yaml")
407
+ trackio.save("models/*.pth")
408
+ ```
409
+ """
410
+ if project is None:
411
+ project = context_vars.current_project.get()
412
+ if project is None:
413
+ raise RuntimeError(
414
+ "No project specified. Either call trackio.init() first or provide a "
415
+ "project parameter to trackio.save()."
416
+ )
417
+
418
+ glob_str = Path(glob_str)
419
+ base_path = Path.cwd().resolve()
420
+
421
+ matched_files = []
422
+ if glob_str.is_file():
423
+ matched_files = [glob_str.resolve()]
424
+ else:
425
+ pattern = str(glob_str)
426
+ if not glob_str.is_absolute():
427
+ pattern = str((Path.cwd() / glob_str).resolve())
428
+ matched_files = [
429
+ Path(f).resolve()
430
+ for f in glob.glob(pattern, recursive=True)
431
+ if Path(f).is_file()
432
+ ]
433
+
434
+ if not matched_files:
435
+ raise ValueError(f"No files found matching pattern: {glob_str}")
436
+
437
+ current_run = context_vars.current_run.get()
438
+ is_local = (
439
+ current_run._is_local
440
+ if current_run is not None
441
+ else (context_vars.current_space_id.get() is None)
442
+ )
443
+
444
+ if is_local:
445
+ for file_path in matched_files:
446
+ try:
447
+ relative_to_base = file_path.relative_to(base_path)
448
+ except ValueError:
449
+ relative_to_base = Path(file_path.name)
450
+
451
+ if current_run is not None:
452
+ current_run._queue_upload(
453
+ file_path,
454
+ step=None,
455
+ relative_path=str(relative_to_base.parent),
456
+ use_run_name=False,
457
+ )
458
+ else:
459
+ media_path = get_project_media_path(
460
+ project=project,
461
+ run=None,
462
+ step=None,
463
+ relative_path=str(relative_to_base),
464
+ )
465
+ shutil.copy(str(file_path), str(media_path))
466
+ else:
467
+ url = context_vars.current_server.get()
468
+
469
+ upload_entries = []
470
+ for file_path in matched_files:
471
+ try:
472
+ relative_to_base = file_path.relative_to(base_path)
473
+ except ValueError:
474
+ relative_to_base = Path(file_path.name)
475
+
476
+ if current_run is not None:
477
+ current_run._queue_upload(
478
+ file_path,
479
+ step=None,
480
+ relative_path=str(relative_to_base.parent),
481
+ use_run_name=False,
482
+ )
483
+ else:
484
+ upload_entry: UploadEntry = {
485
+ "project": project,
486
+ "run": None,
487
+ "step": None,
488
+ "relative_path": str(relative_to_base),
489
+ "uploaded_file": handle_file(file_path),
490
+ }
491
+ upload_entries.append(upload_entry)
492
+
493
+ if upload_entries:
494
+ if url is None:
495
+ raise RuntimeError(
496
+ "No server available. Call trackio.init() before trackio.save() to start the server."
497
+ )
498
+
499
+ try:
500
+ client = Client(url, verbose=False, httpx_kwargs={"timeout": 90})
501
+ client.predict(
502
+ api_name="/bulk_upload_media",
503
+ uploads=upload_entries,
504
+ hf_token=huggingface_hub.utils.get_token(),
505
+ )
506
+ except Exception as e:
507
+ warnings.warn(
508
+ f"Failed to upload files: {e}. "
509
+ "Files may not be available in the dashboard."
510
+ )
511
+
512
+ return str(utils.MEDIA_DIR / project / "files")
513
+
514
+
515
+ def show(
516
+ project: str | None = None,
517
+ *,
518
+ theme: str | ThemeClass | None = None,
519
+ mcp_server: bool | None = None,
520
+ footer: bool = True,
521
+ color_palette: list[str] | None = None,
522
+ open_browser: bool = True,
523
+ block_thread: bool | None = None,
524
+ host: str | None = None,
525
+ ):
526
+ """
527
+ Launches the Trackio dashboard.
528
+
529
+ Args:
530
+ project (`str`, *optional*):
531
+ The name of the project whose runs to show. If not provided, all projects
532
+ will be shown and the user can select one.
533
+ theme (`str` or `ThemeClass`, *optional*):
534
+ A Gradio Theme to use for the dashboard instead of the default Gradio theme,
535
+ can be a built-in theme (e.g. `'soft'`, `'citrus'`), a theme from the Hub
536
+ (e.g. `"gstaff/xkcd"`), or a custom Theme class. If not provided, the
537
+ `TRACKIO_THEME` environment variable will be used, or if that is not set,
538
+ the default Gradio theme will be used.
539
+ mcp_server (`bool`, *optional*):
540
+ If `True`, the Trackio dashboard will be set up as an MCP server and certain
541
+ functions will be added as MCP tools. If `None` (default behavior), then the
542
+ `GRADIO_MCP_SERVER` environment variable will be used to determine if the
543
+ MCP server should be enabled (which is `"True"` on Hugging Face Spaces).
544
+ footer (`bool`, *optional*, defaults to `True`):
545
+ Whether to show the Gradio footer. When `False`, the footer will be hidden.
546
+ This can also be controlled via the `footer` query parameter in the URL.
547
+ color_palette (`list[str]`, *optional*):
548
+ A list of hex color codes to use for plot lines. If not provided, the
549
+ `TRACKIO_COLOR_PALETTE` environment variable will be used (comma-separated
550
+ hex codes), or if that is not set, the default color palette will be used.
551
+ Example: `['#FF0000', '#00FF00', '#0000FF']`
552
+ open_browser (`bool`, *optional*, defaults to `True`):
553
+ If `True` and not in a notebook, a new browser tab will be opened with the
554
+ dashboard. If `False`, the browser will not be opened.
555
+ block_thread (`bool`, *optional*):
556
+ If `True`, the main thread will be blocked until the dashboard is closed.
557
+ If `None` (default behavior), then the main thread will not be blocked if the
558
+ dashboard is launched in a notebook, otherwise the main thread will be blocked.
559
+ host (`str`, *optional*):
560
+ The host to bind the server to. If not provided, defaults to `'127.0.0.1'`
561
+ (localhost only). Set to `'0.0.0.0'` to allow remote access.
562
+
563
+ Returns:
564
+ `app`: The Gradio app object corresponding to the dashboard launched by Trackio.
565
+ `url`: The local URL of the dashboard.
566
+ `share_url`: The public share URL of the dashboard.
567
+ `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).
568
+ """
569
+ demo, CSS, HEAD = _get_demo()
570
+
571
+ if color_palette is not None:
572
+ os.environ["TRACKIO_COLOR_PALETTE"] = ",".join(color_palette)
573
+
574
+ theme = theme or os.environ.get("TRACKIO_THEME")
575
+
576
+ _mcp_server = (
577
+ mcp_server
578
+ if mcp_server is not None
579
+ else os.environ.get("GRADIO_MCP_SERVER", "False") == "True"
580
+ )
581
+
582
+ app, url, share_url = demo.launch(
583
+ css=CSS,
584
+ head=HEAD,
585
+ footer_links=["gradio", "settings"] + (["api"] if _mcp_server else []),
586
+ quiet=True,
587
+ inline=False,
588
+ prevent_thread_lock=True,
589
+ favicon_path=TRACKIO_LOGO_DIR / "trackio_logo_light.png",
590
+ allowed_paths=[TRACKIO_LOGO_DIR, TRACKIO_DIR],
591
+ mcp_server=_mcp_server,
592
+ theme=theme,
593
+ ssr_mode=False,
594
+ server_name=host,
595
+ )
596
+
597
+ base_url = share_url + "/" if share_url else url
598
+ full_url = utils.get_full_url(
599
+ base_url, project=project, write_token=demo.write_token, footer=footer
600
+ )
601
+
602
+ if not utils.is_in_notebook():
603
+ print(f"* Trackio UI launched at: {full_url}")
604
+ if open_browser:
605
+ webbrowser.open(full_url)
606
+ block_thread = block_thread if block_thread is not None else True
607
+ else:
608
+ utils.embed_url_in_notebook(full_url)
609
+ block_thread = block_thread if block_thread is not None else False
610
+
611
+ if block_thread:
612
+ utils.block_main_thread_until_keyboard_interrupt()
613
+ return TupleNoPrint((demo, url, share_url, full_url))
trackio/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (18.7 kB). View file
 
trackio/__pycache__/api.cpython-310.pyc ADDED
Binary file (3.32 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__/markdown.cpython-310.pyc ADDED
Binary file (882 Bytes). View file
 
trackio/__pycache__/run.cpython-310.pyc ADDED
Binary file (14.4 kB). View file
 
trackio/__pycache__/sqlite_storage.cpython-310.pyc ADDED
Binary file (39.9 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.1 kB). View file
 
trackio/api.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 rename(self, new_name: str) -> "Run":
32
+ SQLiteStorage.rename_run(self.project, self.name, new_name)
33
+ self.name = new_name
34
+ return self
35
+
36
+ def __repr__(self) -> str:
37
+ return f"<Run {self.name} in project {self.project}>"
38
+
39
+
40
+ class Runs:
41
+ def __init__(self, project: str):
42
+ self.project = project
43
+ self._runs = None
44
+
45
+ def _load_runs(self):
46
+ if self._runs is None:
47
+ run_names = SQLiteStorage.get_runs(self.project)
48
+ self._runs = [Run(self.project, name) for name in run_names]
49
+
50
+ def __iter__(self) -> Iterator[Run]:
51
+ self._load_runs()
52
+ return iter(self._runs)
53
+
54
+ def __getitem__(self, index: int) -> Run:
55
+ self._load_runs()
56
+ return self._runs[index]
57
+
58
+ def __len__(self) -> int:
59
+ self._load_runs()
60
+ return len(self._runs)
61
+
62
+ def __repr__(self) -> str:
63
+ self._load_runs()
64
+ return f"<Runs project={self.project} count={len(self._runs)}>"
65
+
66
+
67
+ class Api:
68
+ def runs(self, project: str) -> Runs:
69
+ if not SQLiteStorage.get_project_db_path(project).exists():
70
+ raise ValueError(f"Project '{project}' does not exist")
71
+ 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,659 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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.markdown import Markdown
15
+ from trackio.sqlite_storage import SQLiteStorage
16
+ from trackio.ui.main import get_project_summary, get_run_summary
17
+
18
+
19
+ def _handle_status():
20
+ print("Reading local Trackio projects...\n")
21
+ projects = SQLiteStorage.get_projects()
22
+ if not projects:
23
+ print("No Trackio projects found.")
24
+ return
25
+
26
+ local_projects = []
27
+ synced_projects = []
28
+ unsynced_projects = []
29
+
30
+ for project in projects:
31
+ space_id = SQLiteStorage.get_space_id(project)
32
+ if space_id is None:
33
+ local_projects.append(project)
34
+ elif SQLiteStorage.has_pending_data(project):
35
+ unsynced_projects.append(project)
36
+ else:
37
+ synced_projects.append(project)
38
+
39
+ print("Finished reading Trackio projects")
40
+ if local_projects:
41
+ print(f" * {len(local_projects)} local trackio project(s) [OK]")
42
+ if synced_projects:
43
+ print(f" * {len(synced_projects)} trackio project(s) synced to Spaces [OK]")
44
+ if unsynced_projects:
45
+ print(
46
+ f" * {len(unsynced_projects)} trackio project(s) with unsynced changes [WARNING]:"
47
+ )
48
+ for p in unsynced_projects:
49
+ print(f" - {p}")
50
+
51
+ if unsynced_projects:
52
+ print(
53
+ f"\nRun `trackio sync --project {unsynced_projects[0]}` to sync. "
54
+ "Or run `trackio sync --all` to sync all unsynced changes."
55
+ )
56
+
57
+
58
+ def _handle_sync(args):
59
+ from trackio.deploy import sync_incremental
60
+
61
+ if args.sync_all and args.project:
62
+ error_exit("Cannot use --all and --project together.")
63
+ if not args.sync_all and not args.project:
64
+ error_exit("Must provide either --project or --all.")
65
+
66
+ if args.sync_all:
67
+ projects = SQLiteStorage.get_projects()
68
+ synced_any = False
69
+ for project in projects:
70
+ space_id = SQLiteStorage.get_space_id(project)
71
+ if space_id and SQLiteStorage.has_pending_data(project):
72
+ sync_incremental(
73
+ project, space_id, private=args.private, pending_only=True
74
+ )
75
+ synced_any = True
76
+ if not synced_any:
77
+ print("No projects with unsynced data found.")
78
+ else:
79
+ space_id = args.space_id
80
+ if space_id is None:
81
+ space_id = SQLiteStorage.get_space_id(args.project)
82
+ sync(
83
+ project=args.project,
84
+ space_id=space_id,
85
+ private=args.private,
86
+ force=args.force,
87
+ )
88
+
89
+
90
+ def _extract_reports(
91
+ run: str, logs: list[dict], report_name: str | None = None
92
+ ) -> list[dict]:
93
+ reports = []
94
+ for log in logs:
95
+ timestamp = log.get("timestamp")
96
+ step = log.get("step")
97
+ for key, value in log.items():
98
+ if report_name is not None and key != report_name:
99
+ continue
100
+ if isinstance(value, dict) and value.get("_type") == Markdown.TYPE:
101
+ content = value.get("_value")
102
+ if isinstance(content, str):
103
+ reports.append(
104
+ {
105
+ "run": run,
106
+ "report": key,
107
+ "step": step,
108
+ "timestamp": timestamp,
109
+ "content": content,
110
+ }
111
+ )
112
+ return reports
113
+
114
+
115
+ def main():
116
+ parser = argparse.ArgumentParser(description="Trackio CLI")
117
+ subparsers = parser.add_subparsers(dest="command")
118
+
119
+ ui_parser = subparsers.add_parser(
120
+ "show", help="Show the Trackio dashboard UI for a project"
121
+ )
122
+ ui_parser.add_argument(
123
+ "--project", required=False, help="Project name to show in the dashboard"
124
+ )
125
+ ui_parser.add_argument(
126
+ "--theme",
127
+ required=False,
128
+ default="default",
129
+ 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').",
130
+ )
131
+ ui_parser.add_argument(
132
+ "--mcp-server",
133
+ action="store_true",
134
+ 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.",
135
+ )
136
+ ui_parser.add_argument(
137
+ "--footer",
138
+ action="store_true",
139
+ default=True,
140
+ help="Show the Gradio footer. Use --no-footer to hide it.",
141
+ )
142
+ ui_parser.add_argument(
143
+ "--no-footer",
144
+ dest="footer",
145
+ action="store_false",
146
+ help="Hide the Gradio footer.",
147
+ )
148
+ ui_parser.add_argument(
149
+ "--color-palette",
150
+ required=False,
151
+ 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.",
152
+ )
153
+ ui_parser.add_argument(
154
+ "--host",
155
+ required=False,
156
+ 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).",
157
+ )
158
+
159
+ subparsers.add_parser(
160
+ "status",
161
+ help="Show the status of all local Trackio projects, including sync status.",
162
+ )
163
+
164
+ sync_parser = subparsers.add_parser(
165
+ "sync",
166
+ help="Sync a local project's database to a Hugging Face Space. If the Space does not exist, it will be created.",
167
+ )
168
+ sync_parser.add_argument(
169
+ "--project",
170
+ required=False,
171
+ help="The name of the local project.",
172
+ )
173
+ sync_parser.add_argument(
174
+ "--space-id",
175
+ required=False,
176
+ 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.",
177
+ )
178
+ sync_parser.add_argument(
179
+ "--all",
180
+ action="store_true",
181
+ dest="sync_all",
182
+ help="Sync all projects that have unsynced data to their configured Spaces.",
183
+ )
184
+ sync_parser.add_argument(
185
+ "--private",
186
+ action="store_true",
187
+ 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.",
188
+ )
189
+ sync_parser.add_argument(
190
+ "--force",
191
+ action="store_true",
192
+ help="Overwrite the existing database without prompting for confirmation.",
193
+ )
194
+
195
+ list_parser = subparsers.add_parser(
196
+ "list",
197
+ help="List projects, runs, or metrics",
198
+ )
199
+ list_subparsers = list_parser.add_subparsers(dest="list_type", required=True)
200
+
201
+ list_projects_parser = list_subparsers.add_parser(
202
+ "projects",
203
+ help="List all projects",
204
+ )
205
+ list_projects_parser.add_argument(
206
+ "--json",
207
+ action="store_true",
208
+ help="Output in JSON format",
209
+ )
210
+
211
+ list_runs_parser = list_subparsers.add_parser(
212
+ "runs",
213
+ help="List runs for a project",
214
+ )
215
+ list_runs_parser.add_argument(
216
+ "--project",
217
+ required=True,
218
+ help="Project name",
219
+ )
220
+ list_runs_parser.add_argument(
221
+ "--json",
222
+ action="store_true",
223
+ help="Output in JSON format",
224
+ )
225
+
226
+ list_metrics_parser = list_subparsers.add_parser(
227
+ "metrics",
228
+ help="List metrics for a run",
229
+ )
230
+ list_metrics_parser.add_argument(
231
+ "--project",
232
+ required=True,
233
+ help="Project name",
234
+ )
235
+ list_metrics_parser.add_argument(
236
+ "--run",
237
+ required=True,
238
+ help="Run name",
239
+ )
240
+ list_metrics_parser.add_argument(
241
+ "--json",
242
+ action="store_true",
243
+ help="Output in JSON format",
244
+ )
245
+
246
+ list_system_metrics_parser = list_subparsers.add_parser(
247
+ "system-metrics",
248
+ help="List system metrics for a run",
249
+ )
250
+ list_system_metrics_parser.add_argument(
251
+ "--project",
252
+ required=True,
253
+ help="Project name",
254
+ )
255
+ list_system_metrics_parser.add_argument(
256
+ "--run",
257
+ required=True,
258
+ help="Run name",
259
+ )
260
+ list_system_metrics_parser.add_argument(
261
+ "--json",
262
+ action="store_true",
263
+ help="Output in JSON format",
264
+ )
265
+
266
+ list_reports_parser = list_subparsers.add_parser(
267
+ "reports",
268
+ help="List markdown reports for a project or run",
269
+ )
270
+ list_reports_parser.add_argument(
271
+ "--project",
272
+ required=True,
273
+ help="Project name",
274
+ )
275
+ list_reports_parser.add_argument(
276
+ "--run",
277
+ required=False,
278
+ help="Run name (optional)",
279
+ )
280
+ list_reports_parser.add_argument(
281
+ "--json",
282
+ action="store_true",
283
+ help="Output in JSON format",
284
+ )
285
+
286
+ get_parser = subparsers.add_parser(
287
+ "get",
288
+ help="Get project, run, or metric information",
289
+ )
290
+ get_subparsers = get_parser.add_subparsers(dest="get_type", required=True)
291
+
292
+ get_project_parser = get_subparsers.add_parser(
293
+ "project",
294
+ help="Get project summary",
295
+ )
296
+ get_project_parser.add_argument(
297
+ "--project",
298
+ required=True,
299
+ help="Project name",
300
+ )
301
+ get_project_parser.add_argument(
302
+ "--json",
303
+ action="store_true",
304
+ help="Output in JSON format",
305
+ )
306
+
307
+ get_run_parser = get_subparsers.add_parser(
308
+ "run",
309
+ help="Get run summary",
310
+ )
311
+ get_run_parser.add_argument(
312
+ "--project",
313
+ required=True,
314
+ help="Project name",
315
+ )
316
+ get_run_parser.add_argument(
317
+ "--run",
318
+ required=True,
319
+ help="Run name",
320
+ )
321
+ get_run_parser.add_argument(
322
+ "--json",
323
+ action="store_true",
324
+ help="Output in JSON format",
325
+ )
326
+
327
+ get_metric_parser = get_subparsers.add_parser(
328
+ "metric",
329
+ help="Get metric values for a run",
330
+ )
331
+ get_metric_parser.add_argument(
332
+ "--project",
333
+ required=True,
334
+ help="Project name",
335
+ )
336
+ get_metric_parser.add_argument(
337
+ "--run",
338
+ required=True,
339
+ help="Run name",
340
+ )
341
+ get_metric_parser.add_argument(
342
+ "--metric",
343
+ required=True,
344
+ help="Metric name",
345
+ )
346
+ get_metric_parser.add_argument(
347
+ "--json",
348
+ action="store_true",
349
+ help="Output in JSON format",
350
+ )
351
+
352
+ get_system_metric_parser = get_subparsers.add_parser(
353
+ "system-metric",
354
+ help="Get system metric values for a run",
355
+ )
356
+ get_system_metric_parser.add_argument(
357
+ "--project",
358
+ required=True,
359
+ help="Project name",
360
+ )
361
+ get_system_metric_parser.add_argument(
362
+ "--run",
363
+ required=True,
364
+ help="Run name",
365
+ )
366
+ get_system_metric_parser.add_argument(
367
+ "--metric",
368
+ required=False,
369
+ help="System metric name (optional, if not provided returns all system metrics)",
370
+ )
371
+ get_system_metric_parser.add_argument(
372
+ "--json",
373
+ action="store_true",
374
+ help="Output in JSON format",
375
+ )
376
+
377
+ get_report_parser = get_subparsers.add_parser(
378
+ "report",
379
+ help="Get markdown report entries for a run",
380
+ )
381
+ get_report_parser.add_argument(
382
+ "--project",
383
+ required=True,
384
+ help="Project name",
385
+ )
386
+ get_report_parser.add_argument(
387
+ "--run",
388
+ required=True,
389
+ help="Run name",
390
+ )
391
+ get_report_parser.add_argument(
392
+ "--report",
393
+ required=True,
394
+ help="Report metric name",
395
+ )
396
+ get_report_parser.add_argument(
397
+ "--json",
398
+ action="store_true",
399
+ help="Output in JSON format",
400
+ )
401
+
402
+ args = parser.parse_args()
403
+
404
+ if args.command == "show":
405
+ color_palette = None
406
+ if args.color_palette:
407
+ color_palette = [color.strip() for color in args.color_palette.split(",")]
408
+ show(
409
+ project=args.project,
410
+ theme=args.theme,
411
+ mcp_server=args.mcp_server,
412
+ footer=args.footer,
413
+ color_palette=color_palette,
414
+ host=args.host,
415
+ )
416
+ elif args.command == "status":
417
+ _handle_status()
418
+ elif args.command == "sync":
419
+ _handle_sync(args)
420
+ elif args.command == "list":
421
+ if args.list_type == "projects":
422
+ projects = SQLiteStorage.get_projects()
423
+ if args.json:
424
+ print(format_json({"projects": projects}))
425
+ else:
426
+ print(format_list(projects, "Projects"))
427
+ elif args.list_type == "runs":
428
+ db_path = SQLiteStorage.get_project_db_path(args.project)
429
+ if not db_path.exists():
430
+ error_exit(f"Project '{args.project}' not found.")
431
+ runs = SQLiteStorage.get_runs(args.project)
432
+ if args.json:
433
+ print(format_json({"project": args.project, "runs": runs}))
434
+ else:
435
+ print(format_list(runs, f"Runs in '{args.project}'"))
436
+ elif args.list_type == "metrics":
437
+ db_path = SQLiteStorage.get_project_db_path(args.project)
438
+ if not db_path.exists():
439
+ error_exit(f"Project '{args.project}' not found.")
440
+ runs = SQLiteStorage.get_runs(args.project)
441
+ if args.run not in runs:
442
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
443
+ metrics = SQLiteStorage.get_all_metrics_for_run(args.project, args.run)
444
+ if args.json:
445
+ print(
446
+ format_json(
447
+ {"project": args.project, "run": args.run, "metrics": metrics}
448
+ )
449
+ )
450
+ else:
451
+ print(
452
+ format_list(
453
+ metrics, f"Metrics for '{args.run}' in '{args.project}'"
454
+ )
455
+ )
456
+ elif args.list_type == "system-metrics":
457
+ db_path = SQLiteStorage.get_project_db_path(args.project)
458
+ if not db_path.exists():
459
+ error_exit(f"Project '{args.project}' not found.")
460
+ runs = SQLiteStorage.get_runs(args.project)
461
+ if args.run not in runs:
462
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
463
+ system_metrics = SQLiteStorage.get_all_system_metrics_for_run(
464
+ args.project, args.run
465
+ )
466
+ if args.json:
467
+ print(
468
+ format_json(
469
+ {
470
+ "project": args.project,
471
+ "run": args.run,
472
+ "system_metrics": system_metrics,
473
+ }
474
+ )
475
+ )
476
+ else:
477
+ print(format_system_metric_names(system_metrics))
478
+ elif args.list_type == "reports":
479
+ db_path = SQLiteStorage.get_project_db_path(args.project)
480
+ if not db_path.exists():
481
+ error_exit(f"Project '{args.project}' not found.")
482
+ runs = SQLiteStorage.get_runs(args.project)
483
+ if args.run and args.run not in runs:
484
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
485
+
486
+ target_runs = [args.run] if args.run else runs
487
+ all_reports = []
488
+ for run_name in target_runs:
489
+ logs = SQLiteStorage.get_logs(args.project, run_name)
490
+ all_reports.extend(_extract_reports(run_name, logs))
491
+
492
+ if args.json:
493
+ print(
494
+ format_json(
495
+ {
496
+ "project": args.project,
497
+ "run": args.run,
498
+ "reports": all_reports,
499
+ }
500
+ )
501
+ )
502
+ else:
503
+ report_lines = [
504
+ f"{entry['run']} | {entry['report']} | step={entry['step']} | {entry['timestamp']}"
505
+ for entry in all_reports
506
+ ]
507
+ if args.run:
508
+ print(
509
+ format_list(
510
+ report_lines,
511
+ f"Reports for '{args.run}' in '{args.project}'",
512
+ )
513
+ )
514
+ else:
515
+ print(format_list(report_lines, f"Reports in '{args.project}'"))
516
+ elif args.command == "get":
517
+ if args.get_type == "project":
518
+ db_path = SQLiteStorage.get_project_db_path(args.project)
519
+ if not db_path.exists():
520
+ error_exit(f"Project '{args.project}' not found.")
521
+ summary = get_project_summary(args.project)
522
+ if args.json:
523
+ print(format_json(summary))
524
+ else:
525
+ print(format_project_summary(summary))
526
+ elif args.get_type == "run":
527
+ db_path = SQLiteStorage.get_project_db_path(args.project)
528
+ if not db_path.exists():
529
+ error_exit(f"Project '{args.project}' not found.")
530
+ runs = SQLiteStorage.get_runs(args.project)
531
+ if args.run not in runs:
532
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
533
+ summary = get_run_summary(args.project, args.run)
534
+ if args.json:
535
+ print(format_json(summary))
536
+ else:
537
+ print(format_run_summary(summary))
538
+ elif args.get_type == "metric":
539
+ db_path = SQLiteStorage.get_project_db_path(args.project)
540
+ if not db_path.exists():
541
+ error_exit(f"Project '{args.project}' not found.")
542
+ runs = SQLiteStorage.get_runs(args.project)
543
+ if args.run not in runs:
544
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
545
+ metrics = SQLiteStorage.get_all_metrics_for_run(args.project, args.run)
546
+ if args.metric not in metrics:
547
+ error_exit(
548
+ f"Metric '{args.metric}' not found in run '{args.run}' of project '{args.project}'."
549
+ )
550
+ values = SQLiteStorage.get_metric_values(
551
+ args.project, args.run, args.metric
552
+ )
553
+ if args.json:
554
+ print(
555
+ format_json(
556
+ {
557
+ "project": args.project,
558
+ "run": args.run,
559
+ "metric": args.metric,
560
+ "values": values,
561
+ }
562
+ )
563
+ )
564
+ else:
565
+ print(format_metric_values(values))
566
+ elif args.get_type == "system-metric":
567
+ db_path = SQLiteStorage.get_project_db_path(args.project)
568
+ if not db_path.exists():
569
+ error_exit(f"Project '{args.project}' not found.")
570
+ runs = SQLiteStorage.get_runs(args.project)
571
+ if args.run not in runs:
572
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
573
+ if args.metric:
574
+ system_metrics = SQLiteStorage.get_system_logs(args.project, args.run)
575
+ all_system_metric_names = SQLiteStorage.get_all_system_metrics_for_run(
576
+ args.project, args.run
577
+ )
578
+ if args.metric not in all_system_metric_names:
579
+ error_exit(
580
+ f"System metric '{args.metric}' not found in run '{args.run}' of project '{args.project}'."
581
+ )
582
+ filtered_metrics = [
583
+ {
584
+ k: v
585
+ for k, v in entry.items()
586
+ if k == "timestamp" or k == args.metric
587
+ }
588
+ for entry in system_metrics
589
+ if args.metric in entry
590
+ ]
591
+ if args.json:
592
+ print(
593
+ format_json(
594
+ {
595
+ "project": args.project,
596
+ "run": args.run,
597
+ "metric": args.metric,
598
+ "values": filtered_metrics,
599
+ }
600
+ )
601
+ )
602
+ else:
603
+ print(format_system_metrics(filtered_metrics))
604
+ else:
605
+ system_metrics = SQLiteStorage.get_system_logs(args.project, args.run)
606
+ if args.json:
607
+ print(
608
+ format_json(
609
+ {
610
+ "project": args.project,
611
+ "run": args.run,
612
+ "system_metrics": system_metrics,
613
+ }
614
+ )
615
+ )
616
+ else:
617
+ print(format_system_metrics(system_metrics))
618
+ elif args.get_type == "report":
619
+ db_path = SQLiteStorage.get_project_db_path(args.project)
620
+ if not db_path.exists():
621
+ error_exit(f"Project '{args.project}' not found.")
622
+ runs = SQLiteStorage.get_runs(args.project)
623
+ if args.run not in runs:
624
+ error_exit(f"Run '{args.run}' not found in project '{args.project}'.")
625
+
626
+ logs = SQLiteStorage.get_logs(args.project, args.run)
627
+ reports = _extract_reports(args.run, logs, report_name=args.report)
628
+ if not reports:
629
+ error_exit(
630
+ f"Report '{args.report}' not found in run '{args.run}' of project '{args.project}'."
631
+ )
632
+
633
+ if args.json:
634
+ print(
635
+ format_json(
636
+ {
637
+ "project": args.project,
638
+ "run": args.run,
639
+ "report": args.report,
640
+ "values": reports,
641
+ }
642
+ )
643
+ )
644
+ else:
645
+ output = []
646
+ for idx, entry in enumerate(reports, start=1):
647
+ output.append(
648
+ f"Entry {idx} | step={entry['step']} | timestamp={entry['timestamp']}"
649
+ )
650
+ output.append(entry["content"])
651
+ if idx < len(reports):
652
+ output.append("-" * 80)
653
+ print("\n".join(output))
654
+ else:
655
+ parser.print_help()
656
+
657
+
658
+ if __name__ == "__main__":
659
+ 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/markdown.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ class Markdown:
2
+ """
3
+ Markdown report data type for Trackio.
4
+
5
+ Args:
6
+ text (`str`):
7
+ Markdown content to log.
8
+ """
9
+
10
+ TYPE = "trackio.markdown"
11
+
12
+ def __init__(self, text: str = ""):
13
+ if not isinstance(text, str):
14
+ raise ValueError("Markdown text must be a string")
15
+ self.text = text
16
+
17
+ def _to_dict(self) -> dict:
18
+ return {
19
+ "_type": self.TYPE,
20
+ "_value": self.text,
21
+ }
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.17.0",
4
+ "description": "",
5
+ "python": "true"
6
+ }