JafarUruc/orange_cube

1import glob2import json3import os4import sqlite35from datetime import datetime6 7from huggingface_hub import CommitScheduler8 9try:10    from trackio.context_vars import current_scheduler11    from trackio.dummy_commit_scheduler import DummyCommitScheduler12    from trackio.utils import TRACKIO_DIR13except:  # noqa: E72214    from context_vars import current_scheduler15    from dummy_commit_scheduler import DummyCommitScheduler16    from utils import TRACKIO_DIR17 18 19class SQLiteStorage:20    @staticmethod21    def get_project_db_path(project: str) -> str:22        """Get the database path for a specific project."""23        safe_project_name = "".join(24            c for c in project if c.isalnum() or c in ("-", "_")25        ).rstrip()26        if not safe_project_name:27            safe_project_name = "default"28        return os.path.join(TRACKIO_DIR, f"{safe_project_name}.db")29 30    @staticmethod31    def init_db(project: str) -> str:32        """33        Initialize the SQLite database with required tables.34        Returns the database path.35        """36        db_path = SQLiteStorage.get_project_db_path(project)37        os.makedirs(os.path.dirname(db_path), exist_ok=True)38        with SQLiteStorage.get_scheduler().lock:39            with sqlite3.connect(db_path) as conn:40                cursor = conn.cursor()41                cursor.execute("""42                    CREATE TABLE IF NOT EXISTS metrics (43                        id INTEGER PRIMARY KEY AUTOINCREMENT,44                        timestamp TEXT NOT NULL,45                        project_name TEXT NOT NULL,46                        run_name TEXT NOT NULL,47                        step INTEGER NOT NULL,48                        metrics TEXT NOT NULL49                    )50                """)51                conn.commit()52        return db_path53 54    @staticmethod55    def get_scheduler():56        """57        Get the scheduler for the database based on the environment variables.58        This applies to both local and Spaces.59        """60        if current_scheduler.get() is not None:61            return current_scheduler.get()62        hf_token = os.environ.get("HF_TOKEN")63        dataset_id = os.environ.get("TRACKIO_DATASET_ID")64        if dataset_id is None:65            scheduler = DummyCommitScheduler()66        else:67            scheduler = CommitScheduler(68                repo_id=dataset_id,69                repo_type="dataset",70                folder_path=TRACKIO_DIR,71                private=True,72                squash_history=True,73                token=hf_token,74            )75        current_scheduler.set(scheduler)76        return scheduler77 78    @staticmethod79    def log(project: str, run: str, metrics: dict):80        """81        Safely log metrics to the database. Before logging, this method will ensure the database exists82        and is set up with the correct tables. It also uses the scheduler to lock the database so83        that there is no race condition when logging / syncing to the Hugging Face Dataset.84        """85        db_path = SQLiteStorage.init_db(project)86 87        with SQLiteStorage.get_scheduler().lock:88            with sqlite3.connect(db_path) as conn:89                cursor = conn.cursor()90 91                cursor.execute(92                    """93                    SELECT MAX(step) 94                    FROM metrics 95                    WHERE project_name = ? AND run_name = ?96                    """,97                    (project, run),98                )99                last_step = cursor.fetchone()[0]100                current_step = 0 if last_step is None else last_step + 1101 102                current_timestamp = datetime.now().isoformat()103 104                cursor.execute(105                    """106                    INSERT INTO metrics 107                    (timestamp, project_name, run_name, step, metrics)108                    VALUES (?, ?, ?, ?, ?)109                    """,110                    (111                        current_timestamp,112                        project,113                        run,114                        current_step,115                        json.dumps(metrics),116                    ),117                )118                conn.commit()119 120    @staticmethod121    def bulk_log(122        project: str,123        run: str,124        metrics_list: list[dict],125        steps: list[int] | None = None,126        timestamps: list[str] | None = None,127    ):128        """Bulk log metrics to the database with specified steps and timestamps."""129        if not metrics_list:130            return131 132        if steps is None:133            steps = list(range(len(metrics_list)))134 135        if timestamps is None:136            timestamps = [datetime.now().isoformat()] * len(metrics_list)137 138        if len(metrics_list) != len(steps) or len(metrics_list) != len(timestamps):139            raise ValueError(140                "metrics_list, steps, and timestamps must have the same length"141            )142 143        db_path = SQLiteStorage.init_db(project)144        with SQLiteStorage.get_scheduler().lock:145            with sqlite3.connect(db_path) as conn:146                cursor = conn.cursor()147 148                data = []149                for i, metrics in enumerate(metrics_list):150                    data.append(151                        (152                            timestamps[i],153                            project,154                            run,155                            steps[i],156                            json.dumps(metrics),157                        )158                    )159 160                cursor.executemany(161                    """162                    INSERT INTO metrics 163                    (timestamp, project_name, run_name, step, metrics)164                    VALUES (?, ?, ?, ?, ?)165                    """,166                    data,167                )168                conn.commit()169 170    @staticmethod171    def get_metrics(project: str, run: str) -> list[dict]:172        """Retrieve metrics for a specific run. The metrics also include the step count (int) and the timestamp (datetime object)."""173        db_path = SQLiteStorage.get_project_db_path(project)174        if not os.path.exists(db_path):175            return []176 177        with sqlite3.connect(db_path) as conn:178            cursor = conn.cursor()179            cursor.execute(180                """181                SELECT timestamp, step, metrics182                FROM metrics183                WHERE project_name = ? AND run_name = ?184                ORDER BY timestamp185                """,186                (project, run),187            )188            rows = cursor.fetchall()189 190            results = []191            for row in rows:192                timestamp, step, metrics_json = row193                metrics = json.loads(metrics_json)194                metrics["timestamp"] = timestamp195                metrics["step"] = step196                results.append(metrics)197            return results198 199    @staticmethod200    def get_projects() -> list[str]:201        """Get list of all projects by scanning database files."""202        projects = []203        if not os.path.exists(TRACKIO_DIR):204            return projects205 206        db_files = glob.glob(os.path.join(TRACKIO_DIR, "*.db"))207 208        for db_file in db_files:209            try:210                with sqlite3.connect(db_file) as conn:211                    cursor = conn.cursor()212                    cursor.execute(213                        "SELECT name FROM sqlite_master WHERE type='table' AND name='metrics'"214                    )215                    if cursor.fetchone():216                        cursor.execute("SELECT DISTINCT project_name FROM metrics")217                        project_names = [row[0] for row in cursor.fetchall()]218                        projects.extend(project_names)219            except sqlite3.Error:220                continue221 222        return list(set(projects))223 224    @staticmethod225    def get_runs(project: str) -> list[str]:226        """Get list of all runs for a project."""227        db_path = SQLiteStorage.get_project_db_path(project)228        if not os.path.exists(db_path):229            return []230 231        with sqlite3.connect(db_path) as conn:232            cursor = conn.cursor()233            cursor.execute(234                "SELECT DISTINCT run_name FROM metrics WHERE project_name = ?",235                (project,),236            )237            return [row[0] for row in cursor.fetchall()]238 239    def finish(self):240        """Cleanup when run is finished."""241        pass242