abidlabs/trackio-32004

1import glob2import json3import os4import sqlite35 6from huggingface_hub import CommitScheduler7 8try:9    from trackio.dummy_commit_scheduler import DummyCommitScheduler10    from trackio.utils import RESERVED_KEYS, TRACKIO_DIR11except:  # noqa: E72212    from dummy_commit_scheduler import DummyCommitScheduler13    from utils import RESERVED_KEYS, TRACKIO_DIR14 15 16class SQLiteStorage:17    def __init__(18        self, project: str, name: str, config: dict, dataset_id: str | None = None19    ):20        self.project = project21        self.name = name22        self.config = config23        self.db_path = self._get_project_db_path(project)24        self.dataset_id = dataset_id25        self.scheduler = self._get_scheduler()26 27        os.makedirs(TRACKIO_DIR, exist_ok=True)28 29        self._init_db()30        self._save_config()31 32    @staticmethod33    def _get_project_db_path(project: str) -> str:34        """Get the database path for a specific project."""35        safe_project_name = "".join(36            c for c in project if c.isalnum() or c in ("-", "_")37        ).rstrip()38        if not safe_project_name:39            safe_project_name = "default"40        return os.path.join(TRACKIO_DIR, f"{safe_project_name}.db")41 42    def _get_scheduler(self):43        hf_token = os.environ.get(44            "HF_TOKEN"45        )  # Get the token from the environment variable on Spaces46        dataset_id = self.dataset_id or os.environ.get("TRACKIO_DATASET_ID")47        if dataset_id is None:48            scheduler = DummyCommitScheduler()49        else:50            scheduler = CommitScheduler(51                repo_id=dataset_id,52                repo_type="dataset",53                folder_path=TRACKIO_DIR,54                private=True,55                squash_history=True,56                token=hf_token,57            )58        return scheduler59 60    def _init_db(self):61        """Initialize the SQLite database with required tables."""62        with self.scheduler.lock:63            with sqlite3.connect(self.db_path) as conn:64                cursor = conn.cursor()65 66                cursor.execute("""67                    CREATE TABLE IF NOT EXISTS metrics (68                        id INTEGER PRIMARY KEY AUTOINCREMENT,69                        timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,70                        project_name TEXT NOT NULL,71                        run_name TEXT NOT NULL,72                        metrics TEXT NOT NULL73                    )74                """)75 76                cursor.execute("""77                    CREATE TABLE IF NOT EXISTS configs (78                        project_name TEXT NOT NULL,79                        run_name TEXT NOT NULL,80                        config TEXT NOT NULL,81                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,82                        PRIMARY KEY (project_name, run_name)83                    )84                """)85 86                conn.commit()87 88    def _save_config(self):89        """Save the run configuration to the database."""90        with self.scheduler.lock:91            with sqlite3.connect(self.db_path) as conn:92                cursor = conn.cursor()93                cursor.execute(94                    "INSERT OR REPLACE INTO configs (project_name, run_name, config) VALUES (?, ?, ?)",95                    (self.project, self.name, json.dumps(self.config)),96                )97                conn.commit()98 99    def log(self, metrics: dict):100        """Log metrics to the database."""101        for k in metrics.keys():102            if k in RESERVED_KEYS or k.startswith("__"):103                raise ValueError(104                    f"Please do not use this reserved key as a metric: {k}"105                )106 107        with self.scheduler.lock:108            with sqlite3.connect(self.db_path) as conn:109                cursor = conn.cursor()110                cursor.execute(111                    """112                    INSERT INTO metrics 113                    (project_name, run_name, metrics)114                    VALUES (?, ?, ?)115                    """,116                    (self.project, self.name, json.dumps(metrics)),117                )118                conn.commit()119 120    @staticmethod121    def get_metrics(project: str, run: str) -> list[dict]:122        """Retrieve metrics for a specific run."""123        db_path = SQLiteStorage._get_project_db_path(project)124        if not os.path.exists(db_path):125            return []126 127        with sqlite3.connect(db_path) as conn:128            cursor = conn.cursor()129            cursor.execute(130                """131                SELECT timestamp, metrics132                FROM metrics133                WHERE project_name = ? AND run_name = ?134                ORDER BY timestamp135                """,136                (project, run),137            )138            rows = cursor.fetchall()139 140            results = []141            for row in rows:142                timestamp, metrics_json = row143                metrics = json.loads(metrics_json)144                metrics["timestamp"] = timestamp145                results.append(metrics)146 147            return results148 149    @staticmethod150    def get_projects() -> list[str]:151        """Get list of all projects by scanning database files."""152        projects = []153        if not os.path.exists(TRACKIO_DIR):154            return projects155 156        db_files = glob.glob(os.path.join(TRACKIO_DIR, "*.db"))157 158        for db_file in db_files:159            try:160                with sqlite3.connect(db_file) as conn:161                    cursor = conn.cursor()162                    cursor.execute(163                        "SELECT name FROM sqlite_master WHERE type='table' AND name='metrics'"164                    )165                    if cursor.fetchone():166                        cursor.execute("SELECT DISTINCT project_name FROM metrics")167                        project_names = [row[0] for row in cursor.fetchall()]168                        projects.extend(project_names)169            except sqlite3.Error:170                continue171 172        return list(set(projects))173 174    @staticmethod175    def get_runs(project: str) -> list[str]:176        """Get list of all runs for a project."""177        db_path = SQLiteStorage._get_project_db_path(project)178        if not os.path.exists(db_path):179            return []180 181        with sqlite3.connect(db_path) as conn:182            cursor = conn.cursor()183            cursor.execute(184                "SELECT DISTINCT run_name FROM metrics WHERE project_name = ?",185                (project,),186            )187            return [row[0] for row in cursor.fetchall()]188 189    def finish(self):190        """Cleanup when run is finished."""191        pass192