elephantmipt/trackio_stable-diffusion-v1-5_stable-diffusion-v1-5_None
0
1import json2import os3import shutil4import sqlite35from datetime import datetime6from pathlib import Path7from threading import Lock8 9from huggingface_hub import hf_hub_download10from huggingface_hub.errors import EntryNotFoundError11 12try: # absolute imports when installed13 from trackio.commit_scheduler import CommitScheduler14 from trackio.dummy_commit_scheduler import DummyCommitScheduler15 from trackio.utils import TRACKIO_DIR16except Exception: # relative imports for local execution on Spaces17 from commit_scheduler import CommitScheduler18 from dummy_commit_scheduler import DummyCommitScheduler19 from utils import TRACKIO_DIR20 21 22class SQLiteStorage:23 _dataset_import_attempted = False24 _current_scheduler: CommitScheduler | DummyCommitScheduler | None = None25 _scheduler_lock = Lock()26 27 @staticmethod28 def _get_connection(db_path: Path) -> sqlite3.Connection:29 conn = sqlite3.connect(str(db_path))30 conn.row_factory = sqlite3.Row31 return conn32 33 @staticmethod34 def get_project_db_filename(project: str) -> Path:35 """Get the database filename for a specific project."""36 safe_project_name = "".join(37 c for c in project if c.isalnum() or c in ("-", "_")38 ).rstrip()39 if not safe_project_name:40 safe_project_name = "default"41 return f"{safe_project_name}.db"42 43 @staticmethod44 def get_project_db_path(project: str) -> Path:45 """Get the database path for a specific project."""46 filename = SQLiteStorage.get_project_db_filename(project)47 return TRACKIO_DIR / filename48 49 @staticmethod50 def init_db(project: str) -> Path:51 """52 Initialize the SQLite database with required tables.53 If there is a dataset ID provided, copies from that dataset instead.54 Returns the database path.55 """56 db_path = SQLiteStorage.get_project_db_path(project)57 db_path.parent.mkdir(parents=True, exist_ok=True)58 with SQLiteStorage.get_scheduler().lock:59 dataset_id = os.environ.get("TRACKIO_DATASET_ID")60 if dataset_id is not None and not SQLiteStorage._dataset_import_attempted:61 filename = SQLiteStorage.get_project_db_filename(project)62 try:63 downloaded_path = hf_hub_download(64 dataset_id, filename, repo_type="dataset"65 )66 shutil.copy(downloaded_path, db_path)67 except EntryNotFoundError:68 pass69 SQLiteStorage._dataset_import_attempted = True70 71 with sqlite3.connect(db_path) as conn:72 cursor = conn.cursor()73 cursor.execute("""74 CREATE TABLE IF NOT EXISTS metrics (75 id INTEGER PRIMARY KEY AUTOINCREMENT,76 timestamp TEXT NOT NULL,77 run_name TEXT NOT NULL,78 step INTEGER NOT NULL,79 metrics TEXT NOT NULL80 )81 """)82 cursor.execute(83 """84 CREATE INDEX IF NOT EXISTS idx_metrics_run_step85 ON metrics(run_name, step)86 """87 )88 conn.commit()89 return db_path90 91 @staticmethod92 def get_scheduler():93 """94 Get the scheduler for the database based on the environment variables.95 This applies to both local and Spaces.96 """97 with SQLiteStorage._scheduler_lock:98 if SQLiteStorage._current_scheduler is not None:99 return SQLiteStorage._current_scheduler100 hf_token = os.environ.get("HF_TOKEN")101 dataset_id = os.environ.get("TRACKIO_DATASET_ID")102 space_repo_name = os.environ.get("SPACE_REPO_NAME")103 if dataset_id is None or space_repo_name is None:104 scheduler = DummyCommitScheduler()105 else:106 scheduler = CommitScheduler(107 repo_id=dataset_id,108 repo_type="dataset",109 folder_path=TRACKIO_DIR,110 private=True,111 squash_history=True,112 token=hf_token,113 )114 SQLiteStorage._current_scheduler = scheduler115 return scheduler116 117 @staticmethod118 def log(project: str, run: str, metrics: dict):119 """120 Safely log metrics to the database. Before logging, this method will ensure the database exists121 and is set up with the correct tables. It also uses the scheduler to lock the database so122 that there is no race condition when logging / syncing to the Hugging Face Dataset.123 """124 db_path = SQLiteStorage.init_db(project)125 126 with SQLiteStorage.get_scheduler().lock:127 with SQLiteStorage._get_connection(db_path) as conn:128 cursor = conn.cursor()129 130 cursor.execute(131 """132 SELECT MAX(step) 133 FROM metrics 134 WHERE run_name = ?135 """,136 (run,),137 )138 last_step = cursor.fetchone()[0]139 current_step = 0 if last_step is None else last_step + 1140 141 current_timestamp = datetime.now().isoformat()142 143 cursor.execute(144 """145 INSERT INTO metrics146 (timestamp, run_name, step, metrics)147 VALUES (?, ?, ?, ?)148 """,149 (150 current_timestamp,151 run,152 current_step,153 json.dumps(metrics),154 ),155 )156 conn.commit()157 158 @staticmethod159 def bulk_log(160 project: str,161 run: str,162 metrics_list: list[dict],163 steps: list[int] | None = None,164 timestamps: list[str] | None = None,165 ):166 """Bulk log metrics to the database with specified steps and timestamps."""167 if not metrics_list:168 return169 170 if steps is None:171 steps = list(range(len(metrics_list)))172 173 if timestamps is None:174 timestamps = [datetime.now().isoformat()] * len(metrics_list)175 176 if len(metrics_list) != len(steps) or len(metrics_list) != len(timestamps):177 raise ValueError(178 "metrics_list, steps, and timestamps must have the same length"179 )180 181 db_path = SQLiteStorage.init_db(project)182 with SQLiteStorage.get_scheduler().lock:183 with SQLiteStorage._get_connection(db_path) as conn:184 cursor = conn.cursor()185 186 data = []187 for i, metrics in enumerate(metrics_list):188 data.append(189 (190 timestamps[i],191 run,192 steps[i],193 json.dumps(metrics),194 )195 )196 197 cursor.executemany(198 """199 INSERT INTO metrics200 (timestamp, run_name, step, metrics)201 VALUES (?, ?, ?, ?)202 """,203 data,204 )205 conn.commit()206 207 @staticmethod208 def get_metrics(project: str, run: str) -> list[dict]:209 """Retrieve metrics for a specific run. The metrics also include the step count (int) and the timestamp (datetime object)."""210 db_path = SQLiteStorage.get_project_db_path(project)211 if not db_path.exists():212 return []213 214 with SQLiteStorage._get_connection(db_path) as conn:215 cursor = conn.cursor()216 cursor.execute(217 """218 SELECT timestamp, step, metrics219 FROM metrics220 WHERE run_name = ?221 ORDER BY timestamp222 """,223 (run,),224 )225 226 rows = cursor.fetchall()227 results = []228 for row in rows:229 metrics = json.loads(row["metrics"])230 metrics["timestamp"] = row["timestamp"]231 metrics["step"] = row["step"]232 results.append(metrics)233 234 return results235 236 @staticmethod237 def get_projects() -> list[str]:238 """239 Get list of all projects by scanning the database files in the trackio directory.240 """241 projects: set[str] = set()242 if not TRACKIO_DIR.exists():243 return []244 245 for db_file in TRACKIO_DIR.glob("*.db"):246 project_name = db_file.stem247 projects.add(project_name)248 return sorted(projects)249 250 @staticmethod251 def get_runs(project: str) -> list[str]:252 """Get list of all runs for a project."""253 db_path = SQLiteStorage.get_project_db_path(project)254 if not db_path.exists():255 return []256 257 with SQLiteStorage._get_connection(db_path) as conn:258 cursor = conn.cursor()259 cursor.execute(260 "SELECT DISTINCT run_name FROM metrics",261 )262 return [row[0] for row in cursor.fetchall()]263 264 @staticmethod265 def get_max_steps_for_runs(project: str, runs: list[str]) -> dict[str, int]:266 """Efficiently get the maximum step for multiple runs in a single query."""267 db_path = SQLiteStorage.get_project_db_path(project)268 if not db_path.exists():269 return {run: 0 for run in runs}270 271 with SQLiteStorage._get_connection(db_path) as conn:272 cursor = conn.cursor()273 placeholders = ",".join("?" * len(runs))274 cursor.execute(275 f"""276 SELECT run_name, MAX(step) as max_step277 FROM metrics278 WHERE run_name IN ({placeholders})279 GROUP BY run_name280 """,281 runs,282 )283 284 results = {run: 0 for run in runs} # Default to 0 for runs with no data285 for row in cursor.fetchall():286 results[row["run_name"]] = row["max_step"]287 288 return results289 290 def finish(self):291 """Cleanup when run is finished."""292 pass293