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

rapidfireai/utils/constants.py

@@ -0,0 +1,141 @@
+from enum import Enum
+
+MLFLOW_URL = "http://127.0.0.1:5002"
+
+# Shared Memory Constants
+SHM_WARN_THRESHOLD = 80
+SHM_MIN_FREE_SPACE = 1.0
+
+# Logging Constants
+LOG_FILENAME = "rapidfire.log"
+TRAINING_LOG_FILENAME = "training.log"
+
+
+class LogType(Enum):
+    """Enum class for log types"""
+
+    RF_LOG = "rf_log"
+    TRAINING_LOG = "training_log"
+
+
+# Dispatcher Constants
+class DispatcherConfig:
+    """Class to manage the dispatcher configuration"""
+
+    HOST: str = "127.0.0.1"
+    PORT: int = 8080
+
+
+# Database Constants
+class DBConfig:
+    """Class to manage the database configuration for SQLite"""
+
+    # Use user's home directory for database path
+    import os
+
+    DB_PATH: str = os.path.expanduser("~/db/rapidfire.db")
+
+    # Connection settings
+    CONNECTION_TIMEOUT: float = 30.0
+
+    # Performance optimizations
+    CACHE_SIZE: int = 10000
+    MMAP_SIZE: int = 268435456  # 256MB
+    PAGE_SIZE: int = 4096
+    BUSY_TIMEOUT: int = 30000
+
+    # Retry settings
+    DEFAULT_MAX_RETRIES: int = 3
+    DEFAULT_BASE_DELAY: float = 0.1
+    DEFAULT_MAX_DELAY: float = 1.0
+
+
+# Experiment Constants
+class ExperimentTask(Enum):
+    """Enum class for experiment tasks"""
+
+    IDLE = "Idle"
+    CREATE_MODELS = "Create Models"
+    IC_OPS = "Interactive Control Operations"
+    RUN_FIT = "Training and Validation"
+
+
+class ExperimentStatus(Enum):
+    """Enum class for experiment status"""
+
+    RUNNING = "Running"
+    COMPLETED = "Completed"
+    FAILED = "Failed"
+    CANCELLED = "Cancelled"
+
+
+# Task Constants
+class TaskStatus(Enum):
+    """Enum class for task status"""
+
+    SCHEDULED = "Scheduled"
+    IN_PROGRESS = "In Progress"
+    COMPLETED = "Completed"
+    SKIPPED = "Skipped"
+    FAILED = "Failed"
+
+
+class ControllerTask(Enum):
+    """Enum class for ML controller tasks"""
+
+    RUN_FIT = "Run Fit"
+    CREATE_MODELS = "Create Models"
+    IC_DELETE = "Interactive Control Delete"
+    IC_STOP = "Interactive Control Stop"
+    IC_RESUME = "Interactive Control Resume"
+    IC_CLONE_MODIFY = "Interactive Control Clone Modify"
+    IC_CLONE_MODIFY_WARM = "Interactive Control Clone Modify with Warm-Start"
+    EPOCH_BOUNDARY = "Epoch Boundary Task"
+    GET_RUN_METRICS = "Get Run Metrics"
+
+
+class WorkerTask(Enum):
+    """Enum class for ML worker tasks"""
+
+    CREATE_MODELS = "Create Models"
+    TRAIN_VAL = "Train Validation"
+
+
+# Run Constants
+class RunStatus(Enum):
+    """Enum class for run status"""
+
+    NEW = "New"
+    ONGOING = "Ongoing"
+    STOPPED = "Stopped"
+    DELETED = "Deleted"
+    COMPLETED = "Completed"
+    FAILED = "Failed"
+
+
+class RunSource(Enum):
+    """Enum class for how a run was created"""
+
+    SHA = "Successive Halving Algorithm"
+    INITIAL = "Initial"
+    INTERACTIVE_CONTROL = "Interactive Control"
+
+
+class RunEndedBy(Enum):
+    """Enum class for how a run was ended"""
+
+    SHA = "Successive Halving Algorithm"
+    EPOCH_COMPLETED = "Epoch Completed"
+    INTERACTIVE_CONTROL = "Interactive Control"
+    TOLERENCE = "Tolerence Threshold Met"
+
+
+# SHM Model Type Constants
+class SHMObjectType(Enum):
+    """Enum class for model types in shared memory"""
+
+    BASE_MODEL = "base_model"
+    FULL_MODEL = "full_model"
+    REF_FULL_MODEL = "ref_full_model"
+    REF_STATE_DICT = "ref_state_dict"
+    CHECKPOINTS = "checkpoints"

rapidfireai/utils/datapaths.py

@@ -0,0 +1,69 @@
+"""
+This module contains the DatasetLocators class which encapsulates the filepaths for the datasets.
+"""
+
+from pathlib import Path
+
+from rapidfireai.utils.exceptions import DataPathException
+
+
+class DataPath:
+    """Class to set the data paths for ML"""
+
+    @classmethod
+    def initialize(cls, experiment_name: str, experiments_path: str) -> None:
+        """Create directories for the ML paths"""
+
+        try:
+            # set standard paths
+            cls.experiments_path: Path = Path(experiments_path) / f"{experiment_name}"
+            cls.user_code_path: Path = cls.experiments_path / "code"
+            cls.mlflow_path: Path = cls.experiments_path / "mlflow"
+
+            # create directories
+            Path.mkdir(cls.experiments_path, parents=True, exist_ok=True)
+            Path.mkdir(cls.user_code_path, parents=True, exist_ok=True)
+            Path.mkdir(cls.mlflow_path, parents=True, exist_ok=True)
+        except (PermissionError, OSError) as e:
+            raise DataPathException(f"Failed to create required DataPaths directories: {e}") from e
+
+    @classmethod
+    def base_run_path(cls, run_id: str | int) -> Path:
+        """Return the work directory path"""
+        return cls.experiments_path / "runs" / f"{run_id}"
+
+    @classmethod
+    def dataset_path(cls) -> Path:
+        """Return the dataset path"""
+        return cls.experiments_path / "datasets.dill"
+
+    @classmethod
+    def work_dir_path(cls, base_run_path: Path) -> Path:
+        """Return the work directory path"""
+        return base_run_path / "work_dir"
+
+    @classmethod
+    def initial_checkpoint_path(cls, base_run_path: Path) -> Path:
+        """Return the initial checkpoint path"""
+        return base_run_path / "checkpoints" / "initial_checkpoint"
+
+    @classmethod
+    def final_checkpoint_path(cls, base_run_path: Path) -> Path:
+        """Return the final checkpoint path"""
+        return base_run_path / "checkpoints" / "final_checkpoint"
+
+    @classmethod
+    def intermediate_checkpoint_path(cls, base_run_path: Path) -> Path:
+        """Return the intermediate checkpoint path"""
+        return base_run_path / "checkpoints" / "intermediate_checkpoints"
+
+    @classmethod
+    def val_metrics_path(cls, base_run_path: Path) -> Path:
+        """Return the validation metrics path"""
+        return cls.work_dir_path(base_run_path) / "val_metrics.csv"
+
+    @classmethod
+    def ref_model_path(cls, base_run_path: Path) -> Path:
+        """Return the reference model path for DPO training"""
+        return base_run_path / "ref_model"
+        

rapidfireai/utils/exceptions.py

@@ -0,0 +1,82 @@
+"""
+This module contains custom exceptions for RapidFire.
+"""
+
+
+class ExperimentException(Exception):
+    """Custom exception for experiment creation"""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
+class DispatcherException(Exception):
+    """Custom exception for dispatcher"""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class DBException(Exception):
+    """Custom exception for database operations"""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class DataPathException(Exception):
+    """Custom exception for data path operations"""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class NoGPUsFoundException(Exception):
+    """Custom exception for no GPUs found"""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class InitializeRunException(Exception):
+    """Custom exception for initialize run"""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class ControllerException(Exception):
+    """Custom exception for controller"""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class WorkerException(Exception):
+    """Custom exception for worker"""
+
+    def __init__(self, message: str, traceback: str = None):
+        self.message = message
+        super().__init__(self.message)
+
+
+class AutoMLException(Exception):
+    """Custom exception for AutoML"""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class InsufficientSharedMemoryException(Exception):
+    """Custom exception for insufficient shared memory"""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)

rapidfireai/utils/experiment_utils.py

@@ -0,0 +1,370 @@
+"""This module contains utility functions for the experiment."""
+
+import multiprocessing as mp
+import os
+import re
+import signal
+import sys
+import warnings
+from typing import Any
+
+import mlflow
+import pandas as pd
+import torch
+from IPython.display import display
+from tqdm import tqdm
+from transformers import logging as transformers_logging
+
+from rapidfireai.db.rf_db import RfDb
+from rapidfireai.utils.constants import MLFLOW_URL, ExperimentStatus, ExperimentTask
+from rapidfireai.utils.datapaths import DataPath
+from rapidfireai.utils.exceptions import DBException, ExperimentException
+from rapidfireai.utils.logging import RFLogger
+from rapidfireai.utils.mlflow_manager import MLflowManager
+
+
+class ExperimentUtils:
+    """Class to contain utility functions for the experiment."""
+
+    def __init__(self) -> None:
+        # initialize database handler
+        self.db = RfDb()
+
+    def _disable_ml_warnings_display(self) -> None:
+        """Disable notebook display"""
+        tqdm.disable = True
+
+        # Suppress the transformers logging
+        os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+        transformers_logging.set_verbosity_error()
+
+        # Suppress the torch warnings
+        torch.set_warn_always(False)
+
+        # Suppress the FutureWarning
+        warnings.filterwarnings("ignore", message=".*torch.cuda.amp.autocast.*")
+        warnings.filterwarnings("ignore", message=".*generation flags are not valid.*")
+        warnings.filterwarnings("ignore", message=".*decoder-only architecture.*")
+        warnings.filterwarnings("ignore", message=".*attention mask is not set.*")
+
+    def setup_signal_handlers(
+        self,
+        worker_processes: list[mp.Process],
+    ) -> None:
+        """Setup signal handlers for graceful shutdown on the main process."""
+
+        def signal_handler(signum, frame):
+            """Handle SIGINT and SIGTERM signals"""
+            signal_name = "SIGINT" if signum == signal.SIGINT else "SIGTERM"
+            print(f"\nReceived {signal_name}, shutting down gracefully...")
+
+            try:
+                # Cancel current task if any
+                self.cancel_current()
+
+                self.shutdown_workers(worker_processes)
+
+                print("Graceful shutdown completed.")
+                sys.exit(0)
+            except Exception as e:
+                print(f"Error during graceful shutdown: {e}")
+                sys.exit(1)
+
+        # Register signal handlers
+        signal.signal(signal.SIGINT, signal_handler)
+        signal.signal(signal.SIGTERM, signal_handler)
+
+    def create_experiment(self, given_name: str, experiments_path: str) -> tuple[int, str, list[str]]:
+        """Create a new experiment. Returns the experiment id, name, and log messages."""
+        log_messages: list[str] = []
+
+        # disable warnings from notebook
+        self._disable_ml_warnings_display()
+
+        # Clear any existing MLflow context before starting new experiment
+        try:
+            if mlflow.active_run():
+                print("Clearing existing MLflow context before starting new experiment")
+                mlflow.end_run()
+        except Exception as e:
+            print(f"Error clearing existing MLflow context: {e}")
+
+        # check if experiment is already running
+        running_experiment = None
+        try:
+            running_experiment = self.db.get_running_experiment()
+        except DBException:
+            pass
+        if running_experiment:
+            # check if the running experiment is the same as the new experiment
+            if given_name == running_experiment["experiment_name"]:
+                msg = (
+                    f"Experiment {running_experiment['experiment_name']} is currently running."
+                    " Returning the same experiment object."
+                )
+                print(msg)
+                log_messages.append(msg)
+
+                # check if the running experiment is also running a task
+                current_task = self.db.get_experiment_current_task()
+                if current_task != ExperimentTask.IDLE:
+                    msg = f"Task {current_task.value} that was running has been cancelled."
+                    print(msg)
+                    log_messages.append(msg)
+                self.cancel_current(internal=True)
+
+                # get experiment id
+                experiment_id, experiment_name = (
+                    running_experiment["experiment_id"],
+                    running_experiment["experiment_name"],
+                )
+            else:
+                self.end_experiment(internal=True)
+                experiment_id, experiment_name, mlflow_experiment_id = self._create_experiment_internal(
+                    given_name,
+                    experiments_path,
+                )
+                msg = (
+                    f"The previously running experiment {running_experiment['experiment_name']} was forcibly ended."
+                    f" Created a new experiment with name '{experiment_name}' with Experiment ID: {experiment_id}"
+                    f" and MLFlow Experiment ID: {mlflow_experiment_id} saved at {experiments_path}/{experiment_name}"
+                )
+                print(msg)
+                log_messages.append(msg)
+        # check if experiment name already exists
+        elif given_name in self.db.get_all_experiment_names():
+            experiment_id, experiment_name, mlflow_experiment_id = self._create_experiment_internal(
+                given_name,
+                experiments_path,
+            )
+            msg = (
+                "An experiment with the same name already exists."
+                f" Created a new experiment with name '{experiment_name}' with Experiment ID: {experiment_id}"
+                f" and MLFlow Experiment ID: {mlflow_experiment_id} saved at {experiments_path}/{experiment_name}"
+            )
+            print(msg)
+            log_messages.append(msg)
+        else:
+            experiment_id, experiment_name, mlflow_experiment_id = self._create_experiment_internal(
+                given_name,
+                experiments_path,
+            )
+            msg = (
+                f"Experiment {experiment_name} created with Experiment ID: {experiment_id}"
+                f" and MLFlow Experiment ID: {mlflow_experiment_id} saved at {experiments_path}/{experiment_name}"
+            )
+            print(msg)
+            log_messages.append(msg)
+
+        # initialize the data paths and create directories
+        try:
+            DataPath.initialize(experiment_name, experiments_path)
+        except Exception as e:
+            raise ExperimentException(f"Failed to initialize data paths: {e}")
+
+        return experiment_id, experiment_name, log_messages
+
+    def end_experiment(self, internal: bool = False) -> None:
+        """End the experiment"""
+        # check if experiment is running
+        try:
+            current_experiment = self.db.get_running_experiment()
+        except DBException:
+            if not internal:
+                print("No experiment is currently running. Nothing to end.")
+            return
+
+        # create logger
+        experiment_name = current_experiment["experiment_name"]
+        logger = RFLogger().create_logger(experiment_name)
+
+        # cancel current task if there's any
+        self.cancel_current(internal=True)
+
+        # reset DB states
+        self.db.set_experiment_status(current_experiment["experiment_id"], ExperimentStatus.COMPLETED)
+        self.db.reset_all_tables()
+
+        # Clear MLflow context
+        try:
+            if mlflow.active_run():
+                print("Ending active MLflow run before ending experiment")
+                mlflow.end_run()
+
+            # Also clear context through MLflowManager if available
+            try:
+                mlflow_manager = MLflowManager(MLFLOW_URL)
+                mlflow_manager.clear_context()
+            except Exception as e2:
+                print(f"[Error clearing MLflow context through MLflowManager: {e2}")
+        except Exception as e:
+            print(f"Error clearing MLflow context: {e}")
+
+        # print experiment ended message
+        msg = f"Experiment {experiment_name} ended"
+        if not internal:
+            print(msg)
+        logger.info(msg)
+
+    def cancel_current(self, internal: bool = False) -> None:
+        """Cancel the current task"""
+        # check if experiment is running
+        try:
+            current_experiment = self.db.get_running_experiment()
+        except DBException:
+            if not internal:
+                print("No experiment is currently running. Nothing to cancel.")
+            return
+
+        # create logger
+        logger = RFLogger().create_logger(current_experiment["experiment_name"])
+
+        try:
+            current_task = self.db.get_experiment_current_task()
+        except DBException:
+            msg = "No task is currently running. Nothing to cancel."
+            if not internal:
+                print(msg)
+            logger.info(msg)
+            return
+
+        # reset experiment states and set current task to idle
+        self.db.reset_experiment_states()
+        self.db.set_experiment_current_task(ExperimentTask.IDLE)
+        if current_task != ExperimentTask.IDLE:
+            msg = f"Task {current_task.value} cancelled"
+            print(msg)
+            logger.info(msg)
+        logger.debug("Reset experiment states and set current experiment task to idle")
+
+    def shutdown_workers(
+        self,
+        worker_processes: list[mp.Process],
+    ) -> None:
+        """Shutdown the workers"""
+        # stop workers
+        for worker_process in worker_processes:
+            worker_process.terminate()
+        print("Workers stopped")
+
+    def get_runs_info(self) -> pd.DataFrame:
+        """Get the run info"""
+        try:
+            runs = self.db.get_all_runs()
+            runs_info = {}
+            for run_id, run_details in runs.items():
+                new_run_details = {k: v for k, v in run_details.items() if k not in ("flattened_config", "config_leaf")}
+                if "config_leaf" in run_details:
+                    config_leaf = run_details["config_leaf"].copy()
+                    config_leaf.pop("additional_kwargs", None)
+                    new_run_details["config"] = config_leaf
+
+                runs_info[run_id] = new_run_details
+
+            if runs_info:
+                df = pd.DataFrame.from_dict(runs_info, orient="index")
+                df = df.reset_index().rename(columns={"index": "run_id"})
+                cols = ["run_id"] + [col for col in df.columns if col != "run_id"]
+                df = df[cols]
+                return df
+            else:
+                return pd.DataFrame(columns=["run_id"])
+                
+        except DBException as e:
+            raise ExperimentException("Error getting runs info") from e
+
+    def _display_runs_info(self, runs_info: dict[int, dict[str, Any]]) -> pd.DataFrame:
+        """Fetch runs info, display as a pandas DataFrame, and return the DataFrame.
+
+        Returns:
+            pd.DataFrame: A DataFrame containing all runs information with run_id as first column.
+        """
+        try:
+            # Convert the runs info to a pandas DataFrame
+            df = pd.DataFrame.from_dict(runs_info, orient="index")
+
+            # Reset index to make run_id a regular column for better display
+            df = df.reset_index().rename(columns={"index": "run_id"})
+
+            # Reorder columns to put run_id first
+            cols = ["run_id"] + [col for col in df.columns if col != "run_id"]
+            df = df[cols]
+
+            # Set pandas display options for better notebook viewing
+            pd.set_option("display.max_columns", None)
+            pd.set_option("display.max_rows", None)
+            pd.set_option("display.width", None)
+            pd.set_option("display.max_colwidth", 50)
+
+            # Display the results
+            print(f"Total runs: {len(df)}")
+            print("\n" + "=" * 80)
+
+            try:
+                display(df)  # For notebook environments
+            except NameError:
+                print(df.to_string())  # Fallback for non-notebook environments
+
+            return df
+
+        except ExperimentException as e:
+            print(f"Error displaying runs info: {e}")
+            raise
+
+    def _create_experiment_internal(self, given_name: str, experiments_path: str) -> tuple[int, str, str]:
+        """Create new experiment -
+        if given_name already exists - increment suffix and create new experiment
+        if given_name is new - create new experiment with given name
+        """
+        try:
+            given_name = given_name if given_name else "rf-exp"
+            experiment_name = self._generate_unique_experiment_name(given_name, self.db.get_all_experiment_names())
+
+            mlflow_manager = MLflowManager(MLFLOW_URL)
+            mlflow_experiment_id = mlflow_manager.create_experiment(experiment_name)
+            mlflow.tracing.disable_notebook_display()
+
+            # write new experiment details to database
+            experiment_id = self.db.create_experiment(
+                experiment_name,
+                mlflow_experiment_id,
+                config_options={"experiments_path": experiments_path},
+            )
+            return experiment_id, experiment_name, mlflow_experiment_id
+        except mlflow.exceptions.RestException as e:  # pyright: ignore
+            raise ExperimentException(f"Error creating MLFlow experiment: {e}") from e
+
+    def _generate_unique_experiment_name(self, name: str, existing_names: list[str]) -> str:
+        """Increment the suffix of the name after the last '_' till it is unique"""
+        if not name:
+            name = "rf-exp"
+
+        pattern = r"^(.+?)(_(\d+))?$"
+        max_attempts = 1000  # Prevent infinite loops
+        attempts = 0
+
+        new_name = name
+        while new_name in existing_names and attempts < max_attempts:
+            match = re.match(pattern, new_name)
+
+            if match:
+                base_name = match.group(1)
+                current_suffix = match.group(3)
+                if current_suffix:
+                    try:
+                        new_suffix = int(current_suffix) + 1
+                    except ValueError:
+                        # If suffix is not a valid integer, start from 1
+                        new_suffix = 1
+                else:
+                    new_suffix = 1
+                new_name = f"{base_name}_{new_suffix}"
+            else:
+                new_name = f"{new_name}_1"
+
+            attempts += 1
+
+        if attempts >= max_attempts:
+            raise ExperimentException("Could not generate unique experiment name")
+
+        return new_name