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

rapidfireai/db/tables.sql

@@ -0,0 +1,64 @@
+-- Experiments table
+CREATE TABLE IF NOT EXISTS experiments (
+    experiment_id INTEGER PRIMARY KEY AUTOINCREMENT,
+    experiment_name TEXT NOT NULL,
+    mlflow_experiment_id TEXT NOT NULL,
+    config_options TEXT NOT NULL,
+    status TEXT NOT NULL,
+    current_task TEXT NOT NULL,
+    error TEXT DEFAULT ''
+);
+
+-- Runs table
+CREATE TABLE IF NOT EXISTS runs (
+    run_id INTEGER PRIMARY KEY AUTOINCREMENT,
+    status TEXT NOT NULL,
+    mlflow_run_id TEXT,
+    flattened_config TEXT DEFAULT '{}',
+    config_leaf TEXT DEFAULT '{}',
+    completed_steps INTEGER DEFAULT 0,
+    total_steps INTEGER DEFAULT 0,
+    start_chunk_id INTEGER DEFAULT 0,
+    num_chunks_visited_curr_epoch INTEGER DEFAULT 0,
+    num_epochs_completed INTEGER DEFAULT 0,
+    error TEXT DEFAULT '',
+    source TEXT DEFAULT '',
+    ended_by TEXT DEFAULT '',
+    warm_started_from INTEGER DEFAULT NULL
+);
+
+-- Interactive Control table
+CREATE TABLE IF NOT EXISTS interactive_control (
+    task_id INTEGER PRIMARY KEY AUTOINCREMENT,
+    run_id INTEGER NOT NULL,
+    ic_op TEXT NOT NULL,
+    config_leaf TEXT DEFAULT '{}',
+    status TEXT NOT NULL,
+    FOREIGN KEY (run_id) REFERENCES runs (run_id)
+);
+
+-- Worker Task table
+CREATE TABLE IF NOT EXISTS worker_task (
+    task_id INTEGER PRIMARY KEY AUTOINCREMENT,
+    worker_id INTEGER NOT NULL,
+    task_type TEXT NOT NULL,
+    status TEXT NOT NULL,
+    run_id INTEGER NOT NULL,
+    chunk_id INTEGER NOT NULL,
+    config_options TEXT DEFAULT '{}',
+    FOREIGN KEY (run_id) REFERENCES runs (run_id)
+);
+
+-- Controller Progress table
+CREATE TABLE IF NOT EXISTS controller_progress (
+    run_id INTEGER PRIMARY KEY,
+    progress REAL DEFAULT 0.0,
+    FOREIGN KEY (run_id) REFERENCES runs (run_id)
+);
+
+-- Worker Progress table
+CREATE TABLE IF NOT EXISTS worker_progress (
+    run_id INTEGER PRIMARY KEY,
+    subchunk_progress REAL DEFAULT 0.0,
+    FOREIGN KEY (run_id) REFERENCES runs (run_id)
+);

rapidfireai/dispatcher/dispatcher.py

@@ -0,0 +1,391 @@
+"""This module contains functions for the dispatcher module."""
+
+import traceback
+from logging import Logger
+from typing import Any
+
+from flask import Flask, Response, jsonify, request
+from flask_cors import CORS
+
+from rapidfireai.db.rf_db import RfDb
+from rapidfireai.utils.constants import LOG_FILENAME, ControllerTask, DispatcherConfig
+from rapidfireai.utils.exceptions import DispatcherException
+from rapidfireai.utils.logging import RFLogger
+
+CORS_ALLOWED_ORIGINS = ["http://localhost:3000", "http://localhost"]
+
+
+class Dispatcher:
+    """Class to co-ordinate the flow of tasks between the user and Controllers"""
+
+    def __init__(self) -> None:
+        # initialize loggers
+        self.logger_experiment_name: str | None = None
+        self.logger: Logger | None = None
+
+        # create Db handle
+        self.db: RfDb = RfDb()
+
+        # create Flask app
+        self.app: Flask = Flask(__name__)
+
+        # Enable CORS for all routes
+        _ = CORS(self.app, resources={r"/*": {"origins": CORS_ALLOWED_ORIGINS}})
+
+        # register routes
+        self.register_routes()
+
+    def _get_logger(self) -> Logger | None:
+        """Get the latest logger for the dispatcher"""
+        current_experiment_name = self.db.get_running_experiment()["experiment_name"]
+        if self.logger is None or self.logger_experiment_name != current_experiment_name:
+            self.logger = RFLogger().create_logger("dispatcher")
+            self.logger_experiment_name = current_experiment_name
+        return self.logger
+
+    def register_routes(self) -> None:
+        """Register the routes for the dispatcher"""
+
+        try:
+            route_prefix = "/dispatcher"
+
+            # health check route
+            self.app.add_url_rule(f"{route_prefix}/health-check", "health_check", self.health_check, methods=["GET"])
+
+            # UI routes
+            self.app.add_url_rule(f"{route_prefix}/get-all-runs", "get_all_runs", self.get_all_runs, methods=["GET"])
+            self.app.add_url_rule(f"{route_prefix}/get-run", "get_run", self.get_run, methods=["POST"])
+            self.app.add_url_rule(
+                f"{route_prefix}/get-all-experiment-names",
+                "get_all_experiment_names",
+                self.get_all_experiment_names,
+                methods=["GET"],
+            )
+            self.app.add_url_rule(
+                f"{route_prefix}/get-running-experiment",
+                "get_running_experiment",
+                self.get_running_experiment,
+                methods=["GET"],
+            )
+
+            # Interactive Control routes
+            self.app.add_url_rule(
+                f"{route_prefix}/clone-modify-run", "clone_modify_run", self.clone_modify_run, methods=["POST"]
+            )
+            self.app.add_url_rule(f"{route_prefix}/stop-run", "stop_run", self.stop_run, methods=["POST"])
+            self.app.add_url_rule(f"{route_prefix}/resume-run", "resume_run", self.resume_run, methods=["POST"])
+            self.app.add_url_rule(f"{route_prefix}/delete-run", "delete_run", self.delete_run, methods=["POST"])
+            self.app.add_url_rule(
+                f"{route_prefix}/get-ic-logs", "get_ic_ops_logs", self.get_ic_ops_logs, methods=["POST"]
+            )
+            self.app.add_url_rule(
+                f"{route_prefix}/get-experiment-logs", "get_experiment_logs", self.get_experiment_logs, methods=["POST"]
+            )
+        except Exception as e:
+            raise DispatcherException(f"Error while registering routes: {e}") from e
+
+    # Misc routes
+    def health_check(self) -> tuple[Response, int]:
+        """Health check route"""
+        try:
+            return jsonify("Dispatcher is up and running"), 200
+        except Exception as e:
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+    # UI routes
+    def get_all_runs(self) -> tuple[Response, int]:
+        """Get all the runs for the UI"""
+        try:
+            results = self.db.get_all_runs()
+            safe_results: list[dict[str, Any]] = []
+            for run_id, result in results.items():
+                # remove additional_kwargs from config_leaf
+                result["config_leaf"].pop("additional_kwargs", None)
+
+                # remove peft_params.task_type if it exists
+                if "peft_params" in result["config_leaf"]:
+                    result["config_leaf"]["peft_params"].pop("task_type", None)
+
+                # remove model_kwargs.torch_dtype if it exists
+                if "model_kwargs" in result["config_leaf"]:
+                    result["config_leaf"]["model_kwargs"].pop("torch_dtype", None)
+
+                if "reward_funcs" in result["config_leaf"]:
+                    result["config_leaf"].pop("reward_funcs", None)
+
+                safe_results.append(
+                    {
+                        "run_id": run_id,
+                        "status": result["status"].value,
+                        "mlflow_run_id": result["mlflow_run_id"],
+                        "config": result["config_leaf"],
+                        "flattened_config": result["flattened_config"],
+                        "completed_steps": result["completed_steps"],
+                        "total_steps": result["total_steps"],
+                        "num_epochs_completed": result["num_epochs_completed"],
+                        "error": result["error"],
+                        "source": result["source"].value if result["source"] else None,
+                        "ended_by": result["ended_by"].value if result["ended_by"] else None,
+                    }
+                )
+            return jsonify(safe_results), 200
+        except Exception as e:
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+    def get_run(self) -> tuple[Response, int]:
+        """Get a run for the UI"""
+        try:
+            data = request.get_json()
+            result = self.db.get_run(data["run_id"])
+            if not result:
+                return jsonify({"error": "Run not found"}), 404
+
+            # remove additional_kwargs from config_leaf
+            result["config_leaf"].pop("additional_kwargs", None)
+
+            # remove peft_params.task_type if it exists
+            if "peft_params" in result["config_leaf"]:
+                result["config_leaf"]["peft_params"].pop("task_type", None)
+
+            # remove model_kwargs.torch_dtype if it exists
+            if "model_kwargs" in result["config_leaf"]:
+                result["config_leaf"]["model_kwargs"].pop("torch_dtype", None)
+
+            if "reward_funcs" in result["config_leaf"]:
+                result["config_leaf"].pop("reward_funcs", None)
+
+            safe_result = {
+                "run_id": data["run_id"],
+                "status": result["status"].value,
+                "mlflow_run_id": result["mlflow_run_id"],
+                "config": result["config_leaf"],
+                "flattened_config": result["flattened_config"],
+                "completed_steps": result["completed_steps"],
+                "total_steps": result["total_steps"],
+                "num_epochs_completed": result["num_epochs_completed"],
+                "error": result["error"],
+                "source": result["source"].value if result["source"] else None,
+                "ended_by": result["ended_by"].value if result["ended_by"] else None,
+            }
+            return jsonify(safe_result), 200
+        except Exception as e:
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+    def get_all_experiment_names(self) -> tuple[Response, int]:
+        """Get all the experiment names for the UI"""
+        try:
+            results = self.db.get_all_experiment_names()
+            return jsonify(results), 200
+        except Exception as e:
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+    def get_running_experiment(self) -> tuple[Response, int]:
+        """Get the running experiment for the UI"""
+        try:
+            result = self.db.get_running_experiment()
+            return jsonify(result), 200
+        except Exception as e:
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+    # Interactive Control routes
+    def clone_modify_run(self) -> tuple[Response, int]:
+        """Clone and modify a run"""
+        try:
+            data = request.get_json()
+            if not data:
+                return jsonify({"error": "No JSON data received"}), 400
+
+            run_id = data["run_id"]
+            if not run_id:
+                return jsonify({"error": "run_id is required"}), 400
+
+            if "config" not in data or not data["config"]:
+                return jsonify({"error": "config is required"}), 400
+
+            # validate and parse the ML config text
+            # TODO: Implement validate_config
+            # status = validate_config(data["config"])
+
+            # get the ML config
+            config = data["config"]
+
+            # Validate the ML config text
+            # data["warm_start"]: bool indicating if the run should be warm started
+            task = ControllerTask.IC_CLONE_MODIFY_WARM if data["warm_start"] else ControllerTask.IC_CLONE_MODIFY
+
+            # set create models subtask
+            _ = self.db.create_ic_ops_task(
+                run_id=run_id,
+                ic_op=task,
+                config_leaf=config,
+            )
+
+            # log the task
+            logger = self._get_logger()
+            if logger:
+                logger.info(f"Received clone-modify task with warm start {data['warm_start']} for run_id {run_id}")
+            return jsonify({}), 200
+        except ValueError as ve:
+            logger = self._get_logger()
+            if logger:
+                logger.opt(exception=True).error(f"ValueError in clone_modify_run: {ve}")
+            return jsonify({"error": str(ve)}), 400
+        except TypeError as te:
+            logger = self._get_logger()
+            if logger:
+                logger.opt(exception=True).error(f"TypeError in clone_modify_run: {te}")
+            return jsonify({"error": str(te)}), 400
+        except Exception as e:
+            logger = self._get_logger()
+            if logger:
+                logger.opt(exception=True).error(f"Unexpected error in clone_modify_run: {e}", exc_info=True)
+            return jsonify({"error": "An unexpected error occurred", "details": str(e)}), 500
+
+    def stop_run(self) -> tuple[Response, int]:
+        """Stop the run"""
+        try:
+            data = request.get_json()
+            run_id = data["run_id"]
+            task = ControllerTask.IC_STOP
+
+            # get ml config from db
+            config_leaf = self.db.get_run(run_id)["config_leaf"]
+
+            # create ic ops task
+            _ = self.db.create_ic_ops_task(run_id, task, config_leaf)
+
+            # log the task
+            logger = self._get_logger()
+            if logger:
+                logger.info(f"Received stop task for run_id: {run_id}")
+            return jsonify({}), 200
+        except Exception as e:
+            logger = self._get_logger()
+            if logger:
+                logger.opt(exception=True).error(f"Error in stop_run: {e}", exc_info=True)
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+    def resume_run(self) -> tuple[Response, int]:
+        """Resume the run"""
+        try:
+            data = request.get_json()
+            run_id = data["run_id"]
+            task = ControllerTask.IC_RESUME
+
+            # get ml config from db
+            config_leaf = self.db.get_run(run_id)["config_leaf"]
+
+            # set resume run task
+            _ = self.db.create_ic_ops_task(
+                run_id=run_id,
+                ic_op=task,
+                config_leaf=config_leaf,
+            )
+
+            # log the task
+            logger = self._get_logger()
+            if logger:
+                logger.info(f"Received resume task for run_id: {run_id}")
+            return jsonify({}), 200
+        except Exception as e:
+            logger = self._get_logger()
+            if logger:
+                logger.opt(exception=True).error(f"Error in resume_run: {e}", exc_info=True)
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+    def delete_run(self) -> tuple[Response, int]:
+        """Delete the run"""
+        try:
+            data = request.get_json()
+            run_id = data["run_id"]
+            task = ControllerTask.IC_DELETE
+
+            # get ml config from db
+            config_leaf = self.db.get_run(run_id)["config_leaf"]
+
+            # set delete run task
+            _ = self.db.create_ic_ops_task(
+                run_id=run_id,
+                ic_op=task,
+                config_leaf=config_leaf,
+            )
+
+            # log the task
+            logger = self._get_logger()
+            if logger:
+                logger.info(f"Received delete task for run_id: {run_id}")
+            return jsonify({}), 200
+        except Exception as e:
+            logger = self._get_logger()
+            if logger:
+                logger.opt(exception=True).error(f"Error in delete_run: {e}", exc_info=True)
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+    def get_ic_ops_logs(self) -> tuple[Response, int]:
+        """Get the IC ops logs for the given experiment"""
+        try:
+            experiment_name = None
+            if request.is_json:
+                data = request.get_json()
+                if data and data.get("experiment_name"):
+                    experiment_name = data["experiment_name"]
+
+            if not experiment_name:
+                experiment_name = self.db.get_running_experiment()["experiment_name"]
+
+            log_file_path = self.db.get_experiments_path(experiment_name) / LOG_FILENAME
+
+            # Check if the log file exists
+            if not log_file_path.exists():
+                return jsonify({"error": f"Log file not found for experiment: {experiment_name}"}), 404
+
+            # Read and filter logs for interactive-control entries
+            interactive_control_logs = []
+            with open(log_file_path) as f:
+                for line in f:
+                    if f"| {experiment_name} | interactive-control |" in line:
+                        interactive_control_logs.append(line.strip())
+
+            return jsonify(interactive_control_logs), 200
+        except Exception as e:
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+    def get_experiment_logs(self) -> tuple[Response, int]:
+        """Get all the logs for the given experiment"""
+        # TODO: find a way to optimize this, instead of reading the entire log file, we can read the last N lines
+        try:
+            experiment_name = None
+            if request.is_json:
+                data = request.get_json()
+                if data and data.get("experiment_name"):
+                    experiment_name = data["experiment_name"]
+
+            if not experiment_name:
+                experiment_name = self.db.get_running_experiment()["experiment_name"]
+
+            log_file_path = self.db.get_experiments_path(experiment_name) / LOG_FILENAME
+
+            experiment_logs = []
+            with open(log_file_path) as f:
+                for line in f:
+                    if f"| {experiment_name} |" in line:
+                        experiment_logs.append(line.strip())
+                return jsonify(experiment_logs), 200
+        except Exception as e:
+            return jsonify({"error": str(e) + " " + str(traceback.format_exc())}), 500
+
+
+def serve_forever() -> Flask:
+    """start the Dispatcher via Gunicorn"""
+    return Dispatcher().app
+
+
+if __name__ == "__main__":
+    # initialize the database tables
+    rf_db = RfDb()
+    rf_db.create_tables()
+    print("Database tables initialized successfully")
+
+    # start the Dispatcher on local via Flask
+    dispatcher = Dispatcher()
+    dispatcher.app.run(host=DispatcherConfig.HOST, port=DispatcherConfig.PORT)

rapidfireai/dispatcher/gunicorn.conf.py

@@ -0,0 +1,25 @@
+"""Gunicorn configuration file for the RapidFire dispatcher"""
+
+from rapidfireai.db.rf_db import RfDb
+from rapidfireai.utils.constants import DispatcherConfig
+
+# Other Gunicorn settings...
+bind = f"{DispatcherConfig.HOST}:{DispatcherConfig.PORT}"
+workers = 2
+
+wsgi_app = "rapidfireai.dispatcher.dispatcher:serve_forever()"
+
+
+def on_starting(server):
+    """
+    This function is called once before the master process is initialized.
+    We use it to create tables, ensuring this happens only once.
+    """
+    print("Initializing database tables...")
+    try:
+        rf_db = RfDb()
+        rf_db.create_tables()
+        print("Database tables initialized successfully")
+    except Exception as e:
+        print(f"Error initializing database tables: {e}")
+        raise

rapidfireai/experiment.py

@@ -0,0 +1,168 @@
+"""
+This module contains the Experiment class which manages the entire experiment lifecycle.
+"""
+
+import multiprocessing as mp
+import os
+import traceback
+from typing import Any, Callable
+
+import pandas as pd
+from torch.utils.data import Dataset
+
+from rapidfireai.backend.controller import Controller
+from rapidfireai.db.rf_db import RfDb
+from rapidfireai.utils.constants import MLFLOW_URL
+from rapidfireai.utils.exceptions import ExperimentException
+from rapidfireai.utils.experiment_utils import ExperimentUtils
+from rapidfireai.utils.logging import RFLogger
+from rapidfireai.utils.mlflow_manager import MLflowManager
+
+
+class Experiment:
+    """Class to manage the entire experiment lifecycle."""
+
+    def __init__(
+        self,
+        experiment_name: str,
+        experiments_path: str = os.getenv("RF_EXPERIMENT_PATH", "./rapidfire_experiments"),
+    ) -> None:
+        """
+        Args:
+            experiment_name: The name of the experiment.
+            experiments_path: The base path to the experiments directory.
+        """
+        # initialize experiment variables
+        self.experiment_name: str = experiment_name
+        self.experiment_id: int | None = None
+        self.log_server_process: mp.Process | None = None
+        self.worker_processes: list[mp.Process] = []
+
+        # create db tables
+        try:
+            RfDb().create_tables()
+        except Exception as e:
+            raise ExperimentException(f"Error creating db tables: {e}, traceback: {traceback.format_exc()}") from e
+
+        # create experiment utils object
+        self.experiment_utils = ExperimentUtils()
+
+        # create experiment
+        try:
+            self.experiment_id, self.experiment_name, log_messages = self.experiment_utils.create_experiment(
+                given_name=self.experiment_name,
+                experiments_path=os.path.abspath(experiments_path),
+            )
+        except Exception as e:
+            raise ExperimentException(f"Error creating experiment: {e}, traceback: {traceback.format_exc()}") from e
+
+        # create logger
+        try:
+            self.logger = RFLogger().create_logger("experiment")
+            for msg in log_messages:
+                self.logger.info(msg)
+        except Exception as e:
+            raise ExperimentException(f"Error creating logger: {e}, traceback: {traceback.format_exc()}") from e
+
+        # setup signal handlers for graceful shutdown
+        try:
+            self.experiment_utils.setup_signal_handlers(self.worker_processes)
+        except Exception as e:
+            if hasattr(self, "logger"):
+                self.logger.opt(exception=True).error(f"Error setting up signal handlers: {e}")
+            raise ExperimentException(
+                f"Error setting up signal handlers: {e}, traceback: {traceback.format_exc()}"
+            ) from e
+
+    def run_fit(
+        self,
+        param_config: Any,
+        create_model_fn: Callable,
+        train_dataset: Dataset,
+        eval_dataset: Dataset,
+        num_chunks: int,
+        seed: int = 42,
+    ) -> None:
+        """Run the fit"""
+        try:
+            controller = Controller(self.experiment_id, self.experiment_name)
+            controller.run_fit(param_config, create_model_fn, train_dataset, eval_dataset, num_chunks, seed)
+        except Exception as e:
+            if hasattr(self, "logger"):
+                self.logger.opt(exception=True).error(f"Error running fit: {e}")
+            raise ExperimentException(f"Error running fit: {e}, traceback: {traceback.format_exc()}") from e
+
+    def get_results(self) -> pd.DataFrame:
+        """
+        Get the MLflow training metrics for all runs in the experiment.
+        """
+        try:
+            runs_info_df = self.experiment_utils.get_runs_info()
+            mlflow_manager = MLflowManager(MLFLOW_URL)
+
+            metrics_data = []
+
+            for _, run_row in runs_info_df.iterrows():
+                run_id = run_row["run_id"]
+                mlflow_run_id = run_row.get("mlflow_run_id")
+
+                if not mlflow_run_id:
+                    continue
+
+                run_metrics = mlflow_manager.get_run_metrics(mlflow_run_id)
+
+                step_metrics = {}
+                for metric_name, metric_values in run_metrics.items():
+                    for step, value in metric_values:
+                        if step not in step_metrics:
+                            step_metrics[step] = {"run_id": run_id, "step": step}
+                        step_metrics[step][metric_name] = value
+
+                metrics_data.extend(step_metrics.values())
+
+            if metrics_data:
+                return pd.DataFrame(metrics_data).sort_values(["run_id", "step"])
+            else:
+                return pd.DataFrame(columns=["run_id", "step"])
+
+        except Exception as e:
+            if hasattr(self, "logger"):
+                self.logger.opt(exception=True).error(f"Error getting results: {e}")
+            raise ExperimentException(f"Error getting results: {e}, traceback: {traceback.format_exc()}") from e
+
+    def cancel_current(self) -> None:
+        """Cancel the current task"""
+        try:
+            self.experiment_utils.cancel_current(internal=False)
+        except Exception as e:
+            if hasattr(self, "logger"):
+                self.logger.opt(exception=True).error(f"Error canceling current task: {e}")
+            raise ExperimentException(f"Error canceling current task: {e}, traceback: {traceback.format_exc()}") from e
+
+    def get_runs_info(self) -> pd.DataFrame:
+        """Get the run info"""
+        try:
+            return self.experiment_utils.get_runs_info()
+        except Exception as e:
+            if hasattr(self, "logger"):
+                self.logger.opt(exception=True).error(f"Error getting run info: {e}")
+            raise ExperimentException(f"Error getting run info: {e}, traceback: {traceback.format_exc()}") from e
+
+    def end(self) -> None:
+        """End the experiment"""
+        try:
+            self.experiment_utils.end_experiment(internal=False)
+        except Exception as e:
+            if hasattr(self, "logger"):
+                self.logger.opt(exception=True).error(f"Error ending experiment: {e}")
+            raise ExperimentException(f"Error ending experiment: {e}, traceback: {traceback.format_exc()}") from e
+
+        # shutdown all child processes
+        try:
+            self.experiment_utils.shutdown_workers(self.worker_processes)
+        except Exception as e:
+            if hasattr(self, "logger"):
+                self.logger.opt(exception=True).error(f"Error shutting down RapidFire processes: {e}")
+            raise ExperimentException(
+                f"Error shutting down RapidFire processes: {e}, traceback: {traceback.format_exc()}"
+            ) from e