umair894/quickstart-trackio
0
1import os2import platform3import sqlite34import time5from datetime import datetime6from pathlib import Path7from threading import Lock8 9try:10 import fcntl11except ImportError: # fcntl is not available on Windows12 fcntl = None13 14import huggingface_hub as hf15import orjson16import pandas as pd17 18try: # absolute imports when installed from PyPI19 from trackio.commit_scheduler import CommitScheduler20 from trackio.dummy_commit_scheduler import DummyCommitScheduler21 from trackio.utils import (22 TRACKIO_DIR,23 deserialize_values,24 serialize_values,25 )26except ImportError: # relative imports when installed from source on Spaces27 from commit_scheduler import CommitScheduler28 from dummy_commit_scheduler import DummyCommitScheduler29 from utils import TRACKIO_DIR, deserialize_values, serialize_values30 31DB_EXT = ".db"32 33 34class ProcessLock:35 """A file-based lock that works across processes. Is a no-op on Windows."""36 37 def __init__(self, lockfile_path: Path):38 self.lockfile_path = lockfile_path39 self.lockfile = None40 self.is_windows = platform.system() == "Windows"41 42 def __enter__(self):43 """Acquire the lock with retry logic."""44 if self.is_windows:45 return self46 self.lockfile_path.parent.mkdir(parents=True, exist_ok=True)47 self.lockfile = open(self.lockfile_path, "w")48 49 max_retries = 10050 for attempt in range(max_retries):51 try:52 fcntl.flock(self.lockfile.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)53 return self54 except IOError:55 if attempt < max_retries - 1:56 time.sleep(0.1)57 else:58 raise IOError("Could not acquire database lock after 10 seconds")59 60 def __exit__(self, exc_type, exc_val, exc_tb):61 """Release the lock."""62 if self.is_windows:63 return64 65 if self.lockfile:66 fcntl.flock(self.lockfile.fileno(), fcntl.LOCK_UN)67 self.lockfile.close()68 69 70class SQLiteStorage:71 _dataset_import_attempted = False72 _current_scheduler: CommitScheduler | DummyCommitScheduler | None = None73 _scheduler_lock = Lock()74 75 @staticmethod76 def _get_connection(db_path: Path) -> sqlite3.Connection:77 conn = sqlite3.connect(str(db_path), timeout=30.0)78 # Keep WAL for concurrency + performance on many small writes79 conn.execute("PRAGMA journal_mode = WAL")80 # ---- Minimal perf tweaks for many tiny transactions ----81 # NORMAL = fsync at critical points only (safer than OFF, much faster than FULL)82 conn.execute("PRAGMA synchronous = NORMAL")83 # Keep temp data in memory to avoid disk hits during small writes84 conn.execute("PRAGMA temp_store = MEMORY")85 # Give SQLite a bit more room for cache (negative = KB, engine-managed)86 conn.execute("PRAGMA cache_size = -20000")87 # --------------------------------------------------------88 conn.row_factory = sqlite3.Row89 return conn90 91 @staticmethod92 def _get_process_lock(project: str) -> ProcessLock:93 lockfile_path = TRACKIO_DIR / f"{project}.lock"94 return ProcessLock(lockfile_path)95 96 @staticmethod97 def get_project_db_filename(project: str) -> str:98 """Get the database filename for a specific project."""99 safe_project_name = "".join(100 c for c in project if c.isalnum() or c in ("-", "_")101 ).rstrip()102 if not safe_project_name:103 safe_project_name = "default"104 return f"{safe_project_name}{DB_EXT}"105 106 @staticmethod107 def get_project_db_path(project: str) -> Path:108 """Get the database path for a specific project."""109 filename = SQLiteStorage.get_project_db_filename(project)110 return TRACKIO_DIR / filename111 112 @staticmethod113 def init_db(project: str) -> Path:114 """115 Initialize the SQLite database with required tables.116 Returns the database path.117 """118 db_path = SQLiteStorage.get_project_db_path(project)119 db_path.parent.mkdir(parents=True, exist_ok=True)120 with SQLiteStorage._get_process_lock(project):121 with sqlite3.connect(str(db_path), timeout=30.0) as conn:122 conn.execute("PRAGMA journal_mode = WAL")123 conn.execute("PRAGMA synchronous = NORMAL")124 conn.execute("PRAGMA temp_store = MEMORY")125 conn.execute("PRAGMA cache_size = -20000")126 cursor = conn.cursor()127 cursor.execute(128 """129 CREATE TABLE IF NOT EXISTS metrics (130 id INTEGER PRIMARY KEY AUTOINCREMENT,131 timestamp TEXT NOT NULL,132 run_name TEXT NOT NULL,133 step INTEGER NOT NULL,134 metrics TEXT NOT NULL135 )136 """137 )138 cursor.execute(139 """140 CREATE TABLE IF NOT EXISTS configs (141 id INTEGER PRIMARY KEY AUTOINCREMENT,142 run_name TEXT NOT NULL,143 config TEXT NOT NULL,144 created_at TEXT NOT NULL,145 UNIQUE(run_name)146 )147 """148 )149 cursor.execute(150 """151 CREATE INDEX IF NOT EXISTS idx_metrics_run_step152 ON metrics(run_name, step)153 """154 )155 cursor.execute(156 """157 CREATE INDEX IF NOT EXISTS idx_configs_run_name158 ON configs(run_name)159 """160 )161 cursor.execute(162 """163 CREATE INDEX IF NOT EXISTS idx_metrics_run_timestamp164 ON metrics(run_name, timestamp)165 """166 )167 conn.commit()168 return db_path169 170 @staticmethod171 def export_to_parquet():172 """173 Exports all projects' DB files as Parquet under the same path but with extension ".parquet".174 """175 # don't attempt to export (potentially wrong/blank) data before importing for the first time176 if not SQLiteStorage._dataset_import_attempted:177 return178 if not TRACKIO_DIR.exists():179 return180 181 all_paths = os.listdir(TRACKIO_DIR)182 db_names = [f for f in all_paths if f.endswith(DB_EXT)]183 for db_name in db_names:184 db_path = TRACKIO_DIR / db_name185 parquet_path = db_path.with_suffix(".parquet")186 if (not parquet_path.exists()) or (187 db_path.stat().st_mtime > parquet_path.stat().st_mtime188 ):189 with sqlite3.connect(str(db_path)) as conn:190 df = pd.read_sql("SELECT * FROM metrics", conn)191 # break out the single JSON metrics column into individual columns192 metrics = df["metrics"].copy()193 metrics = pd.DataFrame(194 metrics.apply(195 lambda x: deserialize_values(orjson.loads(x))196 ).values.tolist(),197 index=df.index,198 )199 del df["metrics"]200 for col in metrics.columns:201 df[col] = metrics[col]202 203 df.to_parquet(parquet_path)204 205 @staticmethod206 def _cleanup_wal_sidecars(db_path: Path) -> None:207 """Remove leftover -wal/-shm files for a DB basename (prevents disk I/O errors)."""208 for suffix in ("-wal", "-shm"):209 sidecar = Path(str(db_path) + suffix)210 try:211 if sidecar.exists():212 sidecar.unlink()213 except Exception:214 pass215 216 @staticmethod217 def import_from_parquet():218 """219 Imports to all DB files that have matching files under the same path but with extension ".parquet".220 """221 if not TRACKIO_DIR.exists():222 return223 224 all_paths = os.listdir(TRACKIO_DIR)225 parquet_names = [f for f in all_paths if f.endswith(".parquet")]226 for pq_name in parquet_names:227 parquet_path = TRACKIO_DIR / pq_name228 db_path = parquet_path.with_suffix(DB_EXT)229 230 SQLiteStorage._cleanup_wal_sidecars(db_path)231 232 df = pd.read_parquet(parquet_path)233 # fix up df to have a single JSON metrics column234 if "metrics" not in df.columns:235 # separate other columns from metrics236 metrics = df.copy()237 other_cols = ["id", "timestamp", "run_name", "step"]238 df = df[other_cols]239 for col in other_cols:240 del metrics[col]241 # combine them all into a single metrics col242 metrics = orjson.loads(metrics.to_json(orient="records"))243 df["metrics"] = [orjson.dumps(serialize_values(row)) for row in metrics]244 245 with sqlite3.connect(str(db_path), timeout=30.0) as conn:246 df.to_sql("metrics", conn, if_exists="replace", index=False)247 conn.commit()248 249 @staticmethod250 def get_scheduler():251 """252 Get the scheduler for the database based on the environment variables.253 This applies to both local and Spaces.254 """255 with SQLiteStorage._scheduler_lock:256 if SQLiteStorage._current_scheduler is not None:257 return SQLiteStorage._current_scheduler258 hf_token = os.environ.get("HF_TOKEN")259 dataset_id = os.environ.get("TRACKIO_DATASET_ID")260 space_repo_name = os.environ.get("SPACE_REPO_NAME")261 if dataset_id is None or space_repo_name is None:262 scheduler = DummyCommitScheduler()263 else:264 scheduler = CommitScheduler(265 repo_id=dataset_id,266 repo_type="dataset",267 folder_path=TRACKIO_DIR,268 private=True,269 allow_patterns=["*.parquet", "media/**/*"],270 squash_history=True,271 token=hf_token,272 on_before_commit=SQLiteStorage.export_to_parquet,273 )274 SQLiteStorage._current_scheduler = scheduler275 return scheduler276 277 @staticmethod278 def log(project: str, run: str, metrics: dict, step: int | None = None):279 """280 Safely log metrics to the database. Before logging, this method will ensure the database exists281 and is set up with the correct tables. It also uses a cross-process lock to prevent282 database locking errors when multiple processes access the same database.283 284 This method is not used in the latest versions of Trackio (replaced by bulk_log) but285 is kept for backwards compatibility for users who are connecting to a newer version of286 a Trackio Spaces dashboard with an older version of Trackio installed locally.287 """288 db_path = SQLiteStorage.init_db(project)289 with SQLiteStorage._get_process_lock(project):290 with SQLiteStorage._get_connection(db_path) as conn:291 cursor = conn.cursor()292 cursor.execute(293 """294 SELECT MAX(step) 295 FROM metrics 296 WHERE run_name = ?297 """,298 (run,),299 )300 last_step = cursor.fetchone()[0]301 current_step = (302 0303 if step is None and last_step is None304 else (step if step is not None else last_step + 1)305 )306 current_timestamp = datetime.now().isoformat()307 cursor.execute(308 """309 INSERT INTO metrics310 (timestamp, run_name, step, metrics)311 VALUES (?, ?, ?, ?)312 """,313 (314 current_timestamp,315 run,316 current_step,317 orjson.dumps(serialize_values(metrics)),318 ),319 )320 conn.commit()321 322 @staticmethod323 def bulk_log(324 project: str,325 run: str,326 metrics_list: list[dict],327 steps: list[int] | None = None,328 timestamps: list[str] | None = None,329 config: dict | None = None,330 ):331 """332 Safely log bulk metrics to the database. Before logging, this method will ensure the database exists333 and is set up with the correct tables. It also uses a cross-process lock to prevent334 database locking errors when multiple processes access the same database.335 """336 if not metrics_list:337 return338 339 if timestamps is None:340 timestamps = [datetime.now().isoformat()] * len(metrics_list)341 342 db_path = SQLiteStorage.init_db(project)343 with SQLiteStorage._get_process_lock(project):344 with SQLiteStorage._get_connection(db_path) as conn:345 cursor = conn.cursor()346 347 if steps is None:348 steps = list(range(len(metrics_list)))349 elif any(s is None for s in steps):350 cursor.execute(351 "SELECT MAX(step) FROM metrics WHERE run_name = ?", (run,)352 )353 last_step = cursor.fetchone()[0]354 current_step = 0 if last_step is None else last_step + 1355 processed_steps = []356 for step in steps:357 if step is None:358 processed_steps.append(current_step)359 current_step += 1360 else:361 processed_steps.append(step)362 steps = processed_steps363 364 if len(metrics_list) != len(steps) or len(metrics_list) != len(365 timestamps366 ):367 raise ValueError(368 "metrics_list, steps, and timestamps must have the same length"369 )370 371 data = []372 for i, metrics in enumerate(metrics_list):373 data.append(374 (375 timestamps[i],376 run,377 steps[i],378 orjson.dumps(serialize_values(metrics)),379 )380 )381 382 cursor.executemany(383 """384 INSERT INTO metrics385 (timestamp, run_name, step, metrics)386 VALUES (?, ?, ?, ?)387 """,388 data,389 )390 391 if config:392 current_timestamp = datetime.now().isoformat()393 cursor.execute(394 """395 INSERT OR REPLACE INTO configs396 (run_name, config, created_at)397 VALUES (?, ?, ?)398 """,399 (400 run,401 orjson.dumps(serialize_values(config)),402 current_timestamp,403 ),404 )405 406 conn.commit()407 408 @staticmethod409 def get_logs(project: str, run: str) -> list[dict]:410 """Retrieve logs for a specific run. Logs include the step count (int) and the timestamp (datetime object)."""411 db_path = SQLiteStorage.get_project_db_path(project)412 if not db_path.exists():413 return []414 415 with SQLiteStorage._get_connection(db_path) as conn:416 cursor = conn.cursor()417 cursor.execute(418 """419 SELECT timestamp, step, metrics420 FROM metrics421 WHERE run_name = ?422 ORDER BY timestamp423 """,424 (run,),425 )426 427 rows = cursor.fetchall()428 results = []429 for row in rows:430 metrics = orjson.loads(row["metrics"])431 metrics = deserialize_values(metrics)432 metrics["timestamp"] = row["timestamp"]433 metrics["step"] = row["step"]434 results.append(metrics)435 return results436 437 @staticmethod438 def load_from_dataset():439 dataset_id = os.environ.get("TRACKIO_DATASET_ID")440 space_repo_name = os.environ.get("SPACE_REPO_NAME")441 if dataset_id is not None and space_repo_name is not None:442 hfapi = hf.HfApi()443 updated = False444 if not TRACKIO_DIR.exists():445 TRACKIO_DIR.mkdir(parents=True, exist_ok=True)446 with SQLiteStorage.get_scheduler().lock:447 try:448 files = hfapi.list_repo_files(dataset_id, repo_type="dataset")449 for file in files:450 # Download parquet and media assets451 if not (file.endswith(".parquet") or file.startswith("media/")):452 continue453 if (TRACKIO_DIR / file).exists():454 continue455 hf.hf_hub_download(456 dataset_id, file, repo_type="dataset", local_dir=TRACKIO_DIR457 )458 updated = True459 except hf.errors.EntryNotFoundError:460 pass461 except hf.errors.RepositoryNotFoundError:462 pass463 if updated:464 SQLiteStorage.import_from_parquet()465 SQLiteStorage._dataset_import_attempted = True466 467 @staticmethod468 def get_projects() -> list[str]:469 """470 Get list of all projects by scanning the database files in the trackio directory.471 """472 if not SQLiteStorage._dataset_import_attempted:473 SQLiteStorage.load_from_dataset()474 475 projects: set[str] = set()476 if not TRACKIO_DIR.exists():477 return []478 479 for db_file in TRACKIO_DIR.glob(f"*{DB_EXT}"):480 project_name = db_file.stem481 projects.add(project_name)482 return sorted(projects)483 484 @staticmethod485 def get_runs(project: str) -> list[str]:486 """Get list of all runs for a project."""487 db_path = SQLiteStorage.get_project_db_path(project)488 if not db_path.exists():489 return []490 491 with SQLiteStorage._get_connection(db_path) as conn:492 cursor = conn.cursor()493 cursor.execute(494 "SELECT DISTINCT run_name FROM metrics",495 )496 return [row[0] for row in cursor.fetchall()]497 498 @staticmethod499 def get_max_steps_for_runs(project: str) -> dict[str, int]:500 """Get the maximum step for each run in a project."""501 db_path = SQLiteStorage.get_project_db_path(project)502 if not db_path.exists():503 return {}504 505 with SQLiteStorage._get_connection(db_path) as conn:506 cursor = conn.cursor()507 cursor.execute(508 """509 SELECT run_name, MAX(step) as max_step510 FROM metrics511 GROUP BY run_name512 """513 )514 515 results = {}516 for row in cursor.fetchall():517 results[row["run_name"]] = row["max_step"]518 519 return results520 521 @staticmethod522 def store_config(project: str, run: str, config: dict) -> None:523 """Store configuration for a run."""524 db_path = SQLiteStorage.init_db(project)525 526 with SQLiteStorage._get_process_lock(project):527 with SQLiteStorage._get_connection(db_path) as conn:528 cursor = conn.cursor()529 current_timestamp = datetime.now().isoformat()530 531 cursor.execute(532 """533 INSERT OR REPLACE INTO configs534 (run_name, config, created_at)535 VALUES (?, ?, ?)536 """,537 (run, orjson.dumps(serialize_values(config)), current_timestamp),538 )539 conn.commit()540 541 @staticmethod542 def get_run_config(project: str, run: str) -> dict | None:543 """Get configuration for a specific run."""544 db_path = SQLiteStorage.get_project_db_path(project)545 if not db_path.exists():546 return None547 548 with SQLiteStorage._get_connection(db_path) as conn:549 cursor = conn.cursor()550 try:551 cursor.execute(552 """553 SELECT config FROM configs WHERE run_name = ?554 """,555 (run,),556 )557 558 row = cursor.fetchone()559 if row:560 config = orjson.loads(row["config"])561 return deserialize_values(config)562 return None563 except sqlite3.OperationalError as e:564 if "no such table: configs" in str(e):565 return None566 raise567 568 @staticmethod569 def delete_run(project: str, run: str) -> bool:570 """Delete a run from the database (both metrics and config)."""571 db_path = SQLiteStorage.get_project_db_path(project)572 if not db_path.exists():573 return False574 575 with SQLiteStorage._get_process_lock(project):576 with SQLiteStorage._get_connection(db_path) as conn:577 cursor = conn.cursor()578 try:579 cursor.execute("DELETE FROM metrics WHERE run_name = ?", (run,))580 cursor.execute("DELETE FROM configs WHERE run_name = ?", (run,))581 conn.commit()582 return True583 except sqlite3.Error:584 return False585 586 @staticmethod587 def get_all_run_configs(project: str) -> dict[str, dict]:588 """Get configurations for all runs in a project."""589 db_path = SQLiteStorage.get_project_db_path(project)590 if not db_path.exists():591 return {}592 593 with SQLiteStorage._get_connection(db_path) as conn:594 cursor = conn.cursor()595 try:596 cursor.execute(597 """598 SELECT run_name, config FROM configs599 """600 )601 602 results = {}603 for row in cursor.fetchall():604 config = orjson.loads(row["config"])605 results[row["run_name"]] = deserialize_values(config)606 return results607 except sqlite3.OperationalError as e:608 if "no such table: configs" in str(e):609 return {}610 raise611 612 @staticmethod613 def get_metric_values(project: str, run: str, metric_name: str) -> list[dict]:614 """Get all values for a specific metric in a project/run."""615 db_path = SQLiteStorage.get_project_db_path(project)616 if not db_path.exists():617 return []618 619 with SQLiteStorage._get_connection(db_path) as conn:620 cursor = conn.cursor()621 cursor.execute(622 """623 SELECT timestamp, step, metrics624 FROM metrics625 WHERE run_name = ?626 ORDER BY timestamp627 """,628 (run,),629 )630 631 rows = cursor.fetchall()632 results = []633 for row in rows:634 metrics = orjson.loads(row["metrics"])635 metrics = deserialize_values(metrics)636 if metric_name in metrics:637 results.append(638 {639 "timestamp": row["timestamp"],640 "step": row["step"],641 "value": metrics[metric_name],642 }643 )644 return results645 646 @staticmethod647 def get_all_metrics_for_run(project: str, run: str) -> list[str]:648 """Get all metric names for a specific project/run."""649 db_path = SQLiteStorage.get_project_db_path(project)650 if not db_path.exists():651 return []652 653 with SQLiteStorage._get_connection(db_path) as conn:654 cursor = conn.cursor()655 cursor.execute(656 """657 SELECT metrics658 FROM metrics659 WHERE run_name = ?660 ORDER BY timestamp661 """,662 (run,),663 )664 665 rows = cursor.fetchall()666 all_metrics = set()667 for row in rows:668 metrics = orjson.loads(row["metrics"])669 metrics = deserialize_values(metrics)670 for key in metrics.keys():671 if key not in ["timestamp", "step"]:672 all_metrics.add(key)673 return sorted(list(all_metrics))674 675 def finish(self):676 """Cleanup when run is finished."""677 pass678