rapidfireai 0.0.1__py3-none-any.whl → 0.9.10__py3-none-any.whl

rapidfireai/db/rf_db.py

@@ -0,0 +1,694 @@
+"""This module contains the RfDb class which is responsible for handling the database operations."""
+
+import json
+import os
+import sqlite3
+from pathlib import Path
+from typing import Any
+
+from rapidfireai.db.db_interface import DatabaseInterface
+from rapidfireai.utils.constants import (
+    ControllerTask,
+    ExperimentStatus,
+    ExperimentTask,
+    RunEndedBy,
+    RunSource,
+    RunStatus,
+    TaskStatus,
+    WorkerTask,
+)
+from rapidfireai.utils.exceptions import DBException
+from rapidfireai.utils.serialize import decode_db_payload, encode_payload
+
+# TODO: add custom exceptions like - RunNotFoundError, ExperimentNotFoundError, ICOpsOnKilledRuns, etc
+
+
+class RfDb:
+    """Class to handle the database operations"""
+
+    def __init__(self):
+        self.db: DatabaseInterface = DatabaseInterface()
+
+    def create_tables(self):
+        """Create the tables in the database"""
+        try:
+            # Get the directory where this file is located
+            current_dir = os.path.dirname(os.path.abspath(__file__))
+            tables_file = os.path.join(current_dir, "tables.sql")
+
+            # First check if tables exist
+            table_check = self.db.conn.execute(
+                "SELECT name FROM sqlite_master WHERE type='table' AND name='experiments'"
+            ).fetchone()
+
+            if table_check is None:
+                # Tables don't exist, create them
+                with open(tables_file, encoding="utf-8") as f:
+                    sql_content = f.read()
+                _ = self.db.conn.executescript(sql_content)
+
+        except FileNotFoundError as e:
+            raise DBException(f"tables.sql file not found at {tables_file}") from e
+        except sqlite3.Error as e:
+            raise DBException(f"Failed to create tables: {e}") from e
+        except Exception as e:
+            raise DBException(f"Unexpected error creating tables: {e}") from e
+
+    def close(self):
+        """Close the database connection"""
+        self.db.close()
+
+    def reset_all_tables(self, experiments_table: bool = False) -> None:
+        """Truncate tables when an experiment is ended"""
+        # ordering based on foreign key constraints
+        tables = [
+            "controller_progress",
+            "worker_progress",
+            "worker_task",
+            "interactive_control",
+            "runs",
+        ]
+
+        if experiments_table:
+            tables.append("experiments")
+
+        for table in tables:
+            query = f"DELETE FROM {table};"
+            self.db.execute(query, commit=True)
+
+        # Reset auto-increment indices to start from 1
+        for table in tables:
+            query = "DELETE FROM sqlite_sequence WHERE name = ?;"
+            self.db.execute(query, (table,), commit=True)
+
+    def reset_experiment_states(self) -> None:
+        """Reset the experiment states when a running task is cancelled"""
+
+        # mark all scheduled and in-progress worker tasks as failed
+        query = """
+            UPDATE worker_task
+            SET status = ?
+            WHERE status = ? OR status = ?
+        """
+        self.db.execute(
+            query, (TaskStatus.FAILED.value, TaskStatus.IN_PROGRESS.value, TaskStatus.SCHEDULED.value), commit=True
+        )
+
+        # mark ongoing and new Runs as failed
+        all_runs = self.get_runs_by_status([RunStatus.ONGOING, RunStatus.NEW]).keys()
+        for run_id in all_runs:
+            self.set_run_status(run_id, RunStatus.FAILED)
+
+        # reset all interactive control tasks
+        all_ic_ops_tasks = self.get_scheduled_ic_ops_tasks()
+        for task in all_ic_ops_tasks:
+            self.set_ic_ops_task_status(task["task_id"], TaskStatus.SKIPPED)
+
+        # reset all progress tables
+        for table in ["controller_progress", "worker_progress"]:
+            query = f"DELETE FROM {table};"
+            self.db.execute(query, commit=True)
+
+    # Experiments Table
+    def create_experiment(
+        self,
+        experiment_name: str,
+        mlflow_experiment_id: str,
+        config_options: dict[str, Any],
+    ) -> int:
+        """Create a new experiment"""
+        query = """
+            INSERT INTO experiments (experiment_name, mlflow_experiment_id, config_options,
+            status, current_task, error)
+            VALUES (?, ?, ?, ?, ?, ?)
+            RETURNING experiment_id
+        """
+
+        result = self.db.execute(
+            query,
+            (
+                experiment_name,
+                mlflow_experiment_id,
+                encode_payload(config_options),
+                ExperimentStatus.RUNNING.value,
+                ExperimentTask.IDLE.value,
+                "",
+            ),
+            commit=True,
+            fetch=True,
+        )
+        if result:
+            # run db optimizer every time an experiment is created
+            self.db.optimize_periodically()
+
+            # return experiment_id
+            return result[0][0]
+        raise DBException("Failed to create experiment")
+
+    def get_running_experiment(self) -> dict[str, Any]:
+        """Get an experiment's details by its ID"""
+        query = """
+            SELECT experiment_id, experiment_name, status, error, mlflow_experiment_id, config_options
+            FROM experiments
+            WHERE status = ?
+            ORDER BY experiment_id DESC
+            LIMIT 1
+        """
+        experiment_details = self.db.execute(query, (ExperimentStatus.RUNNING.value,), fetch=True)
+
+        if experiment_details:
+            experiment_details = experiment_details[0]
+            experiment_details = {
+                "experiment_id": experiment_details[0],
+                "experiment_name": experiment_details[1],
+                "status": ExperimentStatus(experiment_details[2]),
+                "error": experiment_details[3],
+                "mlflow_experiment_id": experiment_details[4],
+                "config_options": decode_db_payload(experiment_details[5]),
+            }
+            return experiment_details
+        raise DBException("No running experiment found")
+
+    def get_experiment_status(self) -> ExperimentStatus | None:
+        """Get the status of an experiment"""
+        query = """
+            SELECT status
+            FROM experiments
+            ORDER BY experiment_id DESC
+            LIMIT 1
+        """
+        status = self.db.execute(query, fetch=True)
+
+        if status:
+            return ExperimentStatus(status[0][0])
+        raise DBException("No experiment status found")
+
+    def set_experiment_status(self, experiment_id: int, status: ExperimentStatus) -> None:
+        """Set the status of an experiment"""
+        query = """
+            UPDATE experiments
+            SET status = ?
+            WHERE experiment_id = ?
+        """
+        self.db.execute(query, (status.value, experiment_id), commit=True)
+
+    def set_experiment_error(self, error: str) -> None:
+        """Set the error message of an experiment"""
+        query = """
+            UPDATE experiments
+            SET error = ?
+            WHERE status = ?
+        """
+        self.db.execute(query, (error, ExperimentStatus.RUNNING.value), commit=True)
+
+    def get_experiment_error(self) -> str:
+        """Get the error message of an experiment"""
+        query = """
+            SELECT error
+            FROM experiments
+            ORDER BY experiment_id DESC
+            LIMIT 1
+        """
+        error = self.db.execute(query, fetch=True)
+
+        if error:
+            return error[0][0]
+        else:
+            return ""
+
+    def set_experiment_current_task(self, task: ExperimentTask) -> None:
+        """Set the current task of an experiment"""
+        query = """
+            UPDATE experiments
+            SET current_task = ?
+            WHERE status = ?
+        """
+        self.db.execute(query, (task.value, ExperimentStatus.RUNNING.value), commit=True)
+
+    def get_experiment_current_task(self) -> ExperimentTask:
+        """Get the current task of an experiment"""
+        query = """
+            SELECT current_task
+            FROM experiments
+            WHERE status = ?
+            ORDER BY experiment_id DESC
+            LIMIT 1
+        """
+        task = self.db.execute(query, (ExperimentStatus.RUNNING.value,), fetch=True)
+
+        if task:
+            return ExperimentTask(task[0][0])
+        raise DBException("No running experiment found")
+
+    def get_all_experiment_names(self) -> list[str]:
+        """Get all experiment names"""
+        query = """
+            SELECT experiment_name
+            FROM experiments
+        """
+        experiment_names = self.db.execute(query, fetch=True)
+
+        if experiment_names:
+            return [experiment[0] for experiment in experiment_names]
+        else:
+            return []
+
+    def get_experiments_path(self, experiment_name: str) -> Path:
+        """Get the experiments path for a given experiment name"""
+        query = """
+            SELECT config_options
+            FROM experiments
+            WHERE experiment_name = ?
+        """
+        config_options = self.db.execute(query, (experiment_name,), fetch=True)
+        if config_options:
+            config_dict = decode_db_payload(config_options[0][0])
+            return Path(config_dict["experiments_path"])
+        raise DBException("Experiments path not found for running experiment")
+
+    # Runs Table
+    def create_run(
+        self,
+        config_leaf: dict[str, Any],
+        status: RunStatus,
+        mlflow_run_id: str | None = None,
+        flattened_config: dict[str, Any] | None = None,
+        completed_steps: int = 0,
+        total_steps: int = 0,
+        start_chunk_id: int = 0,
+        num_chunks_visited_curr_epoch: int = 0,
+        num_epochs_completed: int = 0,
+        error: str = "",
+        source: RunSource | None = None,
+        ended_by: RunEndedBy | None = None,
+        warm_started_from: int | None = None,
+    ) -> int:
+        """Create a new run"""
+        query = """
+            INSERT INTO runs (status, mlflow_run_id, flattened_config, config_leaf,
+            completed_steps, total_steps, start_chunk_id, num_chunks_visited_curr_epoch, num_epochs_completed, error, source, ended_by, warm_started_from)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+        """
+        self.db.execute(
+            query,
+            (
+                status.value,
+                mlflow_run_id,
+                json.dumps(flattened_config) if flattened_config else "{}",
+                encode_payload(config_leaf) if config_leaf else "{}",
+                completed_steps,
+                total_steps,
+                start_chunk_id,
+                num_chunks_visited_curr_epoch,
+                num_epochs_completed,
+                error,
+                source.value if source else "",
+                ended_by.value if ended_by else "",
+                warm_started_from,
+            ),
+            commit=True,
+        )
+        result = self.db.execute("SELECT last_insert_rowid()", fetch=True)
+        if result:
+            return result[0][0]
+        raise DBException("Failed to create run")
+
+    def set_run_details(
+        self,
+        run_id: int,
+        status: RunStatus | None = None,
+        mlflow_run_id: str | None = None,
+        flattened_config: dict[str, Any] | None = None,
+        config_leaf: dict[str, Any] | None = None,
+        completed_steps: int | None = None,
+        total_steps: int | None = None,
+        start_chunk_id: int | None = None,
+        num_chunks_visited_curr_epoch: int | None = None,
+        num_epochs_completed: int | None = None,
+        error: str | None = None,
+        source: RunSource | None = None,
+        ended_by: RunEndedBy | None = None,
+        warm_started_from: int | None = None,
+    ) -> None:
+        """Set the details of an existing run"""
+        # Initialize a dictionary to hold the column-value pairs
+        columns = {
+            "status": status.value if status else None,
+            "mlflow_run_id": mlflow_run_id,
+            "flattened_config": json.dumps(flattened_config) if flattened_config else None,
+            "config_leaf": encode_payload(config_leaf) if config_leaf else None,
+            "completed_steps": completed_steps,
+            "total_steps": total_steps,
+            "start_chunk_id": start_chunk_id,
+            "num_chunks_visited_curr_epoch": num_chunks_visited_curr_epoch,
+            "num_epochs_completed": num_epochs_completed,
+            "error": error,
+            "source": source.value if source else None,
+            "ended_by": ended_by.value if ended_by else None,
+            "warm_started_from": warm_started_from,
+        }
+
+        # Filter out None values
+        columns = {k: v for k, v in columns.items() if v is not None}
+
+        # If no columns to update, return early
+        if not columns:
+            return
+
+        # Construct the query parts
+        query_parts = [f"{col} = ?" for col in columns]
+        values = list(columns.values())
+
+        # Ensure run_id is always included
+        values.append(run_id)
+
+        # Construct the final query
+        query = f"UPDATE runs SET {', '.join(query_parts)} WHERE run_id = ?"
+
+        # Execute the query
+        self.db.execute(query, tuple(values), commit=True)
+
+    def get_run(self, run_id: int) -> dict[str, Any]:
+        """Get a run's details"""
+        query = """
+            SELECT status, mlflow_run_id, flattened_config, config_leaf, completed_steps, total_steps, start_chunk_id,
+            num_chunks_visited_curr_epoch, num_epochs_completed, error, source, ended_by, warm_started_from
+            FROM runs
+            WHERE run_id = ?
+        """
+        run_details = self.db.execute(query, (run_id,), fetch=True)
+
+        if run_details:
+            run_details = run_details[0]
+            formatted_details = {
+                "status": RunStatus(run_details[0]),
+                "mlflow_run_id": run_details[1],
+                "flattened_config": json.loads(run_details[2]),
+                "config_leaf": decode_db_payload(run_details[3]) if run_details[3] and run_details[3] != "{}" else {},
+                "completed_steps": run_details[4],
+                "total_steps": run_details[5],
+                "start_chunk_id": run_details[6],
+                "num_chunks_visited_curr_epoch": run_details[7],
+                "num_epochs_completed": run_details[8],
+                "error": run_details[9],
+                "source": RunSource(run_details[10]) if run_details[10] else None,
+                "ended_by": RunEndedBy(run_details[11]) if run_details[11] else None,
+                "warm_started_from": run_details[12],
+            }
+            return formatted_details
+        raise DBException("No run found")
+
+    def get_runs_by_status(self, statuses: list[RunStatus]) -> dict[int, dict[str, Any]]:
+        """Get all runs by statuses"""
+        if not statuses:
+            return {}
+
+        # Create placeholders for SQL IN clause
+        placeholders = ",".join(["?"] * len(statuses))
+        query = f"""
+            SELECT run_id, status, mlflow_run_id, flattened_config, config_leaf, completed_steps, total_steps,
+            start_chunk_id, num_chunks_visited_curr_epoch, num_epochs_completed, error, source, ended_by, warm_started_from
+            FROM runs
+            WHERE status IN ({placeholders})
+        """
+        # Extract status values for the query parameters
+        status_values = [status.value for status in statuses]
+        run_details = self.db.execute(query, status_values, fetch=True)
+        formatted_details: dict[int, dict[str, Any]] = {}
+        if run_details:
+            for run in run_details:
+                formatted_details[run[0]] = {
+                    "status": RunStatus(run[1]),
+                    "mlflow_run_id": run[2],
+                    "flattened_config": json.loads(run[3]),
+                    "config_leaf": decode_db_payload(run[4]) if run[4] and run[4] != "{}" else {},
+                    "completed_steps": run[5],
+                    "total_steps": run[6],
+                    "start_chunk_id": run[7],
+                    "num_chunks_visited_curr_epoch": run[8],
+                    "num_epochs_completed": run[9],
+                    "error": run[10],
+                    "source": RunSource(run[11]) if run[11] else None,
+                    "ended_by": RunEndedBy(run[12]) if run[12] else None,
+                    "warm_started_from": run[13],
+                }
+        return formatted_details
+
+    def get_all_runs(self) -> dict[int, dict[str, Any]]:
+        """Get all runs for UI display (ignore all complex fields)"""
+        query = """
+            SELECT run_id, status, mlflow_run_id, flattened_config, config_leaf, completed_steps, total_steps,
+            start_chunk_id, num_chunks_visited_curr_epoch, num_epochs_completed, error, source, ended_by, warm_started_from
+            FROM runs
+        """
+        run_details = self.db.execute(query, fetch=True)
+
+        formatted_details: dict[int, dict[str, Any]] = {}
+        if run_details:
+            for run in run_details:
+                formatted_details[run[0]] = {
+                    "status": RunStatus(run[1]),
+                    "mlflow_run_id": run[2],
+                    "flattened_config": json.loads(run[3]),
+                    "config_leaf": decode_db_payload(run[4]) if run[4] and run[4] != "{}" else {},
+                    "completed_steps": run[5],
+                    "total_steps": run[6],
+                    "start_chunk_id": run[7],
+                    "num_chunks_visited_curr_epoch": run[8],
+                    "num_epochs_completed": run[9],
+                    "error": run[10],
+                    "source": RunSource(run[11]) if run[11] else None,
+                    "ended_by": RunEndedBy(run[12]) if run[12] else None,
+                    "warm_started_from": run[13],
+                }
+        return formatted_details
+
+    def set_run_status(self, run_id: int, status: RunStatus) -> None:
+        """Set the status of a run"""
+        query = """
+            UPDATE runs
+            SET status = ?
+            WHERE run_id = ?
+        """
+        self.db.execute(query, (status.value, run_id), commit=True)
+
+    def set_completed_steps(self, run_id: int, completed_steps: int) -> None:
+        """Set the current completed steps for a run"""
+        query = """
+            UPDATE runs
+            SET completed_steps = ?
+            WHERE run_id = ?
+        """
+        self.db.execute(query, (completed_steps, run_id), commit=True)
+
+    def get_completed_steps(self, run_id: int) -> int:
+        """Get the current completed steps for a run"""
+        query = """
+            SELECT completed_steps
+            FROM runs
+            WHERE run_id = ?
+        """
+        completed_steps = self.db.execute(query, (run_id,), fetch=True)
+        if completed_steps:
+            return completed_steps[0][0]
+        raise DBException("No completed steps found")
+
+    # Interactive Control Table
+    def create_ic_ops_task(self, run_id: int, ic_op: ControllerTask, config_leaf: dict[str, Any]) -> int:
+        """Create a new interactive control task"""
+        query = """
+            INSERT INTO interactive_control (run_id, ic_op, config_leaf, status)
+            VALUES (?, ?, ?, ?)
+            RETURNING task_id
+        """
+        config_leaf_str = encode_payload(config_leaf) if config_leaf else "{}"
+        result = self.db.execute(
+            query,
+            (run_id, ic_op.value, config_leaf_str, TaskStatus.SCHEDULED.value),
+            commit=True,
+            fetch=True,
+        )
+        if result:
+            return result[0][0]
+        raise DBException("Failed to create interactive control task")
+
+    def get_scheduled_ic_ops_tasks(self) -> list[dict[str, Any]]:
+        """Get all scheduled interactive control operations"""
+        query = """
+            SELECT task_id, run_id, ic_op, config_leaf
+            FROM interactive_control
+            WHERE status = ?
+        """
+        tasks = self.db.execute(query, (TaskStatus.SCHEDULED.value,), fetch=True)
+        if not tasks:
+            return []
+        return [
+            {
+                "task_id": task[0],
+                "run_id": task[1],
+                "ic_op": ControllerTask(task[2]),
+                "config_leaf": decode_db_payload(task[3]),
+            }
+            for task in tasks
+        ]
+
+    def set_ic_ops_task_status(self, task_id: int, status: TaskStatus) -> None:
+        """Set the status of an interactive control operation"""
+        query = """
+            UPDATE interactive_control
+            SET status = ?
+            WHERE task_id = ?
+        """
+        self.db.execute(query, (status.value, task_id), commit=True)
+
+    # Worker Task Table
+    def create_worker_task(
+        self,
+        worker_id: int,
+        task_type: WorkerTask,
+        status: TaskStatus,
+        run_id: int,
+        chunk_id: int = -1,
+        config_options: dict[str, Any] | None = None,
+    ) -> int:
+        """Create a worker task"""
+
+        query = """
+            INSERT INTO worker_task (worker_id, task_type, status, run_id, chunk_id, config_options)
+            VALUES (?, ?, ?, ?, ?, ?)
+        """
+        config_options_str = encode_payload(config_options) if config_options else "{}"
+        self.db.execute(
+            query,
+            (
+                worker_id,
+                task_type.value,
+                status.value,
+                run_id,
+                chunk_id,
+                config_options_str,
+            ),
+            commit=True,
+        )
+        result = self.db.execute("SELECT last_insert_rowid()", fetch=True)
+        if result:
+            return result[0][0]
+        raise DBException("Failed to create worker task")
+
+    def get_all_worker_tasks(self) -> dict[int, dict[str, Any]]:
+        """Get the latest task of each worker"""
+        query = """
+            SELECT worker_id, task_id, task_type, status, run_id, chunk_id, config_options
+            FROM worker_task wt1
+            WHERE task_id = (
+                SELECT MAX(task_id)
+                FROM worker_task wt2
+                WHERE wt2.worker_id = wt1.worker_id
+            )
+        """
+        task_details = self.db.execute(query, fetch=True)
+
+        formatted_details: dict[int, dict[str, Any]] = {}
+        if task_details:
+            for task in task_details:
+                formatted_details[task[0]] = {
+                    "task_id": task[1],
+                    "task_type": WorkerTask(task[2]),
+                    "status": TaskStatus(task[3]),
+                    "run_id": task[4],
+                    "chunk_id": task[5],
+                    "config_options": decode_db_payload(task[6]) if task[6] and task[6] != "{}" else {},
+                }
+        return formatted_details
+
+    def get_worker_scheduled_task(self, worker_id: int) -> dict[str, Any]:
+        """Get the latest scheduled task for a worker"""
+        query = """
+            SELECT task_id, task_type, run_id, chunk_id, config_options
+            FROM worker_task
+            WHERE worker_id = ? AND status = ?
+            ORDER BY task_id DESC
+            LIMIT 1
+        """
+        task_details = self.db.execute(query, (worker_id, TaskStatus.SCHEDULED.value), fetch=True)
+
+        if task_details:
+            task_details = task_details[0]
+            formatted_details = {
+                "task_id": task_details[0],
+                "task_type": WorkerTask(task_details[1]),
+                "run_id": task_details[2],
+                "chunk_id": task_details[3],
+                "config_options": decode_db_payload(task_details[4])
+                if task_details[4] and task_details[4] != "{}"
+                else {},
+            }
+            return formatted_details
+        return {}
+
+    def set_worker_task_status(self, worker_id: int, status: TaskStatus) -> None:
+        """Set the status of the latest task of a worker"""
+        query = """
+            UPDATE worker_task
+            SET status = ?
+            WHERE task_id = (
+                SELECT task_id
+                FROM worker_task
+                WHERE worker_id = ?
+                ORDER BY task_id DESC
+                LIMIT 1
+            )
+        """
+        self.db.execute(query, (status.value, worker_id), commit=True)
+
+    # Train Controller Progress Table
+    def set_controller_progress(self, run_id: int, progress: float) -> None:
+        """Set the Train progress for a Train Controller"""
+        query = """
+            INSERT INTO controller_progress (run_id, progress)
+            VALUES (?, ?)
+            ON CONFLICT (run_id)
+            DO UPDATE SET
+                progress = EXCLUDED.progress;
+        """
+        progress_rounded = round(progress, 2)
+        self.db.execute(query, (run_id, progress_rounded), commit=True)
+
+    def get_controller_progress(self, run_id: int) -> float:
+        """Get the train progress for a Controller"""
+        query = """
+            SELECT progress
+            FROM controller_progress
+            WHERE run_id = ?
+        """
+        progress_details = self.db.execute(query, (run_id,), fetch=True)
+
+        if progress_details:
+            return progress_details[0][0]
+        return 0.0
+
+    # Train Worker Progress Table
+    def set_worker_progress(self, run_id: int, subchunk_progress: float) -> None:
+        """Set the progress of a Worker for training"""
+        query = """
+            INSERT INTO worker_progress (run_id, subchunk_progress)
+            VALUES (?, ?)
+            ON CONFLICT (run_id)
+            DO UPDATE SET
+                subchunk_progress = EXCLUDED.subchunk_progress;
+        """
+        progress_rounded = round(subchunk_progress, 2)
+        self.db.execute(query, (run_id, progress_rounded), commit=True)
+
+    def get_worker_progress(self, run_id: int) -> float:
+        """Get the progress of a Worker for training"""
+        query = """
+            SELECT subchunk_progress
+            FROM worker_progress
+            WHERE run_id = ?
+        """
+        progress_details = self.db.execute(query, (run_id,), fetch=True)
+
+        if progress_details:
+            return progress_details[0][0]
+        return 0.0